jax-ml/jax · error · ValueError

TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref

Error message

TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref.shape[0]}

What it means

Raised by async_copy_smem_to_tmem when the TMEM reference's row count (shape[0]) does not equal TMEM_ROWS (128). The tcgen05.cp instruction writes a full 128-row TMEM block per copy, so the destination reference must span all 128 lanes.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:2014

    tmem_ref: TMEMRef,
    swizzle: int,
    collective: bool = False,
) -> None:
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)
  smem_ty = ir.MemRefType(smem_ref.type)
  if (dtype := smem_ty.element_type) != tmem_ref.dtype:
    raise ValueError(f"Incompatible dtypes: SMEM has {dtype}, TMEM has {tmem_ref.dtype}")
  if swizzle not in {16, 32, 64, 128}:
    raise ValueError(f"Unsupported swizzle, expected 16, 32, 64 or 128, but got: {swizzle}")
  bitwidth = utils.bitwidth(dtype)
  if tmem_ref.packing != 32 // bitwidth:
    raise ValueError(
        "tcgen05.cp only supports fully packed TMEM references"
        f" (packing={32 // bitwidth}), but got packing={tmem_ref.packing}"
    )
  if tmem_ref.shape[0] != TMEM_ROWS:
    raise ValueError(
        f"TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref.shape[0]}"
    )
  if tmem_ref.layout != tmem_default_layout(packing=tmem_ref.packing):
    raise ValueError(
        f"Only standard TMEM layout is supported, got: {tmem_ref.layout}"
    )
  swizzle_elems = 8 * swizzle // bitwidth
  expected_smem_shape = utils.tile_shape(tmem_ref.shape, (8, swizzle_elems))
  smem_shape = tuple(smem_ty.shape)
  if smem_shape != expected_smem_shape:
    raise ValueError(
        f"SMEM has shape {smem_shape}, but expected {expected_smem_shape} for"
        f" TMEM shape {tmem_ref.shape} with swizzle={swizzle}"
    )
  strides, _ = smem_ty.get_strides_and_offset()
  row_tile_stride, col_tile_stride, inner_row_stride, inner_col_stride = strides
  if inner_col_stride != 1 or inner_row_stride != swizzle_elems:
    raise ValueError("The SMEM tiles must be contiguous")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate/copy into a full 128-row TMEM reference, then slice per-warpgroup afterwards
2. Check the TMEM_ROWS constant (128) against your tmem_alloc shape[0]
3. Redesign the kernel so the SMEM staging tile covers all 128 rows in one copy

Example fix

// before
tmem = tmem_alloc(dtype, (64, n))
async_copy_smem_to_tmem(smem, tmem, swizzle=32)
// after
tmem = tmem_alloc(dtype, (128, n))
async_copy_smem_to_tmem(smem, tmem, swizzle=32)

assert tmem_ref.shape[0] == 128, 'tcgen05.cp requires full 128-row TMEM'

def is_full_tmem(tmem_ref, rows=128):
    return tmem_ref.shape[0] == rows

When it happens

Trigger: Passing a tmem_ref that was sliced to fewer than 128 rows (e.g. tmem_ref.slice or a per-warpgroup 64-row view), or allocating TMEM with an M dimension other than 128, to async_copy_smem_to_tmem.

Common situations: Splitting TMEM across warpgroups or warps (each getting 32/64 rows) and then attempting a bulk SMEM->TMEM copy on the sub-view; using a matmul_shape with M != 128.

Related errors

• Stored array has shape {value.shape}, but TMEM has shape {se
• A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}),
• B scale shape[0] must be a multiple of 128 and >= N={n * num
• B scale shape mismatch: expected ({b_scale.shape[0]}, {k_sca
• A sparse metadata shape mismatch: expected {(m, expected_met

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