fix(inference): fused_sample_probs differs by 1 ULP at temperature 1.0 on sm_70
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Found while implementing #1544 (part of epic #1536). Not caused by that change, and out of its scope. Same pattern as #1558: a pre-existing failure that surfaced only because a wider test scope finally ran.
## What happens
On a Tesla V100 (sm_70, CUDA 12.9.41):
```
cargo test -p mlxcel-core --release --features cuda --lib -- --test-threads=1
```
gives **1672 passed, 1 failed**. The failure is `sampling::tests::temperature_one_support_unchanged`, on a 1 ULP difference in **5 of 64** `fused_sample_probs` entries at `T = 1.0`.
## It predates the change that found it
Verified by reverting `src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu` to `c2e54939`, rebuilding, and re-running: the failure reproduces **byte-identically**. The grouped-GEMM work is not implicated, and neither is anything else in this epic's merged units, which touch `qmv`, `qmm_naive` tile selection, and the grouped-GEMM arch tag rather than the sampler.
## Why it surfaced now
This was the first full library-suite run in the epic. #1559 ran only a filtered subset, and the earlier units narrowed scope to the modules they touched, because a full `cargo test --lib` exceeds the sub-agent stream-idle watchdog on this host. The test itself is not new and is not skipped; it simply had not been executed on this hardware before.
## Scope
1. Determine whether the 1 ULP divergence at `T = 1.0` is a genuine defect in `fused_sample_probs` on sm_70 or an over-tight assertion. The test's name says it checks that temperature one leaves the support unchanged, so the question is whether a 1 ULP move can change the sampled support at all, or whether the comparison should be on support membership rather than bitwise equality.
2. Establish whether sm_80 and later are affected. The kernel is shared, so this may not be Volta-specific, and nobody has run this suite on an sm_121 host as part of this work.
3. If the assertion is correct and the kernel is wrong, fix the kernel. If the kernel is right and the assertion is too strict for a fused path, fix the assertion and say what tolerance is defensible and why.
## Acceptance criteria
- [ ] A written determination of whether this is a kernel defect or a test defect, with the mechanism.
- [ ] The same question answered for sm_80 and later, or an explicit statement that it could not be tested and why.
- [ ] `cargo test -p mlxcel-core --release --features cuda --lib` green on sm_70.
- [ ] If the kernel changes: sampled output unchanged at temperature 0, and no throughput regression on the decode path.
## Related
- #1558, quantized prefill not bitwise reproducible on sm_70, found the same way and still open.
- Both are arguments for running the full library suite on this host at least once per epic rather than only filtered subsets. The watchdog constraint that prevents sub-agents from doing it is real, but the orchestrator can run it in the background, where the watchdog does not apply.
Contributor guide
Research direction
Start with sampling::tests::temperature_one_support_unchanged and the fused_sample_probs entry point, then run the provided full CUDA library test on sm_70. Compare support membership and bitwise values at T = 1.0 across sm_70, sm_80, and later hardware where available. Done means documenting whether the kernel or assertion is wrong, making the appropriate fix, and satisfying the listed temperature-zero and regression checks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, performance, testing
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100