jax-ml/jax · error · ValueError

Only M=128 and M=64 are supported for MMA, but got M={m}

Error message

Only M=128 and M=64 are supported for MMA, but got M={m}

What it means

The Blackwell tcgen05 MMA instruction only supports M=128 and M=64 operand shapes. Any other leading dimension of the A operand is rejected up front with this ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:281

    # 2CTA M=128 instruction uses a different TMEM layout than 1CTA M=64.
    expected_d_layout = _infer_tmem_layout(d.shape, collective, packing=1)
    if d.layout != expected_d_layout:
      raise ValueError(
          f"Accumulator layout mismatch: expected {expected_d_layout}, got {d.layout}"
      )
    if collective:
      n_lane_groups = 1
    else:
      n_lane_groups = 2
      # We can't split N into groups if we would partition it below the tile size.
      # TODO: We only need to check this if N is the minormost dim in B.
      if 8 * b_swizzle // utils.bitwidth(a_element_type) > n // n_lane_groups:
        raise ValueError(
            f"Swizzle={b_swizzle} is too big for MMA with M=64. Try"
            " lowering it."
        )
  else:
    raise ValueError(f"Only M=128 and M=64 are supported for MMA, but got M={m}")
  f32 = ir.F32Type.get()
  f16 = ir.F16Type.get()
  s32 = ir.IntegerType.get_signless(32)
  elem_type_str = (
      f"{a_element_type}"
      if a_element_type == b_element_type
      else f"({a_element_type}, {b_element_type})"
  )
  if a_element_type == f32 or a_element_type == ir.BF16Type.get():
    if a_element_type == f32 and is_sparse:
      raise NotImplementedError("Sparse MMA unsupported for f32")
    if is_scaled:
      raise ValueError(
          f"MMA with element type {elem_type_str} does not support block scaling"
      )
    if d.dtype != f32:
      raise ValueError(
          f"MMA with element type {elem_type_str} only supports accumulators"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set the A operand's leading dimension to 128 or 64 (pad the tile if the natural M differs)
2. If M is dynamic, assert or clamp it to 128/64 before building the MMA
3. Check for slicing mistakes that produced an unexpected M

Example fix

# before
m = lhs.shape[0]  # e.g. 32
tcgen05.mma(lhs, rhs, d)
# after
assert lhs.shape[0] in (64, 128), f"unsupported M={lhs.shape[0]}"
tcgen05.mma(lhs, rhs, d)

assert lhs.shape[0] in (64, 128), f'M must be 64 or 128, got {lhs.shape[0]}'

When it happens

Trigger: Calling tcgen05.mma where lhs's first dimension is anything other than 128 or 64 (e.g. 256, 32, or a symbolic/dynamic M).

Common situations: Using tile sizes inherited from Hopper wgmma kernels (e.g. M=64 works but M=32 or M=256 does not); computing M from a dynamic batch dimension that resolves to an unsupported value; off-by-one/slicing bugs that shrink M.

Related errors

• Sparse MMA not supported for M=64
• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block
• MMA with element type {elem_type_str} only supports accumula

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