jax-ml/jax · error · ValueError

WGMMA requires m and n to be multiples of 64 and 8, got {m}

Error message

WGMMA requires m and n to be multiples of 64 and 8, got {m} and {n}

What it means

WGMMAAccumulator.zero (wgmma.py:69) validates WGMMA tile geometry: the Hopper wgmma.mma_async instruction requires the M dimension to be a multiple of 64 and N a multiple of 8. Non-conforming shapes cannot be executed by the tensor cores.

Source

Thrown at jax/experimental/mosaic/gpu/wgmma.py:69

      self,
      *,
      _value: fa.FragmentedArray,
      _original_layout: fa.FragmentedLayout,
      _sync: bool = True,
  ):
    self._original_layout = _original_layout
    self._value = _value
    if _sync:
      self._value = wgmma_fence(_value)

  @property
  def value(self) -> fa.FragmentedArray:
    return self._value.to_layout(self._original_layout)

  @classmethod
  def zero(cls, m, n, dtype=None, *, is_signed: bool | None = None):
    if m % 64 or n % 8:
      raise ValueError("WGMMA requires m and n to be multiples of 64 and 8, "
                       f"got {m} and {n}")
    if is_signed is False:
      raise TypeError("PTX does not support unsigned WGMMA accumulators")
    f32 = ir.F32Type.get()
    if dtype is None:
      dtype = f32
    if isinstance(dtype, ir.IntegerType):
      zero = arith.constant(dtype, ir.IntegerAttr.get(dtype, 0))
    else:
      zero = arith.constant(dtype, ir.FloatAttr.get(dtype, 0.0))
    return cls.from_registers(
        fa.FragmentedArray.splat(
            zero, (m, n), fa.WGMMA_LAYOUT, is_signed=is_signed
        )
    )

  @classmethod
  def from_registers(cls, registers, sync=True):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Round m up to a multiple of 64 and n up to a multiple of 8 (padding as needed)
2. Restructure the kernel loop so each wgmma call operates on conformant tiles
3. Add a shape assertion early so the failure is caught at kernel-definition time

Example fix

# before
acc = wgmma.WGMMAAccumulator.zero(m=48, n=128)
# after
m = (m + 63) // 64 * 64
n = (n + 7) // 8 * 8
acc = wgmma.WGMMAAccumulator.zero(m=m, n=n)

assert m % 64 == 0 and n % 8 == 0, f'WGMMA tiles must satisfy m%64==0, n%8==0 (got {m}, {n})'

When it happens

Trigger: Calling wgmma.WGMMAAccumulator.zero(m, n) with m not divisible by 64 or n not divisible by 8 (e.g. zero(64, 12) or zero(32, 16)).

Common situations: Padding attention/GEMM tiles incorrectly; deriving m/n from head_dim or batch sizes without rounding up; porting TMA layouts whose tile shape is not WGMMA-conformant.

Related errors

• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Accumulator shape {inner.shape} does not match value shape {
• Accumulator dtype {inner.dtype} does not match value dtype {
• Unsupported bitwidth: {bitwidth}

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