jax-ml/jax · error · NotImplementedError

Sparse MMA unsupported for f32

Error message

Sparse MMA unsupported for f32

What it means

Sparse MMA (tcgen05 sparse tensor cores) is only defined for fp16 and fp8 input types; the hardware has no fp32 sparse instruction, so requesting is_sparse with f32 A operands raises NotImplementedError.

Source

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

      # TODO: We only need to check this if N is the minormost dim in B.
      if 8 * b_swizzle // utils.bitwidth(a_element_type) > n // n_lane_groups:
        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}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use fp16 or fp8 (Float8E5M2/E4M3FN) for the A operand when doing sparse MMA
2. Keep the A/B operands in f32 only for the dense path (disable sparse metadata)
3. If numeric range is the concern, use block-scaled fp8 instead of f32

Example fix

# before
lhs_f32 = lhs.astype(jnp.float32)
tcgen05.mma(lhs_f32, rhs, d, sparse_metadata=meta)  # raises
# after
lhs_f16 = lhs.astype(jnp.float16)
tcgen05.mma(lhs_f16, rhs_f16, d, sparse_metadata=meta)

if is_sparse:
    assert a_element_type not in (ir.F32Type.get(),), 'Sparse MMA unsupported for f32'

When it happens

Trigger: Calling tcgen05.mma with lhs of dtype f32 while supplying sparse metadata / enabling the sparse path.

Common situations: Enabling sparsity on a mixed-precision pipeline still configured with f32 activations; migrating a sparse fp16 kernel to f32 for accumulation-accuracy experiments without disabling sparsity.

Related errors

• Sparse MMA not supported for M=64
• MMA with element type {elem_type_str} does not support block
• MMA with element type {elem_type_str} only supports accumula
• MMA with element type {elem_type_str} only supports accumula
• 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/009ee6465bb6fb13. Report an issue: GitHub.