jax-ml/jax · error · ValueError

Only f16 WGMMA supports transposes

Error message

Only f16 WGMMA supports transposes

What it means

The WGMMA instruction can only transpose operands when the element type is 16 bits wide (f16/bf16); narrower types have no transposed layout encoding. It fires when a_transpose or b_transpose is set with a sub-16-bit element type.

Solutions

1. Remove a_transpose/b_transpose for non-16-bit dtypes and physically transpose the data instead (swap index math or pre-transpose in SMEM)
2. Convert operands to f16/bf16 if transposes are essential
3. Use transpose via TMA layout rather than the wgmma flags

Example fix

# before
acc = wgmma.wgmma(a_i8, b_i8, acc, b_transpose=True)
# after
b_t = utils.transpose_smem(b)  # or load B pre-transposed
acc = wgmma.wgmma(a_i8, b_t, acc)

if bytewidth(b_element_type) != 2:
    assert not a_transpose and not b_transpose, 'transpose requires 16-bit operands'

When it happens

Trigger: Calling wgmma.wgmma(..., a_transpose=True or b_transpose=True) with i8/f8/s32 operands.

Common situations: Porting an f16 attention kernel to int8 quantized operands and keeping the transpose flags; enabling transpose on low-precision B stored in SMEM.

Related errors

• Unsupported A register array dtype
• Unsupported wgmma types
• Unsupported wgmma types
• WGMMA instruction only supports f32, f16 and s32 out
• Accumulator dtype does not match value dtype

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

Appendix: source

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

  if not _supported_wgmma_types(out_ty, b_element_type):
    raise ValueError(f"Unsupported wgmma types {(out_ty, b_element_type)=}")
  if n % 8:
    raise ValueError

  bf16 = ir.BF16Type.get()
  f16 = ir.F16Type.get()
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)
  i64 = ir.IntegerType.get_signless(64)
  f8e5m2 = ir.Float8E5M2Type.get()
  f8e4m3fn = ir.Float8E4M3FNType.get()
  if b_k_stride % 16:
    raise ValueError
  assert bytewidth(a_element_type) == bytewidth(b_element_type)
  # Only 16-bit types support transposes
  supports_transpose = bytewidth(b_element_type) == 2
  if not supports_transpose and (a_transpose or b_transpose):
    raise ValueError("Only f16 WGMMA supports transposes")
  if a_in_regs := isinstance(a, fa.FragmentedArray):
    if a.mlir_dtype not in {bf16, f16, i8, f8e5m2, f8e4m3fn}:
      raise ValueError(f"Unsupported A register array dtype: {a.mlir_dtype}")
    # Column count must be equal to swizzle // bytewidth.
    elt_bytewidth = utils.bytewidth(a_element_type)
    swizzle_elems = swizzle // elt_bytewidth
    if a.shape != (64, swizzle_elems):
      raise ValueError("Unsupported A register array shape")
    if a.layout not in {fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_8BIT}:
      raise ValueError("Unsupported A register array layout")
    if a_k_stride is not None or a_transpose is not None:
      raise ValueError("Unsupported WGMMA features with A in registers")
  else:
    if a_k_stride is None or a_k_stride % 16:
      raise ValueError
    if a_transpose is None:
      raise ValueError

View on GitHub (pinned to 1e1c6a8fc0)