jax-ml/jax · error · NotImplementedError

Only s32 accumulator supported for integer operands.

Error message

Only s32 accumulator supported for integer operands.

What it means

For integer MMA operands (i8/i4), the accumulator must be 32-bit integers (s32); PTX integer MMA only produces i32 results. Any other accumulator dtype raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/mma.py:214

    raise ValueError(f"K mismatch: {k} != {k2}")

  # todo(cperivol): A tile shape can have dimensions that are higher
  # multiples of the mma op size as long as those dimensions are not
  # sharded across warps.
  i4 = ir.IntegerType.get_signless(4)
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)
  bf16 = ir.BF16Type.get()
  f16 = ir.F16Type.get()
  f8e4m3fn = ir.Float8E4M3FNType.get()
  f8e5m2 = ir.Float8E5M2Type.get()
  if (element_type := a.mlir_dtype) != b.mlir_dtype:
    raise ValueError(f"Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}")
  if element_type not in (bf16, f16, f8e4m3fn, f8e5m2, i8, i4):
    raise NotImplementedError(f"Unsupported operand type: {element_type}")
  if isinstance(element_type, ir.IntegerType):
    if acc.mlir_dtype != i32:
      raise NotImplementedError("Only s32 accumulator supported for integer operands.")
    if not acc.is_signed:
      raise ValueError("Only signed accumulator supported for integer operands.")
  elif acc.mlir_dtype != ir.F32Type.get():
    raise NotImplementedError("Only f32 accumulator supported for floating operands.")

  can_infer_from_acc_layout = (
      isinstance(acc.layout, fa.TiledLayout)
      and len(acc.layout.base_tile_shape) == 2
      and acc.layout.base_tile_shape[0] % 16 == 0
  )
  if not can_infer_from_acc_layout:
    raise ValueError("Expected MMALayouts.acc for acc")
  m_warps = acc.layout.base_tile_shape[0] // 16  # type: ignore
  layouts = MMALayouts(element_type, m_warps=m_warps)
  if layouts.lhs != a.layout:
    raise ValueError("Expected MMALayouts.lhs layout for A")
  if layouts.rhs != b.layout:
    raise ValueError("Expected MMALayouts.rhs layout for B")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the accumulator as jnp.int32 (signed) for integer operands
2. Keep f32 accumulators only for floating-point operand types

Example fix

// before
acc = fa.from_tensor(jnp.zeros((m, n), jnp.float32))
acc = mma.mma(a_i8, b_i8, acc)
// after
acc = fa.from_tensor(jnp.zeros((m, n), jnp.int32))
acc = mma.mma(a_i8, b_i8, acc)

if isinstance(a.mlir_dtype, ir.IntegerType):
    assert acc.mlir_dtype == ir.IntegerType.get_signless(32)

When it happens

Trigger: Calling mma() with i8 operands but an f32 or i8 accumulator.

Common situations: Copy-pasting an fp16 kernel's f32 accumulator setup into an int8 quantized kernel.

Related errors

• Only signed accumulator supported for integer operands.
• Acc ref must be at least 2D, got shape {shape}
• m_warps must be 1, 2, or 4, but got {m_warps=}
• is_signed must be specified for integer types
• M mismatch: {m} != {m2}

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