jax-ml/jax · error · NotImplementedError

B tiling too small. Increase swizzle or transpose the input.

Error message

B tiling too small. Increase swizzle or transpose the input.

What it means

In sparse MMA, a B swizzle of 32 bytes combined with K-fastest layout is unsupported because the sparse tensor core needs larger tiles. The fix is a bigger swizzle or a transposed B.

Source

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

      swizzle=b_swizzle,
      group_size=(k_group_elems, n_group_elems),
      logical_k_major=True,
      mma_bytewidth_k=64 if is_sparse else 32,
      split_const=True,
  )

  if is_scaled and utils.bitwidth(mma_a_element_type) == 4:
    if a_fastest != mma_utils.Dim.K:
      raise ValueError(
          "4-bit block scaled MMA only supports K-fastest operands, but A is M-fastest"
      )
    if b_fastest != mma_utils.Dim.K:
      raise ValueError(
          "4-bit block scaled MMA only supports K-fastest operands, but B is N-fastest"
      )
  if is_sparse:
    if b_swizzle == 32 and b_fastest == mma_utils.Dim.K:
      raise NotImplementedError(
          "B tiling too small. Increase swizzle or transpose the input."
      )

  # Step 4. Issue the instructions.
  true = arith.constant(ir.IntegerType.get_signless(1), 1)
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Increase b_swizzle to 64 or 128
2. Or transpose B so it is N-fastest (if the element type allows)

Example fix

# before
tcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=32)
# after
tcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=64)

if is_sparse and b_swizzle == 32 and b_fastest == mma_utils.Dim.K:
    b_swizzle = 64

try:
    tcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=32)
except NotImplementedError:
    tcgen05.mma(a, b, d, a_sparse_metadata=meta, b_swizzle=64)

When it happens

Trigger: Calling mma() with is_sparse=True, b_swizzle==32, and b_fastest == Dim.K.

Common situations: Reusing dense kernel swizzle settings (32B) for sparse 2:4 GEMMs; auto-tuning picking the smallest swizzle.

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 sparse metadata address calculation for multiple tiles
• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• Unsupported element type for block scaling: {a_element_type}
• A sparse metadata shape mismatch: expected {(m, expected_met
• A sparse metadata dtype mismatch: expected i2, got {a_sparse

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