{"record":{"id":"a709878e0c47bf6c","repo":"jax-ml/jax","slug":"invalid-wgmma-accumulator-dtype-for-s8-i8-wgmma-e","errorCode":null,"errorMessage":"Invalid WGMMA accumulator dtype for s8/i8 WGMMA. Expected signed integer, but got {aval.dtype}.","messagePattern":"Invalid WGMMA accumulator dtype for s8/i8 WGMMA\\. Expected signed integer, but got (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3594,"sourceCode":"  if collective_axes and collective_axes != (wg_axis,):\n    raise ValueError(\n        \"Only thread-collective allocations are supported in run_scoped.\"\n    )\n  if is_multithreaded and is_thread_collective:\n    gpu_dialect.barrier()\n  with contextlib.ExitStack() as alloc_stack:\n    for v in jaxpr.invars:\n      aval = cast(ShapedAbstractValue, v.aval)\n      if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):\n        if collective_axes:\n          raise ValueError(\n              \"WGMMA accumulators can only be allocated non-collectively. Hint:\"\n              \" remove collective_axes from run_scoped. If other allocations\"\n              \" are performed as well, split the run_scoped into two.\"\n          )\n        is_signed = mgpu_utils.is_signed(aval.dtype)\n        if is_signed is not None and not is_signed:\n          raise ValueError(\n              \"Invalid WGMMA accumulator dtype for s8/i8 WGMMA. \"\n              f\"Expected signed integer, but got {aval.dtype}.\"\n          )\n\n        dtype = mlir.dtype_to_ir_type(aval.dtype)\n        if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n          input_refs.append(\n              mgpu.WGMMAAccumulator.zero(\n                  *aval.shape, dtype=dtype, is_signed=is_signed\n              )\n          )\n        else:\n          zero = _ir_constant(0, dtype)\n          acc_type = ir.VectorType.get(aval.shape, dtype)\n          acc = vector_dialect.broadcast(acc_type, zero)\n          acc = mgpu.dialect.optimization_barrier([acc])\n          nvvm_dialect.wgmma_fence_aligned()\n          input_refs.append(acc)","sourceCodeStart":3576,"sourceCodeEnd":3612,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3576-L3612","documentation":"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.","triggerScenarios":"Allocating a WGMMA accumulator via run_scoped with dtype jnp.uint8 in a kernel that performs int8 WGMMA.","commonSituations":"Quantized/int8 matmul kernels where the accumulator dtype was declared uint8 (e.g. copied from unsigned activation dtypes); mixing u8 activations with s8 accumulators.","solutions":["Declare the accumulator with jnp.int8 instead of jnp.uint8","Check mgpu_utils.is_signed(aval.dtype) in your allocation helper","Keep activation dtype and accumulator dtype separately configured so unsigned activations don't leak into the accumulator spec"],"exampleFix":"// before\nacc_ref.dtype = jnp.uint8\n// after\nacc_ref.dtype = jnp.int8","handlingStrategy":"type-guard","validationCode":"assert aval.dtype == jnp.int8 or mgpu_utils.is_signed(aval.dtype) is not False","typeGuard":"def is_signed_int8(dt) -> bool:\n    return dt == jnp.int8","tryCatchPattern":null,"preventionTips":["Use jnp.int8 for int8 WGMMA accumulators","Separate activation dtype from accumulator dtype configs"],"tags":["jax","pallas","wgmma","int8","dtype"],"backgroundTag":"wrong-accumulator-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}