lablup / lablup/mlxcel

fix(cuda/quant): quantized prefill is not bitwise reproducible on sm_70, and its determinism test never runs

Open
#1,558 0 comments 0 reactions 0 assignees View on GitHub
area:core platform:linux priority:medium status:investigation type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

Found while implementing #1539 (part of epic #1536). Not caused by that change, and out of its scope. Same class as #910, which fixed a `cp.async` race in `qmm_sm80` on sm_121.

## What happens

`tests/cuda_qmm_determinism.rs` fails on a Tesla V100 (sm_70, CUDA 12.9.41). The divergence is at the **prefill** step, not decode.

This was confirmed to predate #1539 rather than be introduced by it: reverting `src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu`, rebuilding `origin/main`, and rerunning reproduces the same divergence at step 0. The two builds also agree on iteration 0's prefill logits hash, so the accumulator change is not implicated.

## Why `qmv` is not the culprit

Prefill runs at `M = 64`, which `quantized.cpp` routes to `qmm_naive`. `qmv` only serves `M * B < 8` and never runs at prefill shapes, so the kernel #1539 touched is not on this path at all. The suspect is `qmm_naive` on sm_70, which is the `UniversalFMA` scalar path from `mlx/backend/cuda/device/gemm_sm70.cuh`.

## Why nobody noticed

The test has never actually executed in CI. Its default checkpoint `models/llama-3.2-1b-4bit` is absent from the runners, and the test **skips rather than fails** when the model is missing, so it has been reporting green without running.

That is arguably the more important half of this report: a determinism gate that silently no-ops is worse than no gate, because it is counted as coverage.

## Scope

1. Determine whether prefill non-determinism on sm_70 is a real race in `qmm_naive` or an artifact of the test's tolerance and hashing. Check the shared-memory staging and the k-residue path first, by analogy with #910.
2. Establish whether sm_80 and later are affected. The audit that produced #1536 did not test this, and the kernel is shared, so this may not be Volta-specific.
3. Make the test fail loudly when its checkpoint is missing, or point it at a checkpoint the runners actually have. A skip that reads as a pass should not survive this issue.

## Acceptance criteria

- [ ] A written determination of whether `qmm_naive` prefill is non-deterministic on sm_70, with the mechanism if so.
- [ ] The same question answered for sm_80 and later, or an explicit statement that it could not be tested and why.
- [ ] `tests/cuda_qmm_determinism.rs` no longer reports success when it did not run.
- [ ] If a race is confirmed and fixed: repeated greedy prefill byte-identical run to run on sm_70, and no throughput regression.

## Evidence

Reported in `docs/benchmark_results/qmv-float-accum-v100-2026-08-31.md` and in the technical report for PR #1557. Three greedy 32-token decodes give one distinct output per build and both builds emit the same tokens, so the decode path is reproducible; only prefill is not.

Contributor guide

Open the contributing guide

Research direction

Start with tests/cuda_qmm_determinism.rs and verify whether the missing models/llama-3.2-1b-4bit checkpoint causes a successful skip. Then inspect quantized.cpp, mlx/backend/cuda/device/gemm_sm70.cuh, and the qmm_naive shared-memory and k-residue paths, comparing sm_70 with sm_80 and later. Done means the test cannot report success without running and the prefill determinism question is documented, with a fix and repeated byte-identical runs if a race is confirmed.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, rust
Domain
backend, performance, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.