jax-ml/jax · error · NotImplementedError

Splat layout does not support multimem

Error message

Splat layout does not support multimem

What it means

The WGSplatFragLayout (warpgroup splat layout, where every lane holds the same value) cannot store through MultimemRef references. Multimem loads/stores interact with the memory descriptor layout in a way splat fragments don't model, so Mosaic rejects the combination outright.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:3645

      fmt_str = fmt.format(f"[{idx_fmt}]: {{}}")
      utils.debug_print(fmt_str, *idx, val, uniform=False)

  def store_untiled(
      self,
      ref: ir.Value | utils.MultimemRef,
      *,
      swizzle: int = 16,
      optimized: bool = True,
      atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None,
  ) -> None:
    index = ir.IndexType.get()
    i64 = ir.IntegerType.get_signless(64)
    if not isinstance(ref.type, ir.MemRefType):
      raise ValueError(ref)
    match self.layout:
      case WGSplatFragLayout():
        if isinstance(ref, utils.MultimemRef):
          raise NotImplementedError("Splat layout does not support multimem")
        if atomic is not None:
          raise NotImplementedError(
              "Atomic stores not supported for splat layout"
          )
        # All values are the same so swizzle does not affect anything here.
        self._store_untiled_splat(ref)
      case WGStridedFragLayout():
        if swizzle != 16:
          raise ValueError("Only TiledLayouts support swizzling")
        assert isinstance(self.layout, WGStridedFragLayout)
        vec_size = self.layout.vec_size
        bitwidth = utils.bitwidth(self.mlir_dtype)
        total_bits = vec_size * bitwidth
        if total_bits % 8 != 0:
          raise NotImplementedError("Vector length should be a multiple of byte size")
        # pyrefly: ignore[bad-argument-type]
        for get, _update, transfer_ref, idx in self.transfer_strided(ref, vec_size):
          if isinstance(transfer_ref, utils.MultimemRef):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the array to a storable layout before storing: use fa.to_layout / re-fragment into a TiledLayout or WGStridedFragLayout, then store.
2. Store via a plain memref instead of the MultimemRef when the destination is regular global/shared memory.
3. If the value is truly scalar, write it with a single scalar store (e.g. via a small tiled array of one element) rather than a splat store.

Example fix

# before
fa_splat.store(multimem_ref)  # WGSplatFragLayout + MultimemRef

# after
fa_tiled = fa_splat.to_layout(TiledLayout(1, 1))
fa_tiled.store(memref)  # plain memref, tiled layout

from jax.experimental.mosaic.gpu.fragmented_array import WGSplatFragLayout
from jax.experimental.mosaic.gpu import utils
if isinstance(fa.layout, WGSplatFragLayout) and isinstance(ref, utils.MultimemRef):
    raise SystemExit('splat layout cannot store to multimem; convert layout first')

def splat_multimem_conflict(fa, ref):
    return isinstance(fa.layout, WGSplatFragLayout) and isinstance(ref, utils.MultimemRef)

When it happens

Trigger: fa.store(ref) where fa.layout is WGSplatFragLayout and ref is a utils.MultimemRef (e.g. a TMA/multimem descriptor obtained from warp-specialized memory ops).

Common situations: Kernels that splat a scalar (e.g. an accumulator init or alpha coefficient) and try to write it back through a TMA descriptor; adapting shared-memory store examples to multimem without changing the array's layout.

Related errors

• copy_gmem_to_smem with a barrier is only supported Hopper an
• Cannot broadcast shape {self.shape} to layout {o.layout}
• Atomic stores not supported for splat layout
• Arrays with the splat layout can only be stored when they ha
• 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/7882ce2aa155d535. Report an issue: GitHub.