{"record":{"id":"f7a12e9604f70154","repo":"jax-ml/jax","slug":"mma-rhs-tiling-does-not-fit-swizzle-rhs-tiling","errorCode":null,"errorMessage":"MMA rhs tiling does not fit swizzle {rhs_tiling=} expected={(8, swizzle_elems)}","messagePattern":"MMA rhs tiling does not fit swizzle (.+?) expected=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2761,"sourceCode":"  match b_transforms:\n    case (\n        gpu_core.UnswizzleRef(rhs_swizzle),\n        gpu_core.UntilingTransform(rhs_tiling),\n    ):\n      rhs_transpose = False\n    case (\n        gpu_core.UnswizzleRef(rhs_swizzle),\n        gpu_core.UntilingTransform(rhs_tiling),\n        state_types.TransposeTransform((1, 0)),\n    ):\n      rhs_transpose = True\n    case _:\n      raise NotImplementedError(\n          f\"Unsupported transforms for RHS: {b_transforms}.\"\n      )\n  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:","sourceCodeStart":2743,"sourceCodeEnd":2779,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2743-L2779","documentation":"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.","triggerScenarios":"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).","commonSituations":"Switching B's dtype between fp16 and fp8 without recomputing the tiling; custom copy kernels that tile SMEM differently from the reference implementation.","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"],"exampleFix":"# before\nb_smem = manual_copy(b, tiling=(8, wrong_width))\ntcgen05.mma(a, b_smem, acc, k_dim=k)\n# after\nsw = 128\nb_smem = load_to_smem(b, swizzle=sw)  # tiling auto-derived (8, 8*sw//bits)\ntcgen05.mma(a, b_smem, acc, k_dim=k)","handlingStrategy":"validation","validationCode":"from jax._src import dtypes\nswizzle_elems = 8 * rhs_swizzle // dtypes.itemsize_bits(b_dtype)\nassert rhs_tiling == (8, swizzle_elems), (rhs_tiling, (8, swizzle_elems))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive RHS tiling from swizzle and dtype, don't hardcode","Keep swizzle width consistent between load and MMA"],"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"}