jax-ml/jax · error · NotImplementedError

Sparse metadata format not implemented for {operand_dtype=}

Error message

Sparse metadata format not implemented for {operand_dtype=}

What it means

format_tcgen05_sparse_metadata implements tiling patterns only for 8-bit and 16-bit operand dtypes. Other bit sizes (e.g. 32-bit floats/ints) have no tcgen05 sparse metadata packing defined, so NotImplementedError is raised.

Source

Thrown at jax/_src/pallas/mosaic_gpu/helpers.py:220

  if meta.ndim != 3:
    raise ValueError(
        "Expected metadata to be 3-dimensional (M, K // 4, 2), but it is"
        f" {meta.ndim}D"
    )
  m, k, _2 = meta.shape
  if _2 != 2:
    raise ValueError(
        "Expected the trailing dimension of the metadata to be 2, got:"
        f" {meta.shape[-1]}"
    )
  k *= 2
  bitsize = dtypes.itemsize_bits(operand_dtype)
  if bitsize == 8:
    meta_tiled = meta.reshape(m // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)
  elif bitsize == 16:
    meta_tiled = meta.reshape(m // 128, 8, 2, 8, k // 64, 4, 2, 8).transpose(0, 4, 1, 6, 3, 5, 2, 7)
  else:
    raise NotImplementedError(
        f"Sparse metadata format not implemented for {operand_dtype=}"
    )
  return meta_tiled.reshape(m // 128, k // 64, 128, 64)


def find_swizzle(minor_dim_bits: int, what: str = ""):
  """Returns the largest swizzle that can be applied to a memory region.

  Swizzling is usually necessary when dealing with 2D data in SMEM, especially
  if the reference is used as an MMA operand. The returned swizzle is usually
  applied as ``plgpu`` transform:

    transforms = (
        plgpu.TilingTransform((8, 8 * swizzle // elem_bits)),
        plgpu.SwizzleTransform(swizzle))
    )

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use an 8- or 16-bit operand dtype such as float8_e4m3fn, float8_e5m2, float16, or bfloat16
2. Pass the actual MMA operand dtype, not the accumulator dtype (accumulators are usually fp32)
3. If you need wider dtypes, use the dense (non-sparse) path

Example fix

// before
meta_tiled = format_tcgen05_sparse_metadata(meta, jnp.float32)

// after
meta_tiled = format_tcgen05_sparse_metadata(meta, jnp.float8_e4m3fn)

from jax import dtypes
bits = dtypes.itemsize_bits(operand_dtype)
assert bits in (8, 16), f'sparse metadata unsupported for {bits}-bit operands'

def sparse_dtype_supported(operand_dtype) -> bool:
    from jax import dtypes
    return dtypes.itemsize_bits(operand_dtype) in (8, 16)

try:
    return format_tcgen05_sparse_metadata(meta, operand_dtype)
except NotImplementedError:
    # fall back to dense operand dtype or dense matmul
    return format_tcgen05_sparse_metadata(meta, jnp.float8_e4m3fn)

When it happens

Trigger: Calling format_tcgen05_sparse_metadata(meta, operand_dtype) with operand_dtype of itemsize other than 8 or 16 bits, e.g. jnp.float32 or jnp.bfloat16 is 16-bit ok but jnp.int32 is not.

Common situations: Attempting sparse matmul with fp32 accumulators passed as the operand dtype by mistake; exploring sparsity with dtypes the hardware sparse path doesn't support.

Related errors

• Expected metadata dtype to be uint2, got: {meta.dtype}
• Unsupported dtype: {ref.dtype}
• Expected metadata to be 3-dimensional (M, K // 4, 2), but it
• Expected the trailing dimension of the metadata to be 2, got
• Data type bitcast is only supported from i8 to other types.

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