jax-ml/jax · error · ValueError

A sparse metadata dtype mismatch: expected i2, got {a_sparse

Error message

A sparse metadata dtype mismatch: expected i2, got {a_sparse_metadata.dtype}

What it means

Sparse MMA requires the A sparse metadata tensor to be i2 (signless 2-bit integer), matching the tcgen05 sparse metadata encoding. Any other dtype raises this error.

Source

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

          f"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},"
          f" got {b_scale.shape[0]}"
      )
    if b_scale.shape[1] != k_scales:
      raise ValueError(
          f"B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}),"
          f" got {b_scale.shape}"
      )
  if is_sparse:
    sparse_group_elems = 8 if utils.bitwidth(a_element_type) == 4 else 4
    # Each sparse group has 2 entries.
    expected_meta_k = k // sparse_group_elems * 2
    if a_sparse_metadata.shape != (m, expected_meta_k):
      raise ValueError(
          f"A sparse metadata shape mismatch: expected {(m, expected_meta_k)},"
          f" got {a_sparse_metadata.shape}"
      )
    if a_sparse_metadata.dtype != ir.IntegerType.get_signless(2):
      raise ValueError(
          "A sparse metadata dtype mismatch: expected i2, got"
          f" {a_sparse_metadata.dtype}"
      )

  # Step 3. Compute the operand descriptors.
  if not isinstance(a, TMEMRef):
    # Both dense and sparse matmul consume A with a K bytewidth of 32, only
    # the group size is halved when it's sparse.
    (
        (a_desc_base, a_k_instr_strides),
        (a_m_group_stride, a_k_group_stride),
        a_fastest,
    ) = mma_utils.create_descriptor(
        a,
        swizzle=a_swizzle,
        group_size=(m_group_elems, k_group_elems // (1 + is_sparse)),
        logical_k_major=False,
        mma_bytewidth_k=32,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pack the metadata into 2-bit signless integers (4 values per byte) and view the buffer as i2
2. Check a_sparse_metadata.dtype == ir.IntegerType.get_signless(2) before calling mma

Example fix

# before
meta = memref (m, meta_k) of i8
tcgen05.mma(..., a_sparse_metadata=meta)
# after
meta_i2 = packed_i2_metadata(m, meta_k)  # ir.IntegerType.get_signless(2)
tcgen05.mma(..., a_sparse_metadata=meta_i2)

assert a_sparse_metadata.dtype == ir.IntegerType.get_signless(2)

def is_i2_meta(t) -> bool:
    return t.dtype == ir.IntegerType.get_signless(2)

When it happens

Trigger: Passing a_sparse_metadata with dtype i8, i32, or signed i2 variants instead of ir.IntegerType.get_signless(2).

Common situations: Metadata produced by JAX/torch as uint8 or int32 without repacking into packed i2; loading metadata from files with default integer dtypes.

Related errors

• A sparse metadata shape mismatch: expected {(m, expected_met
• Unsupported accumulator dtype: {acc_dtype}
• Unsupported input dtype: {ty}
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block

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