jax-ml/jax · error · ValueError

Block size is not a multiple of {scope_size}

Error message

Block size is not a multiple of {scope_size}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/profiler.py:217

      )
    else:
      self.dump_path = dump_path
    if trace_scope not in (ThreadSubset.WARP, ThreadSubset.WARPGROUP):
      raise ValueError(f"Unsupported trace scope: {trace_scope}")
    self.trace_scope = trace_scope

  def _num_traces(
      self, grid: tuple[int, ...], block: tuple[int, ...]
  ) -> int:
    if self.trace_scope == ThreadSubset.WARP:
      scope_size = WARP_SIZE
    elif self.trace_scope == ThreadSubset.WARPGROUP:
      scope_size = WARPGROUP_SIZE
    else:
      raise NotImplementedError(f"Scope {self.trace_scope} not supported")

    if math.prod(block) % scope_size:
      raise ValueError(f"Block size is not a multiple of {scope_size}")
    return math.prod(grid) * math.prod(block) // scope_size

  def mlir_buffer_type(
      self, grid: tuple[int, ...], block: tuple[int, ...]
  ) -> ir.MemRefType:
    return ir.MemRefType.get(
        (self._num_traces(grid, block) * self.entries_per_warpgroup,),
        ir.IntegerType.get_signless(32),
    )

  def jax_buffer_type(
      self, grid: tuple[int, ...], block: tuple[int, ...]
  ) -> jax.ShapeDtypeStruct:
    return jax.ShapeDtypeStruct(
        (self._num_traces(grid, block) * self.entries_per_warpgroup,),
        jnp.uint32,
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Round the block size up to a multiple of 32 (WARP) or 128 (WARPGROUP)
2. Verify math.prod(block) % 32 == 0 (or % 128) before launching

Example fix

# before
spec = ProfilerSpec(trace_scope=ThreadSubset.WARPGROUP)
block = (100, 1, 1)

# after
block = (128, 1, 1)

import math
scope_size = 32 if spec.trace_scope == ThreadSubset.WARP else 128
assert math.prod(block) % scope_size == 0, f'block must be multiple of {scope_size}'

When it happens

Trigger: Launching a profiled kernel with a block like (7,1,1) with WARP scope, or block size 130 with WARPGROUP scope (130 % 128 != 0).

Common situations: Hand-tuned launch configs; templated grid/block sizes not checked against 32/128 alignment.

Related errors

• {iterations=} must be positive
• Unsupported trace scope: {trace_scope}
• Scope {self.trace_scope} not supported
• Allocated too many names
• Insufficient space to capture a full trace

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/8299c16c3ee3e460. Report an issue: GitHub.