fix(cuda): wrong values from sorted gather_mm and first graph run #2462

Problem / Background

Two CUDA-backend defects on GB10 (sm_121, MLX 0.32, driver 580.178.04) produce silently wrong values. Each currently has a model-level workaround, so nothing is visibly broken, but the root causes are open and other models can hit them. Implement both items in one PR, run the verification pass once at the end, and keep pushes to a minimum.

A1. Sorted gather_mm over f32 activations and bf16 weights depends on allocation history

Found in PR #2460 (OpenAI Privacy Filter, 128-expert top-4 MoE, f32 activations, bf16 expert weights). The bf16 forward is bit-exact in a fresh process and in the order bf16, 8-bit, bf16. After the 8-bit model's workload ran first in the same process, the tokens of 2024-05-17 drifted from private_date (log-prob -0.003) toward private_phone (-0.36 and -0.53). The drift is deterministic per predecessor workload.

Bisection (env toggles, since removed): unsorted gather is exact; sorted gather with bf16 activations is exact; evaluating per layer is still wrong; sorted gather_qmm (the 8-bit path, also GPT-OSS's path) is exact in every order.

A2. First execution of a freshly captured large CUDA graph intermittently returns wrong values (DeBERTa-v2 encoder)

Found in PR #2461 (GLiNER2.5). With the 12-layer encoder captured as one CUDA graph, on a 306-token input about half of the processes saw the first encode differ from every later replay by 1e-2 to 9e-1. Replays were always correct. MLX_USE_CUDA_GRAPHS=0 gave 0/80 failures; eval per layer gave 0/192. Removing the disentangled position terms makes it stable; softmax is not involved; materializing int32 indices changes nothing; general gather, 1-D take plus a strided view, or a contiguous index are all worse. The trigger is the position-embedding gathers inside one large graph. Related but distinct: #2174 (Llama fresh-process divergence that persists with graphs off, attributed to non-deterministic SDPA).

Current Behavior

• A1: src/lib/mlxcelverse/overlays/cuda/mlx/backend/cuda/matmul.cpp:355, GatherMM::eval_gpu, branch M == 1 && right_sorted_ == true calls gather_mm_rhs (defined at line 142). Workaround: src/models/openai_privacy_filter_layers.rs:285 sets sort_routed true only when all three expert planes are ExpertLinear::Quantized, so dense planes gather unsorted (about 1.5x the 8-bit forward on 7.5k tokens). The sorted dispatch is used at line 325.
• A1 exposure: every model that calls SwitchLinear::forward(.., sorted = true) (src/models/switch_layers.rs:463) with non-quantized weights and f32 activations can reach the branch. Callers of gather_mm or the sorted switch path in src/models/: deepseek, deepseek_v2, deepseek_v32, ernie4_5_moe, exaone_moe, glm4_moe, glm4_moe_lite, gpt_oss_common, hunyuan_moe, llama4, nemotron_h, qwen3_next, afmoe, bailing_moe, bailing_moe_linear, klear, jamba, phixtral, openai_privacy_filter_layers. Most run bf16 activations, which was exact; the issue must record which of them can feed f32 activations (check dtype at each sorted = true call) and list them in the PR.
• A2: src/models/deberta_v2.rs:293 sets graph_break_per_layer: mlxcel_core::cuda_is_available(), used at line 412, ending a CUDA graph after every layer (Metal untouched); rationale is in DebertaV2Layer::forward. Guarded by an encode == replay bit-exact assertion and a 259-token GLiNER determinism test. The GLiNER head's own take calls were observed clean (3 to 5 processes x 4 rounds), not proven.

Proposed Solution

A1: find the root cause in gather_mm_rhs (suspects: a kernel reading uninitialized or stale scratch for the f32-activation case, an alignment or stride assumption on the pooled allocation, or a missing stream dependency, since the failure depends on allocator history). Fix it in the overlay. If no fix is found within reasonable effort, disable the gather_mm_rhs branch when the activation dtype is float32 and weights are bfloat16, falling back to the generic path, and record the measured cost. If the fix lands, set sort_routed back to true for dense planes in Privacy Filter (remove the quantized-only condition at openai_privacy_filter_layers.rs:285).

A2: determine whether the cause is graph capture of gather/take with dependent indices, a missing dependency edge in the capture, or an MLX upstream bug. If fixable in the overlay or by an mlxcel-side change, remove graph_break_per_layer. If not, keep the split, document the reason in the code, file or link the upstream MLX report, and grep the other encoders for the same pattern (position-embedding gathers inside a single large graph) and list the hits.

Scope

In scope: the CUDA overlay matmul.cpp, src/models/openai_privacy_filter_layers.rs, src/models/deberta_v2.rs, and new CUDA tests.

Out of scope: the Llama divergence in #2174, Metal.

Implementation Notes

• Reuse: the existing CUDA test conventions under tests/ and src/models/*_tests.rs; SwitchLinear::forward_indexed for building the sorted call.
• Constraints: do not change Metal behavior. Runs must be on a quiet host (no concurrent GPU jobs) and use MLX_ENABLE_TF32=0.
• Edge cases: A1 needs the order effect, so the test must run a predecessor workload with different shapes (for example an 8-bit gather_qmm over several thousand tokens) in the same process before the sorted f32 x bf16 gather_mm with M == 1. A2 is process-level: use N fresh processes (N >= 20) because roughly half fail, so a single process proves nothing.
• Error handling: not applicable; the defect is silent.

Acceptance Criteria

• A1: a CUDA test that runs a predecessor workload, then a sorted gather_mm (f32 activations x bf16 weights, M == 1), and compares bit for bit against the unsorted result. It fails on current main.
• A1: the root cause is fixed in the overlay, or the branch is disabled for f32 x bf16 with the measured cost recorded in the PR. The PR lists which models in src/models/ can take the path.
• A1: if the fix allows, Privacy Filter dense planes use the sorted path again, and the long-text and order-swap checks from PR feat(models): support OpenAI Privacy Filter PII tagger in mlxcel tag #2460 (bf16, 8-bit, bf16 in one process) are rerun and match.
• A2: a CUDA test captures the full 12-layer encoder as one graph (with graph_break_per_layer forced off) and asserts first execution equals replay across N fresh processes. It fails on current main.
• A2: the root cause is identified. Either the per-layer split is removed and the test passes, or the split stays with a documented reason, an upstream report link, and a list of other encoders with the same pattern.
• Integrated into the real code flow: the Privacy Filter and GLiNER2.5 paths run through the fix end to end.

Verification

export MLX_CUDA_ARCHITECTURES=121 MLX_ENABLE_TF32=0
cargo build --profile test-fast --features cuda --tests
gpu-lock run --tag cuda-a1 -- <test-binary> <a1_test_name> --test-threads=1
for i in $(seq 1 20); do gpu-lock run --tag cuda-a2 -- <test-binary> <a2_test_name> --test-threads=1 || echo "FAIL run $i"; done
cargo fmt --all -- --check
cargo clippy --profile test-fast --features cuda -- -D warnings

A pass is 0 failures in all 20 fresh processes and bit-identical A1 comparisons. Do not wrap the cargo build in gpu-lock; lock only the test binary. Run the full lib suite with --test-threads=1.