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/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/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..4b3bac6e 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -355,11 +355,97 @@ 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. +/// +/// 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 { 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); 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(); 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,