jax-ml/jax · error · ValueError

Unsupported TMA index shape {shape}

Error message

Unsupported TMA index shape {shape}

What it means

For async TMA load/store/prefetch ops, vector-valued indices must have a first dimension divisible by 16 (TMA_INDICES_LAYOUT) or by 4 (TMA_INDICES_4_LAYOUT). If neither divides shape[0], no register layout exists for the index vector and a ValueError is raised.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2199

      base_operand_index = 3
    case mgpu.AsyncStoreOp():
      base_operand_index = 2
    case mgpu.AsyncPrefetchOp():
      base_operand_index = 1

  for i, idx in enumerate(op.indices):
    if isinstance(idx.type, ir.VectorType):
      value_site = ValueSite(op, VariableType.OPERAND, base_operand_index + i)
      value_site_var = cs.Variable(value_site)
      shape = tuple(idx.type.shape)

      allowed_layouts = []
      if shape[0] % 16 == 0:
        allowed_layouts.append(cs.RegisterLayout(value=fa.TMA_INDICES_LAYOUT))
      if shape[0] % 4 == 0:
        allowed_layouts.append(cs.RegisterLayout(value=fa.TMA_INDICES_4_LAYOUT))
      if not allowed_layouts:
        raise ValueError(f"Unsupported TMA index shape {shape}")
      values_sites[value_site_var] = [value_site]
      constraints.append(cs.OneOf(value_site_var, tuple(allowed_layouts)))
  return values_sites, constraints


@_add_constraint_system_derivation_rule(mgpu.AsyncLoadOp)
@_add_constraint_system_derivation_rule(mgpu.AsyncStoreOp)
def _async_load_store_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncLoadOp | mgpu.AsyncStoreOp,
) -> ConstraintSystemDerivationRuleResult:
  # We only support 2D gathers/scatters along the leading dimension. Tiling
  # either keeps the gather/scatter dimension leading or allows
  # collapsing leading dimensions to maintain contiguity without
  # transforming global memory.
  tiling_multiple = []
  for i, (size, index) in enumerate(zip(op.slice_lengths, op.indices, strict=True)):
    if size == -1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or reshape the index vector so its first dimension is a multiple of 4 (or 16) — pad with dummy/unused indices and mask them out
2. Choose gather batch sizes that are multiples of 16 to get the more efficient TMA_INDICES_LAYOUT
3. Split the gather into chunks whose index counts are multiples of 4/16

Example fix

# before: idx shape (6, ...) -> ValueError
mgpu.async_load(source, smem, indices=idx6, slice_lengths=(1, w), ...)

# after: pad to multiple of 4 (mask or ignore extra rows)
idx8 = pad_to_multiple(idx6, multiple=4)
mgpu.async_load(source, smem, indices=idx8, slice_lengths=(1, w), ...)

n = indices.shape[0]
if n % 16 and n % 4:
    indices = pad_indices_to_multiple(indices, 4)  # and mask extras

def is_valid_tma_index_shape(shape) -> bool:
    return len(shape) >= 1 and (shape[0] % 16 == 0 or shape[0] % 4 == 0)

When it happens

Trigger: Passing an index vector whose leading dimension is not a multiple of 4 — e.g. shape[0] == 6 or 10 — as the indices operand of mgpu.async_load/async_store/async_prefetch in gather/scatter mode.

Common situations: Gather kernels with an arbitrary number of rows gathered per program; changing the gather batch size to a value not divisible by 4/16 after tuning; padding index vectors inconsistently with the TMA hardware constraints.

Related errors

• Only leading gather dimensions allowed.
• Only 2D gathers/scatters for async load/store are supported.
• copy_gmem_to_smem with a barrier is only supported Hopper an
• Splat layout does not support multimem
• async_copy requires all GMEM strides except the last one to

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