jax-ml/jax · error · NotImplementedError

cannot cast to `{dtype_name}`

Error message

cannot cast to `{dtype_name}`

What it means

Same _ir_cast check as the source-dtype variant, but here the destination element type is a dtype unsupported at the GPU's compute capability. The lowering refuses to emit a Triton cast to that type because the hardware cannot execute it natively.

Source

Thrown at jax/_src/pallas/triton/lowering.py:1688

    src_type = ir.RankedTensorType(src.type)
    dst_type = ir.RankedTensorType.get(
        src_type.shape,
        dst_type,
        src_type.encoding,
    )
  if src.type == dst_type:
    return src

  src_element_type = _element_type(src.type)
  dst_element_type = _element_type(dst_type)

  for dtype, dtype_name, is_supported in _UNSUPPORTED_CAST_DTYPES:
    if isinstance(src_element_type, dtype):
      if not is_supported(compute_capability):
        raise NotImplementedError(f"cannot cast from `{dtype_name}`")
    if isinstance(dst_element_type, dtype):
      if not is_supported(compute_capability):
        raise NotImplementedError(f"cannot cast to `{dtype_name}`")

  if isinstance(src_element_type, (ir.F16Type, ir.BF16Type)) and not isinstance(
      dst_element_type, ir.F32Type
  ):
    return _ir_cast(
        _ir_cast(src, ir.F32Type.get(), signed=False),
        dst_type, signed=False, dst_signed=dst_signed
    )

  if isinstance(src_element_type, ir.FloatType) and isinstance(
      dst_element_type, ir.FloatType
  ):
    return _float_float_cast(src, dst_type)

  if isinstance(src_element_type, ir.IntegerType) and isinstance(
      dst_element_type, ir.IntegerType
  ):
    return _int_int_cast(src, dst_type, signed=signed)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to the dtype outside the kernel (on the XLA side) or after the store, keeping the kernel in f16/f32/bf16
2. Move execution to a GPU with sufficient compute capability for the target dtype
3. Upgrade JAX to a version with wider fp8 cast support in the Triton lowering
4. Restructure the kernel to store the wider type and bit-cast later if semantics allow

Example fix

// before
out_ref[...] = x.astype(jnp.float8_e4m3fn)  # inside kernel on sm80

// after
out_ref[...] = x.astype(jnp.float16)  # cast to fp8 outside the kernel

dst = jnp.float8_e4m3fn
cc = float(jax.devices()[0].compute_capability)
if dst in (jnp.float8_e4m3fn, jnp.float8_e5m2) and cc < 8.9:
    dst = jnp.float16  # degrade gracefully inside the kernel

def output_dtype_ok(dst_dtype, cc: float) -> bool:
    return not (dst_dtype in (jnp.float8_e4m3fn, jnp.float8_e5m2) and cc < 8.9)

try:
    out_ref[...] = x.astype(fp8_dtype)
except NotImplementedError:
    out_ref[...] = x.astype(jnp.float16)  # fallback store

When it happens

Trigger: Casting a value TO an unsupported dtype (e.g. float8_e4m3fn, float8_e5m2) inside a Triton Pallas kernel on a GPU below the required compute capability; reached via _cast, _load with `other`, masked loads/stores, or offset computation casting to the index type.

Common situations: Quantization-style kernels writing fp8 outputs on Ampere or older GPUs; kernels ported from Triton-lang code that assume newer hardware; CI on older GPUs failing while dev machines with H100 pass.

Related errors

• cannot cast from `{dtype_name}`
• Unsigned integer dtype {aval.dtype} is not supported for dot
• Only positional arguments are supported by debug_print on Pa
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
• cannot cast {src} to {dst_type}

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