{"record":{"id":"e8d02de0c1e7250b","repo":"jax-ml/jax","slug":"columns-must-be-a-multiple-of-16-got-columns","errorCode":null,"errorMessage":"Columns must be a multiple of 16, got: {columns}","messagePattern":"Columns must be a multiple of 16, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1101,"sourceCode":"\ndef tmem_default_layout(packing: int = 1) -> TMEMLayout:\n  \"\"\"A TMEM layout used for 1CTA MMA with M=128 and 2CTA MMA with M=256.\"\"\"\n  if packing.bit_count() != 1:\n    raise ValueError(f\"Packing must be a power of 2, got: {packing}\")\n  return TMEMLayout(\n      fa.Tiling(((TMEM_ROWS, packing), (fa.WARP_SIZE, packing))),\n      warp_dims=(-4,),\n      lane_dims=(-2,),\n      vector_dim=-1,\n  )\n\n\ndef tmem_half_lane_layout(columns, packing: int = 1) -> TMEMLayout:\n  \"\"\"A TMEM layout used for 1CTA MMA with M=64.\"\"\"\n  if packing > (columns // 2) or packing.bit_count() != 1:\n    raise ValueError(f\"Packing must be <= 8 and a power of 2, got: {packing}\")\n  if columns % 16:\n    raise ValueError(f\"Columns must be a multiple of 16, got: {columns}\")\n  return TMEMLayout(\n      fa.Tiling((\n          (TMEM_ROWS // 2, columns),\n          (fa.WARP_SIZE // 2, columns // 2),\n          (packing,),\n      )),\n      warp_dims=(-5,),\n      lane_dims=(-4, -3),\n      vector_dim=-1,\n  )\n\n\ndef tmem_m64_collective_layout(columns: int, packing: int = 1) -> TMEMLayout:\n  \"\"\"A TMEM layout used for 2CTA MMA with M=128.\"\"\"\n  if packing > 8 or packing.bit_count() != 1:\n    raise ValueError(f\"Packing must be <= 8 and a power of 2, got: {packing}\")\n  if columns % 16:\n    raise ValueError(f\"Columns must be a multiple of 16, got: {columns}\")","sourceCodeStart":1083,"sourceCodeEnd":1119,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1083-L1119","documentation":"tmem_half_lane_layout requires the column count to be a multiple of 16 because the underlying TMEM tiling for the 64-row layout uses 16-column granularity; a non-multiple would leave a partial tile that the hardware layout cannot express.","triggerScenarios":"Calling tmem_half_lane_layout(columns=10) or any columns % 16 != 0; indirectly from tcgen05.load/store or infer_tmem_layout on a 64-row TMEMRef whose second dimension isn't a multiple of 16.","commonSituations":"MMA with M=64 and an N dimension like 8, 24, or 40 (common with small tiles or f8 dtypes) resulting in a TMEM shape violating the 16-column alignment.","solutions":["Pad N (columns) up to the next multiple of 16","Choose tile sizes where N % 16 == 0 (e.g. 16, 32, 64)","If partial results are needed, load the full padded region and mask in registers"],"exampleFix":"# before\nlayout = tcgen05.tmem_half_lane_layout(columns=24)\n# after\nlayout = tcgen05.tmem_half_lane_layout(columns=32)","handlingStrategy":"validation","validationCode":"assert columns % 16 == 0, f'columns must be multiple of 16, got {columns}'","typeGuard":"def columns_ok(columns: int) -> bool:\n    return columns % 16 == 0","tryCatchPattern":null,"preventionTips":["Pick N tile sizes from {16, 32, 64, 128}","Validate all TMEM shapes against the 16-column rule in one place"],"tags":["jax","mosaic","tmem","layout","alignment"],"backgroundTag":"shape-alignment-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}