tile-ai / tile-ai/tilelang

[BUG][Fuzzer][wrong-code] `T.cast` to `float8_e5m2` silently maps `+inf`/over-range `float32` to the max-finite value `57344` instead of `inf`

Open
#2,944 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.4k
Forks
745
Avg merge
1d 1h
Merged PRs (30d)
104

Description

### Required prerequisites

- [x] I have read the documentation .
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) that this hasn't already been reported.

### What version of TileLang are you using?

0.1.13 (latest release).

### System information

NVIDIA L40S (sm_89), CUDA 12.x, PyTorch 2.13. The conversion is emitted in the target-independent CUDA codegen / the shared fp8 template; not arch-gated, reproduced on sm_89.

### Problem description

Casting `+inf` (or an over-range `float32`) to `float8_e5m2` silently returns the largest finite value `57344` (bits `0x7B`) instead of `inf` (bits `0x7C`), with no error. `float8_e5m2` **has an `inf`/`nan` encoding** (unlike the finite-only `float8_e4m3fn`), and `torch.float8_e5m2` maps both `inf` and over-range values to `inf`; TileLang instead clamps `±inf` and every `float32` above the round-to-inf threshold to `±57344`.

The clearest violation is `±inf` itself: `inf` is exactly representable in `e5m2`, needs no rounding, and is nonetheless turned into a finite number.

```
input 65504.0 1e5 1e6 +inf -1e5 -inf 57344.0
TileLang 57344.0 57344 57344 57344 -57344 -57344 57344 (bits 0x7B/0xFB)
torch inf inf inf inf -inf -inf 57344 (bits 0x7C/0xFC)
^ max normal, both correct
```

This is not specific to `float32`: casting from **`float16` and `bfloat16`** to `float8_e5m2` drops `inf` the same way (measured — `+inf`/`-inf` come back as `0x7B`/`0xFB` = ±57344 instead of `0x7C`/`0xFC`), because all of them go through the same saturating conversion. `nan` is preserved on every path; only `inf`/over-range is affected. Both the scalar and the vectorized cast paths behave this way (`n_mismatch 7/8` in each), and finite in-range values round-trip bit-exactly — so the cast is otherwise correct, only the overflow/`inf` region diverges. Not a regression — see Provenance.

### Reproducible example code

```python
import tilelang, tilelang.language as T, torch

N = 8

@tilelang.jit(out_idx=[1])
def kernel(n):
@T.prim_func
def main(A: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float8_e5m2")):
with T.Kernel(1, threads=128) as _:
for i in T.Parallel(n):
C[i] = T.cast(A[i], "float8_e5m2")
return main

# over-range fp32 (max e5m2 normal = 57344) plus actual +/-inf
vals = torch.tensor([65504.0, 1e5, 1e6, 120000.0, float("inf"),
-1e5, float("-inf"), 57344.0], dtype=torch.float32)

got = kernel(N)(vals.cuda()).view(torch.uint8).cpu()
ref = vals.to(torch.float8_e5m2).view(torch.uint8) # trusted reference
print("got", [hex(int(b)) for b in got]) # -> 0x7B ... (all saturated to 57344)
print("ref", [hex(int(b)) for b in ref]) # -> 0x7C ... (inf), 57344 only for the last
print("ORACLE:", "PASS" if torch.equal(got, ref)
else f"FAIL {int((got!=ref).sum())}/{N} differ") # -> FAIL 7/8 differ

# CONTROL: same op to the finite-only float8_e4m3fn, which has NO inf encoding,
# so saturating IS correct and TileLang matches torch bit-exactly.
@tilelang.jit(out_idx=[1])
def kernel_e4m3(n):
@T.prim_func
def main(A: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float8_e4m3")):
with T.Kernel(1, threads=128) as _:
for i in T.Parallel(n):
C[i] = T.cast(A[i], "float8_e4m3")
return main

v4 = torch.tensor([500.0, 1e5, 1e6, 449.0, 448.0, -1e5, -449.0, 100.0], dtype=torch.float32)
g4 = kernel_e4m3(N)(v4.cuda()).view(torch.uint8).cpu()
r4 = v4.to(torch.float8_e4m3fn).view(torch.uint8)
print("CONTROL:", "PASS" if torch.equal(g4, r4) else "FAIL") # -> PASS
```

### Traceback

No traceback — the kernel compiles and runs to completion; the result is silently wrong and deterministic.

### Expected behavior

Two parts, with different levels of certainty:

- **`±inf` input → `inf` (`0x7C`/`0xFC`): this one is unambiguous.** `inf` is exactly representable in `e5m2` and needs no rounding or saturation decision; turning an exactly-representable `inf` into a finite number loses information for no reason, and matches neither `torch.float8_e5m2` nor IEEE conversion. This is the part that should clearly change.

- **Over-range *finite* input → `inf` vs saturate: the intuitive default, but a discussable policy.** For a format that has an `inf` encoding, the natural/IEEE behaviour is round-to-nearest, which sends values past the round-to-inf threshold to `inf` (this is what `torch` does), so producing `inf` is the more predictable default. That said, saturating over-range finite values to the max normal is a choice some ML stacks make deliberately (to keep `inf` out of downstream math), so this half is arguably a policy call rather than an outright defect — worth deciding explicitly rather than inheriting it by accident.

The `e4m3fn` control shows the current saturation policy is exactly right for a finite-only format (it has no `inf` to produce); the question is only whether that same policy should be applied wholesale to `e5m2`, which can represent `inf`.

### Additional context

**Root cause.** The `float32 → float8_e5m2` conversion always saturates to the max finite value rather than emitting the `inf` encoding the format supports. The vectorized path (`lanes` 2/4/8) hard-codes the saturating flag when it builds the conversion call ([`extra_args = ", __NV_SATFINITE, " + type_suffix`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/cuda/codegen/codegen_cuda.cc#L1848)); the scalar path (`lanes == 1`) is not handled there and falls through to the default C++ cast `(fp8_e5_t)x`, where [`fp8_e5_t = tl::float_e5m2_t`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/tl_templates/cuda/cuda_fp8.h#L8) (CUTLASS), whose converting constructor also saturates. `__NV_SATFINITE` is the correct choice for `e4m3fn` but discards `e5m2`'s `inf`; the CUDA conversion API also offers `__NV_NOSAT`, and a round-to-inf path exists for `e5m2`.

generated CUDA for the scalar path

```cuda
C[((int)threadIdx.x)] = ((fp8_e5_t)A[((int)threadIdx.x)]);
```
The `(fp8_e5_t)` constructor saturates over-range and `inf` to `0x7B`.

**Disposition.** Whether over-range *finite* inputs should round to `inf` or saturate is a policy choice some ML stacks make deliberately; the unambiguous part is that an exactly-representable `±inf` input is silently turned into a finite number. The evidence suggests `e5m2` should preserve `inf` (torch does; the format has the encoding; the finite-only sibling `e4m3fn` is already handled correctly). Which region to change — and whether to expose the saturation choice — is a call for the maintainers.

**Suggested fix.** For the `e5m2` conversion, use the CUDA API's `__NV_NOSAT` (round-to-inf) mode on the vectorized path instead of the hard-coded `__NV_SATFINITE`, and route the scalar path through a conversion that preserves `inf` rather than falling through to the saturating CUTLASS `float_e5m2_t` constructor — leaving `__NV_SATFINITE` in place for the finite-only `e4m3fn`. This touches both cast paths (vectorized codegen and the scalar fallback), so it is more than a one-line change (not verified end-to-end).

**Provenance.** The saturating conversion has shipped since fp8 cast support was added; the vectorized `__NV_SATFINITE` line was last relocated by the NFC namespace refactor [#2297](https://github.com/tile-ai/tilelang/pull/2297) (2026-06-02), not introduced by it. The scalar path's saturation comes from the CUTLASS `float_e5m2_t` constructor. Exact introducing commit not pinned.

**Dedup.** I searched the open and closed tracker and found no existing report of this defect. Distinct from #2670 (`T.infinity("float8_e5m2")` aborts compilation — a different symptom on the same dtype) and #2702 (e4m3→f16 *dequant* mis-decode).

**Reach.** The trigger is a documented, typed operation: `T.cast(x, "float8_e5m2")` on a `float32`, with `float8_e5m2` an exposed dtype. `float8_e5m2` appears in 12 `examples/`/`testing/` sites (`grep`), including the fp8 GEMM benchmarks and `examples/dequantize_gemm/quantize/quantization.py`; the GEMM examples supply e5m2 tensors from small-integer test data and the quantize path scales inputs, so none feed an over-range/`inf` value into the cast and none trip this. The bug fires whenever a `float32` at or above the round-to-inf threshold (~`61440`), or an `±inf`, reaches the cast — e.g. an un-clamped activation or an `inf` produced upstream.

Contributor guide

Open the contributing guide

Research direction

Start in src/cuda/codegen/codegen_cuda.cc around the vectorized conversion's extra_args and src/tl_templates/cuda/cuda_fp8.h for the scalar float_e5m2_t path. Run the supplied T.cast reproduction and compare scalar and vector results with torch.float8_e5m2, including the e4m3 control. Done means the agreed e5m2 policy is applied consistently to both paths, with ±inf preserved and regression coverage for the affected cases.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend-api-design, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.