jax-ml/jax · error · ValueError

Only standard TMEM layout is supported, got: {tmem_ref.layou

Error message

Only standard TMEM layout is supported, got: {tmem_ref.layout}

What it means

Raised by async_copy_smem_to_tmem when the TMEM reference's layout does not match tmem_default_layout(packing=...). tcgen05.cp only emits correct addresses for the standard (default) TMEM layout; custom blocked/interleaved TMEM layouts are not supported for this copy instruction.

Source

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

  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")
  # Make sure strides are a multiple of the byte packing for narrow types.
  byte_packing = max(8 // bitwidth, 1)
  assert row_tile_stride % byte_packing == 0
  assert col_tile_stride % byte_packing == 0

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit the layout argument in tmem_alloc so the default layout is used, or pass tmem_default_layout(packing=32//bitwidth(dtype)) explicitly
2. Copy into a fresh default-layout TMEM reference and move data with TMEM-to-TMEM or register ops if a custom layout is required downstream

Example fix

// before
tmem = tmem_alloc(dtype, shape, layout=custom_layout)
// after
from jax.experimental.mosaic.gpu import tcgen05
tmem = tmem_alloc(dtype, shape, layout=tcgen05.tmem_default_layout(packing=32//bitwidth(dtype)))

expected = tcgen05.tmem_default_layout(packing=tmem_ref.packing)
assert tmem_ref.layout == expected, 'non-default TMEM layout cannot be a tcgen05.cp destination'

When it happens

Trigger: Passing a tmem_ref created with a custom layout argument to tmem_alloc (or one produced by tmem slice ops carrying a non-default layout) into async_copy_smem_to_tmem.

Common situations: Experimenting with TMEM layouts for accumulator reuse across MMAs, then feeding the same reference into a tcgen05.cp; defaulting to a legacy layout after a JAX/Mosaic version change altered default layouts.

Related errors

• Loading multiple row tiles
• Loads from TMEM layout {self.layout} to register layout {lay
• Stored array has layout {value.layout}, but TMEM stores expe
• Storing from register layout {value.layout} to TMEM layout {
• A scale layout {a_scale.layout} is not supported

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