jax-ml/jax · error · ValueError

Only WGMMA layouts supported in WGMMAAccumulator

Error message

Only WGMMA layouts supported in WGMMAAccumulator

What it means

WGMMAAccumulator.from_registers (wgmma.py:90) only accepts FragmentedArrays laid out as fa.WGMMA_LAYOUT or fa.WGMMA_LAYOUT_ACC_32BIT, because the accumulator must present registers exactly as the wgmma instruction produces them.

Source

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

      raise TypeError("PTX does not support unsigned WGMMA accumulators")
    f32 = ir.F32Type.get()
    if dtype is None:
      dtype = f32
    if isinstance(dtype, ir.IntegerType):
      zero = arith.constant(dtype, ir.IntegerAttr.get(dtype, 0))
    else:
      zero = arith.constant(dtype, ir.FloatAttr.get(dtype, 0.0))
    return cls.from_registers(
        fa.FragmentedArray.splat(
            zero, (m, n), fa.WGMMA_LAYOUT, is_signed=is_signed
        )
    )

  @classmethod
  def from_registers(cls, registers, sync=True):
    original_layout = registers.layout
    if registers.layout != fa.WGMMA_LAYOUT and registers.layout != fa.WGMMA_LAYOUT_ACC_32BIT:
      raise ValueError("Only WGMMA layouts supported in WGMMAAccumulator")
    if utils.bitwidth(registers.mlir_dtype) == 32:
      registers = registers.to_layout(fa.WGMMA_LAYOUT_ACC_32BIT)
    return cls(_value=registers, _original_layout=original_layout, _sync=sync)

  def tree_flatten(self):
    return (self._value,), (self._original_layout,)

  @classmethod
  def tree_unflatten(cls, aux, value):
    return cls(_value=value[0], _original_layout=aux[0], _sync=False)


def _supported_wgmma_types(dtype, abtype) -> bool:
  input_types_are = lambda ty: isinstance(abtype, ty)
  f16_acc_types = (ir.F16Type, ir.Float8E5M2Type, ir.Float8E4M3FNType)
  if isinstance(dtype, ir.F32Type):
    return any(input_types_are(ty) for ty in (ir.FloatTF32Type, ir.BF16Type, *f16_acc_types))
  elif isinstance(dtype, ir.F16Type):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert registers first: regs.to_layout(fa.WGMMA_LAYOUT_ACC_32BIT if bitwidth==32 else fa.WGMMA_LAYOUT)
2. Ensure you pass the direct result of wgmma.wgmma (or wgmma_m64) without intermediate relayouts
3. Check fa.WGMMA_LAYOUT constants exist in your JAX version; upgrade or use the renamed constants

Example fix

# before
acc = wgmma.WGMMAAccumulator.from_registers(regs)  # regs in row-major
# after
regs = regs.to_layout(fa.WGMMA_LAYOUT_ACC_32BIT)
acc = wgmma.WGMMAAccumulator.from_registers(regs)

assert regs.layout in (fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_ACC_32BIT), f'bad layout {regs.layout}'

def is_wgmma_layout(a): return a.layout in (fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_ACC_32BIT)

When it happens

Trigger: Calling wgmma.WGMMAAccumulator.from_registers(regs) where regs is a FragmentedArray in a row-major/blocked/register-tensor layout instead of a WGMMA layout; also hit via lowering rules (_mgpu_wgmma_op_lowering_rule, accumulator store lowering, kv_loop in FlashAttention examples).

Common situations: Using wgmma(...) output that was already relayouted (e.g. after arithmetic in a different layout); passing a tensor created via fa.make_tensor and forgetting .to_layout(fa.WGMMA_LAYOUT_ACC_32BIT); version changes that renamed layouts.

Related errors

• Unsupported A register array layout
• Source and destination layouts aren't compatible for a broad
• packed cannot be specified if layout is specified.
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Swizzle {self.swizzle} requires the trailing dimension to be

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