[Bug] [Relax] `reduce_max` / `reduce_min` silently drop NaN on CUDA — CPU and CUDA produce different results
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
## Expected behavior
`R.max()` and `R.min()` with NaN inputs should propagate NaN in the output, consistent with IEEE 754 semantics: any comparison involving NaN returns NaN. This is the behavior on CPU (`llvm` target).
## Actual behavior
On CUDA, `R.max()` and `R.min()` silently ignore NaN values and return the max/min of only the non-NaN elements. CPU correctly propagates NaN.
`R.sum()` and `R.mean()` handle NaN correctly on both CPU and CUDA (propagate NaN), so this is specific to `reduce_max` and `reduce_min`.
## Reproduction
```python
import numpy as np
import tvm
from tvm import relax
import tvm.relax.op as R
from tvm.s_tir import dlight
def build_reduce(shape, axis, reduce_fn):
bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorStructInfo(shape, "float32"))
with bb.function("main", [x]):
with bb.dataflow():
out = bb.emit(reduce_fn(x, axis=axis, keepdims=False))
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return bb.get()
def run_on_target(mod, x_np, target):
if target == "llvm":
pipeline = tvm.ir.transform.Sequential([relax.transform.LegalizeOps()])
mod_l = pipeline(mod)
exe = relax.build(mod_l, target="llvm")
dev = tvm.cpu()
else:
pipeline = tvm.ir.transform.Sequential([
relax.transform.LegalizeOps(),
dlight.ApplyDefaultSchedule(dlight.gpu.Fallback()),
])
with tvm.target.Target("cuda"):
mod_l = pipeline(mod)
exe = relax.build(mod_l, target="cuda")
dev = tvm.cuda()
vm = relax.VirtualMachine(exe, device=dev)
return vm["main"](tvm.runtime.tensor(x_np, device=dev)).numpy()
np.random.seed(42)
shape = (4, 8)
x = np.random.randn(*shape).astype("float32")
x[1, 3] = float("nan") # inject NaN
# reduce_max
mod = build_reduce(shape, axis=1, reduce_fn=R.max)
cpu_out = run_on_target(mod, x, "llvm")
cuda_out = run_on_target(mod, x, "cuda")
print(f"Input row 1: {x[1]}")
print(f"CPU max row 1: {cpu_out[1]} (nan = correct, IEEE 754)")
print(f"CUDA max row 1: {cuda_out[1]} (finite value = WRONG)")
# reduce_min
mod_min = build_reduce(shape, axis=1, reduce_fn=R.min)
cpu_min = run_on_target(mod_min, x, "llvm")
cuda_min = run_on_target(mod_min, x, "cuda")
print(f"CPU min row 1: {cpu_min[1]} (nan = correct)")
print(f"CUDA min row 1: {cuda_min[1]} (finite value = WRONG)")
# reduce_sum — NOT affected (both propagate NaN)
mod_sum = build_reduce(shape, axis=1, reduce_fn=R.sum)
cpu_sum = run_on_target(mod_sum, x, "llvm")
cuda_sum = run_on_target(mod_sum, x, "cuda")
print(f"CPU sum row 1: {cpu_sum[1]} (nan = correct)")
print(f"CUDA sum row 1: {cuda_sum[1]} (nan = correct)")
```
## Expected output
```
CPU max row 1: nan
CUDA max row 1: nan <-- should match CPU
```
## Actual output
```
CPU max row 1: nan
CUDA max row 1: 1.579 <-- NaN silently dropped, returns max of non-NaN values
```
## Root cause
The CUDA codegen for `reduce_max` / `reduce_min` uses a tree reduction with `tir.max()` / `tir.min()`, which on CUDA maps to `fmaxf()` / `fminf()`. Per the CUDA math API, `fmaxf(x, NaN) = x` — NaN is treated as missing rather than propagated. The CPU backend uses a comparison-based reduction that naturally propagates NaN per IEEE 754.
A correct NaN-propagating max would be: `result = (isnan(a) || isnan(b)) ? NaN : max(a, b)`, or equivalently using the `__hmax_nan` / raw comparison approach.
## Impact
Any model using `reduce_max` or `reduce_min` on data that may contain NaN (e.g., attention masks with `-inf`/`NaN`, loss values, data with missing entries) will **silently produce wrong results** on CUDA while appearing correct on CPU. This is a correctness bug, not a precision issue.
## Environment
- TVM: main branch
- Target: `cuda` (tested on Tesla T4, sm_75)
- Python: 3.11
- OS: Ubuntu Linux
cc @junrushao
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by tracing the CUDA codegen path for Relax reduce_max and reduce_min from the R.max/R.min reproduction, focusing on the tree reduction using tir.max() and tir.min(). Compare the CUDA behavior with the CPU llvm path and verify completion with NaN-containing reductions on both targets, including the unaffected sum and mean cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100