jax-ml/jax · error · ValueError

MMA lhs tiling does not fit swizzle. {lhs_tiling=} expected=

Error message

MMA lhs tiling does not fit swizzle. {lhs_tiling=} expected={(8, swizzle_elems)}

What it means

For an SMEM LHS in tcgen05.mma, the tiling transform must equal (8, swizzle_elems) where swizzle_elems = 8 * swizzle_bits / dtype_itemsize_bits. The MMA lowering requires this exact 8-row micro-tile tied to the swizzle.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:2731

      ):
        lhs_transpose = False
      case (
          gpu_core.UnswizzleRef(lhs_swizzle),
          gpu_core.UntilingTransform(lhs_tiling),
          state_types.TransposeTransform((1, 0)),
      ):
        lhs_transpose = True
      case () if isinstance(a_ref, tcgen05.TMEMRef):
        lhs_tiling = None
      case _:
        raise NotImplementedError(
            f"Unsupported transforms for LHS: {a_transforms}."
        )
    if not isinstance(a_ref, tcgen05.TMEMRef):
      assert lhs_swizzle is not None
      swizzle_elems = 8 * lhs_swizzle // dtypes.itemsize_bits(a_dtype)
      if lhs_tiling != (8, swizzle_elems):
        raise ValueError("MMA lhs tiling does not fit swizzle. "
                        f"{lhs_tiling=} expected={(8, swizzle_elems)}")

  assert b_transforms_tree is not None
  b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)
  b_out_ty = state_types.transform_type(b_transforms, b_aval)
  assert isinstance(b_out_ty, state_types.AbstractRef)
  b_dtype = b_out_ty.dtype
  b_transform_avals = b_transforms_tree.unflatten(b_transforms_leaves_avals)
  b_ref, _, b_transforms = lowering._handle_transforms(
      ctx, b_aval, b_ref, b_transform_avals, b_transforms, handle_transposes=False,
      handle_reshapes=True)
  match b_transforms:
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
    ):
      rhs_transpose = False
    case (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the standard load/swizzle utility so the tiling is derived as (8, 8*swizzle//itemsize_bits)
2. Recompute the tiling after changing dtype or swizzle width
3. Match the pattern used in Mosaic GPU matmul reference kernels

Example fix

# before
# custom tiling (16, w) with swizzle=128 -> mismatch
tcgen05.mma(a_custom_tiled, b, acc, k_dim=k)
# after
sw = 128
tile = (8, 8*sw // dtypes.itemsize_bits(a.dtype))
a_smem = load_to_smem(a, swizzle=sw, tiling=tile)
tcgen05.mma(a_smem, b, acc, k_dim=k)

from jax._src import dtypes
swizzle_elems = 8 * swizzle // dtypes.itemsize_bits(a_dtype)
assert lhs_tiling == (8, swizzle_elems), (lhs_tiling, (8, swizzle_elems))

When it happens

Trigger: Loading the A tile with a custom tiling (e.g. (16, x) or unsuitable vectorization) while specifying swizzle=128, producing lhs_tiling != (8, swizzle_elems).

Common situations: Hand-rolled SMEM layouts that deviate from the standard swizzle-compatible micro-tile; changing the element dtype (fp8 vs fp16) without adjusting the tiling width.

Related errors

• MMA rhs tiling does not fit swizzle {rhs_tiling=} expected={
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Swizzle {self.swizzle} requires the trailing dimension to be
• MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_
• packed cannot be specified if layout is specified.

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