Skip to content

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

Open

Description

@inureyes

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 bug: Gemma 4 26B A4B temp-0 decode is non-deterministic run-to-run and corrupts near-tied tokens independently of the fused MoE kernel #910.
  2. Establish whether sm_80 and later are affected. The audit that produced epic: NVIDIA Volta (sm_70) inference acceleration program #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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificpriority:mediumMedium prioritystatus:investigationFeasibility spike / under investigationtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions