jax-ml/jax · error · NotImplementedError

Unsupported accumulator dtype: {acc_dtype}

Error message

Unsupported accumulator dtype: {acc_dtype}

What it means

create_instr_descriptor encodes the MMA accumulator (D) dtype into a 2-bit hardware instruction descriptor field. Only f16 (0), f32 (1) and i32 (2) have encodings; anything else (e.g. f32[i8-packed], bf16) raises NotImplementedError.

Source

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

    sparsity_selector: int | None = None,
) -> ir.Value:
  f16 = ir.F16Type.get()
  f32 = ir.F32Type.get()
  i32 = ir.IntegerType.get_signless(32)

  desc = 0
  if sparsity_selector is not None:
    assert 0 <= sparsity_selector < 3
    desc |= sparsity_selector
    desc |= 1 << 2  # Enable sparsity
  if acc_dtype == f16:
    d_type_val = 0
  elif acc_dtype == f32:
    d_type_val = 1
  elif acc_dtype == i32:
    d_type_val = 2
  else:
    raise NotImplementedError(f"Unsupported accumulator dtype: {acc_dtype}")
  desc |= (d_type_val << 4)  # D type, bits 4-5
  # Bit 6 is reserved
  def get_input_encoding(ty):
    if ty == f16:
      assert acc_dtype in {f16, f32}
      return 0
    elif ty == ir.BF16Type.get():
      assert acc_dtype == f32
      return 1
    elif ty == ir.Float8E4M3FNType.get():
      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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use f32 (most common), f16, or i32 accumulators
2. For integer matmuls, ensure the accumulator is i32
3. Check the JAX version — newer releases may add encodings

Example fix

# before
acc = tmem.alloc(dtype=ir.BFloat16Type.get(), shape=...)

# after
acc = tmem.alloc(dtype=ir.F32Type.get(), shape=...)

supported = {ir.F16Type.get(), ir.F32Type.get(), ir.IntegerType.get_signless(32)}
assert acc_dtype in supported, f'unsupported accumulator {acc_dtype}'

def is_supported_acc_dtype(dt) -> bool:
    return dt in (ir.F16Type.get(), ir.F32Type.get(), ir.IntegerType.get_signless(32))

When it happens

Trigger: Calling mma()/tcgen05 matmul helpers with an accumulator memref of dtype other than f16/f32/i32, e.g. an f32 tensor-core accumulate replaced by bf16.

Common situations: Porting Hopper/Warpgroup matmul code to tcgen05 with non-standard accumulator dtypes; using packed f16x2 accumulators before support landed.

Related errors

• Unsupported input dtype: {ty}
• N must be a multiple of 8 and <= 256, got: {n}
• M must be a multiple of 16 and <= 256, got: {m}
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block

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