{"record":{"id":"9e1db71cfbb81b01","repo":"jax-ml/jax","slug":"swizzle-b-swizzle-is-too-big-for-mma-with-m-64","errorCode":null,"errorMessage":"Swizzle={b_swizzle} is too big for MMA with M=64. Try lowering it.","messagePattern":"Swizzle=(.+?) is too big for MMA with M=64\\. Try lowering it\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":276,"sourceCode":"    if is_scaled and not collective:\n      raise NotImplementedError(\"MMA with block scaling is not supported for 1CTA M=64\")\n    if is_sparse:\n      raise NotImplementedError(\"Sparse MMA not supported for M=64\")\n    # Watch out: this layout must be consistent with A's layout (up to packing).\n    # 2CTA M=128 instruction uses a different TMEM layout than 1CTA M=64.\n    expected_d_layout = _infer_tmem_layout(d.shape, collective, packing=1)\n    if d.layout != expected_d_layout:\n      raise ValueError(\n          f\"Accumulator layout mismatch: expected {expected_d_layout}, got {d.layout}\"\n      )\n    if collective:\n      n_lane_groups = 1\n    else:\n      n_lane_groups = 2\n      # We can't split N into groups if we would partition it below the tile size.\n      # TODO: We only need to check this if N is the minormost dim in B.\n      if 8 * b_swizzle // utils.bitwidth(a_element_type) > n // n_lane_groups:\n        raise ValueError(\n            f\"Swizzle={b_swizzle} is too big for MMA with M=64. Try\"\n            \" lowering it.\"\n        )\n  else:\n    raise ValueError(f\"Only M=128 and M=64 are supported for MMA, but got M={m}\")\n  f32 = ir.F32Type.get()\n  f16 = ir.F16Type.get()\n  s32 = ir.IntegerType.get_signless(32)\n  elem_type_str = (\n      f\"{a_element_type}\"\n      if a_element_type == b_element_type\n      else f\"({a_element_type}, {b_element_type})\"\n  )\n  if a_element_type == f32 or a_element_type == ir.BF16Type.get():\n    if a_element_type == f32 and is_sparse:\n      raise NotImplementedError(\"Sparse MMA unsupported for f32\")\n    if is_scaled:\n      raise ValueError(","sourceCodeStart":258,"sourceCodeEnd":294,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L258-L294","documentation":"For 1CTA M=64 MMAs the N dimension is split into 2 lane groups, so each group must still hold a full tile column. The check 8 * b_swizzle // bitwidth(a) > n // 2 fires when the B operand's swizzle atom is wider than the per-group N slice.","triggerScenarios":"Calling tcgen05.mma with m=64 (non-collective), a small n (e.g. n < 2 * 8 * b_swizzle / bitwidth), and a large b_swizzle such as 128-byte swizzling on fp16 with narrow N.","commonSituations":"Tuning swizzle values copied from an M=128 kernel; using 128B swizzle with fp8/fp16 operands and narrow N tiles; forgetting that M=64 halves the available N per lane group.","solutions":["Lower b_swizzle (e.g. from 128 to 64 or 32 bytes) on the B operand","Increase n so that n // 2 >= 8 * b_swizzle // bitwidth(a_element_type)","Use M=128 (possibly collective) where the N-split constraint differs"],"exampleFix":"# before\nb_tiled = ... swizzle=128 ...  # with m=64, small n -> raises\ntcgen05.mma(a, b_tiled, d)\n# after\nb_tiled = ... swizzle=32 ...\ntcgen05.mma(a, b_tiled, d)","handlingStrategy":"validation","validationCode":"if m == 64 and not collective:\n    max_swizzle = n * 2 * utils.bitwidth(a_dtype) // 8\n    assert b_swizzle <= max_swizzle, f'b_swizzle={b_swizzle} too large for M=64, n={n}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Parameterize swizzle by (m, n, dtype) rather than hardcoding 128","Validate swizzle against tile shape in kernel config helpers"],"tags":["jax","mosaic","gpu","mma","swizzle","tile-size","tcgen05"],"backgroundTag":"tile-shape-config-invalid","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}