{"record":{"id":"0f9fd87be716eb6c","repo":"jax-ml/jax","slug":"d-address-calculation-for-multiple-m-tiles","errorCode":null,"errorMessage":"D address calculation for multiple M tiles","messagePattern":"D address calculation for multiple M tiles","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":623,"sourceCode":"      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),\n      )\n      b_scale_addr = arith.addi(\n          b_scale_addr_base,\n          utils.c(ki * k_scales_per_group * b_scale_n_stride, i32)\n      )\n    else:\n      a_scale_addr = b_scale_addr = None\n    acc = accumulate if ki == 0 else true\n    ni_lane_group, ni_col = ni // n_col_groups, ni % n_col_groups\n    d_offset = (\n        ((ni_lane_group * lanes_per_n_group) << 16)\n        + ni_col * n_collective_group_elems\n    )\n    if m_groups != 1:\n      raise NotImplementedError(\"D address calculation for multiple M tiles\")\n    _do_mma(\n        arith.addi(d.address, arith.constant(i32, d_offset)),\n        a_mk,\n        b_nk,\n        d_type=d.dtype,\n        m=m_group_elems,\n        n=n_group_elems,\n        k=k_group_elems,\n        collective=collective,\n        a_transpose=a_fastest != mma_utils.Dim.K,\n        b_transpose=b_fastest != mma_utils.Dim.K,\n        a_k_strides=a_k_instr_strides,\n        b_k_strides=b_k_instr_strides,\n        a_scale_addr=a_scale_addr,\n        b_scale_addr=b_scale_addr,\n        a_scale_m_stride=a_scale_m_stride,\n        b_scale_n_stride=b_scale_n_stride,\n        a_sparse_addr=a_sparse_addr,","sourceCodeStart":605,"sourceCodeEnd":641,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L605-L641","documentation":"Computing the TMEM address offset of the accumulator D across multiple M tiles is not implemented in tcgen05.mma; the D path requires m_groups == 1.","triggerScenarios":"Calling mma() with a D tensor and tiling that yields m_groups > 1.","commonSituations":"D (accumulator in TMEM) covering M larger than a single instruction tile; porting kernels that relied on automatic M tiling.","solutions":["Issue one mma per M tile, slicing D accordingly","Verify m_group_elems matches D's M extent so only one group is needed"],"exampleFix":"# before\ntcgen05.mma(a, b, d, m=256)  # m_groups=2\n# after\ntcgen05.mma(a.slice(0,128), b, d.slice(0,128), m=128)\ntcgen05.mma(a.slice(128,128), b, d.slice(128,128), m=128)","handlingStrategy":"validation","validationCode":"assert m_groups == 1  # D TMEM addressing supports single M tile","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Issue one mma per M tile and slice D accordingly","Keep D tile M equal to instruction m"],"tags":["gpu","mosaic","tcgen05","tmem","not-implemented","tiling"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}