jax-ml/jax · error · ValueError

Invalid WGMMA accumulator dtype for s8/i8 WGMMA. Expected si

Error message

Invalid WGMMA accumulator dtype for s8/i8 WGMMA. Expected signed integer, but got {aval.dtype}.

What it means

When a WGMMA accumulator ref has an 8-bit integer dtype, it must be signed (s8). If mgpu_utils.is_signed reports the dtype as unsigned (e.g. uint8), the run_scoped lowering raises ValueError because s8/i8 WGMMA requires a signed accumulator.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3594

  if collective_axes and collective_axes != (wg_axis,):
    raise ValueError(
        "Only thread-collective allocations are supported in run_scoped."
    )
  if is_multithreaded and is_thread_collective:
    gpu_dialect.barrier()
  with contextlib.ExitStack() as alloc_stack:
    for v in jaxpr.invars:
      aval = cast(ShapedAbstractValue, v.aval)
      if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):
        if collective_axes:
          raise ValueError(
              "WGMMA accumulators can only be allocated non-collectively. Hint:"
              " remove collective_axes from run_scoped. If other allocations"
              " are performed as well, split the run_scoped into two."
          )
        is_signed = mgpu_utils.is_signed(aval.dtype)
        if is_signed is not None and not is_signed:
          raise ValueError(
              "Invalid WGMMA accumulator dtype for s8/i8 WGMMA. "
              f"Expected signed integer, but got {aval.dtype}."
          )

        dtype = mlir.dtype_to_ir_type(aval.dtype)
        if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
          input_refs.append(
              mgpu.WGMMAAccumulator.zero(
                  *aval.shape, dtype=dtype, is_signed=is_signed
              )
          )
        else:
          zero = _ir_constant(0, dtype)
          acc_type = ir.VectorType.get(aval.shape, dtype)
          acc = vector_dialect.broadcast(acc_type, zero)
          acc = mgpu.dialect.optimization_barrier([acc])
          nvvm_dialect.wgmma_fence_aligned()
          input_refs.append(acc)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Declare the accumulator with jnp.int8 instead of jnp.uint8
2. Check mgpu_utils.is_signed(aval.dtype) in your allocation helper
3. Keep activation dtype and accumulator dtype separately configured so unsigned activations don't leak into the accumulator spec

Example fix

// before
acc_ref.dtype = jnp.uint8
// after
acc_ref.dtype = jnp.int8

assert aval.dtype == jnp.int8 or mgpu_utils.is_signed(aval.dtype) is not False

def is_signed_int8(dt) -> bool:
    return dt == jnp.int8

When it happens

Trigger: Allocating a WGMMA accumulator via run_scoped with dtype jnp.uint8 in a kernel that performs int8 WGMMA.

Common situations: Quantized/int8 matmul kernels where the accumulator dtype was declared uint8 (e.g. copied from unsigned activation dtypes); mixing u8 activations with s8 accumulators.

Related errors

• Index map function {debug_info.func_src_info} for {origin} m
• bitcast_convert_type with different bitwidths not supported
• Acc ref dtype must be float32 or int32, got {dtype}
• {uninitialized_memory} + {str(dtype)}
• masked swap with non-32-bit data

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