jax-ml/jax · error · ValueError

a_scale must be a TMEM Ref

Error message

a_scale must be a TMEM Ref

What it means

In a block-scaled tcgen05.mma (scaled=True), the a_scale operand must be a TMEM ref; SMEM/GMEM scale refs are rejected.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:2595

    # Figure out a way to fix this.
    if isinstance(acc, gpu_core.AbstractTMEMRef) and not acc.collective:
      raise ValueError(
          "Accumulator Ref must be collective if collective_axis is set.")
    if isinstance(a, gpu_core.AbstractTMEMRef) and not a.collective:
      raise ValueError(
          "LHS Ref must be collective if collective_axis is set.")

  scales_and_transforms_leaves = barrier_scales_and_transforms_leaves
  if arrive:
    barrier, *scales_and_transforms_leaves = barrier_scales_and_transforms_leaves
    orders_tensor_core = getattr(
        barrier.inner_aval.dtype, "orders_tensor_core", False)
    if not orders_tensor_core:
      raise ValueError("MMA barrier must have orders_tensor_core set to True.")
  if scaled:
    a_scale, b_scale = scales_and_transforms_leaves[:2]
    if a_scale.memory_space != gpu_core.TMEM:
      raise ValueError("a_scale must be a TMEM Ref")
    if b_scale.memory_space != gpu_core.TMEM:
      raise ValueError("b_scale must be a TMEM Ref")

  return [], {gpu_core._memory_effect}


@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWG_SEMANTICS)
@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWARP_SEMANTICS)
def _tcgen05_mma_lowering(
    ctx: lowering.LoweringRuleContext,
    acc: tcgen05.TMEMRef,
    a_ref,
    b_ref,
    accumulate: bool | ir.Value,
    *barrier_scales_and_transforms_leaves,
    acc_transforms_tree,
    a_transforms_tree,
    b_transforms_tree,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate a_scale in TMEM
2. Follow the scaled-MMA example's scale allocation pattern exactly
3. Verify a_scale.memory_space == TMEM before calling

Example fix

# before
a_scale = allocate(SMEM, scale_shape, jnp.uint8)
tcgen05.mma(a, b, acc, k_dim=k, a_scale=a_scale, b_scale=b_scale)
# after
a_scale = allocate(TMEM, scale_shape, jnp.uint8)
tcgen05.mma(a, b, acc, k_dim=k, a_scale=a_scale, b_scale=b_scale)

if scaled:
    assert a_scale.memory_space == gpu_core.TMEM

def scale_is_tmem(ref):
    return getattr(ref, 'memory_space', None) == gpu_core.TMEM

When it happens

Trigger: Passing a_scale allocated in SMEM or as a raw GMEM input to a scaled tcgen05.mma call.

Common situations: Allocating scales next to the operand tiles in SMEM instead of TMEM when implementing MXFP8 block scaling.

Related errors

• b_scale must be a TMEM Ref
• TMEM aliasing only supported for Refs with the same first di
• Unsupported TMEM ref {ref}.
• Stores to TMEM are asynchronous operations and cannot be per
• a_scale and b_scale must both be present or absent.

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