jax-ml/jax · error · ValueError

Unsupported swizzle, expected 16, 32, 64 or 128, but got: {s

Error message

Unsupported swizzle, expected 16, 32, 64 or 128, but got: {swizzle}

What it means

Raised by async_copy_smem_to_tmem when the swizzle parameter is not one of the hardware-supported swizzle sizes. NVIDIA tcgen05.cp only implements 16-, 32-, 64- and 128-byte swizzle modes for the SMEM-to-TMEM path; anything else (including 0 for no swizzle) cannot be lowered.

Source

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

    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,
    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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use one of the supported swizzle values: 16, 32, 64, or 128
2. Match the swizzle to the SMEM layout you allocated (the SMEM shape check later enforces swizzle_elems = 8*swizzle/bitwidth consistency)
3. If you truly need unswizzled copies, use a different copy path (e.g. regular async_copy) instead of tcgen05.cp

Example fix

// before
async_copy_smem_to_tmem(smem, tmem, swizzle=0)
// after
async_copy_smem_to_tmem(smem, tmem, swizzle=128)

assert swizzle in {16, 32, 64, 128}, f'swizzle {swizzle} unsupported for tcgen05.cp'

def is_valid_swizzle(s):
    return s in {16, 32, 64, 128}

When it happens

Trigger: Calling tcgen05.async_copy_smem_to_tmem(..., swizzle=N) with N not in {16, 32, 64, 128} — e.g. swizzle=0 for unswizzled access, or swizzle=256 copied from a TMA descriptor config.

Common situations: Porting kernels that previously used unswizzled SMEM loads, or reusing TMA descriptor swizzle values (which also only allow 32/64/128) incorrectly; misreading swizzle as bits vs bytes.

Related errors

• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• B tiling too small. Increase swizzle or transpose the input.
• SMEM has shape {smem_shape}, but expected {expected_smem_sha
• Found multiple SwizzleTransformAttr
• Expected both or neither of scales to be specified.

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