{"record":{"id":"8299c16c3ee3e460","repo":"jax-ml/jax","slug":"block-size-is-not-a-multiple-of-scope-size","errorCode":null,"errorMessage":"Block size is not a multiple of {scope_size}","messagePattern":"Block size is not a multiple of (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/profiler.py","lineNumber":217,"sourceCode":"      )\n    else:\n      self.dump_path = dump_path\n    if trace_scope not in (ThreadSubset.WARP, ThreadSubset.WARPGROUP):\n      raise ValueError(f\"Unsupported trace scope: {trace_scope}\")\n    self.trace_scope = trace_scope\n\n  def _num_traces(\n      self, grid: tuple[int, ...], block: tuple[int, ...]\n  ) -> int:\n    if self.trace_scope == ThreadSubset.WARP:\n      scope_size = WARP_SIZE\n    elif self.trace_scope == ThreadSubset.WARPGROUP:\n      scope_size = WARPGROUP_SIZE\n    else:\n      raise NotImplementedError(f\"Scope {self.trace_scope} not supported\")\n\n    if math.prod(block) % scope_size:\n      raise ValueError(f\"Block size is not a multiple of {scope_size}\")\n    return math.prod(grid) * math.prod(block) // scope_size\n\n  def mlir_buffer_type(\n      self, grid: tuple[int, ...], block: tuple[int, ...]\n  ) -> ir.MemRefType:\n    return ir.MemRefType.get(\n        (self._num_traces(grid, block) * self.entries_per_warpgroup,),\n        ir.IntegerType.get_signless(32),\n    )\n\n  def jax_buffer_type(\n      self, grid: tuple[int, ...], block: tuple[int, ...]\n  ) -> jax.ShapeDtypeStruct:\n    return jax.ShapeDtypeStruct(\n        (self._num_traces(grid, block) * self.entries_per_warpgroup,),\n        jnp.uint32,\n    )\n","sourceCodeStart":199,"sourceCodeEnd":235,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/profiler.py#L199-L235","documentation":"The trace buffer allocates one trace region per WARP (32 threads) or WARPGROUP (128 threads), so the CUDA block dimensions must be an exact multiple of that scope size. _num_traces raises when math.prod(block) % scope_size != 0.","triggerScenarios":"Launching a profiled kernel with a block like (7,1,1) with WARP scope, or block size 130 with WARPGROUP scope (130 % 128 != 0).","commonSituations":"Hand-tuned launch configs; templated grid/block sizes not checked against 32/128 alignment.","solutions":["Round the block size up to a multiple of 32 (WARP) or 128 (WARPGROUP)","Verify math.prod(block) % 32 == 0 (or % 128) before launching"],"exampleFix":"# before\nspec = ProfilerSpec(trace_scope=ThreadSubset.WARPGROUP)\nblock = (100, 1, 1)\n\n# after\nblock = (128, 1, 1)","handlingStrategy":"validation","validationCode":"import math\nscope_size = 32 if spec.trace_scope == ThreadSubset.WARP else 128\nassert math.prod(block) % scope_size == 0, f'block must be multiple of {scope_size}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always size CUDA blocks as multiples of 128 to satisfy both scopes","Add a launch-config sanity check before profiled launches"],"tags":["profiling","mosaic","gpu","cuda-launch-config"],"backgroundTag":"cuda-block-size-alignment","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}