jax-ml/jax · error · NotImplementedError

No swizzle is not supported

Error message

No swizzle is not supported

What it means

tcgen05 mma() rejects swizzle mode 16 (the 'no swizzle'/32-byte-atomic layout) because that shared-memory layout is not supported by the tcgen05 tensor-core path. Only 32B/64B/128B swizzles work.

Source

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

      raise NotImplementedError(f"Unsupported input dtype: {ty}")
  return _create_scaled_instr_descriptor(get_input_encoding, *args, **kwargs)


def mma(
    d: TMEMRef,
    a: ir.Value | TMEMRef,
    b: ir.Value,
    *,
    a_swizzle: int = 128,
    b_swizzle: int = 128,
    a_scale: TMEMRef | None = None,
    b_scale: TMEMRef | None = None,
    a_sparse_metadata: TMEMRef | None = None,
    accumulate: ir.Value | bool = True,
    collective: bool = False,
) -> None:
  if a_swizzle == 16 or b_swizzle == 16:
    raise NotImplementedError("No swizzle is not supported")
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)
  if isinstance(accumulate, bool):
    accumulate = arith.constant(ir.IntegerType.get_signless(1), accumulate)
  num_cta = 2 if collective else 1
  if (is_scaled := a_scale is not None) != (b_scale is not None):
    raise ValueError("Either none or both scales should be provided")
  is_sparse = a_sparse_metadata is not None
  if is_scaled and is_sparse:
    if isinstance(a, TMEMRef):
      raise NotImplementedError(
          "A in TMEM unsupported for block-scaled sparse matmuls"
      )

  # 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}")
  (k, n), b_element_type = mma_utils.tiled_memref_shape(b)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a_swizzle/b_swizzle of 32, 64, or 128
2. Re-layout shared memory operands with a supported swizzle mode before the mma

Example fix

# before
mma(acc, a, b, a_swizzle=16, b_swizzle=16)

# after
mma(acc, a, b, a_swizzle=128, b_swizzle=128)

assert a_swizzle != 16 and b_swizzle != 16, 'swizzle 16 unsupported; use 32/64/128'

When it happens

Trigger: Calling mma(a_swizzle=16) or mma(b_swizzle=16).

Common situations: Porting Hopper WGMMA kernels that used the 16-byte (none) swizzle mode, or defaulting swizzle parameters to 16 from older code.

Related errors

• No valid out swizzle{what}: minor dimension has {minor_dim_b
• SMEM has shape {smem_shape}, but expected {expected_smem_sha
• MMA lhs tiling does not fit swizzle. {lhs_tiling=} expected=
• MMA rhs tiling does not fit swizzle {rhs_tiling=} expected={
• MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_

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