jax-ml/jax · error · NotImplementedError

cannot cast to

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.

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
• Unsigned integer dtype
• a and b must have the same element type, but got
• cannot cast to
• cmp and val must have identical dtypes and shapes

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

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