{"record":{"id":"20ae10b238f172ed","repo":"jax-ml/jax","slug":"a-sparse-metadata-address-calculation-for-multiple","errorCode":null,"errorMessage":"A sparse metadata address calculation for multiple tiles","messagePattern":"A sparse metadata address calculation for multiple tiles","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":590,"sourceCode":"    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)\n  else:\n    a_scale_m_stride = b_scale_n_stride = None\n  for mi, ni, ki in np.ndindex(m_groups, n_groups, k_groups):\n    if isinstance(a, TMEMRef):\n      if m_groups != 1:\n        raise NotImplementedError(\"A address calculation for multiple M tiles\")\n      a_k_group_elems = k_group_elems // (1 + is_sparse)\n      a_mk = a.slice(slice(None), utils.ds(ki * a_k_group_elems, a_k_group_elems)).address\n    else:\n      assert a_desc_base is not None\n      a_offset = mi * a_m_group_stride + ki * a_k_group_stride\n      a_mk = (a_desc_base[0], a_desc_base[1] + mma_utils.encode_addr(a_offset))\n    b_offset = ni * b_n_group_stride + ki * b_k_group_stride\n    b_nk = (b_desc_base[0], b_desc_base[1] + mma_utils.encode_addr(b_offset))\n    if a_sparse_addr_base is not None:\n      if n_groups != 1 or m_groups != 1:\n        raise NotImplementedError(\"A sparse metadata address calculation for multiple tiles\")\n      sparse_group_elems = 8 if utils.bitwidth(mma_a_element_type) == 4 else 4\n      # Each sparse group has 2 entries, each TMEM column holds 16 i2 entries.\n      cols_per_k_group = k_group_elems // sparse_group_elems * 2 // 16\n      a_sparse_addr = arith.addi(a_sparse_addr_base, utils.c(ki * cols_per_k_group, i32))\n    else:\n      a_sparse_addr = None\n    if a_scale_addr_base is not None and b_scale_addr_base is not None:\n      if m_groups != 1:\n        raise NotImplementedError(\"A scale address calculation for multiple M tiles\")\n      if n_groups != 1:\n        raise NotImplementedError(\"B scale address calculation for multiple N tiles\")\n      assert scale_block is not None  # For type checkers.\n      assert k_group_elems % (scale_block * 4) == 0\n      assert m_group_elems % 32 == 0 and n_group_elems % (8 * num_cta) == 0\n      k_scales_per_group = k_group_elems // (scale_block * 4)\n      a_scale_addr = arith.addi(\n          a_scale_addr_base,\n          utils.c(ki * k_scales_per_group * a_scale_m_stride, i32),","sourceCodeStart":572,"sourceCodeEnd":608,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L572-L608","documentation":"Sparse metadata TMEM addressing across multiple M or N tiles is not implemented in tcgen05.mma; the sparse path requires m_groups == 1 and n_groups == 1.","triggerScenarios":"Calling sparse mma() where the tile loop produces n_groups != 1 or m_groups != 1 while a_sparse_addr_base is set.","commonSituations":"Sparse GEMM with N or M larger than one tile, e.g. N=256 with per-tile n=128.","solutions":["Restructure the kernel so each mma call covers exactly one M and one N tile (loop manually outside mma)","Keep K as the only tiled dimension inside a single mma call"],"exampleFix":"# before\ntcgen05.mma(a, b_big, d, a_sparse_metadata=meta)  # n_groups=2\n# after\nfor ni in range(n_tiles):\n  tcgen05.mma(a, b.slice(ni), d.slice(ni), a_sparse_metadata=meta)","handlingStrategy":"validation","validationCode":"assert m_groups == 1 and n_groups == 1  # required for sparse metadata addressing","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Manual M/N loops outside mma for sparse GEMMs","Only tile K inside a single sparse mma call"],"tags":["gpu","mosaic","tcgen05","sparsity","not-implemented","tiling"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}