From deec95b898064333443667cc0ef19387ed634ae7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 7 Jul 2026 21:27:32 +0900 Subject: [PATCH 1/3] fix(gemma4): disable CUDA graphs for decode graph-race collapse (#688) On CUDA, gemma-4 greedy decode collapsed to after 1-3 tokens and, decisively, varied run to run for the same model, prompt, and greedy settings, which is a race signature rather than a deterministic numerical bug. Prefill and the first token were already correct; the collapse was purely in the incremental single-query KV-cache decode. Localization against an mlx_lm gemma4_text oracle: disabling the head_dim-256 sdpa_vector overlay kernel (MLXCEL_SDPA_VECTOR_LARGE_D=0), the maskless single-query path (MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS=1), or mx.compile (MLX_DISABLE_COMPILE=1) each only lengthened the coherent prefix and stayed nondeterministic. Only MLX_USE_CUDA_GRAPHS=0 fully fixed it (5/5 identical, coherent 110-token decodes matching the reference). An eager unit test shows the sdpa_vector kernel is numerically correct, so the fault appears only under CUDA graph capture/replay. Root cause: mlxcel's gemma-4 incremental decode has a CUDA-graph read-before-write hazard (missing ordering between the in-place KV-cache write and its reads under graph capture), with two contributors: the head_dim-256 sdpa_vector kernel reading the rotating sliding-window cache slice (dominant; also the source of gemma-3's latent nondeterminism) plus a residual gemma-4-specific hazard that is not in the SDPA dispatch and not in any compiled kernel. Apple mlx_lm's gemma4_text decode is graph-safe on the same hardware. The 8-bit-MLP "secondary GEMV factor" was only a sensitivity amplifier, not a separate bug. Fix: force MLX_USE_CUDA_GRAPHS=0 for gemma-4 model types (Gemma4/Gemma4VLM/Gemma4Unified) in the shared model-load path, before the first GPU eval, honoring an explicit operator override. Non-gemma-4 families keep graph capture on and are unchanged. Validation (GB10, release, greedy temp 0): gemma-4-12b-it-4bit-uniform and the shipped 8-bit-MLP variant both decode deterministically and coherently for 110+ tokens; gemma-4-31b-it-4bit coherent; explicit MLX_USE_CUDA_GRAPHS=1 reproduces the collapse. gemma-3-4b, qwen3.5-4b, qwen3-4b unchanged. Steady-state decode cost ~2.3% (13.51 vs 13.83 tok/s). Numbers in benchmarks/cuda_gb10_gemma4_decode_688_2026-07-07.csv. --- ...cuda_gb10_gemma4_decode_688_2026-07-07.csv | 8 ++ src/lib/mlxcel-core/src/layers.rs | 81 +++++++++++++++++++ src/loading/mod.rs | 51 ++++++++++++ 3 files changed, 140 insertions(+) create mode 100644 benchmarks/cuda_gb10_gemma4_decode_688_2026-07-07.csv diff --git a/benchmarks/cuda_gb10_gemma4_decode_688_2026-07-07.csv b/benchmarks/cuda_gb10_gemma4_decode_688_2026-07-07.csv new file mode 100644 index 00000000..c0d63f52 --- /dev/null +++ b/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 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" diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 767fac0d..1b539609 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/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 { + 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::() / 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 `` 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. diff --git a/src/loading/mod.rs b/src/loading/mod.rs index 66d1c26e..e20b4a3e 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -355,11 +355,58 @@ pub fn context_window_from_config(config: &serde_json::Value) -> Option { 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 `` 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. +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 only + // under a concurrent getenv/setenv on another thread. This runs inside the + // model load, on the single loading thread, before any weights are realised + // (the first GPU eval) and therefore before MLX's `use_cuda_graphs` static is + // first read and before any generation worker thread is spawned. It mirrors + // the established one-shot startup env write in + // `mlxcel_core::configure_ptx_cache`. + 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." + ); +} + /// 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 +513,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 +556,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 +575,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 { 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); From 3f3c87a1f8059aee5beead128e62b777306431a7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 7 Jul 2026 22:59:42 +0900 Subject: [PATCH 2/3] fix(gemma4): hoist CUDA-graph disable to main thread and cover pipeline-parallel load (#688) Move the gemma-4 MLX_USE_CUDA_GRAPHS=0 env write to the main thread before the generation worker is spawned (new path-based helper called from server startup), so the set_var is not raced against a concurrent getenv on a request-handler thread under --no-warmup, matching the ensure_persistent_ptx_cache precedent. Also apply the guard on the in-process pipeline-parallel load path, which bypassed the four backend-seam entry points and otherwise served gemma-4 with graphs still on. The per-load-site calls remain as idempotent defense-in-depth. Correct the SAFETY comment's ordering rationale and its stale symbol reference. --- .../pipeline/stage_executor/mod.rs | 10 ++++ src/loading/mod.rs | 53 ++++++++++++++++--- src/server/startup.rs | 14 +++++ 3 files changed, 70 insertions(+), 7 deletions(-) diff --git a/src/distributed/pipeline/stage_executor/mod.rs b/src/distributed/pipeline/stage_executor/mod.rs index 248e9118..f78dc7bb 100644 --- a/src/distributed/pipeline/stage_executor/mod.rs +++ b/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 = diff --git a/src/loading/mod.rs b/src/loading/mod.rs index e20b4a3e..4b3bac6e 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -374,6 +374,20 @@ pub fn context_window_from_config(config: &serde_json::Value) -> Option { /// 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, @@ -386,13 +400,22 @@ fn maybe_disable_cuda_graphs_for_gemma4(model_type: ModelType) { if std::env::var_os("MLX_USE_CUDA_GRAPHS").is_some() { return; } - // SAFETY: set_var mutates the process-global environment and is unsound only - // under a concurrent getenv/setenv on another thread. This runs inside the - // model load, on the single loading thread, before any weights are realised - // (the first GPU eval) and therefore before MLX's `use_cuda_graphs` static is - // first read and before any generation worker thread is spawned. It mirrors - // the established one-shot startup env write in - // `mlxcel_core::configure_ptx_cache`. + // 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 \ @@ -401,6 +424,22 @@ fn maybe_disable_cuda_graphs_for_gemma4(model_type: ModelType) { ); } +/// 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); diff --git a/src/server/startup.rs b/src/server/startup.rs index 09657a91..084cb900 100644 --- a/src/server/startup.rs +++ b/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, From 082951f502ee655add5a6c5816157d47500466c4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 7 Jul 2026 23:28:43 +0900 Subject: [PATCH 3/3] style(gemma4): drop redundant i32 cast tripping the PP-modules clippy gate (#688) The pipeline-parallel CI's `clippy -p mlxcel --lib --tests -D warnings` step is path-gated to PRs touching the distributed modules, so this pre-existing unnecessary_cast in gemma4_tests.rs (from #679) only surfaced now that this branch adds the pipeline-parallel guard. Trivial lint fix, no behavior change. --- src/models/gemma4_tests.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/gemma4_tests.rs b/src/models/gemma4_tests.rs index b28e275d..44a07743 100644 --- a/src/models/gemma4_tests.rs +++ b/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 = (0..24).map(|i| (i % 7) as i32).collect(); + let prompt: Vec = (0..24).map(|i| i % 7).collect(); let single = build(); let mut caches = single.make_caches();