jax-ml/jax · error · TypeError

PTX does not support unsigned WGMMA accumulators

Error message

PTX does not support unsigned WGMMA accumulators

What it means

WGMMAAccumulator.zero (wgmma.py:72) rejects is_signed=False. The PTX wgmma instruction only defines signed integer accumulation for s32 accumulators; there is no unsigned accumulator variant.

Source

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

      _original_layout: fa.FragmentedLayout,
      _sync: bool = True,
  ):
    self._original_layout = _original_layout
    self._value = _value
    if _sync:
      self._value = wgmma_fence(_value)

  @property
  def value(self) -> fa.FragmentedArray:
    return self._value.to_layout(self._original_layout)

  @classmethod
  def zero(cls, m, n, dtype=None, *, is_signed: bool | None = None):
    if m % 64 or n % 8:
      raise ValueError("WGMMA requires m and n to be multiples of 64 and 8, "
                       f"got {m} and {n}")
    if is_signed is False:
      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")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a signed s32 (or f32/f16) accumulator even when operands are unsigned
2. Pass is_signed=None or True when using an integer accumulator
3. If unsigned accumulation is needed, accumulate in s32 and convert to unsigned after the wgmma loop

Example fix

# before
acc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.IntegerType.get_unsigned(32), is_signed=False)
# after
acc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.IntegerType.get_signless(32))
out = acc.to_unsigned()  # convert after accumulation if needed

assert not (dtype and isinstance(dtype, ir.IntegerType) and dtype.is_unsigned), 'WGMMA accumulator must be signed'

def wgmma_accumulator_dtype(dt):
    if isinstance(dt, ir.IntegerType) and dt.is_unsigned:
        return ir.IntegerType.get_signless(dt.width)
    return dt

When it happens

Trigger: Calling wgmma.WGMMAAccumulator.zero(m, n, dtype=ir.IntegerType.get_unsigned(32), is_signed=False) or passing an unsigned integer dtype with is_signed=False.

Common situations: Feeding u8 operands (e.g. quantized weights) and assuming the accumulator can be unsigned too; porting int8 GEMM code that reuses the operand's signedness for the accumulator.

Related errors

• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Accumulator shape {inner.shape} does not match value shape {
• Accumulator dtype {inner.dtype} does not match value dtype {
• Unsupported bitwidth: {bitwidth}
• Cannot convert from {self.layout} to {new_layout}

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