[GPU] Heap corruption during autotune-fusion-emitters on log-sum-exp with is_finite/select guard
## Summary
A small (34-op, ~3 KB) StableHLO module containing a log-sum-exp / soft-maximum pattern (with a `is_finite(max) → select` numerical-stability guard) deterministically aborts the process during XLA GPU JIT compile. Glibc reports `double free or corruption (!prev)` (SIGABRT, exit 134) or sometimes `corrupted size vs. prev_size while consolidating` / SIGSEGV (exit 139). The HLO dump pinpoints the crash inside the `autotune-fusion-emitters` pass.
The fragment is a generic numerical-stability pattern (the same shape `jax.nn.logsumexp` lowers to). Reduced via op-level ddmin from ~10 MB of production IR down to the 34-op fragment below.
## Environment
| | |
|---|---|
| GPU | NVIDIA L40S (46 GiB), CC 8.9 |
| Driver | 595.64 (NVIDIA Open Kernel Module, AWS build) |
| OS | Ubuntu 24.04.4 LTS, kernel 6.17.0-1015-aws |
| Python | 3.12.3 |
| JAX | 0.10.0 (also reproduced on 0.10.1) |
| jaxlib | 0.10.0 |
| CUDA pip | `nvidia-cuda-runtime-cu12==12.9.79`, `nvidia-cudnn-cu12==9.19.0.56` |
A cosmetic warning fires on every run:
```
E ... cuda_executor.cc:1526] Could not get kernel mode driver version:
(INVALID_ARGUMENT: Version does not match the format X.Y.Z)
```
We confirmed this is unrelated — the affected code logs and ignores the parse failure, leaving `kernel_mode_driver_version` unset. The crash happens later, in the fusion-emitter autotuner.
## Reproducer
Install (any JAX 0.10.x):
```sh
pip install 'jax[cuda12]==0.10.0'
```
`repro.mlir`:
```mlir
module @repro attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<7xf32> {
%cst = stablehlo.constant dense<0.000000e+00> : tensor<3xf32>
%cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<1xf32>
%cst_1 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%cst_2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%cst_3 = stablehlo.constant dense<9.99999974E-5> : tensor<f32>
%0 = stablehlo.multiply %arg1, %cst_2 : tensor<f32>
%1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32>
%2 = stablehlo.concatenate %cst_0, %cst_0, %1, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3xf32>
%3 = stablehlo.divide %2, %cst : tensor<3xf32>
%4 = stablehlo.reduce(%3 init: %cst_1) applies stablehlo.maximum across dimensions = [0] : (tensor<3xf32>, tensor<f32>) -> tensor<f32>
%5 = stablehlo.maximum %cst_1, %4 : tensor<f32>
%6 = stablehlo.is_finite %5 : (tensor<f32>) -> tensor<i1>
%7 = stablehlo.select %6, %5, %cst_2 : tensor<i1>, tensor<f32>
%8 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor<f32>) -> tensor<1xf32>
%9 = stablehlo.broadcast_in_dim %8, dims = [0] : (tensor<1xf32>) -> tensor<3xf32>
%10 = stablehlo.subtract %3, %9 : tensor<3xf32>
%11 = stablehlo.reduce(%10 init: %cst_2) applies stablehlo.add across dimensions = [0] : (tensor<3xf32>, tensor<f32>) -> tensor<f32>
%12 = stablehlo.log %11 : tensor<f32>
%13 = stablehlo.add %12, %7 : tensor<f32>
%14 = stablehlo.multiply %cst_2, %13 : tensor<f32>
%15 = stablehlo.compare LT, %arg0, %cst_3, FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1>
%16 = stablehlo.select %15, %cst_2, %14 : tensor<i1>, tensor<f32>
%17 = stablehlo.broadcast_in_dim %16, dims = [] : (tensor<f32>) -> tensor<1xf32>
%18 = stablehlo.concatenate %17, %cst_0, %cst_0, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3xf32>
%19 = stablehlo.divide %18, %cst : tensor<3xf32>
%20 = stablehlo.subtract %19, %cst : tensor<3xf32>
%21 = stablehlo.reduce(%20 init: %cst_2) applies stablehlo.add across dimensions = [0] : (tensor<3xf32>, tensor<f32>) -> tensor<f32>
%22 = stablehlo.log %21 : tensor<f32>
%23 = stablehlo.add %22, %cst_2 : tensor<f32>
%24 = stablehlo.multiply %cst_2, %23 : tensor<f32>
%25 = stablehlo.maximum %16, %cst_2 : tensor<f32>
%26 = stablehlo.select %15, %25, %24 : tensor<i1>, tensor<f32>
%27 = stablehlo.add %26, %cst_2 : tensor<f32>
%28 = stablehlo.add %27, %cst_2 : tensor<f32>
%29 = stablehlo.add %28, %cst_2 : tensor<f32>
%30 = stablehlo.add %29, %cst_2 : tensor<f32>
%31 = stablehlo.add %30, %cst_2 : tensor<f32>
%32 = stablehlo.broadcast_in_dim %31, dims = [] : (tensor<f32>) -> tensor<1xf32>
%33 = stablehlo.concatenate %cst_0, %cst_0, %32, %cst_0, %cst_0, %cst_0, %cst_0, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<7xf32>
return %33 : tensor<7xf32>
}
}
```
`repro.py`:
```python
import sys
import jax
from jaxlib.xla_client import CompileOptions
assert jax.default_backend() == "gpu", "GPU backend required"
client = jax.devices("gpu")[0].client
with open(sys.argv[1]) as f:
module = f.read()
print(f"compiling {len(module) / 1e6:.2f} MB module...")
client.compile_and_load(module, client.devices(), CompileOptions())
print("OK")
```
Run:
```
$ python repro.py repro.mlir
compiling 0.00 MB module...
free(): double free detected in tcache 2
Aborted (core dumped)
```
## Expected
Module compiles and `repro.py` prints `OK`. The module is a trivial pure-numerical computation with no asynchronous ops, no collectives, no custom calls.
## Observed
Process aborts. Exit code is `134` (SIGABRT, glibc heap-corruption detection) or occasionally `139` (SIGSEGV). The glibc message varies between runs:
- `free(): double free detected in tcache 2`
- `corrupted size vs. prev_size while consolidating`
- `corrupted double-linked list`
## Locating the failing pass
Setting `XLA_FLAGS="--xla_dump_to=/tmp/xla_dump --xla_dump_hlo_pass_re=.*"` and running the original (un-reduced) IR, the last HLO file written before the abort is named:
```
module_NNNN.jit_scan.0107.autotune-fusion-emitters.after_pipeline-start.before_fusion-wrapper.txt
```
So the crash occurs inside `autotune-fusion-emitters`, specifically between `pipeline-start` and `fusion-wrapper` for the autotuned fusion emitters.
## Bisection notes (FYI)
While reducing, I found:
- `--xla_disable_hlo_passes=priority-fusion` makes the crash go away (but kernels are then unfused, ~17× slower runtime — not a real workaround).
- `--xla_gpu_autotune_max_solutions=1 --xla_gpu_deterministic_ops=true` also makes the crash go away on this reduced fragment, but is wildly slow on the original module.
- The following had no effect: `--xla_gpu_autotune_level=0`, `--xla_gpu_enable_triton_gemm=false`, `--xla_gpu_experimental_disable_binary_libraries=true`, `--xla_gpu_enable_command_buffer=`, `--xla_gpu_enable_cublaslt=false`, `--xla_disable_hlo_passes=autotune-fusion-emitters` (does not appear to actually disable the pass by that name).
- Replacing the single `stablehlo.is_finite` op in the fragment above with a constant `true` removes the crash for the minimized fragment. On the un-reduced ~5 MB IR, removing all 38 `is_finite` ops removes the crash for one captured module but NOT for a slightly larger one captured at a different time, so `is_finite` is part of the trigger but not the whole story.
- The same module compiles and runs to completion on the CPU backend in ~95s.
## Origin
This pattern is `jax.nn.logsumexp` lowered through JAX — the `where(isfinite(amax), amax, 0)` guard against an all-`-inf` input. We hit it because our gradient-descent margin optimizer uses `soft_maximum(a, b, width) = w * logsumexp([a, b] / w)` inside the scan body. Any JAX user calling `jax.nn.logsumexp` on GPU inside a sufficiently large jitted function on this driver/jaxlib combo would likely trip the same bug.
## Notes for triage
- The minimum-size reproducer I could produce; further reduction via `mlir-reduce` would be welcome (we don't have a build of `mlir-reduce` with the `stablehlo` dialect registered).
- Happy to provide larger captures, the full HLO pass dump, or `ASAN`/`MSAN` traces if useful.
1 条评论