jax-ml/jax · error · NotImplementedError

Sparse MMA not supported for M=64

Error message

Sparse MMA not supported for M=64

What it means

Raised by Mosaic GPU's tcgen05 mma op when a sparse MMA is requested with M=64. The underlying tcgen05 sparse tensor-core instruction only exists for the M=128 shape, so the library explicitly rejects the combination rather than emitting invalid MLIR.

Source

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

        "MMA requires A and B to have the same element type, except that "
        "FP8 types (f8E4M3FN and f8E5M2) may be mixed; got: "
        f"{a_element_type} and {b_element_type}"
    )
  if d.shape != (m, n * num_cta):
    raise ValueError(
        f"Accumulator shape mismatch: expected {(m, n * num_cta)}, got {d.shape}"
    )
  if m == 128:
    if d.layout != (expected_d_layout := tmem_default_layout(packing=1)):
      raise ValueError(
          f"Accumulator layout mismatch: expected {expected_d_layout}, got {d.layout}"
      )
    n_lane_groups = 1
  elif m == 64:
    if is_scaled and not collective:
      raise NotImplementedError("MMA with block scaling is not supported for 1CTA M=64")
    if is_sparse:
      raise NotImplementedError("Sparse MMA not supported for M=64")
    # Watch out: this layout must be consistent with A's layout (up to packing).
    # 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."
        )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change the M dimension of the A operand/accumulator to 128 (pad the tile if necessary)
2. If M must stay 64, drop the sparse path and use a dense MMA instead
3. Split your workload so sparse MMAs are issued only on 128-row tiles and dense MMAs handle the remainder

Example fix

# before
acc = tcgen05.mma(lhs_64, rhs, acc, sparse_metadata=meta)  # M=64 -> raises
# after
lhs_128 = pad_tile_to_m128(lhs_64)
acc = tcgen05.mma(lhs_128, rhs, acc, sparse_metadata=meta)

m = lhs.shape[0]
assert not (is_sparse and m == 64), 'Sparse tcgen05 MMA requires M=128'

When it happens

Trigger: Calling jax.experimental.mosaic.gpu.tcgen05.mma with a lhs whose leading (M) dimension is 64 while passing a sparse metadata / is_sparse path (e.g. building a sparse warp-level MMA pipeline on Blackwell).

Common situations: Porting a sparse kernel originally written for M=128 tiles to smaller 64-wide tiles; reusing a sparse GEMM template and changing only the tile size; targeting sm_100/sm_103 with truncated M.

Related errors

• Sparse MMA unsupported for f32
• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• Only M=128 and M=64 are supported for MMA, but got M={m}
• 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/fbf99367adea7c51. Report an issue: GitHub.