jax-ml/jax · error · ValueError

N must be a multiple of 8 and <= 256, got: {n}

Error message

N must be a multiple of 8 and <= 256, got: {n}

What it means

The tcgen05 MMA instruction descriptor packs N into bits 17-22 as n>>3, requiring N to be a multiple of 8 and at most 256. Violating either makes the descriptor unrepresentable, so create_instr_descriptor raises ValueError.

Source

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

      assert acc_dtype in {f16, f32}
      return 0
    elif ty == ir.Float8E5M2Type.get():
      assert acc_dtype in {f16, f32}
      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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad N up to the next multiple of 8 (typically 64/128/256)
2. Choose tile sizes from supported set: 8, 16, ..., 256
3. Validate shapes before building the descriptor

Example fix

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

# after
mma(acc, a, b, m=128, n=104)  # pad to multiple of 8

assert n % 8 == 0 and n <= 256, f'N={n} must be multiple of 8 and <= 256'

When it happens

Trigger: Calling the tcgen05 mma path with N like 100 (not multiple of 8) or 264 (> 256).

Common situations: Using arbitrary matrix shapes from a model config; padding shapes to 128/256 everywhere except N.

Related errors

• M must be a multiple of 16 and <= 256, got: {m}
• 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/59b3f80b416686ce. Report an issue: GitHub.