jax-ml/jax · error · ValueError

MMA rhs tiling does not fit swizzle {rhs_tiling=} expected={

Error message

MMA rhs tiling does not fit swizzle {rhs_tiling=} expected={(8, swizzle_elems)}

What it means

For the SMEM RHS in tcgen05.mma, the untiling transform must be exactly (8, swizzle_elems) with swizzle_elems = 8 * rhs_swizzle / itemsize_bits(b_dtype), the same 8-row micro-tile constraint as the LHS.

Source

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

  match b_transforms:
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
    ):
      rhs_transpose = False
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
        state_types.TransposeTransform((1, 0)),
    ):
      rhs_transpose = True
    case _:
      raise NotImplementedError(
          f"Unsupported transforms for RHS: {b_transforms}."
      )
  swizzle_elems = 8 * rhs_swizzle // dtypes.itemsize_bits(b_dtype)
  if rhs_tiling != (8, swizzle_elems):
    raise ValueError(
        "MMA rhs tiling does not fit swizzle"
        f" {rhs_tiling=} expected={(8, swizzle_elems)}"
    )

  if barrier_transforms_tree is not None and barrier_ref is not None:
    barrier_transforms = barrier_transforms_tree.unflatten(
        barrier_transforms_leaves
    )
    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
    if base_index is not None:
      barrier_ref = barrier_ref[base_index]

  if lhs_swizzle is None:
    lhs_swizzle = rhs_swizzle
  elif rhs_swizzle != lhs_swizzle:
    raise ValueError("MMA rhs swizzle must match lhs swizzle."
                      f" {lhs_swizzle=} {rhs_swizzle=}")
  if lhs_transpose:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Derive rhs tiling as (8, 8*swizzle // dtypes.itemsize_bits(b.dtype))
Use the provided load-to-SMEM utilities instead of manual tiling
Keep swizzle width consistent between the load and the MMA call

Example fix

# before
b_smem = manual_copy(b, tiling=(8, wrong_width))
tcgen05.mma(a, b_smem, acc, k_dim=k)
# after
sw = 128
b_smem = load_to_smem(b, swizzle=sw)  # tiling auto-derived (8, 8*sw//bits)
tcgen05.mma(a, b_smem, acc, k_dim=k)

Defensive patterns

Strategy: validation

Validate before calling

from jax._src import dtypes
swizzle_elems = 8 * rhs_swizzle // dtypes.itemsize_bits(b_dtype)
assert rhs_tiling == (8, swizzle_elems), (rhs_tiling, (8, swizzle_elems))

Prevention

Derive RHS tiling from swizzle and dtype, don't hardcode
Keep swizzle width consistent between load and MMA

When it happens

Trigger: Loading B with a tiling whose row count is not 8 or whose width doesn't match the swizzle width for the chosen dtype (e.g. fp8 B with fp16-derived tiling).

Common situations: Switching B's dtype between fp16 and fp8 without recomputing the tiling; custom copy kernels that tile SMEM differently from the reference implementation.

Related errors

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