jax-ml/jax · error · NotImplementedError

A address calculation for multiple M tiles

Error message

A address calculation for multiple M tiles

What it means

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.

Source

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

  n_collective_group_elems = n_group_elems * num_cta
  n_col_groups = n_groups // n_lane_groups
  assert d.layout.base_tile_shape[0] % 4 == 0
  lanes_per_n_group = d.layout.base_tile_shape[0] // 4
  a_sparse_addr_base = a_sparse_metadata.address if is_sparse else None
  a_scale_addr_base = a_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]
  b_scale_addr_base = b_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]
  # B scales are padded when N is short, so it can't be derived from n_collective_group_elems.
  # Same for A scales when M is short.
  if is_scaled:
    assert isinstance(a_scale, TMEMRef) and isinstance(b_scale, TMEMRef)
    a_scale_m_stride = a_scale.layout.cols_in_shape((a_scale.shape[0], 4), bitwidth=8)
    b_scale_n_stride = b_scale.layout.cols_in_shape((b_scale.shape[0], 4), bitwidth=8)
  else:
    a_scale_m_stride = b_scale_n_stride = None
  for mi, ni, ki in np.ndindex(m_groups, n_groups, k_groups):
    if isinstance(a, TMEMRef):
      if m_groups != 1:
        raise NotImplementedError("A address calculation for multiple M tiles")
      a_k_group_elems = k_group_elems // (1 + is_sparse)
      a_mk = a.slice(slice(None), utils.ds(ki * a_k_group_elems, a_k_group_elems)).address
    else:
      assert a_desc_base is not None
      a_offset = mi * a_m_group_stride + ki * a_k_group_stride
      a_mk = (a_desc_base[0], a_desc_base[1] + mma_utils.encode_addr(a_offset))
    b_offset = ni * b_n_group_stride + ki * b_k_group_stride
    b_nk = (b_desc_base[0], b_desc_base[1] + mma_utils.encode_addr(b_offset))
    if a_sparse_addr_base is not None:
      if n_groups != 1 or m_groups != 1:
        raise NotImplementedError("A sparse metadata address calculation for multiple tiles")
      sparse_group_elems = 8 if utils.bitwidth(mma_a_element_type) == 4 else 4
      # Each sparse group has 2 entries, each TMEM column holds 16 i2 entries.
      cols_per_k_group = k_group_elems // sparse_group_elems * 2 // 16
      a_sparse_addr = arith.addi(a_sparse_addr_base, utils.c(ki * cols_per_k_group, i32))
    else:
      a_sparse_addr = None
    if a_scale_addr_base is not None and b_scale_addr_base is not None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set m_groups=1 (make the per-call M tile cover the full M)
2. Keep A in shared memory (SMEM descriptor) instead of TMEM when multiple M tiles are needed

Example fix

# before
a_tmem = tmem.alloc((m*2, k))
tcgen05.mma(a_tmem, b, d, m=m*2 ... )  # loop m_groups=2
# after
d1, d2 = tmem.alloc(...), tmem.alloc(...)
tcgen05.mma(a_tmem.slice(0, m), b, d1, ...)
tcgen05.mma(a_tmem.slice(m, 2*m), b, d2, ...)

assert not (isinstance(a, TMEMRef) and m_groups > 1)

When it happens

Trigger: Looping over np.ndindex(m_groups, n_groups, k_groups) with m_groups > 1 while a is a TMEMRef.

Common situations: 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.

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

• D address calculation for multiple M tiles
• Expected both or neither of scales to be specified.
• Sparse MMA not supported for M=64
• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• Only M=128 and M=64 are supported for MMA, but got M={m}

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