jax-ml/jax · error · ValueError

Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn,

Error message

Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got {scale_element_type}

What it means

tcgen05.mma block-scaled MMA only supports scale tensors of type f8e8m0fnu (MX-style exponent scales) or f8e4m3fn. This error fires when the scale operands have any other dtype.

Source

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

            f" {scale_block}"
        )
    elif isinstance(a_element_type, ir.Float4E2M1FNType):
      if isinstance(scale_element_type, ir.Float8E4M3FNType):
        if base_scale_block != 16:
          expected = 32 if is_sparse else 16
          raise ValueError(
              f"Scale block size mismatch: expected {expected}, got"
              f" {scale_block}"
          )
      elif isinstance(scale_element_type, ir.Float8E8M0FNUType):
        if base_scale_block not in (16, 32):
          expected = "32 or 64" if is_sparse else "16 or 32"
          raise ValueError(
              f"Scale block size mismatch: expected {expected}, got"
              f" {scale_block}"
          )
      else:
        raise ValueError(
            "Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got"
            f" {scale_element_type}"
        )
    else:
      raise NotImplementedError(
          f"Unsupported element type for block scaling: {a_element_type}"
      )
    k_scales = k // scale_block
    if a_scale.shape != (TMEM_ROWS, k_scales):
      raise ValueError(
          f"A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got"
          f" {a_scale.shape}"
      )
    if a_scale.layout != scales_layout():
      raise ValueError(f"A scale layout {a_scale.layout} is not supported")
    if collective and m == 64:
      if b_scale.layout != b_scales_m64_collective_layout():
        raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast scale tensors to f8e8m0fnu (preferred MX format) or f8e4m3fn before calling mma
2. Ensure the scale buffers are allocated with the matching dtype at kernel setup

Example fix

# before
a_scale = memref.cast(...)  # f32 scales
tcgen05.mma(..., a_scale=a_scale, scale_block=32)
# after
a_scale = memref.cast(..., f8e8m0fnu)
tcgen05.mma(..., a_scale=a_scale, scale_block=32)

from jaxlib.mlir import ir
assert a_scale.dtype in (ir.Float8E8M0FNUType.get(), ir.Float8E4M3FNType.get())

def is_supported_scale_dtype(t) -> bool:
    return isinstance(t, (ir.Float8E8M0FNUType, ir.Float8E4M3FNType))

When it happens

Trigger: Passing a_scale/b_scale whose element type is e.g. Float32, Float8E4M3FN, or an integer type to mma() with block scaling enabled.

Common situations: Using fp32 scales from a reference implementation; dequantization code that upcast scales; scale tensors created with the wrong MLIR type when building a kernel.

Related errors

• MMA with element type {elem_type_str} does not support block
• Unsupported accumulator dtype: {acc_dtype}
• Unsupported input dtype: {ty}
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} only supports accumula

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