jax-ml/jax · error · ValueError

WGMMA instruction only supports f32, f16 and s32 out (got {o

Error message

WGMMA instruction only supports f32, f16 and s32 out (got {out_ty})

What it means

At wgmma.py:188, after the accumulator dtype branch, any out_ty other than f32, f16, or s32 (i32) is rejected — the wgmma instruction's accumulator register file only holds those types.

Source

Thrown at jax/experimental/mosaic/gpu/wgmma.py:188

    if a_transpose is None:
      raise ValueError

  if isinstance(out_ty, ir.F32Type) or out_ty == i32:
    num_acc_regs = n // 2
    out_ty_field = ir.VectorType.get((1,), out_ty)
    acc_regs = list(acc.flat)
    assert acc_regs[0].type == ir.VectorType.get((1,), out_ty)
    to_acc_vec_regs = lambda regs: np.array(regs).reshape(acc.shape)
    acc_constraint = "r" if isinstance(out_ty, ir.IntegerType) else "f"
  elif isinstance(out_ty, ir.F16Type):
    num_acc_regs = n // 4
    out_ty_field = i32
    acc_regs = [_as_i32_reg(reg) for reg in acc.flat]
    vec_ty = ir.VectorType(acc.flat[0].type)
    to_acc_vec_regs = lambda regs: np.array([_unpack_i32(vec_ty, reg) for reg in regs]).reshape(acc.shape)
    acc_constraint = "r"
  else:
    raise ValueError(
        f"WGMMA instruction only supports f32, f16 and s32 out (got {out_ty})")

  if supports_transpose:
    num_imm_regs = 4
  elif out_ty == i32:
    num_imm_regs = 0
  else:
    num_imm_regs = 2

  if a_in_regs:
    a_reg_constraints = ["r"] * 4  # 4x (b)f16x2 or s8x4 registers
    if supports_transpose:
      num_imm_regs -= 1  # transpose not supported for a in registers
  else:
    a_reg_constraints = ["l"]  # descriptor
  # Reference for i/o aliasing: https://gcc.gnu.org/onlinedocs/gcc/Extended-Asm.html
  # Seems like it's not actually documented in LLVM IR docs.
  reg_constraints_list = (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the accumulator as f32 (default), f16, or signless i32 (WGMMAAccumulator.zero does this correctly)
2. Convert results to the desired dtype after the wgmma loop
3. Check that acc was built from WGMMAAccumulator.zero or a prior wgmma result

Example fix

# before
acc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.BF16Type.get())
# after
acc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.F32Type.get())
result_bf16 = acc.value.to_dtype(ir.BF16Type.get())

assert str(out_ty) in ('f32', 'f16', 'i32'), f'bad accumulator dtype {out_ty}'

def wgmma_out_ok(dt): return str(dt) in ('f32', 'f16', 'i32')

When it happens

Trigger: Calling wgmma.wgmma with an accumulator FragmentedArray of dtype f64, bf16, s64, etc. (this check complements _supported_wgmma_types).

Common situations: Creating the accumulator with a dtype matched to the operands instead of the output; bf16 accumulation attempts; custom vector types on acc.flat.

Related errors

• Unsupported wgmma types {(out_ty, a_element_type)=}
• Unsupported wgmma types {(out_ty, b_element_type)=}
• Only f16 WGMMA supports transposes
• Unsupported A register array dtype: {a.mlir_dtype}
• Accumulator aval mismatch: expected {aval}, got {acc.aval}

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