jax-ml/jax · error · ValueError

The SMEM tiles must be contiguous

Error message

The SMEM tiles must be contiguous

What it means

Raised by async_copy_smem_to_tmem when, after the shape check passes, the SMEM memref strides are not the contiguous strides of the tiled shape: the innermost (column) stride must be 1 and the inner row stride must equal swizzle_elems. This catches views (slices/transposes) of a correctly-shaped buffer that broke contiguity of the 8x(swizzle_elems) core matrices.

Source

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

    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.
  # We use a tiling of 8, so it is simply the tile stride.
  leading_byte_offset = 16
  stride_byte_offset = row_tile_stride * bitwidth // 8
  assert tmem_ref.shape[1] * bitwidth // 8 >= 16
  if swizzle == 16:
    cp_shape = nvvm.Tcgen05CpShape.SHAPE_128x128b
    cp_cols_bytes = 16  # 128 bit = 16 bytes

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate a dedicated contiguous SMEM buffer with the exact tiled shape instead of a strided subview
2. If subview is needed for double buffering, make the tile stride change only the outer (row_tile/col_tile) strides, keeping inner strides (1, swizzle_elems) intact
3. Copy/reshape the data into a contiguous buffer before the tcgen05.cp

Example fix

# before
smem_view = memref.subview(big_smem, offsets, strides=[1, 2, 2, 1], ...)
# after
stage_smem = smem_alloc(f32, smem_shape)  # contiguous, own buffer per stage
# store into stage_smem, then copy from it

bw = utils.bitwidth(dtype)
swizzle_elems = 8 * swizzle // bw
strides, _ = ir.MemRefType(smem_ref.type).get_strides_and_offset()
assert strides[3] == 1 and strides[2] == swizzle_elems, 'SMEM tiles not contiguous'

When it happens

Trigger: Passing a transposed, sliced, or padded SMEM memref to async_copy_smem_to_tmem — e.g. smem[:, ::2] or a subview with inner_col_stride != 1, or an inner row stride that no longer equals 8*swizzle/bitwidth.

Common situations: Reusing one big SMEM buffer with strided sub-allocations per pipeline stage; creating the SMEM ref via memref.subview with offsets/strides for double buffering.

Related errors

• K tile stride must be a multiple of 16
• SMEM has shape {smem_shape}, but expected {expected_smem_sha
• `strides` must contain only 1s.
• Strides {strides} have lower rank than tiling {tiling}
• Can not tile strides when tiled dimensions have been transpo

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