jax-ml/jax · error · NotImplementedError

D address calculation for multiple M tiles

Error message

D address calculation for multiple M tiles

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:623

      k_scales_per_group = k_group_elems // (scale_block * 4)
      a_scale_addr = arith.addi(
          a_scale_addr_base,
          utils.c(ki * k_scales_per_group * a_scale_m_stride, i32),
      )
      b_scale_addr = arith.addi(
          b_scale_addr_base,
          utils.c(ki * k_scales_per_group * b_scale_n_stride, i32)
      )
    else:
      a_scale_addr = b_scale_addr = None
    acc = accumulate if ki == 0 else true
    ni_lane_group, ni_col = ni // n_col_groups, ni % n_col_groups
    d_offset = (
        ((ni_lane_group * lanes_per_n_group) << 16)
        + ni_col * n_collective_group_elems
    )
    if m_groups != 1:
      raise NotImplementedError("D address calculation for multiple M tiles")
    _do_mma(
        arith.addi(d.address, arith.constant(i32, d_offset)),
        a_mk,
        b_nk,
        d_type=d.dtype,
        m=m_group_elems,
        n=n_group_elems,
        k=k_group_elems,
        collective=collective,
        a_transpose=a_fastest != mma_utils.Dim.K,
        b_transpose=b_fastest != mma_utils.Dim.K,
        a_k_strides=a_k_instr_strides,
        b_k_strides=b_k_instr_strides,
        a_scale_addr=a_scale_addr,
        b_scale_addr=b_scale_addr,
        a_scale_m_stride=a_scale_m_stride,
        b_scale_n_stride=b_scale_n_stride,
        a_sparse_addr=a_sparse_addr,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Issue one mma per M tile, slicing D accordingly
2. Verify m_group_elems matches D's M extent so only one group is needed

Example fix

# before
tcgen05.mma(a, b, d, m=256)  # m_groups=2
# after
tcgen05.mma(a.slice(0,128), b, d.slice(0,128), m=128)
tcgen05.mma(a.slice(128,128), b, d.slice(128,128), m=128)

assert m_groups == 1  # D TMEM addressing supports single M tile

When it happens

Trigger: Calling mma() with a D tensor and tiling that yields m_groups > 1.

Common situations: D (accumulator in TMEM) covering M larger than a single instruction tile; porting kernels that relied on automatic M tiling.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• A address calculation for multiple M tiles
• A sparse metadata address calculation for multiple tiles
• A scale address calculation for multiple M tiles
• B scale address calculation for multiple N tiles
• Unsupported element type for block scaling: {a_element_type}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/0f9fd87be716eb6c. Report an issue: GitHub.