jax-ml/jax · error · NotImplementedError

Unsupported scale type: {scale_type}

Error message

Unsupported scale type: {scale_type}

What it means

For block-scaled MMA, the scale matrix dtype is encoded in descriptor bit 23: only Float8E8M0 (encoding 1, the MX standard) and Float8E4M3FN (encoding 0) are supported hardware scale formats. Other scale dtypes raise NotImplementedError.

Source

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

  desc |= sparse << 2  # Sparsity, bit 2
  # Bit 3 is reserved
  assert 0 <= b_scale_idx < 4
  desc |= b_scale_idx << 4  # B scale factor data ID, bits 4-5
  # Bit 6 is reserved
  desc |= get_input_encoding(a_type) << 7  # A dtype, bits 7-9
  desc |= get_input_encoding(b_type) << 10  # B dtype, bits 10-12
  # We ignore negate bits 13-14
  desc |= transpose_a << 15  # Transpose A
  desc |= transpose_b << 16  # Transpose B
  if n % 8 or n > 256:
    raise ValueError(f"N must be a multiple of 8 and <= 256, got: {n}")
  desc |= (n >> 3) << 17  # N, bits 17-22
  if scale_type == ir.Float8E8M0FNUType.get():
    scale_encoding = 1
  elif scale_type == ir.Float8E4M3FNType.get():
    scale_encoding = 0
  else:
    raise NotImplementedError(f"Unsupported scale type: {scale_type}")
  desc |= scale_encoding << 23  # Scale matrix type
  # Bits 24-26 are reserved
  if m % 128 or m > 256:
    raise ValueError(f"M must be a multiple of 16 and <= 256, got: {m}")
  desc |= (m >> 7) << 27  # M >> 7, bits 27-28
  desc |= a_scale_idx << 29  # A scale factor data ID, bits 29-30
  # Bit 31 is reserved
  return arith.constant(ir.IntegerType.get_signless(32), desc)


def create_scaled_f8f6f4_instr_descriptor(*args, **kwargs) -> ir.Value:
  def get_input_encoding(ty):
    if ty == ir.Float8E4M3FNType.get():
      return 0
    elif ty == ir.Float8E5M2Type.get():
      return 1
    else:
      raise NotImplementedError(f"Unsupported input dtype: {ty}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use Float8E8M0FNUType scales (the MX standard)
2. Or Float8E4M3FNType if that variant is intended
3. Quantize/convert f32 scales to e8m0 before the MMA

Example fix

# before
scale = tmem.alloc(dtype=ir.F32Type.get(), ...)

# after
scale = tmem.alloc(dtype=ir.Float8E8M0FNUType.get(), ...)

ok = (ir.Float8E8M0FNUType.get(), ir.Float8E4M3FNType.get())
assert scale_type in ok, f'unsupported scale type {scale_type}'

def is_supported_scale_type(dt) -> bool:
    return dt in (ir.Float8E8M0FNUType.get(), ir.Float8E4M3FNType.get())

When it happens

Trigger: Passing scale_type other than f8e8m0 (e4m3 is accepted) — e.g. f32 scales or e5m2 scales — to the scaled descriptor builders.

Common situations: Writing custom scaled matmuls with FP32 scales instead of the MX-standard UE8M0; mixing up scale element formats.

Related errors

• Unsupported accumulator dtype: {acc_dtype}
• Unsupported input dtype: {ty}
• Either none or both scales should be provided
• Only support preferred_element_type in (f32, bf16, f16), but
• {ctx.avals_out[0].dtype}

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