{"record":{"id":"34d45deac8b57967","repo":"jax-ml/jax","slug":"k-tile-stride-must-be-a-multiple-of-16","errorCode":null,"errorMessage":"K tile stride must be a multiple of 16","messagePattern":"K tile stride must be a multiple of 16","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1978,"sourceCode":"  if tmem_ref.shape[1] % 64:\n    raise ValueError(f\"TMEM reference must have a multiple of 64 colums, but got {tmem_ref.shape[1]}\")\n  if tmem_ref.layout != sparse_meta_layout():\n    raise ValueError(f\"TMEM layout {tmem_ref.layout} is not supported\")\n  smem_shape = tuple(smem_ty.shape)\n  expected_smem_shape = (tmem_ref.shape[0] // 128, tmem_ref.shape[1] // 64, 128, 64)\n  if smem_shape != expected_smem_shape:\n    raise NotImplementedError(\n        f\"SMEM has {smem_shape}, but expected {expected_smem_shape} for TMEM\"\n        f\" ref shape {tmem_ref.shape}\"\n    )\n  strides, _ = smem_ty.get_strides_and_offset()\n  if strides != utils.get_contiguous_strides(smem_shape):\n    raise ValueError(\"Only copies from contiguous SMEM references are supported\")\n  if expected_smem_shape[0] != 1:\n    raise NotImplementedError(\"Only M=128 supported\")\n  k_tile_stride = strides[1]\n  if k_tile_stride % 16:\n    raise ValueError(\"K tile stride must be a multiple of 16\")\n  k_tile_byte_stride = k_tile_stride // 4\n  for k_tile in range(expected_smem_shape[1]):\n    load_ptr = utils.getelementptr(\n        utils.memref_ptr(smem_ref), [k_tile * k_tile_byte_stride], i8\n    )\n    store_ptr = arith.addi(tmem_ref.address, arith.constant(i32, 4 * k_tile))\n    # The \"core matrix\" here is the same as in MMA: 8x(16 bytes).\n    desc = mma_utils.encode_descriptor(load_ptr, 0, 8 * 16, swizzle=None)\n    ptr = _tmem_addr_to_ptr(store_ptr)\n    nvvm.tcgen05_cp(\n        nvvm.Tcgen05CpShape.SHAPE_128x128b, ptr, desc,\n        group=nvvm.CTAGroupKind.CTA_2 if collective else nvvm.CTAGroupKind.CTA_1\n    )\n\n\ndef async_copy_smem_to_tmem(\n    smem_ref: ir.Value,\n    tmem_ref: TMEMRef,","sourceCodeStart":1960,"sourceCodeEnd":1996,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1960-L1996","documentation":"Raised by async_copy_sparse_metadata_smem_to_tmem in Mosaic GPU's tcgen05 module when copying sparse tensor metadata from shared memory (SMEM) to tensor memory (TMEM). The K tile stride (strides[1] of the SMEM memref) must be divisible by 16 because the tcgen05 hardware instruction that loads sparse metadata requires 16-element-aligned tile spacing. The stride is derived from the SMEM layout you allocated, so an oddly-strided buffer fails this check.","triggerScenarios":"Calling tcgen05.async_copy_sparse_metadata_smem_to_tmem with an SMEM memref whose second-dimension stride is not a multiple of 16 — e.g. a non-contiguous or manually strided smem allocation, or a shape where K tiles end up spaced at e.g. 12 or 24 elements.","commonSituations":"Hand-writing Mosaic kernels with sparse tensors on Blackwell (tcgen05) where the SMEM buffer for the sparse metadata (e.g. 128x16 tiles) was sliced, padded, or laid out with non-standard strides; also after refactors that reshape the metadata buffer.","solutions":["Reallocate or reshape the SMEM buffer so the K-tile dimension has a stride that is a multiple of 16 (e.g. pad the trailing dimension)","Ensure the SMEM reference is contiguous by allocating with a shape matching utils.get_contiguous_strides, since only contiguous SMEM is supported anyway","If using tmem_alloc / smem allocation helpers, pass a shape whose innermost tile size divides 16 evenly (e.g. K tiles of 16)"],"exampleFix":"// before\nsmem = smem_alloc(i32, (1, 10))  # stride 10 -> not multiple of 16\n// after\nsmem = smem_alloc(i32, (1, 16))  # K tile stride 16 -> passes check","handlingStrategy":"validation","validationCode":"strides, _ = ir.MemRefType(smem_ref.type).get_strides_and_offset()\nif strides[1] % 16:\n    raise ValueError(f'K tile stride {strides[1]} not multiple of 16; reallocate SMEM')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always allocate sparse-metadata SMEM buffers with contiguous 16-aligned K tiles","Assert stride divisibility in kernel setup code rather than relying on the lowering to fail"],"tags":["gpu","mosaic","tcgen05","sparse","shared-memory","alignment"],"backgroundTag":"gpu-memory-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}