jax-ml/jax · error · ValueError

tcgen05.cp only supports fully packed TMEM references (packi

Error message

tcgen05.cp only supports fully packed TMEM references (packing={32 // bitwidth}), but got packing={tmem_ref.packing}

What it means

Raised by async_copy_smem_to_tmem when the TMEM reference's packing factor does not equal 32/bitwidth. tcgen05.cp requires narrow types (f16, bf16, f8, i8...) to be fully packed into 32-bit TMEM lanes; a partially packed or unpacked TMEM reference cannot be a destination of this instruction.

Source

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

    )


def async_copy_smem_to_tmem(
    smem_ref: ir.Value,
    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}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the TMEM reference with fully packed layout: packing = 32 // bitwidth(dtype) (e.g. 2 for f16/bf16, 4 for i8/f8)
2. Use tmem_default_layout/packing helpers rather than hand-computing packing
3. Verify no intermediate op (slice, cast) reset packing to 1

Example fix

// before
tmem = tmem_alloc(f16, shape, packing=1)
// after
tmem = tmem_alloc(f16, shape, packing=32 // 16)  # packing=2

bw = utils.bitwidth(dtype)
required_packing = 32 // bw
assert tmem_ref.packing == required_packing, f'need packing={required_packing}'

When it happens

Trigger: Calling tcgen05.async_copy_smem_to_tmem with a tmem_ref whose packing attribute (from tmem_alloc or a tmem slice) is not 32/bitwidth of the element dtype — e.g. packing=1 with f16 elements (requires packing=2), or a tmem ref refined via tmem_ref with a non-default packing.

Common situations: Allocating TMEM without specifying packing for sub-32-bit types, or slicing/re-laying-out a TMEM reference with explicit_packing and passing the result to the SMEM-to-TMEM copy.

Related errors

• Sparse meta layout loads unsupported.
• Loading multiple row tiles
• Loads from TMEM layout {self.layout} to register layout {lay
• TMEM stores expect a FragmentedArray, got: {value}
• Stored array has shape {value.shape}, but TMEM has shape {se

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