jax-ml/jax · error · NotImplementedError

No swizzle is not supported

Error message

No swizzle is not supported

What it means

At wgmma.py:335, the wgmma wrapper rejects swizzle == 16 (i.e. 'no swizzle'): the shared-memory descriptor path in Mosaic only supports the 32/64/128-byte swizzle modes.

Source

Thrown at jax/experimental/mosaic/gpu/wgmma.py:335

    a: fa.FragmentedArray | ir.Value,
    b: ir.Value,
    *,
    swizzle: int = 128,
):
  """Perform acc += a @ b using the WGMMA instruction.

  `a` may be passed in registers, or as a memref. `b` must be a memref.

  The expected (logical) memref shapes are:
    a: (m // tile_m, k // tile_k, tile_m, tile_k)
    b: (k // tile_k, n // tile_n, tile_k, tile_n).

  While the shapes may be physically transposed, when considering the row-major
  physical shape, the tile dimensions must be the two minor dimensions and must
  have the shape (8, S) where S = swizzle // bytewidth(element_type).
  """
  if swizzle == 16:
    raise NotImplementedError("No swizzle is not supported")
  # Step 1. Establish the shape and element type of the operation.
  if not isinstance(b.type, ir.MemRefType):
    raise ValueError(f"B must be a memref, got: {b.type}")
  bf16 = ir.BF16Type.get()
  f32 = ir.F32Type.get()
  f16 = ir.F16Type.get()
  i32 = ir.IntegerType.get_signless(32)
  i8 = ir.IntegerType.get_signless(8)
  f8e5m2 = ir.Float8E5M2Type.get()
  f8e4m3fn = ir.Float8E4M3FNType.get()
  (k, n), element_type = mma_utils.tiled_memref_shape(b)
  if a_in_regs := isinstance(a, fa.FragmentedArray):
    m, k2 = a.shape
    element_type2 = a.mlir_dtype
    if element_type2 not in {f16, bf16, i8, f8e5m2, f8e4m3fn}:
      raise ValueError(
          "Only f16, bf16, i8, f8e5m2, f8e4m3fn are supported for A "
          f"in registers, got {element_type2}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use utils.swizzle(32/64/128) for the B operand's SMEM allocation
2. Re-allocate/re-layout SMEM with a supported swizzle before the wgmma call
3. Check utils.swizzle constants in your JAX version and pick from the supported set

Example fix

# before
layout = utils.swizzle(16)
b = utils.memref_alloc(..., layout)
acc = wgmma.wgmma(a, b, acc, swizzle=16)
# after
layout = utils.swizzle(32)
acc = wgmma.wgmma(a, b, acc, swizzle=32)

assert swizzle in (32, 64, 128), f'swizzle 16 unsupported, got {swizzle}'

try:
    wgmma.wgmma(a, b, acc, swizzle=swizzle)
except NotImplementedError:
    wgmma.wgmma(a, b, acc, swizzle=32)

When it happens

Trigger: Calling wgmma.wgmma(..., swizzle=16); hit from lowering rules or the FlashAttention-style kernels (compute_qk, compute_pv, etc.) when the SMEM layout was allocated with 16-byte swizzle.

Common situations: Passing utils.swizzle(16) as the swizzle mode when allocating B; porting TMA descriptors configured for no swizzle; assuming all power-of-two swizzles are allowed.

Related errors

• Unsupported A register array shape
• packed, collective and layout arguments are only supported f
• All aliased Refs must have the same memory space (SMEM or TM
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• Swizzle {self.swizzle} requires the trailing dimension to be

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