{"record":{"id":"e7246a89fa302fa8","repo":"jax-ml/jax","slug":"mma-lhs-tiling-does-not-fit-swizzle-lhs-tiling","errorCode":null,"errorMessage":"MMA lhs tiling does not fit swizzle. {lhs_tiling=} expected={(8, swizzle_elems)}","messagePattern":"MMA lhs tiling does not fit swizzle\\. (.+?) expected=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2731,"sourceCode":"      ):\n        lhs_transpose = False\n      case (\n          gpu_core.UnswizzleRef(lhs_swizzle),\n          gpu_core.UntilingTransform(lhs_tiling),\n          state_types.TransposeTransform((1, 0)),\n      ):\n        lhs_transpose = True\n      case () if isinstance(a_ref, tcgen05.TMEMRef):\n        lhs_tiling = None\n      case _:\n        raise NotImplementedError(\n            f\"Unsupported transforms for LHS: {a_transforms}.\"\n        )\n    if not isinstance(a_ref, tcgen05.TMEMRef):\n      assert lhs_swizzle is not None\n      swizzle_elems = 8 * lhs_swizzle // dtypes.itemsize_bits(a_dtype)\n      if lhs_tiling != (8, swizzle_elems):\n        raise ValueError(\"MMA lhs tiling does not fit swizzle. \"\n                        f\"{lhs_tiling=} expected={(8, swizzle_elems)}\")\n\n  assert b_transforms_tree is not None\n  b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)\n  b_out_ty = state_types.transform_type(b_transforms, b_aval)\n  assert isinstance(b_out_ty, state_types.AbstractRef)\n  b_dtype = b_out_ty.dtype\n  b_transform_avals = b_transforms_tree.unflatten(b_transforms_leaves_avals)\n  b_ref, _, b_transforms = lowering._handle_transforms(\n      ctx, b_aval, b_ref, b_transform_avals, b_transforms, handle_transposes=False,\n      handle_reshapes=True)\n  match b_transforms:\n    case (\n        gpu_core.UnswizzleRef(rhs_swizzle),\n        gpu_core.UntilingTransform(rhs_tiling),\n    ):\n      rhs_transpose = False\n    case (","sourceCodeStart":2713,"sourceCodeEnd":2749,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2713-L2749","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Use the standard load/swizzle utility so the tiling is derived as (8, 8*swizzle//itemsize_bits)","Recompute the tiling after changing dtype or swizzle width","Match the pattern used in Mosaic GPU matmul reference kernels"],"exampleFix":"# before\n# custom tiling (16, w) with swizzle=128 -> mismatch\ntcgen05.mma(a_custom_tiled, b, acc, k_dim=k)\n# after\nsw = 128\ntile = (8, 8*sw // dtypes.itemsize_bits(a.dtype))\na_smem = load_to_smem(a, swizzle=sw, tiling=tile)\ntcgen05.mma(a_smem, b, acc, k_dim=k)","handlingStrategy":"validation","validationCode":"from jax._src import dtypes\nswizzle_elems = 8 * swizzle // dtypes.itemsize_bits(a_dtype)\nassert lhs_tiling == (8, swizzle_elems), (lhs_tiling, (8, swizzle_elems))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive tiling from swizzle and dtype programmatically","Recheck layout math whenever dtype or swizzle changes"],"tags":["jax","pallas","tcgen05","swizzle","tiling","layout"],"backgroundTag":"gpu-shared-memory-layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}