jax-ml/jax · error · ValueError

SMEM has shape {smem_shape}, but expected {expected_smem_sha

Error message

SMEM has shape {smem_shape}, but expected {expected_smem_shape} for TMEM shape {tmem_ref.shape} with swizzle={swizzle}

What it means

Raised by async_copy_smem_to_tmem when the SMEM memref's shape does not equal the tile shape implied by the TMEM reference and swizzle: expected_smem_shape = tile_shape(tmem_ref.shape, (8, swizzle_elems)) where swizzle_elems = 8*swizzle/bitwidth. The instruction walks SMEM as 8x(swizzle_elems) core matrices, so the buffer must be exactly that tiled shape.

Source

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

  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

  # Figure out the matrix descriptor parameters (LBO/SBO)
  # The copy happens using the usual "core matrix" structure: a memory region
  # describing a 8x128bit matrix. LBO describes how far apart from each other
  # are consecutive matrices along the minor dimension (in our case the minor
  # dim is contiguous, so exactly 128 bit = 16 bytes apart). SBO describes how
  # far apart is the beginning of the next matrix along the major dimension.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate SMEM with shape = tile_shape(tmem_shape, (8, 8*swizzle//bitwidth(dtype))) — i.e. reshape your data into 8x(swizzle_elems) core-matrix tiles
2. Recompute the expected shape in Python with utils.tile_shape and assert your smem allocation matches before lowering
3. If dtype changes, keep swizzle fixed and let the tile shape change accordingly

Example fix

# before
swizzle = 128; dtype_bits = 32
smem = smem_alloc(f32, (128, 64))
# after
from jax.experimental.mosaic.gpu import utils
se = 8 * swizzle // dtype_bits
smem_shape = utils.tile_shape(tmem_ref.shape, (8, se))
smem = smem_alloc(f32, smem_shape)

bw = utils.bitwidth(dtype)
expected = utils.tile_shape(tmem_ref.shape, (8, 8 * swizzle // bw))
assert tuple(ir.MemRefType(smem_ref.type).shape) == expected, f'smem shape must be {expected}'

When it happens

Trigger: Allocating SMEM as a flat (rows, cols) array instead of the 4D tiled shape (row_tiles, col_tiles, 8, swizzle_elems), or passing a swizzle whose byte size mismatches the SMEM tile width, e.g. smem shape (128, 64) with swizzle=32 and f32 elements (swizzle_elems=8, so tiles are 8x8 and total shape must reflect that).

Common situations: Porting wgmma/legacy mma SMEM staging layouts to tcgen05; changing element dtype without re-deriving swizzle_elems (swizzle is in bytes, shape is in elements).

Related errors

• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• 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/d0d2620b654692c1. Report an issue: GitHub.