jax-ml/jax · error · ValueError

Scale element type mismatch: expected f8e8m0fnu or…

Error message

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

What it means

Block-scaled MMA scale factors must be one of the UE8M0 or UE4M3 scale formats the hardware decodes. It fires when scale_element_type is neither Float8E8M0FNU nor Float8E4M3FN.

Solutions

Cast scale tensors to f8e8m0fnu (preferred MX format) or f8e4m3fn before calling mma
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)

Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

Centralize scale tensor creation with a fixed dtype helper
Fail fast on scale dtype at buffer allocation time

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

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

Appendix: 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)