jax-ml/jax · error · NotImplementedError

A sparse metadata address calculation for multiple tiles

Error message

A sparse metadata address calculation for multiple tiles

What it means

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.

Source

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

    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:
      if m_groups != 1:
        raise NotImplementedError("A scale address calculation for multiple M tiles")
      if n_groups != 1:
        raise NotImplementedError("B scale address calculation for multiple N tiles")
      assert scale_block is not None  # For type checkers.
      assert k_group_elems % (scale_block * 4) == 0
      assert m_group_elems % 32 == 0 and n_group_elems % (8 * num_cta) == 0
      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),

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
tcgen05.mma(a, b_big, d, a_sparse_metadata=meta)  # n_groups=2
# after
for ni in range(n_tiles):
  tcgen05.mma(a, b.slice(ni), d.slice(ni), a_sparse_metadata=meta)

Defensive patterns

Strategy: validation

Validate before calling

assert m_groups == 1 and n_groups == 1  # required for sparse metadata addressing

Prevention

Manual M/N loops outside mma for sparse GEMMs
Only tile K inside a single sparse mma call

When it happens

Trigger: Calling sparse mma() where the tile loop produces n_groups != 1 or m_groups != 1 while a_sparse_addr_base is set.

Common situations: Sparse GEMM with N or M larger than one tile, e.g. N=256 with per-tile n=128.

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

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