jax-ml/jax · error · ValueError

Cannot assign layout to async load with gather indices since

Error message

Cannot assign layout to async load with gather indices since minor dim={slice_lengths[-1]} is not divisible by {divisor=} bits.

What it means

Mosaic GPU layout inference enforces SMEM alignment for async loads/stores with gather indices: each transferred row must be 256-bit aligned, so the minor dimension length of the slice must be divisible by (256 // element_bitwidth). If the last non-(-1) slice length isn't divisible by that divisor, no valid layout can be assigned and a ValueError is raised.

Source

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

      tiling_multiple.append(size)
      continue
    tiling_multiple.append(dynamic_gcd(size, index))

  operand_index = 1 if isinstance(op, mgpu.AsyncLoadOp) else 0
  operand = ValueSite(op, VariableType.OPERAND, operand_index)
  var = ctx.producer_ref(operand)
  constraints: list[cs.Constraint] = [
      cs.Divides(expr=var, tiling_multiple=tuple(tiling_multiple))
  ]
  if any(isinstance(idx.type, ir.VectorType) for idx in op.indices):
    element_bitwidth = utils.bitwidth(op.source.type.element_type)
    # This constraint enforces sufficient SMEM-alignment.
    # The transfer chunk needs to be 1024 bit-aligned. For each write in the
    # lowering we transfer 4 rows, so each row must be 256 bit-aligned.
    divisor = (1024 // 4) // element_bitwidth
    slice_lengths = [s for s in op.slice_lengths if s != -1]
    if slice_lengths and (slice_lengths[-1] % divisor):
      raise ValueError(
          "Cannot assign layout to async load with gather indices since"
          f" minor dim={slice_lengths[-1]} is not divisible by {divisor=}"
          " bits."
      )
    constraints.append(cs.MinorDimDivisibleBy(expr=var, divisor=divisor))

  value_sites_for_variable = {var: [operand]}
  value_sites, extra_constraints = _vector_value_sites_and_constraints_for_async_ops(op)
  value_sites_for_variable.update(value_sites)
  constraints.extend(extra_constraints)
  return cs.ConstraintSystem(constraints=constraints), value_sites_for_variable


@_add_constraint_system_derivation_rule(mgpu.AsyncPrefetchOp)
def _async_prefetch_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncPrefetchOp,
) -> ConstraintSystemDerivationRuleResult:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pad the innermost (minor) dimension of the gathered slice so its length is divisible by 256//element_bitwidth (e.g. multiples of 16 for 16-bit types)
Use a wider element type or reshape so the minor dim satisfies the alignment
Avoid gather indices: use a regular async load which doesn't hit this constraint

Example fix

// before (bf16, minor dim 8 -> divisor 16)
vals = mgpu.async_load(..., slice_lengths=(..., 8), ...)
// after
vals = mgpu.async_load(..., slice_lengths=(..., 16), ...)  # pad minor dim to 16

Defensive patterns

Strategy: validation

Validate before calling

bitwidth = 16  # e.g. bf16
divisor = 256 // bitwidth
assert slice_lengths[-1] % divisor == 0, f"minor dim {slice_lengths[-1]} not divisible by {divisor}"

Prevention

Always size the innermost gathered dimension as a multiple of 16 elements for 16-bit types
Compute 256//bitwidth and validate before issuing async_load with gather indices

When it happens

Trigger: Calling mgpu.async_load (or async store) with gather/slice indices whose innermost slice length (excluding -1 dims) is not a multiple of 256//bitwidth, e.g. minor dim 4 with 16-bit elements (divisor 16).

Common situations: Hand-written Mosaic kernels doing TMA/tensor gather loads with narrow inner dimensions (fp16/bf16 with small tiles), or padding vectors to non-multiple-of-16 element counts.

Related errors

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