jax-ml/jax · error · ValueError

Unsupported WGMMA features with A in registers

Error message

Unsupported WGMMA features with A in registers

What it means

At wgmma.py:166, the A-in-registers form of wgmma does not support a_k_stride or a_transpose — those parameters only apply when A is read from shared memory.

Source

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

  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

  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set a_k_stride=None and a_transpose=None when A is in registers
2. Physically pre-transpose or re-stride A before loading into registers
3. Use the SMEM path if strides/transposes are required

Example fix

# before
acc = wgmma.wgmma(a_regs, b, acc, a_k_stride=64, a_transpose=False)
# after
acc = wgmma.wgmma(a_regs, b, acc, a_k_stride=None, a_transpose=None)

if isinstance(a, fa.FragmentedArray):
    assert a_k_stride is None and a_transpose is None

When it happens

Trigger: Calling wgmma.wgmma with a as a FragmentedArray AND specifying a_k_stride=... or a_transpose=True/False.

Common situations: Copy-pasting a memory-operand wgmma call and swapping in a register array; leaving explicit a_transpose=False (not None) set from earlier config.

Related errors

• Unsupported A register array dtype: {a.mlir_dtype}
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Preinitialized WGMMAAccumulatorRef only supported in pl.run_
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

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