jax-ml/jax · error · ValueError

K tile stride must be a multiple of 16

Error message

K tile stride must be a multiple of 16

What it means

Raised by async_copy_sparse_metadata_smem_to_tmem in Mosaic GPU's tcgen05 module when copying sparse tensor metadata from shared memory (SMEM) to tensor memory (TMEM). The K tile stride (strides[1] of the SMEM memref) must be divisible by 16 because the tcgen05 hardware instruction that loads sparse metadata requires 16-element-aligned tile spacing. The stride is derived from the SMEM layout you allocated, so an oddly-strided buffer fails this check.

Source

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

  if tmem_ref.shape[1] % 64:
    raise ValueError(f"TMEM reference must have a multiple of 64 colums, but got {tmem_ref.shape[1]}")
  if tmem_ref.layout != sparse_meta_layout():
    raise ValueError(f"TMEM layout {tmem_ref.layout} is not supported")
  smem_shape = tuple(smem_ty.shape)
  expected_smem_shape = (tmem_ref.shape[0] // 128, tmem_ref.shape[1] // 64, 128, 64)
  if smem_shape != expected_smem_shape:
    raise NotImplementedError(
        f"SMEM has {smem_shape}, but expected {expected_smem_shape} for TMEM"
        f" ref shape {tmem_ref.shape}"
    )
  strides, _ = smem_ty.get_strides_and_offset()
  if strides != utils.get_contiguous_strides(smem_shape):
    raise ValueError("Only copies from contiguous SMEM references are supported")
  if expected_smem_shape[0] != 1:
    raise NotImplementedError("Only M=128 supported")
  k_tile_stride = strides[1]
  if k_tile_stride % 16:
    raise ValueError("K tile stride must be a multiple of 16")
  k_tile_byte_stride = k_tile_stride // 4
  for k_tile in range(expected_smem_shape[1]):
    load_ptr = utils.getelementptr(
        utils.memref_ptr(smem_ref), [k_tile * k_tile_byte_stride], i8
    )
    store_ptr = arith.addi(tmem_ref.address, arith.constant(i32, 4 * k_tile))
    # The "core matrix" here is the same as in MMA: 8x(16 bytes).
    desc = mma_utils.encode_descriptor(load_ptr, 0, 8 * 16, swizzle=None)
    ptr = _tmem_addr_to_ptr(store_ptr)
    nvvm.tcgen05_cp(
        nvvm.Tcgen05CpShape.SHAPE_128x128b, ptr, desc,
        group=nvvm.CTAGroupKind.CTA_2 if collective else nvvm.CTAGroupKind.CTA_1
    )


def async_copy_smem_to_tmem(
    smem_ref: ir.Value,
    tmem_ref: TMEMRef,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reallocate or reshape the SMEM buffer so the K-tile dimension has a stride that is a multiple of 16 (e.g. pad the trailing dimension)
2. Ensure the SMEM reference is contiguous by allocating with a shape matching utils.get_contiguous_strides, since only contiguous SMEM is supported anyway
3. If using tmem_alloc / smem allocation helpers, pass a shape whose innermost tile size divides 16 evenly (e.g. K tiles of 16)

Example fix

// before
smem = smem_alloc(i32, (1, 10))  # stride 10 -> not multiple of 16
// after
smem = smem_alloc(i32, (1, 16))  # K tile stride 16 -> passes check

strides, _ = ir.MemRefType(smem_ref.type).get_strides_and_offset()
if strides[1] % 16:
    raise ValueError(f'K tile stride {strides[1]} not multiple of 16; reallocate SMEM')

When it happens

Trigger: Calling tcgen05.async_copy_sparse_metadata_smem_to_tmem with an SMEM memref whose second-dimension stride is not a multiple of 16 — e.g. a non-contiguous or manually strided smem allocation, or a shape where K tiles end up spaced at e.g. 12 or 24 elements.

Common situations: Hand-writing Mosaic kernels with sparse tensors on Blackwell (tcgen05) where the SMEM buffer for the sparse metadata (e.g. 128x16 tiles) was sliced, padded, or laid out with non-standard strides; also after refactors that reshape the metadata buffer.

Related errors

• Sparse MMA not supported for M=64
• Sparse MMA unsupported for f32
• B scale shape[0] must be a multiple of 128 and >= N={n * num
• Sparse meta layout loads unsupported.
• Sparse meta layout stores unsupported.

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