jax-ml/jax · error · ValueError

Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri

Error message

Swizzle={b_swizzle} is too big for MMA with M=64. Try lowering it.

What it means

For 1CTA M=64 MMAs the N dimension is split into 2 lane groups, so each group must still hold a full tile column. The check 8 * b_swizzle // bitwidth(a) > n // 2 fires when the B operand's swizzle atom is wider than the per-group N slice.

Source

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

    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."
        )
  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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Lower b_swizzle (e.g. from 128 to 64 or 32 bytes) on the B operand
Increase n so that n // 2 >= 8 * b_swizzle // bitwidth(a_element_type)
Use M=128 (possibly collective) where the N-split constraint differs

Example fix

# before
b_tiled = ... swizzle=128 ...  # with m=64, small n -> raises
tcgen05.mma(a, b_tiled, d)
# after
b_tiled = ... swizzle=32 ...
tcgen05.mma(a, b_tiled, d)

Defensive patterns

Strategy: validation

Validate before calling

if m == 64 and not collective:
    max_swizzle = n * 2 * utils.bitwidth(a_dtype) // 8
    assert b_swizzle <= max_swizzle, f'b_swizzle={b_swizzle} too large for M=64, n={n}'

Prevention

Parameterize swizzle by (m, n, dtype) rather than hardcoding 128
Validate swizzle against tile shape in kernel config helpers

When it happens

Trigger: Calling tcgen05.mma with m=64 (non-collective), a small n (e.g. n < 2 * 8 * b_swizzle / bitwidth), and a large b_swizzle such as 128-byte swizzling on fp16 with narrow N.

Common situations: Tuning swizzle values copied from an M=128 kernel; using 128B swizzle with fp8/fp16 operands and narrow N tiles; forgetting that M=64 halves the available N per lane group.

Related errors

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