jax-ml/jax · error · ValueError

MMA with element type {elem_type_str} does not support block

Error message

MMA with element type {elem_type_str} does not support block scaling

What it means

Block-scaled MMA ( tcgen05 with scale factors, used for MX formats) is only supported for fp8 (and f16 in some modes) operand types. Requesting is_scaled with f32 or bf16 A operands has no hardware instruction, so it is rejected.

Source

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

        raise ValueError(
            f"Swizzle={b_swizzle} is too big for MMA with M=64. Try"
            " lowering it."
        )
  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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Switch A/B operands to fp8 (Float8E5M2 or Float8E4M3FN) to use block scaling
2. If you need bf16/f32 inputs, drop the scale operands and run unscaled MMA
3. For dynamic-range issues in bf16, consider tmem-based normalization outside the MMA instead of block scaling

Example fix

# before
tcgen05.mma(a_bf16, b_bf16, d, scales=(sa, sb))  # raises
# after
a8, b8 = a_bf16.astype(mxfp8), b_bf16.astype(mxfp8)
tcgen05.mma(a8, b8, d, scales=(sa, sb))

if is_scaled:
    ok = isinstance(a_element_type, (ir.Float8E5M2Type, ir.Float8E4M3FNType)) or a_element_type == ir.F16Type.get()
    assert ok, f'block scaling unsupported for {a_element_type}'

When it happens

Trigger: Calling tcgen05.mma with a_element_type f32 or bf16 while passing scale operands / is_scaled=True.

Common situations: Trying to apply microscaling (MXFP) formats to bf16 weights; enabling scaling flags copied from an fp8 MX kernel onto a bf16 GEMM; experimenting with scaled accumulation precision on higher-precision inputs.

Related errors

• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} only supports accumula
• MMA with element type {elem_type_str} only supports accumula
• 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/de6c7068595f7741. Report an issue: GitHub.