jax-ml/jax · error · ValueError

MMA with element type {elem_type_str} only supports accumula

Error message

MMA with element type {elem_type_str} only supports accumulators of type f32, but got: {d.dtype}

What it means

When the A operand is f32 or bf16, the tcgen05 MMA instruction accumulates in f32 only. If the provided accumulator tensor d has any other dtype (e.g. f16), the op raises this ValueError before emitting MLIR.

Source

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

  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(
          f"MMA with element type {elem_type_str} does not support block scaling"
      )
    if d.dtype != f32:
      raise ValueError(
          f"MMA with element type {elem_type_str} only supports accumulators"
          f" of type f32, but got: {d.dtype}"
      )
  elif a_element_type == f16:
    if is_scaled:
      raise ValueError(
          f"MMA with element type {elem_type_str} does not support block scaling"
      )
    if d.dtype != f16 and d.dtype != f32:
      raise ValueError(
          f"MMA with element type {elem_type_str} only supports accumulators of"
          f" type f32 or f16, but got: {d.dtype}"
      )
  elif any(
      isinstance(a_element_type, t)
      for t in {ir.Float8E5M2Type, ir.Float8E4M3FNType}
  ):
    if d.dtype != f16 and d.dtype != f32:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate/convert the accumulator d to f32 (jnp.float32 / ir.F32Type)
2. Match d's dtype to f32 whenever A is f32 or bf16
3. If f16 accumulation is required, switch the operands themselves to f16 instead

Example fix

# before
d = tmem.alloc((m, n), dtype=jnp.float16)
tcgen05.mma(a_bf16, b_bf16, d)  # raises
# after
d = tmem.alloc((m, n), dtype=jnp.float32)
tcgen05.mma(a_bf16, b_bf16, d)

import jax.numpy as jnp
assert d.dtype == jnp.float32, f'f32/bf16 operands require f32 accumulator, got {d.dtype}'

When it happens

Trigger: Calling tcgen05.mma with a/b of type f32 or bf16 and an accumulator d whose dtype is not f32 (commonly f16).

Common situations: Reusing an f16 accumulator allocation from an fp16 kernel in a bf16 kernel; allocating TMEM with a default dtype that doesn't match the operand type; mixed-precision experiments with f16 accumulation.

Related errors

• MMA with element type {elem_type_str} only supports accumula
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block
• Sparse MMA not supported for M=64
• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri

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