jax-ml/jax · error · ValueError

MMA rhs swizzle must match lhs swizzle.

Error message

MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_swizzle=}

What it means

In the Mosaic GPU (Pallas) tcgen05 MMA lowering, the LHS and RHS operands of a Blackwell tcgen05 matrix-multiply must share the same swizzle pattern. If only one side specifies a swizzle it is inherited from the other, but if both are given and differ, the lowering refuses to continue because the Tensor Core memory layout would be inconsistent.

Solutions

1. Set both lhs_swizzle and rhs_swizzle to the same value (or leave one as None so it inherits the other)
2. Build both operands from the same memref layout / copy pipeline so swizzles are derived identically
3. Check the swizzle mode reported in the message and adjust your mgpu.memref with the matching swizzle attribute

Example fix

// before
tcgen05_mma(lhs, rhs, acc, lhs_swizzle=SwizzleMode.W32, rhs_swizzle=SwizzleMode.W64)
// after
tcgen05_mma(lhs, rhs, acc, lhs_swizzle=SwizzleMode.W32, rhs_swizzle=SwizzleMode.W32)

assert lhs_swizzle is None or rhs_swizzle is None or lhs_swizzle == rhs_swizzle, 'swizzle mismatch'

When it happens

Trigger: Calling tcgen05_mma (or the tcgen05 MMA Pallas primitive) where the lhs Swizzle (e.g. SwizzleMode.W32_ATTRIBUTE etc.) differs from the rhs Swizzle, e.g. lhs built with a 32B swizzle and rhs with 64B or no swizzle.

Common situations: Mixing SMEM layouts when one operand comes from TMEM/SMEM with one tiling and the other from an async copy with a different tiling; upgrading JAX versions where swizzle defaults changed; hand-writing a tiled GEMM kernel.

Related errors

• Accumulator and LHS have incompatible shapes. Expected LHS…
• Accumulator and RHS have incompatible shapes. Expected RHS…
• async_store_smem requires a tiled and swizzled ref
• Can't transpose the swizzled dimension.
• Expected metadata dtype to be uint2, got

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

Appendix: source

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

  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:
    if isinstance(a_ref, tcgen05.TMEMRef):
      raise ValueError("TMEM transpose not allowed.")
    a_ref = mgpu.memref_transpose(a_ref, (1, 0, 3, 2))
  if rhs_transpose:
    b_ref = mgpu.memref_transpose(b_ref, (1, 0, 3, 2))
  if isinstance(accumulate, bool):
    accumulate = mgpu.c(accumulate, ir.IntegerType.get_signless(1))
  elif isinstance(accumulate, mgpu.FragmentedArray):
    accumulate = accumulate.registers.item()
    assert isinstance(accumulate, ir.Value)

  if a_scale_ref is not None and a_scale_transforms_tree is not None:
    assert isinstance(a_scale_ref_aval, state.AbstractRef)
    a_scale_transforms = a_scale_transforms_tree.unflatten(
        a_scale_transforms_leaves
    )

View on GitHub (pinned to 1e1c6a8fc0)