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 or f16, but got: {d.dtype}

What it means

For f16 A operands, the tcgen05 MMA can accumulate in either f16 or f32. Any other accumulator dtype (e.g. bf16 or an integer type) is invalid and rejected here.

Source

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

      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:
      raise ValueError(
          f"MMA with element type {elem_type_str} only supports accumulators of"
          f" type f32 or f16, but got: {d.dtype}"
      )
    if is_scaled and d.dtype != f32:
      raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set the accumulator d dtype to f16 or f32
2. Prefer f32 accumulation unless you specifically want f16 for TMEM bandwidth reasons
3. Centralize accumulator allocation so dtype is derived from the operand type

Example fix

# before
d = tmem.alloc((m, n), dtype=jnp.bfloat16)
tcgen05.mma(a_f16, b_f16, d)  # raises
# after
d = tmem.alloc((m, n), dtype=jnp.float32)
tcgen05.mma(a_f16, b_f16, d)

import jax.numpy as jnp
assert d.dtype in (jnp.float16, jnp.float32), f'f16 operands need f16/f32 accumulator, got {d.dtype}'

When it happens

Trigger: Calling tcgen05.mma with f16 operands and an accumulator d typed as bf16, fp32-variant other than F32, or any non f16/f32 type.

Common situations: Using bf16 accumulators because they work elsewhere in the pipeline; allocating TMEM accumulators with a global default dtype; dtype typos when constructing MLIR types manually.

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/0bdbe265bc0d2f07. Report an issue: GitHub.