jax-ml/jax · error · ValueError

M must be a multiple of 16 and <= 256, got: {m}

Error message

M must be a multiple of 16 and <= 256, got: {m}

What it means

M is packed into descriptor bits 24-28 as m>>4, so M must be a multiple of 16 and no larger than 256. create_instr_descriptor rejects other M values with ValueError.

Source

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

      return 1
    elif ty == ir.IntegerType.get_signless(8):  # Only s8 for now.
      assert acc_dtype == i32
      return 1
    else:
      raise NotImplementedError(f"Unsupported input dtype: {ty}")
  a_type_val = get_input_encoding(a_dtype)
  b_type_val = get_input_encoding(b_dtype)
  desc |= (a_type_val << 7)   # A dtype, bits 7-9
  desc |= (b_type_val << 10)  # B dtype, bits 10-12
  # We ignore negate bits 13-14
  desc |= transpose_a << 15  # Transpose A
  desc |= transpose_b << 16  # Transpose B
  if n % 8 or n > 256:
    raise ValueError(f"N must be a multiple of 8 and <= 256, got: {n}")
  desc |= (n >> 3) << 17  # N, bits 17-22
  # Bit 23 is reserved
  if m % 16 or m > 256:
    raise ValueError(f"M must be a multiple of 16 and <= 256, got: {m}")
  desc |= (m >> 4) << 24  # M >> 4, bits 24-28
  # Bit 29 is reserved
  # We ignore max shift under .ws, bits 30-31
  return arith.constant(ir.IntegerType.get_signless(32), desc)


def _create_scaled_instr_descriptor(
    get_input_encoding: Callable[[ir.Type], int],
    m: int,
    n: int,
    a_type: ir.Type,
    b_type: ir.Type,
    a_scale_idx: int,
    b_scale_idx: int,
    transpose_a: bool,
    transpose_b: bool,
    scale_type: ir.Type,
    sparse: bool = False,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use M in {16, 32, ..., 256}
2. Split larger M across multiple MMA instructions/tiles

Example fix

# before
mma(acc, a, b, m=264, n=128)

# after
# split into two m=132? no — use m=128 twice
mma(acc1, a1, b, m=128, n=128)
mma(acc2, a2, b, m=128, n=128)

assert m % 16 == 0 and m <= 256, f'M={m} must be multiple of 16 and <= 256'

When it happens

Trigger: Calling tcgen05 mma with M=100 (not multiple of 16) or M=512 (> 256).

Common situations: Tile-size tuning sweeps that step M by 8; assuming M can be as large as N.

Related errors

• N must be a multiple of 8 and <= 256, got: {n}
• Unsupported accumulator dtype: {acc_dtype}
• Unsupported input dtype: {ty}
• Only M=128 and M=64 are supported for MMA, but got M={m}
• Expected both or neither of scales to be specified.

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