apache / apache/tvm

[Bug] [Relax] `reduce_max` / `reduce_min` silently drop NaN on CUDA — CPU and CUDA produce different results

Open
#19,518 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.