jax-ml/jax · error · ValueError

Only f16 WGMMA supports transposes

Error message

Only f16 WGMMA supports transposes

What it means

Raised at wgmma.py:154 when a_transpose or b_transpose is requested but the operand bytewidth is not 2 — the hardware only supports transposed WGMMA operands for 16-bit (f16/bf16) types.

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)

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 wgmma types {(out_ty, a_element_type)=}
• Unsupported wgmma types {(out_ty, b_element_type)=}
• Unsupported A register array dtype: {a.mlir_dtype}
• WGMMA instruction only supports f32, f16 and s32 out (got {o
• Unsupported dtype: {ref.dtype}

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