{"record":{"id":"581fb4256b151559","repo":"jax-ml/jax","slug":"mma-rhs-swizzle-must-match-lhs-swizzle-lhs-swizz","errorCode":null,"errorMessage":"MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_swizzle=}","messagePattern":"MMA rhs swizzle must match lhs swizzle\\. (.+?) (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2777,"sourceCode":"  swizzle_elems = 8 * rhs_swizzle // dtypes.itemsize_bits(b_dtype)\n  if rhs_tiling != (8, swizzle_elems):\n    raise ValueError(\n        \"MMA rhs tiling does not fit swizzle\"\n        f\" {rhs_tiling=} expected={(8, swizzle_elems)}\"\n    )\n\n  if barrier_transforms_tree is not None and barrier_ref is not None:\n    barrier_transforms = barrier_transforms_tree.unflatten(\n        barrier_transforms_leaves\n    )\n    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)\n    if base_index is not None:\n      barrier_ref = barrier_ref[base_index]\n\n  if lhs_swizzle is None:\n    lhs_swizzle = rhs_swizzle\n  elif rhs_swizzle != lhs_swizzle:\n    raise ValueError(\"MMA rhs swizzle must match lhs swizzle.\"\n                      f\" {lhs_swizzle=} {rhs_swizzle=}\")\n  if lhs_transpose:\n    if isinstance(a_ref, tcgen05.TMEMRef):\n      raise ValueError(\"TMEM transpose not allowed.\")\n    a_ref = mgpu.memref_transpose(a_ref, (1, 0, 3, 2))\n  if rhs_transpose:\n    b_ref = mgpu.memref_transpose(b_ref, (1, 0, 3, 2))\n  if isinstance(accumulate, bool):\n    accumulate = mgpu.c(accumulate, ir.IntegerType.get_signless(1))\n  elif isinstance(accumulate, mgpu.FragmentedArray):\n    accumulate = accumulate.registers.item()\n    assert isinstance(accumulate, ir.Value)\n\n  if a_scale_ref is not None and a_scale_transforms_tree is not None:\n    assert isinstance(a_scale_ref_aval, state.AbstractRef)\n    a_scale_transforms = a_scale_transforms_tree.unflatten(\n        a_scale_transforms_leaves\n    )","sourceCodeStart":2759,"sourceCodeEnd":2795,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2759-L2795","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Set both lhs_swizzle and rhs_swizzle to the same value (or leave one as None so it inherits the other)","Build both operands from the same memref layout / copy pipeline so swizzles are derived identically","Check the swizzle mode reported in the message and adjust your mgpu.memref with the matching swizzle attribute"],"exampleFix":"// before\ntcgen05_mma(lhs, rhs, acc, lhs_swizzle=SwizzleMode.W32, rhs_swizzle=SwizzleMode.W64)\n// after\ntcgen05_mma(lhs, rhs, acc, lhs_swizzle=SwizzleMode.W32, rhs_swizzle=SwizzleMode.W32)","handlingStrategy":"validation","validationCode":"assert lhs_swizzle is None or rhs_swizzle is None or lhs_swizzle == rhs_swizzle, 'swizzle mismatch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive both operand swizzles from the same memref layout helper","Log swizzle modes of both operands before the MMA call"],"tags":["jax","pallas","mosaic-gpu","tcgen05","mma","swizzle"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}