{"record":{"id":"82b0bd74dd006556","repo":"jax-ml/jax","slug":"a-address-calculation-for-multiple-m-tiles","errorCode":null,"errorMessage":"A address calculation for multiple M tiles","messagePattern":"A address calculation for multiple M tiles","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":579,"sourceCode":"  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)\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:","sourceCodeStart":561,"sourceCodeEnd":597,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L561-L597","documentation":"When the A operand of mma() resides in Tensor Memory (TMEMRef), address offsets for multiple M tiles are not implemented; only a single M tile (m_groups == 1) is supported.","triggerScenarios":"Looping over np.ndindex(m_groups, n_groups, k_groups) with m_groups > 1 while a is a TMEMRef.","commonSituations":"Keeping the accumulator or A in TMEM across a large GEMM split into several M tiles; converting an SMEM-A kernel to TMEM-A without splitting the loop.","solutions":["Set m_groups=1 (make the per-call M tile cover the full M)","Keep A in shared memory (SMEM descriptor) instead of TMEM when multiple M tiles are needed"],"exampleFix":"# before\na_tmem = tmem.alloc((m*2, k))\ntcgen05.mma(a_tmem, b, d, m=m*2 ... )  # loop m_groups=2\n# after\nd1, d2 = tmem.alloc(...), tmem.alloc(...)\ntcgen05.mma(a_tmem.slice(0, m), b, d1, ...)\ntcgen05.mma(a_tmem.slice(m, 2*m), b, d2, ...)","handlingStrategy":"validation","validationCode":"assert not (isinstance(a, TMEMRef) and m_groups > 1)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep A in SMEM when tiling M","Slice TMEM A manually per M tile"],"tags":["gpu","mosaic","tcgen05","tmem","not-implemented","mma"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}