fix(gemma4): disable CUDA graphs for decode graph-race collapse (#688)

[inureyes]

Symptom

On CUDA (GB10), gemma-4 (both the shipped gemma-4-12b-it-4bit with its 8-bit MLP and the uniform-4bit requant) greedy (temperature 0) decode collapsed to <pad> (id 0) spam after only 1-3 generated tokens, even on a byte-identical chat prompt that Apple mlx_lm decodes coherently. Prefill and the first token were correct (the merged chat-template fix from #686/#687 already produces a correct first token); the collapse was purely in the incremental single-query KV-cache decode. A decisive tell that the earlier reports missed: the collapse point varied run to run for the SAME model, prompt, and greedy settings (e.g. "The sky", then "The", then "The sky is blue because" across repeated runs), which is the signature of a race, not a deterministic numerical bug.

Localization trail

Working against an mlx_lm gemma4_text oracle (custom loader that strips vision and quantizes, cache-based greedy decode that exercises the same incremental path) that decodes the sky-is-blue prompt coherently and deterministically, I A/B'd the mlxcel decode with the existing env levers and measured coherent-token-count before collapse on the uniform-4bit model: baseline collapses at ~2 tokens; MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS=1 reaches ~15-24 then still collapses; MLXCEL_SDPA_VECTOR_LARGE_D=0 (routes the head_dim-256 sliding decode off the #681 sdpa_vector overlay kernel to the materializing SDPA) reaches ~8-40 but STILL collapses and STILL varies run to run; MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS=1 + MLXCEL_SDPA_VECTOR_LARGE_D=0 together still collapse nondeterministically; MLX_DISABLE_COMPILE=1 does not fix it either. Only MLX_USE_CUDA_GRAPHS=0 fully fixes it: 4/4 and then 5/5 identical, coherent 110-token decodes that match the mlx_lm reference trajectory. An eager unit test constructing gemma-4-shaped inputs (16/8 GQA, head_dim 256, RMS-normed q/k, scale 1.0) shows the sdpa_vector kernel is numerically correct EAGERLY (maskless == masked within 2e-3), so the kernel math is fine; the fault only appears under CUDA graph capture/replay. Instrumenting the sdpa_vector launch site confirmed correct, fresh per-step params (kL grows 24->25->26->27, strides consistent) and forcing a contiguous K/V copy did NOT fix it, ruling out stale params and the strided KV read. Separately, gemma-3 (also head_dim-256, also a rotating sliding-window cache) is ALSO nondeterministic under graphs with the sdpa_vector kernel but stays coherent, and qwen3.5 (plain cache) is deterministic; MLXCEL_SDPA_VECTOR_LARGE_D=0 makes gemma-3 deterministic but does not fix gemma-4, proving the sdpa_vector kernel is only one contributor and a second, non-SDPA CUDA-graph hazard remains in gemma-4's decode.

Root cause

mlxcel's gemma-4 incremental decode has a CUDA-graph read-before-write hazard (a missing ordering edge between the in-place KV-cache write and its downstream reads under graph capture), with at least two contributors: the head_dim-256 sdpa_vector decode kernel (the #681 overlay extension) reading the rotating sliding-window cache slice (dominant, and also the cause of gemma-3's latent nondeterminism), plus a residual gemma-4-specific hazard that is not in the SDPA dispatch and not in any mx.compiled kernel. Apple mlx_lm's gemma4_text decode is graph-safe on the same hardware, so this is an mlxcel-implementation hazard, not a shared-model or checkpoint defect. The 8-bit-MLP "secondary GEMV factor" flagged in the issue was not a separate bug: it only made the model sit closer to the collapse cliff, so the race tipped it over sooner; with the race removed, both the 8-bit-MLP and uniform-4bit variants decode coherently.

The fix

Disable CUDA graph capture for gemma-4 checkpoints (ModelType::Gemma4 / Gemma4VLM / Gemma4Unified) by writing MLX_USE_CUDA_GRAPHS=0 in the shared model-load path (src/loading/mod.rs), before the first GPU eval, so MLX's read-once use_cuda_graphs() static picks it up. An explicit operator MLX_USE_CUDA_GRAPHS value is respected in either direction; all non-gemma-4 families keep graph capture on and are byte-for-byte unchanged. This is the minimal Rust-side dispatch fix that mirrors how mlx_lm avoids the hazard, and it is applied at all four load entry points (direct, tensor-parallel, adapter, from-weights).

Reference-parity validation (CUDA, GB10, release, greedy temp 0)

gemma-4-12b-it-4bit-uniform: 5/5 identical, coherent 110-token decodes matching the mlx_lm reference near-exactly ("The sky is blue because Earth's atmosphere scatters shorter blue wavelengths of sunlight more than longer red wavelengths. This phenomenon, known as Rayleigh scattering, occurs as sunlight interacts with ..."), with only a tiny near-tie FP drift after ~40 tokens ("gas molecules and small particles" vs the oracle's "the gases and particles"). gemma-4-12b-it-4bit (shipped 8-bit MLP): 5/5 identical, coherent 110 tokens. gemma-4-31b-it-4bit: coherent. Explicit MLX_USE_CUDA_GRAPHS=1 reproduces the nondeterministic collapse, confirming the fix is exactly the graph disable and that the override is honored. Coherent-token-count before/after: 1-3 (nondeterministic) -> >=110 (deterministic). A new eager unit test (single_query_maskless_matches_masked_gemma4_large_scale) pins the head_dim-256 vector-kernel numerics at scale 1.0. Numbers recorded in benchmarks/cuda_gb10_gemma4_decode_688_2026-07-07.csv.

Blast radius

Touches only the model-load dispatch and is scoped to gemma-4 model types; no shared decode-attention code, KV-cache, or kernel behavior changes, so gemma-3/3n, other head_dim-256 or MQA models, and the sliding-window decode path are unaffected. Controls confirmed unchanged: gemma-3-4b, qwen3.5-4b, qwen3-4b all decode coherently with graph capture still on. Decode throughput cost for gemma-4 is ~2.3% steady-state (200-token run: 13.51 vs 13.83 tok/s), within a few percent of baseline; gemma-4 got no measured benefit from CUDA graph capture beyond that. Two follow-ups discovered during this work, out of scope here: (1) gemma-3's head_dim-256 sdpa_vector decode is separately nondeterministic under graphs (latent, non-collapsing) via the same sdpa_vector/rotating-cache contributor; (2) the residual non-SDPA CUDA-graph ordering hazard in mlxcel's gemma-4 decode should be root-caused so graph capture can be safely re-enabled for the throughput.

Review hardening (M1/M2)

The env write is now hoisted to the main thread before the generation worker is spawned (a path-based helper called from server startup), so set_var is not raced against a concurrent getenv on a request-handler thread under --no-warmup; this matches the ensure_persistent_ptx_cache precedent. The in-process pipeline-parallel gemma-4 load path, which bypassed the four backend-seam entry points, now also applies the guard. The per-load-site calls remain as idempotent defense-in-depth. SAFETY comment ordering rationale and a stale symbol reference are corrected.

Follow-up hardening (M1/M2 from security review)

Two soundness and coverage gaps found in review, fixed on this branch without changing the graph-disable behavior.

M1 (set_var soundness): on the server path the per-load-site set_var("MLX_USE_CUDA_GRAPHS=0") runs inside the spawned generation worker thread, which is UB under Rust 2024 if a request-handler thread calls getenv concurrently (reachable with --no-warmup, where the listener accepts before load finishes). The gemma-4 detection and env write are now hoisted to the main startup thread in start_server (src/server/startup.rs:1557), before any worker is spawned and before the first GPU eval latches MLX's read-once use_cuda_graphs static, matching the ensure_persistent_ptx_cache / apply_metal_ops_per_buffer_default / TurboKvCacheArgs::apply_to_environment precedents. On the server path the worker-thread load-site call now short-circuits at the var_os guard because the main thread already wrote the var, so no set_var executes off the main thread. A new one-line path helper maybe_disable_cuda_graphs_for_gemma4_for_path (src/loading/mod.rs) is the single home reused by the hoist and the defense-in-depth sites; the four load-site calls stay as idempotent belt-and-suspenders. The CLI generate/run path already loads on the main thread, so it was already sound. The SAFETY comment was rewritten for the new ordering and its stale symbol reference (configure_ptx_cache, which does not exist) corrected to ensure_persistent_ptx_cache.

M2 (pipeline-parallel coverage): the in-process and remote pipeline-parallel loaders do not go through the four backend-seam entry points, so gemma-4 served via pipeline parallel previously had no graph disable. The start_server hoist is upstream of both the in-process pipeline worker spawn and the serve_remote_pipeline_stage branch (both load startup.model_path), so it now covers them. A defense-in-depth guard was also added at the shared pipeline stage-load chokepoint LoadedStageExecutor::load_with_adapter (src/distributed/pipeline/stage_executor/mod.rs), so the env is set before a stage's first GPU eval even if a future pipeline entry bypasses the hoist.

L1: documented that the env-based lever assumes single-model-per-process, which holds today (no in-process hot-swap in start_server).

Re-verification (CUDA, GB10, after the refactor): CLI gemma-4-12b-it-4bit greedy temp-0 decode is coherent ("Rayleigh scattering ...") and byte-identical across two runs. mlxcel-server on gemma-4-12b-it-4bit answered /v1/chat/completions coherently (HTTP 200), and the "Gemma 4 detected: disabling CUDA graph capture" log now fires from the main-thread hoist (module mlxcel::loading) about 3 seconds before the model worker thread starts. A non-gemma control (llama-3.2-1b-4bit, server and CLI) does not log the graph disable and decodes coherently, confirming graph capture stays on for non-gemma families. Touched-module unit tests pass (271 passed, 0 failed).

Closes #688