{"record":{"id":"6d9d9ad746368c50","repo":"jax-ml/jax","slug":"b-tiling-too-small-increase-swizzle-or-transpose","errorCode":null,"errorMessage":"B tiling too small. Increase swizzle or transpose the input.","messagePattern":"B tiling too small\\. Increase swizzle or transpose the input\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":555,"sourceCode":"      swizzle=b_swizzle,\n      group_size=(k_group_elems, n_group_elems),\n      logical_k_major=True,\n      mma_bytewidth_k=64 if is_sparse else 32,\n      split_const=True,\n  )\n\n  if is_scaled and utils.bitwidth(mma_a_element_type) == 4:\n    if a_fastest != mma_utils.Dim.K:\n      raise ValueError(\n          \"4-bit block scaled MMA only supports K-fastest operands, but A is M-fastest\"\n      )\n    if b_fastest != mma_utils.Dim.K:\n      raise ValueError(\n          \"4-bit block scaled MMA only supports K-fastest operands, but B is N-fastest\"\n      )\n  if is_sparse:\n    if b_swizzle == 32 and b_fastest == mma_utils.Dim.K:\n      raise NotImplementedError(\n          \"B tiling too small. Increase swizzle or transpose the input.\"\n      )\n\n  # Step 4. Issue the instructions.\n  true = arith.constant(ir.IntegerType.get_signless(1), 1)\n  n_collective_group_elems = n_group_elems * num_cta\n  n_col_groups = n_groups // n_lane_groups\n  assert d.layout.base_tile_shape[0] % 4 == 0\n  lanes_per_n_group = d.layout.base_tile_shape[0] // 4\n  a_sparse_addr_base = a_sparse_metadata.address if is_sparse else None\n  a_scale_addr_base = a_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]\n  b_scale_addr_base = b_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]\n  # B scales are padded when N is short, so it can't be derived from n_collective_group_elems.\n  # Same for A scales when M is short.\n  if is_scaled:\n    assert isinstance(a_scale, TMEMRef) and isinstance(b_scale, TMEMRef)\n    a_scale_m_stride = a_scale.layout.cols_in_shape((a_scale.shape[0], 4), bitwidth=8)\n    b_scale_n_stride = b_scale.layout.cols_in_shape((b_scale.shape[0], 4), bitwidth=8)","sourceCodeStart":537,"sourceCodeEnd":573,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L537-L573","documentation":"In sparse MMA, a B swizzle of 32 bytes combined with K-fastest layout is unsupported because the sparse tensor core needs larger tiles. The fix is a bigger swizzle or a transposed B.","triggerScenarios":"Calling mma() with is_sparse=True, b_swizzle==32, and b_fastest == Dim.K.","commonSituations":"Reusing dense kernel swizzle settings (32B) for sparse 2:4 GEMMs; auto-tuning picking the smallest swizzle.","solutions":["Increase b_swizzle to 64 or 128","Or transpose B so it is N-fastest (if the element type allows)"],"exampleFix":"# before\ntcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=32)\n# after\ntcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=64)","handlingStrategy":"fallback","validationCode":"if is_sparse and b_swizzle == 32 and b_fastest == mma_utils.Dim.K:\n    b_swizzle = 64","typeGuard":null,"tryCatchPattern":"try:\n    tcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=32)\nexcept NotImplementedError:\n    tcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=64)","preventionTips":["Exclude 32B swizzle from sparse autotuning search spaces","Default sparse kernels to swizzle >= 64"],"tags":["gpu","mosaic","tcgen05","sparsity","swizzle","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}