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

benchmarks/cuda_gb10_gemma4_decode_688_2026-07-07.csv

@@ -0,0 +1,8 @@
+model,variant,cuda_graphs,coherent_tokens,deterministic,decode_tok_s,date,hardware,mlx_version,build_type,notes
+gemma-4-12b-it-4bit,shipped-8bit-mlp,on,1-3,no,13.83,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,"baseline (broken): greedy temp0 collapses to <pad> after 1-3 tokens, nondeterministic run-to-run"
+gemma-4-12b-it-4bit,shipped-8bit-mlp,off,>=200,yes,13.51,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,"fix #688 (default for gemma-4): deterministic 5/5, coherent, ~2.3% slower steady-state decode vs graphs-on"
+gemma-4-12b-it-4bit-uniform,uniform-4bit,on,2,no,~21,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,"baseline (broken): collapse after 2 tokens, nondeterministic"
+gemma-4-12b-it-4bit-uniform,uniform-4bit,off,>=110,yes,19.29,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,"fix #688: deterministic 5/5, matches mlx_lm reference trajectory near-exactly (tiny FP drift ~token 40)"
+gemma-4-31b-it-4bit,shipped-4bit,off,>=40,yes,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,"fix #688: coherent decode on the larger 31b variant"
+gemma-3-4b-it-4bit,control,on,>=30,n/a,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,"control (non-gemma4): unchanged, graphs stay on; note: gemma-3 hd-256 sdpa_vector decode is separately nondeterministic under graphs (latent, non-collapsing) -> follow-up"
+qwen3.5-4b-4bit,control,on,>=30,yes,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,"control (non-gemma4): unchanged, graphs stay on, deterministic"

src/distributed/pipeline/stage_executor/mod.rs

@@ -168,6 +168,16 @@ impl LoadedStageExecutor {
             stage.stage_index
         );
 
+        // Issue #688 (M2 hardening): defense-in-depth CUDA-graph disable for a
+        // Gemma 4 pipeline stage, mirroring the non-PP backend-seam load sites.
+        // Both the in-process and the remote-process pipeline loaders reach this
+        // common stage-load chokepoint before any stage weights are realised. The
+        // authoritative env write already happens on the main startup thread in
+        // `start_server`; this idempotent call (guarded by `var_os`) only takes
+        // effect if a future pipeline entry ever bypasses that hoist, so the env is
+        // set before the stage's first GPU eval regardless.
+        crate::loading::maybe_disable_cuda_graphs_for_gemma4_for_path(model_dir);
+
         let filter = LayerFilter::from_stage(stage);
         let family = resolve_stage_family(model_dir)?;
         let backend =

src/lib/mlxcel-core/src/layers.rs

@@ -4226,6 +4226,87 @@ mod tests {
         );
     }
 
+    /// Fill a `[rows, head_dim]` buffer with a deterministic pseudo-random
+    /// tensor whose every length-`head_dim` row is RMS-normalized (unit RMS,
+    /// so each element is O(1) and the row L2 norm is `sqrt(head_dim)`). This
+    /// mimics the post `q_norm` / `k_norm` distribution of Gemma 4, where the
+    /// attention scale is 1.0 and pre-softmax scores therefore reach O(head_dim).
+    fn rms_normed_rows(rows: usize, head_dim: usize, seed: u64) -> Vec<f32> {
+        let mut state = seed.wrapping_add(0x9E37_79B9_7F4A_7C15);
+        let mut next = || {
+            // SplitMix64 -> uniform in [-1, 1).
+            state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
+            let mut z = state;
+            z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+            z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+            z ^= z >> 31;
+            ((z >> 11) as f32 / (1u64 << 53) as f32) * 2.0 - 1.0
+        };
+        let mut out = vec![0.0_f32; rows * head_dim];
+        for r in 0..rows {
+            let row = &mut out[r * head_dim..(r + 1) * head_dim];
+            for x in row.iter_mut() {
+                *x = next();
+            }
+            let ms: f32 = row.iter().map(|v| v * v).sum::<f32>() / head_dim as f32;
+            let inv_rms = 1.0 / ms.sqrt().max(1e-6);
+            for x in row.iter_mut() {
+                *x *= inv_rms;
+            }
+        }
+        out
+    }
+
+    /// Numerical parity guard for the head_dim-256 single-query decode path
+    /// investigated under issue #688: the maskless single-query path (which
+    /// routes head_dim-256 attention to the CUDA `sdpa_vector` kernel) must match
+    /// the explicit-mask materializing SDPA for a Gemma-4-shaped decode step,
+    /// INCLUDING the `scale = 1.0` regime where pre-softmax scores reach
+    /// O(head_dim). Gemma 4 uses `scale = 1.0` (its `q_norm`/`k_norm` absorb the
+    /// usual `1/sqrt(d)`), unlike gemma-3 / qwen3.5 (`scale = 1/sqrt(d)`). This
+    /// test confirms the vector kernel is numerically correct EAGERLY at that
+    /// scale; the actual #688 `<pad>` collapse was a CUDA-graph read-before-write
+    /// hazard in Gemma 4's decode (not the kernel math), fixed by disabling CUDA
+    /// graph capture for Gemma 4 in `crate`-external model loading. Keeping this
+    /// guard prevents a future regression of the eager vector-kernel numerics.
+    #[test]
+    fn single_query_maskless_matches_masked_gemma4_large_scale() {
+        let n_heads = 16_i32;
+        let n_kv = 8_i32;
+        let head_dim = 256_i32;
+        let k_len = 25_i32;
+
+        let q_data = rms_normed_rows(n_heads as usize, head_dim as usize, 1);
+        let k_data = rms_normed_rows((n_kv * k_len) as usize, head_dim as usize, 2);
+        let v_data = rms_normed_rows((n_kv * k_len) as usize, head_dim as usize, 3);
+
+        let q = crate::from_slice_f32(&q_data, &[1, n_heads, 1, head_dim]);
+        let k = crate::from_slice_f32(&k_data, &[1, n_kv, k_len, head_dim]);
+        let v = crate::from_slice_f32(&v_data, &[1, n_kv, k_len, head_dim]);
+
+        for &scale in &[1.0_f32, 1.0 / (head_dim as f32).sqrt()] {
+            // Maskless single-query path: `mask = None` routes to the fused
+            // decode kernel (the sdpa_vector kernel for head_dim 256 on CUDA).
+            let out_maskless = attention(&q, &k, &v, scale, None, 0.0, 0);
+            // Explicit all-attend mask: `create_causal_mask(1, k_len - 1)` is
+            // all-zero for a single query, forcing the materializing SDPA path.
+            let mask = crate::utils::create_causal_mask(1, k_len - 1);
+            let out_masked = attention(&q, &k, &v, scale, Some(mask.as_ref().unwrap()), 0.0, 0);
+
+            let diff =
+                crate::subtract(out_maskless.as_ref().unwrap(), out_masked.as_ref().unwrap());
+            let diff_abs = crate::abs(&diff);
+            let max_diff_arr = crate::max_all(&diff_abs);
+            crate::eval(&max_diff_arr);
+            let max_diff = crate::item_f32(&max_diff_arr);
+            assert!(
+                max_diff.is_finite() && max_diff < 2e-3,
+                "maskless single-query decode must match the masked reference at \
+                 scale={scale}; max abs diff = {max_diff}"
+            );
+        }
+    }
+
     /// Verify that the detection condition uses the canonical NA check
     /// (has_neural_accelerator && macos_supports_na) rather than just
     /// metal_version >= 4.

src/loading/mod.rs

@@ -355,11 +355,97 @@ pub fn context_window_from_config(config: &serde_json::Value) -> Option<usize> {
     None
 }
 
+/// Force CUDA graph capture off for Gemma 4 checkpoints (issue #688).
+///
+/// mlxcel's Gemma 4 incremental (single-query, KV-cache) decode has a CUDA-graph
+/// read-before-write hazard: greedy (temperature 0) decode is nondeterministic
+/// run to run and collapses to `<pad>` after 1-3 tokens on CUDA, while Apple
+/// mlx_lm decodes the same model and prompt coherently. Prefill and the first
+/// token are correct; the collapse is in the incremental decode graph. The
+/// hazard is not a single op: disabling the head_dim-256 `sdpa_vector` decode
+/// kernel (`MLXCEL_SDPA_VECTOR_LARGE_D=0`), the maskless single-query attention
+/// path (`MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS=1`), or `mx.compile`
+/// (`MLX_DISABLE_COMPILE=1`) each only lengthens the coherent prefix; only
+/// turning off CUDA graph capture removes it entirely, restoring deterministic,
+/// coherent, mlx_lm-matching decode. Gemma 3, Qwen 3.5 and other families do not
+/// exhibit the collapse and keep graph capture on.
+///
+/// This writes `MLX_USE_CUDA_GRAPHS=0` before the first GPU eval so MLX's
+/// process-wide `use_cuda_graphs()` static (read once, on the first graph-eligible
+/// eval) picks it up. An explicit operator value (either direction) is respected.
+/// Non-Gemma-4 loads are untouched.
+///
+/// The sound home for this write is the main startup thread, upstream of any
+/// worker spawn: the server path calls
+/// [`maybe_disable_cuda_graphs_for_gemma4_for_path`] from `start_server` before the
+/// generation or pipeline worker is created (issue #688 M1/M2 hardening), and the
+/// CLI `generate`/`run` path loads on the main thread. The per-load-site calls that
+/// invoke this helper remain as idempotent defense-in-depth.
+///
+/// Invariant: this env-based lever assumes one model per process. The server is
+/// single-model-per-process today (no in-process hot-swap in `start_server`), so a
+/// non-Gemma eval never latches `use_cuda_graphs = true` before a later Gemma 4
+/// load. If in-process model hot-swap is ever added, a non-Gemma model loaded first
+/// would make a subsequent Gemma 4 `set_var` a silent no-op (the static is already
+/// latched) and this approach would need revisiting.
+fn maybe_disable_cuda_graphs_for_gemma4(model_type: ModelType) {
+    let is_gemma4 = matches!(
+        model_type,
+        ModelType::Gemma4 | ModelType::Gemma4VLM | ModelType::Gemma4Unified
+    );
+    if !is_gemma4 {
+        return;
+    }
+    // Respect an explicit operator override (either direction).
+    if std::env::var_os("MLX_USE_CUDA_GRAPHS").is_some() {
+        return;
+    }
+    // SAFETY: set_var mutates the process-global environment and is unsound under
+    // Rust 2024 only when another thread runs getenv/setenv concurrently. The
+    // authoritative write is hoisted to the main startup thread, upstream of any
+    // worker spawn and of the first GPU eval that latches MLX's `use_cuda_graphs`
+    // static: `maybe_disable_cuda_graphs_for_gemma4_for_path` runs in `start_server`
+    // for every serve path (batched, legacy, tensor-parallel, XLA, in-process and
+    // remote pipeline-parallel), and the CLI `generate`/`run` path loads on the main
+    // thread. On the server path this per-load-site call runs inside the spawned
+    // worker, but by then the main-thread write has already set `MLX_USE_CUDA_GRAPHS`,
+    // so the `var_os` guard above returns early and no set_var executes off the main
+    // thread. This write is therefore reached only on a main-thread load (CLI,
+    // benches) or if a future entry bypasses the hoist. It mirrors the established
+    // one-shot startup env writes `mlxcel_core::ensure_persistent_ptx_cache`,
+    // `hardware::apply_metal_ops_per_buffer_default`, and
+    // `TurboKvCacheArgs::apply_to_environment`, all written upstream of runtime
+    // bring-up.
+    unsafe { std::env::set_var("MLX_USE_CUDA_GRAPHS", "0") };
+    tracing::info!(
+        "Gemma 4 detected: disabling CUDA graph capture (MLX_USE_CUDA_GRAPHS=0) to \
+         avoid the decode-path graph hazard from issue #688. Set MLX_USE_CUDA_GRAPHS \
+         explicitly to override."
+    );
+}
+
+/// Path-based entry to [`maybe_disable_cuda_graphs_for_gemma4`] for the main-thread
+/// hoist (issue #688 M1/M2 hardening).
+///
+/// Peeks the on-disk model type and, for a Gemma 4 checkpoint, performs the
+/// `MLX_USE_CUDA_GRAPHS=0` write on the calling (main startup) thread, before any
+/// generation or pipeline worker is spawned and before the first GPU eval. A
+/// detection error (unreadable or absent `config.json`) is a no-op, which preserves
+/// the baseline for non-model paths and defers to the per-load-site calls. Kept as a
+/// thin wrapper so the Gemma-4 detection, operator-override guard, idempotency, and
+/// logging all stay in the single [`maybe_disable_cuda_graphs_for_gemma4`] helper.
+pub(crate) fn maybe_disable_cuda_graphs_for_gemma4_for_path(model_path: &Path) {
+    if let Ok(model_type) = get_model_type(model_path) {
+        maybe_disable_cuda_graphs_for_gemma4(model_type);
+    }
+}
+
 /// Load a model from a directory (or file — parent directory will be used)
 pub fn load_model(model_path: &Path) -> Result<(LoadedModel, MlxcelTokenizer)> {
     let model_path = resolve_model_dir(model_path);
     let model_path = model_path.as_path();
     let model_type = get_model_type(model_path)?;
+    maybe_disable_cuda_graphs_for_gemma4(model_type);
 
     // Whisper is an encoder-decoder ASR model, not a text generator. It is
     // wired into the speech-to-text audio endpoints at server startup rather
@@ -466,6 +552,7 @@ pub fn load_model_with_tensor_parallel(
     let model_path = resolve_model_dir(model_path);
     let model_path = model_path.as_path();
     let support = validate_supported_runtime(model_path, shard_config.clone(), adapter_path)?;
+    maybe_disable_cuda_graphs_for_gemma4(support.summary.model_type);
     let model = match support.summary.model_type {
         ModelType::Llama | ModelType::Qwen2 => LoadedModel::TensorParallelLlama(
             TensorParallelLlamaModel::from_model_dir(model_path, shard_config.clone())?,
@@ -508,6 +595,8 @@ pub fn load_model_with_adapter(
 ) -> Result<(LoadedModel, MlxcelTokenizer)> {
     let model_path = resolve_model_dir(model_path);
     let model_path = model_path.as_path();
+    // Disable CUDA graphs for Gemma 4 before any weight realisation (issue #688).
+    maybe_disable_cuda_graphs_for_gemma4(get_model_type(model_path)?);
     // Load base weights
     let base_weights = mlxcel_core::weights::load_weights_from_dir(model_path)
         .map_err(|e| anyhow::anyhow!("{}", e))?;
@@ -525,6 +614,7 @@ pub fn load_model_with_adapter(
 /// Build a model from pre-loaded weights (used by adapter loading)
 fn load_model_from_weights(model_path: &Path, weights: &mut WeightMap) -> Result<LoadedModel> {
     let model_type = get_model_type(model_path)?;
+    maybe_disable_cuda_graphs_for_gemma4(model_type);
     let config_path = model_path.join("config.json");
     let config_str = std::fs::read_to_string(&config_path)?;
     let config_str = sanitize_config_json(&config_str);

src/models/gemma4_tests.rs

@@ -273,7 +273,7 @@ mod cache_isolation {
                 let weights = make_test_gemma4_weight_map();
                 Gemma4Wrapper::new(Gemma4Model::from_weights(&weights, &args).unwrap())
             };
-            let prompt: Vec<i32> = (0..24).map(|i| (i % 7) as i32).collect();
+            let prompt: Vec<i32> = (0..24).map(|i| i % 7).collect();
 
             let single = build();
             let mut caches = single.make_caches();

src/server/startup.rs

@@ -1542,6 +1542,20 @@ fn install_surgery_pipeline_for_server(startup: &ServerStartupConfig) -> Result<
 /// Shared entry point used by both `mlxcel serve` and `mlxcel-server`.
 pub async fn start_server(mut startup: ServerStartupConfig) -> Result<()> {
     initialize_server_logging(&startup)?;
+
+    // Issue #688 (M1/M2 hardening): disable CUDA graph capture for Gemma 4 here,
+    // on the main startup thread, before any generation or pipeline worker is
+    // spawned and before the first GPU eval latches MLX's process-wide
+    // `use_cuda_graphs` static. Performing the env write on this thread (rather than
+    // only at the per-load-site calls, which run inside the spawned worker) keeps it
+    // sound under Rust 2024's concurrent-getenv rule. This single chokepoint covers
+    // every serve path that flows through `start_server`: the batched, legacy,
+    // tensor-parallel and XLA workers, the in-process pipeline-parallel worker, and
+    // the remote pipeline stage worker (the `serve_remote_pipeline_stage` branch
+    // below) — all load `startup.model_path`. Non-Gemma-4 models and an explicit
+    // `MLX_USE_CUDA_GRAPHS` operator override are left untouched.
+    crate::loading::maybe_disable_cuda_graphs_for_gemma4_for_path(&startup.model_path);
+
     super::media::configure_image_input_limits(super::media::ImageInputLimits {
         max_payload_bytes: startup.max_image_payload_size,
         max_images_per_request: startup.max_images_per_request,