jax-ml/jax · error · ValueError

b_scale must be a TMEM Ref

Error message

b_scale must be a TMEM Ref

What it means

In a block-scaled tcgen05.mma, the b_scale operand must be a TMEM ref, mirroring the a_scale requirement.

Source

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

      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,
    barrier_transforms_tree,
    a_scale_transforms_tree,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate b_scale in TMEM
2. Keep both scale refs' memory spaces consistent (TMEM)
3. Add a pre-call assert on memory_space for both scales

Example fix

# before
b_scale = allocate(SMEM, scale_shape, jnp.uint8)
# after
b_scale = allocate(TMEM, scale_shape, jnp.uint8)

if scaled:
    assert b_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 b_scale in SMEM/GMEM to a scaled tcgen05.mma.

Common situations: Same as a_scale: scales staged in SMEM alongside B instead of TMEM.

Related errors

• a_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/75fe1ac979ab9bff. Report an issue: GitHub.