{"record":{"id":"7d4d4a77d98f73a7","repo":"jax-ml/jax","slug":"ptx-does-not-support-unsigned-wgmma-accumulators","errorCode":null,"errorMessage":"PTX does not support unsigned WGMMA accumulators","messagePattern":"PTX does not support unsigned WGMMA accumulators","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":72,"sourceCode":"      _original_layout: fa.FragmentedLayout,\n      _sync: bool = True,\n  ):\n    self._original_layout = _original_layout\n    self._value = _value\n    if _sync:\n      self._value = wgmma_fence(_value)\n\n  @property\n  def value(self) -> fa.FragmentedArray:\n    return self._value.to_layout(self._original_layout)\n\n  @classmethod\n  def zero(cls, m, n, dtype=None, *, is_signed: bool | None = None):\n    if m % 64 or n % 8:\n      raise ValueError(\"WGMMA requires m and n to be multiples of 64 and 8, \"\n                       f\"got {m} and {n}\")\n    if is_signed is False:\n      raise TypeError(\"PTX does not support unsigned WGMMA accumulators\")\n    f32 = ir.F32Type.get()\n    if dtype is None:\n      dtype = f32\n    if isinstance(dtype, ir.IntegerType):\n      zero = arith.constant(dtype, ir.IntegerAttr.get(dtype, 0))\n    else:\n      zero = arith.constant(dtype, ir.FloatAttr.get(dtype, 0.0))\n    return cls.from_registers(\n        fa.FragmentedArray.splat(\n            zero, (m, n), fa.WGMMA_LAYOUT, is_signed=is_signed\n        )\n    )\n\n  @classmethod\n  def from_registers(cls, registers, sync=True):\n    original_layout = registers.layout\n    if registers.layout != fa.WGMMA_LAYOUT and registers.layout != fa.WGMMA_LAYOUT_ACC_32BIT:\n      raise ValueError(\"Only WGMMA layouts supported in WGMMAAccumulator\")","sourceCodeStart":54,"sourceCodeEnd":90,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L54-L90","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use a signed s32 (or f32/f16) accumulator even when operands are unsigned","Pass is_signed=None or True when using an integer accumulator","If unsigned accumulation is needed, accumulate in s32 and convert to unsigned after the wgmma loop"],"exampleFix":"# before\nacc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.IntegerType.get_unsigned(32), is_signed=False)\n# after\nacc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.IntegerType.get_signless(32))\nout = acc.to_unsigned()  # convert after accumulation if needed","handlingStrategy":"validation","validationCode":"assert not (dtype and isinstance(dtype, ir.IntegerType) and dtype.is_unsigned), 'WGMMA accumulator must be signed'","typeGuard":"def wgmma_accumulator_dtype(dt):\n    if isinstance(dt, ir.IntegerType) and dt.is_unsigned:\n        return ir.IntegerType.get_signless(dt.width)\n    return dt","tryCatchPattern":null,"preventionTips":["Default to f32 accumulators","Use signless i32 for integer accumulation"],"tags":["jax","mosaic-gpu","wgmma","ptx","signedness"],"backgroundTag":"unsupported-type-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}