diff --git a/docs/METAL_PARITY.md b/docs/METAL_PARITY.md new file mode 100644 index 0000000..266d809 --- /dev/null +++ b/docs/METAL_PARITY.md @@ -0,0 +1,120 @@ +# Apple Metal parity for issue #78 + +## Frozen comparison + +- DS4 oracle: `8db89fe083ae4d17c9a2428ccd29803d3ae8f577` +- DS4Server pre-issue baseline: `79468c65b68c28e1539acf0a2a14bd42e000e365` +- DeepSeek V4 Flash 0731 IQ2XXS/w2Q2K model SHA-256: + `ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0` +- DSpark support-model SHA-256: + `7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360` +- Machine: MacBook Pro `Mac17,7`, Apple M5 Max, 128 GiB unified memory, + macOS 26.6.2, power 100% +- Context and prefill allocation: 4096 tokens; quality, steering, speculative + decoding, and DSpark disabled for the throughput gate +- SSD mode: 4096 cached experts, 4096 popularity-preloaded experts, ordinary + (not cold) streaming +- Prompt: `Count from one to two hundred, spelling out every number.` rendered + to 16 tokens; 128 greedy non-EOS tokens are measured + +The longer prompt is deliberate: all 128 generated tokens are reachable before +EOS and are identical between DS4 and DS4Server in each mode. A short-response +fixture reached EOS early and made a forced post-EOS speed run compare +unreachable, mode-dependent states. + +## DS4 performance-change matrix + +The audited range is +`efdadd41e20134af4f3381e1ed90e96fe4faef6f..8db89fe083ae4d17c9a2428ccd29803d3ae8f577`. +Documentation-only, parser/server-only, CUDA, ROCm, tensor-parallel, model +download, and merge commits do not affect the Rust-owned single-session Metal +graph. The applicable performance/runtime commits are accounted for below. + +| DS4 commit(s) | DS4 call site or behavior | DS4Server result | +| --- | --- | --- | +| `3ce6777`, `8e0f997` | Native streamed-expert cache pinning and safe fallback | Active in the byte-matched native Metal cache; Rust preserves the cache budget and fallback. | +| `df04887` | `ds4_gpu_attention_indexed_mixed_batch_heads_tensor` | Active through `encode_layer_with_cache_rows`; the long-context fixture crosses the indexed boundary. | +| `74e8f11`, `4591cb1`, `a51e6ec`, `f284165`, `af80694`, `7fb2830`, `0e89a0e`, `769a8ba`, `8d683d6` | DSpark block commit, deterministic scheduling, acceptance replay/direct commit, greedy identity, and exact/opportunistic sampling | Active in `Dspark::propose` and the verifier commit paths; greedy and exact stochastic fixtures cover the frozen support model. | +| `96c3ba4`, `d69a017`, `81f8534`, `ac5cb2c`, `48392e3`, `222b2cb`, `532ec8b` | Metal 4 Q4 projection, compact routed-MoE, aligned Q8, indexed-attention, and routed prefill kernels | Active through the existing `encode_batch_layer` Q4/indexed/routed entry points; no second prefill graph was added. | +| `24903f3` | `ds4_gpu_head_rms_norm_rope_tail_tensor` | Existing Rust Q-head norm/RoPE fusion is retained and used. | +| `4893e0c`, `d14ce35` | Mapped-prefill expert seeding and consistent expert/non-expert maps | Rust seeds with `ds4_gpu_stream_expert_cache_seed_experts_gpu_copy`, falls back to CPU/pread, and restores one cached static decode map with the same non-routed spans. | +| `0dd0d36` | Release speed regression gate | Rust has an ignored release gate with retained raw CSV, DS4 reference threshold, baseline threshold, p50/p95, cache metrics, and token hashes. | +| `4e401a1`, `0e125e6` | Automatic pre-M5/M5 admission for exact decode fusions | Rust binds both device predicates and mirrors per-generation rollback variables without a UI preference. | +| `71d8c2a` | Q/KV pair plus ratio-4/128 compressor projection/store | `try_qkv_pair_compressor_fusion` calls the native combined entry point only for DS4's resident full-decode shapes. | +| `b64d3e0` | Emit-path compressor finalizer | `update_compression_pair` defers both updates only when the resident finalizer is eligible, then requires exact success. | +| `3826f52`, `ad4d05c` | Exact compressor softmax/product and ratio-4 pool | Native kernels are active; Rust's compressor ABI now passes the upstream `decode_one_token` and `defer_finalize` flags exactly. | +| `92d83e6` | HC norm/mix producer clustering | `decode_hc_mix` uses split-norm producer fusion, with the established norm/mix/split sequence as fallback. | +| `1be41c8` | Packed exact split-K attention reduction | Active in the vendored native attention kernel reached by Rust's decode attention calls. | +| `0e8fb54`, `e6a6ab5` | Router/select, shared gate/up, QKV norm/KV RoPE/store, parallel full FFN, inverse-attention RoPE, and pre-M5 ports | All applicable entry points are called with DS4 shape, tensor, quality, steering, profile, SSD, and device guards. Parallel FFN aborts on partial failure; inverse RoPE falls back when the attention kernel does not consume it. | +| `0658dac` | Long-context prefill/decode correctness | The 4096-crossing ignored fixture covers compressed/indexed state and the next decode token. | +| `023614e`, `0ad494e`, `42033ee` | Remove stream-512 experiment; make exact paths automatic; verify DFlash pipeline | No stream-512 or opt-in experiment remains. Feature rollbacks retain identical token hashes. | +| `24fa85e` | Keep successfully loaded disk KV checkpoints | Existing Rust checkpoint load/round-trip behavior is verified; only invalid checkpoints follow deletion policy. | +| `7694112` | Reuse a shorter exact GLM live prefix | `live_prefix_rewind_target` rewinds to `prompt_len - 1`, clears pending speculation, and reevaluates the last prompt token. | +| `5186e2c` | Keep Metal decode changes portable | Device checks and unavailable returns preserve non-M5 and unfused execution. | + +The Metal MXFP4 commits (`725b084`, `1f862bd`, `7bec128`, `bcd1e57`, +`f75a065`, and their merge/format work) are intentionally excluded: MXFP4 is +issue #79. `ds4_gpu_set_decode_pipeline_fast_lookup` is reset to disabled, +matching DS4 because that cache is enabled only for the pre-M5 MXFP4 path. + +## Results + +Three implementations were run sequentially, three times per mode, with a +five-second cooldown. Medians are from +[`metal-parity-issue-78.csv`](metal-parity-issue-78.csv). + +| Implementation | Mode | Prefill tok/s | TTFT ms | First decode ms | Stable tok/s | p50 ms | p95 ms | +| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| DS4Server baseline | resident | 69.63 | 256.07 | 26.37 | 39.53 | 24.88 | 26.61 | +| DS4 | resident | 56.49 | 305.94 | 22.71 | 45.60 | 21.91 | 22.32 | +| DS4Server issue #78 | resident | 69.76 | 252.01 | 22.66 | 45.60 | 21.94 | 22.33 | +| DS4Server baseline | SSD | 6.81 | 2476.94 | 125.30 | 8.63 | 114.93 | 124.28 | +| DS4 | SSD | 11.88 | 1393.54 | 46.75 | 26.11 | 35.62 | 46.86 | +| DS4Server issue #78 | SSD | 11.49 | 1440.79 | 48.84 | 25.07 | 37.70 | 48.85 | + +TTFT is the per-run prompt prefill duration plus first decode latency; its +median is reported separately from the first decode kernel latency. + +Resident stable decode improves 15.4% over the baseline and matches DS4. +SSD stable decode improves 190.3% over the baseline and is 4.0% below +DS4, inside the required 5% boundary. SSD cache accounting is identical to +DS4 for the measured token sequence: 4096 entries, 32,958 hits, 4,194 misses, +and 29,684,662,272 pread bytes. The resident and SSD token SHA-256 values are, +respectively, `ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d` +and `2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052`; +DS4 emitted the same 128-token sequences in the corresponding modes. + +Representative `/usr/bin/time -l` measurements report: + +| Implementation | Mode | Maximum RSS bytes | Peak footprint bytes | Process swaps | +| --- | --- | ---: | ---: | ---: | +| DS4 | resident | 78,790,656 | 4,746,511,440 | 0 | +| DS4Server issue #78 | resident | 129,024,000 | 104,071,816 | 0 | +| DS4 | SSD | 29,066,280,960 | 29,595,620,664 | 0 | +| DS4Server issue #78 | SSD | 29,098,393,600 | 104,350,344 | 0 | + +The post-run system memory-pressure reading was 92% free. Every measured +process reported zero swaps; the post-run host counters were 0 swap-ins and 20 +historical swap-outs. + +## Reproduction + +DS4 uses the exact 16 Rust-rendered prompt token IDs with its speed-benchmark +format. DS4Server uses the installed-model fixture: + +```sh +DS4_REFERENCE_TPS=45.60 \ +DS4SERVER_BASELINE_TPS=39.526849 \ +cargo test --release --all-features flash_0731_m5_decode_performance_gate \ + -- --ignored --nocapture --test-threads=1 + +DS4SERVER_BENCH_SSD=1 \ +DS4_REFERENCE_TPS=26.11 \ +DS4SERVER_BASELINE_TPS=8.634659 \ +cargo test --release --all-features flash_0731_m5_decode_performance_gate \ + -- --ignored --nocapture --test-threads=1 +``` + +The gate fails unless throughput improves over the frozen DS4Server baseline +and remains at least 95% of DS4. `DS4SERVER_BENCH_TOKENS=1` prints the prompt +and generated sequences for direct DS4 comparison. diff --git a/docs/metal-parity-issue-78.csv b/docs/metal-parity-issue-78.csv new file mode 100644 index 0000000..b27672d --- /dev/null +++ b/docs/metal-parity-issue-78.csv @@ -0,0 +1,19 @@ +implementation,mode,run,prompt_tokens,decode_tokens,prefill_tps,first_token_ms,steady_tps,p50_ms,p95_ms,cache_entries,cache_hits,cache_misses,pread_bytes,token_sha256 +ds4server-baseline-79468c65,resident,1,16,128,69.627264,26.270667,39.534586,24.875959,26.607167,0,0,0,0, +ds4server-baseline-79468c65,resident,2,16,128,69.799133,26.365458,39.505867,24.905750,26.624667,0,0,0,0, +ds4server-baseline-79468c65,resident,3,16,128,69.618099,26.454750,39.526849,24.877584,26.585875,0,0,0,0, +ds4server-baseline-79468c65,ssd,1,16,128,6.824429,122.398208,8.638560,114.928000,124.375958,4096,34040,7208,51017416704, +ds4server-baseline-79468c65,ssd,2,16,128,6.805883,126.033250,8.634659,114.667833,124.220583,4096,34040,7208,51017416704, +ds4server-baseline-79468c65,ssd,3,16,128,6.726701,125.304417,8.532966,116.431375,124.278709,4096,34040,7208,51017416704, +ds4-8db89fe0,resident,1,16,128,56.49,22.706,45.67,21.905,22.297,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d +ds4-8db89fe0,resident,2,16,128,56.49,22.719,45.59,21.908,22.320,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d +ds4-8db89fe0,resident,3,16,128,56.51,22.638,45.60,21.890,22.342,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d +ds4-8db89fe0,ssd,1,16,128,12.07,46.746,25.73,36.461,46.861,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052 +ds4-8db89fe0,ssd,2,16,128,11.88,46.737,26.11,35.622,46.982,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052 +ds4-8db89fe0,ssd,3,16,128,11.16,47.092,26.20,35.513,46.554,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052 +ds4server-issue-78,resident,1,16,128,69.930426,22.547583,45.683038,21.889833,22.275042,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d +ds4server-issue-78,resident,2,16,128,69.760230,22.657125,45.602401,21.936042,22.331375,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d +ds4server-issue-78,resident,3,16,128,69.727346,22.663208,45.361642,22.052833,22.459084,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d +ds4server-issue-78,ssd,1,16,128,11.279689,48.890291,24.849723,37.847292,48.853791,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052 +ds4server-issue-78,ssd,2,16,128,11.520460,48.567417,25.076916,37.696334,48.779625,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052 +ds4server-issue-78,ssd,3,16,128,11.494703,48.844875,25.066058,37.616542,48.911333,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052 diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 6f258e8..777434c 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -8,7 +8,7 @@ use glm::GlmExecutor; use gpu::*; use profile::ExpertProfile; -use super::gguf::{F16, F32, Gguf, IQ2_XXS, Q4_K, Q8_0, Tensor as GgufTensor}; +use super::gguf::{F16, F32, Gguf, IQ2_XXS, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor}; use super::validation::{DsparkConfig, SupportKind, dspark_config}; use super::{Model, ModelFamily, Rng, exact_delta_sample}; use crate::model::ModelChoice; @@ -17,7 +17,7 @@ use crate::settings::{ }; use sha2::{Digest, Sha256}; use std::env; -use std::ffi::c_void; +use std::ffi::{CStr, c_char, c_void}; use std::fs::{self, File}; use std::io::{Read, Write}; use std::path::Path; @@ -29,6 +29,16 @@ const CHECKPOINT_VERSION: u32 = 2; const CHECKPOINT_IO_CHUNK: usize = 8 * 1024 * 1024; const DEFAULT_PREFILL_CHUNK: u32 = 4096; +unsafe extern "C" { + fn getenv(name: *const c_char) -> *mut c_char; +} + +fn environment_present(name: &CStr) -> bool { + // SAFETY: every caller passes a static, NUL-terminated C string and only + // checks whether the process environment contains it. + !unsafe { getenv(name.as_ptr()) }.is_null() +} + const SOURCES: [(&str, &str); 19] = [ ("DS4_METAL_FLASH_ATTN_SOURCE", "flash_attn.metal"), ("DS4_METAL_DENSE_SOURCE", "dense.metal"), @@ -1639,11 +1649,14 @@ struct SsdPlan { per_expert_bytes: u64, admission_bytes: u64, preload_experts: u32, + preload_by_layer: Vec>, + preload_seeded: Vec, + static_decode_map_current: std::sync::atomic::AtomicBool, cold: bool, loader: Option, - selected_requests: std::sync::atomic::AtomicU64, - selected_experts: std::sync::atomic::AtomicU64, - selected_wait_ns: std::sync::atomic::AtomicU64, + selected_requests: std::cell::Cell, + selected_experts: std::cell::Cell, + selected_wait_ns: std::cell::Cell, } struct DeepSeekModelSpans { @@ -1668,21 +1681,21 @@ struct SelectedLoadJob { } struct SelectedLoadResult { - ids: Vec, + ids: [i32; 8], loaded: bool, job: SelectedLoadJob, } struct SelectedLoadWorker { - jobs: std::sync::mpsc::Sender>, + jobs: std::sync::mpsc::SyncSender>, results: std::sync::mpsc::Receiver>, thread: Option>, } impl SelectedLoadWorker { fn new() -> Self { - let (jobs, incoming) = std::sync::mpsc::channel::>(); - let (outgoing, results) = std::sync::mpsc::channel(); + let (jobs, incoming) = std::sync::mpsc::sync_channel::>(1); + let (outgoing, results) = std::sync::mpsc::sync_channel(1); let thread = std::thread::spawn(move || { unsafe { ds4_gpu_stream_expert_cache_note_service_thread() }; while let Ok(Some(job)) = incoming.recv() { @@ -1700,6 +1713,9 @@ impl SelectedLoadWorker { } fn begin(&self, selected: &Buffer, table: StreamExpertTable, count: u32) -> Result<(), String> { + if count == 0 || count > 8 { + return Err("SSD expert loading supports one to eight selected experts".into()); + } let mut event = 0; call( unsafe { ds4_gpu_signal_selected_readback_ready(&mut event) }, @@ -1734,24 +1750,23 @@ impl SelectedLoadWorker { .map_err(|_| "SSD expert loader stopped unexpectedly".to_string())??; if !result.loaded { let table = result.job.table(); + let count = result.job.count as usize; call( unsafe { ds4_gpu_stream_expert_cache_begin_selected_load( &table, result.ids.as_ptr(), - result.ids.len() as u32, + count as u32, ) }, "retrying selected SSD expert loading", )?; } if set_override { + let count = result.job.count as usize; call( unsafe { - ds4_gpu_routed_moe_set_selected_override( - result.ids.as_ptr(), - result.ids.len() as u32, - ) + ds4_gpu_routed_moe_set_selected_override(result.ids.as_ptr(), count as u32) }, "selecting streamed experts", )?; @@ -1779,7 +1794,8 @@ fn load_selected_experts(job: SelectedLoadJob) -> Result Result= job.total_experts) { @@ -1801,9 +1817,9 @@ fn load_selected_experts(job: SelectedLoadJob) -> Result Result<(), String> { + fn configure(&mut self, model: &Model, _weights: &Weights) -> Result<(), String> { unsafe { ds4_gpu_set_streaming_expert_cache_expert_bytes(self.per_expert_bytes); ds4_gpu_set_streaming_expert_cache_budget(self.cache_experts); } - if self.cold || self.preload_experts == 0 { + if self.cold + || self.preload_experts == 0 + || env::var_os("DS4_METAL_DISABLE_STREAMING_EXPERT_HOTLIST").is_some() + { self.loader = Some(SelectedLoadWorker::new()); return Ok(()); } - let mut by_layer = vec![Vec::<(i32, u32)>::new(); model.shape.layers as usize]; let mut loaded = 0_u32; let hotlist = match model.shape.model { ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Flash0731 => hotlist::FLASH, @@ -2002,34 +2025,55 @@ impl SsdPlan { if u32::from(layer) >= model.shape.layers || u64::from(expert) >= model.shape.experts { continue; } - by_layer[layer as usize].push((i32::from(expert), self.preload_experts - loaded)); + self.preload_by_layer[layer as usize] + .push((i32::from(expert), self.preload_experts - loaded)); loaded += 1; } - for (layer_index, entries) in by_layer.iter().enumerate() { - if entries.is_empty() { - continue; - } - let ids = entries.iter().map(|entry| entry.0).collect::>(); - let priorities = entries.iter().map(|entry| entry.1).collect::>(); - let count = ids.len() as u32; - let layer = &weights.layers[layer_index]; - let table = stream_expert_table(model, layer, layer_index as u32, model.shape.experts); - call( - unsafe { - ds4_gpu_stream_expert_cache_seed_experts( - &table, - ids.as_ptr(), - priorities.as_ptr(), - count, - ) - }, - "preloading SSD experts", - )?; - } self.loader = Some(SelectedLoadWorker::new()); Ok(()) } + fn seed_mapped_layer( + &self, + model: &Model, + layer: &Layer, + index: usize, + gpu_copy: bool, + ) -> Result { + if self.preload_seeded[index].load(std::sync::atomic::Ordering::Acquire) { + return Ok(true); + } + let entries = &self.preload_by_layer[index]; + if entries.is_empty() { + self.preload_seeded[index].store(true, std::sync::atomic::Ordering::Release); + return Ok(true); + } + let ids = entries.iter().map(|entry| entry.0).collect::>(); + let priorities = entries.iter().map(|entry| entry.1).collect::>(); + let table = stream_expert_table(model, layer, index as u32, model.shape.experts); + let seeded = unsafe { + if gpu_copy { + ds4_gpu_stream_expert_cache_seed_experts_gpu_copy( + &table, + ids.as_ptr(), + priorities.as_ptr(), + ids.len() as u32, + ) + } else { + ds4_gpu_stream_expert_cache_seed_experts( + &table, + ids.as_ptr(), + priorities.as_ptr(), + ids.len() as u32, + ) + } + } != 0; + if seeded { + self.preload_seeded[index].store(true, std::sync::atomic::Ordering::Release); + } + Ok(seeded) + } + fn begin_selected( &self, selected: &Buffer, @@ -2037,9 +2081,9 @@ impl SsdPlan { count: u32, ) -> Result<(), String> { self.selected_requests - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + .set(self.selected_requests.get().saturating_add(1)); self.selected_experts - .fetch_add(u64::from(count), std::sync::atomic::Ordering::Relaxed); + .set(self.selected_experts.get().saturating_add(u64::from(count))); self.loader .as_ref() .ok_or("SSD expert loader was not initialized")? @@ -2053,9 +2097,10 @@ impl SsdPlan { .as_ref() .ok_or("SSD expert loader was not initialized")? .finish(set_override); - self.selected_wait_ns.fetch_add( - u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX), - std::sync::atomic::Ordering::Relaxed, + self.selected_wait_ns.set( + self.selected_wait_ns + .get() + .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)), ); result } @@ -2377,6 +2422,8 @@ struct Scratch { kv: Buffer, compressed_kv: Buffer, compressed_score: Buffer, + index_compressed_kv: Buffer, + index_compressed_score: Buffer, compressed_stage: Buffer, indexer_q: Buffer, indexer_weights: Buffer, @@ -2528,6 +2575,8 @@ impl Scratch { kv: Buffer::floats(shape.head_dim)?, compressed_kv: Buffer::floats(2 * shape.head_dim)?, compressed_score: Buffer::floats(2 * shape.head_dim)?, + index_compressed_kv: Buffer::floats(2 * shape.indexer_head_dim)?, + index_compressed_score: Buffer::floats(2 * shape.indexer_head_dim)?, compressed_stage: Buffer::floats(shape.head_dim)?, indexer_q: Buffer::floats(shape.indexer_heads * shape.indexer_head_dim)?, indexer_weights: Buffer::floats(shape.indexer_heads)?, @@ -2827,7 +2876,11 @@ impl DeepSeekExecutor { .enabled .then(|| SsdPlan::new(&model, &weights, ssd, context, prefill_chunk)) .transpose()?; - let spans = ssd_plan.as_ref().map(|plan| plan.model_spans.as_slice()); + let initial_ssd_spans = ssd_plan + .as_ref() + .map(|_| deepseek_token_model_spans(&model).map(|spans| spans.ranges)) + .transpose()?; + let spans = initial_ssd_spans.as_deref(); let admission = if let Some(plan) = &ssd_plan { plan.admission_bytes } else { @@ -2932,8 +2985,10 @@ impl DeepSeekExecutor { pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { self.eval_target(token)?; - if self.dspark.is_some() && self.ssd.is_some() { + if let (Some(_), Some(ssd)) = (&self.dspark, &self.ssd) { install_speculative_model_maps(&self.model, "DSpark support mapping")?; + ssd.static_decode_map_current + .store(false, std::sync::atomic::Ordering::Release); } if let (Some(dspark), Some(support)) = (&mut self.dspark, self.model.support.as_ref()) { dspark.seed_current_cache( @@ -3682,6 +3737,8 @@ impl DeepSeekExecutor { None, None, Some(mtp.raw_rows), + true, + false, )?; mtp.next_hc.copy_from( 0, @@ -3833,7 +3890,9 @@ impl DeepSeekExecutor { .collect(), }) .collect::>(); - if self.ssd.is_some() { + if let Some(ssd) = &self.ssd { + ssd.static_decode_map_current + .store(false, std::sync::atomic::Ordering::Release); install_deepseek_model_spans( &self.model, &deepseek_token_model_spans(&self.model)?, @@ -3924,7 +3983,23 @@ impl DeepSeekExecutor { shape, )?; } + let seeded_from_map = match &self.ssd { + Some(ssd) => ssd.seed_mapped_layer(&self.model, weights, index, true)?, + None => true, + }; commands.finish()?; + if !seeded_from_map { + let seeded = self + .ssd + .as_ref() + .expect("SSD preload fallback lost its plan") + .seed_mapped_layer(&self.model, weights, index, false)?; + if !seeded { + return Err(format!( + "Metal could not preload SSD experts for layer {index}" + )); + } + } throttle( &mut self.prefill_layer_average[index], started.elapsed(), @@ -4012,9 +4087,7 @@ impl DeepSeekExecutor { if let Some(ssd) = &self.ssd { let mut native = StreamExpertCacheStats::default(); unsafe { ds4_gpu_stream_expert_cache_get_stats(&mut native) }; - let selected = ssd - .selected_experts - .load(std::sync::atomic::Ordering::Relaxed); + let selected = ssd.selected_experts.get(); stats.ssd_enabled = true; stats.ssd_resident_bytes = ssd.resident_bytes; stats.ssd_cache_experts = u64::from(ssd.cache_experts); @@ -4033,14 +4106,9 @@ impl DeepSeekExecutor { stats.ssd_pread_ms = native.pread_ms.max(0.0) as u64; stats.ssd_evict_advise_bytes = native.evict_advise_bytes; stats.ssd_willneed_advise_bytes = native.willneed_advise_bytes; - stats.ssd_selected_requests = ssd - .selected_requests - .load(std::sync::atomic::Ordering::Relaxed); + stats.ssd_selected_requests = ssd.selected_requests.get(); stats.ssd_requested_bytes = ssd.per_expert_bytes.saturating_mul(selected); - stats.ssd_wait_ms = ssd - .selected_wait_ns - .load(std::sync::atomic::Ordering::Relaxed) - / 1_000_000; + stats.ssd_wait_ms = ssd.selected_wait_ns.get() / 1_000_000; } stats } @@ -4161,6 +4229,12 @@ impl DeepSeekExecutor { self.encode_token_embedding(token)?; for index in 0..self.weights.layers.len() { self.encode_token_layer(index, token)?; + if index == 3 && index + 1 < self.weights.layers.len() { + call( + unsafe { ds4_gpu_flush_commands() }, + "flushing the DeepSeek decode graph", + )?; + } } encode_output( &self.session.scratch, @@ -4172,43 +4246,50 @@ impl DeepSeekExecutor { } fn encode_streaming_token(&mut self, token: u32) -> Result<(), String> { - let slab_bytes = self + let ssd = self .ssd .as_ref() - .expect("streaming token lost its SSD plan") - .per_expert_bytes; - install_deepseek_model_spans( - &self.model, - &deepseek_token_model_spans(&self.model)?, - "DeepSeek token embedding mapping", - )?; - let commands = Commands::begin()?; - self.encode_token_embedding(token)?; - commands.finish()?; - - for index in 0..self.weights.layers.len() { + .expect("streaming token lost its SSD plan"); + if !ssd + .static_decode_map_current + .load(std::sync::atomic::Ordering::Acquire) + { install_deepseek_model_spans( &self.model, - &deepseek_layer_model_spans( - &self.model, - &self.weights.layers[index], - index as u32, - true, - slab_bytes, - )?, - "DeepSeek decode layer mapping", + &DeepSeekModelSpans { + ranges: ssd.model_spans.clone(), + max_tensor_bytes: self.model.main.max_tensor_bytes(), + }, + "DeepSeek static decode mapping", )?; - let commands = Commands::begin()?; - self.encode_token_layer(index, token)?; - commands.finish()?; + ssd.static_decode_map_current + .store(true, std::sync::atomic::Ordering::Release); } - install_deepseek_model_spans( - &self.model, - &deepseek_output_model_spans(&self.model)?, - "DeepSeek output mapping", - )?; - let commands = Commands::begin()?; + let layer_batch = self.profile.is_none(); + let mut commands = Some(Commands::begin()?); + self.encode_token_embedding(token)?; + if !layer_batch { + commands + .take() + .expect("decode command batch disappeared") + .finish()?; + } + for index in 0..self.weights.layers.len() { + if !layer_batch { + commands = Some(Commands::begin()?); + } + self.encode_token_layer(index, token)?; + if !layer_batch { + commands + .take() + .expect("decode command batch disappeared") + .finish()?; + } + } + if !layer_batch { + commands = Some(Commands::begin()?); + } encode_output( &self.session.scratch, &self.weights, @@ -4216,7 +4297,10 @@ impl DeepSeekExecutor { self.model.main.map_ptr().cast(), self.model.main.len(), )?; - commands.finish() + commands + .take() + .expect("decode command batch disappeared") + .finish() } fn encode_token_embedding(&mut self, token: u32) -> Result<(), String> { @@ -4257,6 +4341,8 @@ impl DeepSeekExecutor { self.session.raw_cap, self.steering.as_ref(), self.ssd.as_ref(), + self.quality, + self.profile.is_some(), )?; if let Some(profile) = &mut self.profile { profile.record( @@ -4751,6 +4837,8 @@ fn compress_attention_batch( shape.rope_beta_slow, shape.rms_epsilon, false, + false, + false, ) }, "updating batched attention compression", @@ -4950,6 +5038,8 @@ fn compress_index_batch( shape.rope_beta_slow, shape.rms_epsilon, false, + false, + false, ) }, "updating batched indexer compression", @@ -5910,12 +6000,277 @@ fn encode_layer( raw_cap: u32, steering: Option<&Steering>, ssd: Option<&SsdPlan>, + quality: bool, + profile_active: bool, ) -> Result<(), String> { encode_layer_with_cache_rows( - s, state, w, shape, map, size, layer, pos, token, raw_cap, steering, ssd, None, + s, + state, + w, + shape, + map, + size, + layer, + pos, + token, + raw_cap, + steering, + ssd, + None, + quality, + profile_active, ) } +fn decode_feature_allowed( + pre_m5: bool, + m5: bool, + pre_m5_disabled: bool, + m5_disabled: bool, +) -> bool { + (pre_m5 && !pre_m5_disabled) || (m5 && !m5_disabled) +} + +fn decode_feature_enabled(pre_m5_disable: &CStr, m5_disable: &CStr) -> bool { + let pre_m5 = unsafe { ds4_gpu_device_is_pre_m5_apple_silicon() } != 0; + let m5 = unsafe { ds4_gpu_device_is_m5_apple_silicon() } != 0; + decode_feature_allowed( + pre_m5, + m5, + environment_present(pre_m5_disable), + environment_present(m5_disable), + ) +} + +#[allow(clippy::too_many_arguments)] +fn decode_hc_mix( + s: &Scratch, + residual: &Buffer, + mix_weight: Weight, + scale: Weight, + base: Weight, + norm_weight: Weight, + shape: super::Shape, + map: *const c_void, + size: u64, + purpose: &'static str, +) -> Result<(), String> { + let hc_dim = shape.hc * shape.embd; + let mix_hc = 2 * shape.hc + shape.hc * shape.hc; + let fused = hc_dim == 16_384 + && mix_hc == 24 + && mix_weight.kind == F16 + && !environment_present(c"DS4_METAL_DISABLE_PRE_M5_HC_NORM_MIX_FUSE") + && decode_feature_enabled( + c"DS4_METAL_DISABLE_PRE_M5_HC_PRODUCER_PRE_NORM_FUSE", + c"DS4_METAL_DISABLE_M5_HC_PRODUCER_PRE_NORM_FUSE", + ) + && unsafe { ds4_gpu_hc_rms_norm_mix_f16_available() } != 0; + if fused { + let result = unsafe { + ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( + s.hc_mix.raw(), + s.current.raw(), + s.norm.raw(), + s.hc_split.raw(), + residual.raw(), + map, + size, + mix_weight.offset, + scale.offset, + base.offset, + norm_weight.offset, + hc_dim as u32, + mix_hc as u32, + shape.embd as u32, + shape.hc as u32, + shape.hc_sinkhorn as u32, + shape.rms_epsilon, + shape.hc_epsilon, + shape.rms_epsilon, + ) + }; + if result < 0 { + return Err(format!("Metal failed while {purpose}")); + } + if result > 0 { + return Ok(()); + } + call( + unsafe { + ds4_gpu_hc_rms_norm_mix_f16_tensor( + s.hc_mix.raw(), + residual.raw(), + map, + size, + mix_weight.offset, + hc_dim as u32, + mix_hc as u32, + shape.rms_epsilon, + ) + }, + purpose, + )?; + } else { + call( + unsafe { + ds4_gpu_rms_norm_plain_tensor( + s.flat_hc.raw(), + residual.raw(), + hc_dim as u32, + shape.rms_epsilon, + ) + }, + purpose, + )?; + matmul(&s.hc_mix, mix_weight, hc_dim, mix_hc, &s.flat_hc, map, size)?; + } + call( + unsafe { + ds4_gpu_hc_split_weighted_sum_norm_tensor( + s.current.raw(), + s.norm.raw(), + s.hc_split.raw(), + s.hc_mix.raw(), + residual.raw(), + map, + size, + scale.offset, + base.offset, + norm_weight.offset, + shape.embd as u32, + shape.hc as u32, + shape.hc_sinkhorn as u32, + shape.hc_epsilon, + shape.rms_epsilon, + ) + }, + purpose, + ) +} + +#[allow(clippy::too_many_arguments)] +fn try_qkv_pair_compressor_fusion( + s: &Scratch, + state: &LayerState, + w: &Layer, + shape: super::Shape, + map: *const c_void, + size: u64, + layer: u32, + pos: u32, + resident: bool, +) -> Result { + let ratio = compression_ratio(shape, layer); + let Some(attn_weights) = w.attn_compressor else { + return Ok(false); + }; + let Some(attn_state) = state.compression.as_ref() else { + return Ok(false); + }; + if !resident + || !matches!(ratio, 4 | 128) + || w.attn_q_a.kind != Q8_0 + || w.attn_kv.kind != Q8_0 + || attn_weights.kv.kind != F16 + || attn_weights.gate.kind != F16 + || attn_weights.kv.dims + != [ + shape.embd, + (if ratio == 4 { 2 } else { 1 }) * shape.head_dim, + 1, + ] + || attn_weights.gate.dims != attn_weights.kv.dims + { + return Ok(false); + } + let enabled = if ratio == 4 { + decode_feature_enabled( + c"DS4_METAL_DISABLE_PRE_M5_QKV_PAIR_QUAD_FUSE", + c"DS4_METAL_DISABLE_M5_QKV_PAIR_QUAD_FUSE", + ) + } else { + decode_feature_enabled( + c"DS4_METAL_DISABLE_PRE_M5_QKV_PAIR_COMPRESSOR_FUSE", + c"DS4_METAL_DISABLE_M5_QKV_PAIR_COMPRESSOR_FUSE", + ) + }; + if !enabled { + return Ok(false); + } + let (out1_kv, out1_score, state1_kv, state1_score, weights1, width1) = if ratio == 4 { + let Some(index_weights) = w.indexer.map(|weights| weights.compressor) else { + return Ok(false); + }; + let Some(index_state) = state.indexer.as_ref() else { + return Ok(false); + }; + if index_weights.kv.kind != F16 + || index_weights.gate.kind != F16 + || index_weights.kv.dims != [shape.embd, 2 * shape.indexer_head_dim, 1] + || index_weights.gate.dims != index_weights.kv.dims + { + return Ok(false); + } + ( + &s.index_compressed_kv, + &s.index_compressed_score, + &index_state.state_kv, + &index_state.state_score, + index_weights, + (2 * shape.indexer_head_dim) as u32, + ) + } else { + ( + &s.compressed_kv, + &s.compressed_score, + &attn_state.state_kv, + &attn_state.state_score, + attn_weights, + 0, + ) + }; + let result = unsafe { + ds4_gpu_qkv_pair_quad_compressor_store_tensor( + s.q_rank.raw(), + s.kv_raw.raw(), + s.compressed_kv.raw(), + s.compressed_score.raw(), + out1_kv.raw(), + out1_score.raw(), + attn_state.state_kv.raw(), + attn_state.state_score.raw(), + state1_kv.raw(), + state1_score.raw(), + map, + size, + w.attn_q_a.offset, + w.attn_kv.offset, + attn_weights.kv.offset, + attn_weights.gate.offset, + weights1.kv.offset, + weights1.gate.offset, + attn_weights.ape.offset, + attn_weights.ape.kind, + weights1.ape.offset, + weights1.ape.kind, + shape.embd as u32, + shape.lora_q as u32, + shape.head_dim as u32, + (u64::from(if ratio == 4 { 2_u32 } else { 1 }) * shape.head_dim) as u32, + width1, + s.norm.raw(), + ratio, + pos, + ) + }; + if result < 0 { + Err("Metal failed while fusing Q/KV and compressor projections".into()) + } else { + Ok(result > 0) + } +} + #[allow(clippy::too_many_arguments)] fn encode_layer_with_cache_rows( s: &Scratch, @@ -5931,9 +6286,9 @@ fn encode_layer_with_cache_rows( steering: Option<&Steering>, ssd: Option<&SsdPlan>, cache_rows: Option, + quality: bool, + profile_active: bool, ) -> Result<(), String> { - let hc_dim = shape.hc * shape.embd; - let mix_hc = 2 * shape.hc + shape.hc * shape.hc; let compressed = layer >= 2; let freq_base = if compressed { shape.compress_rope_base @@ -5961,85 +6316,111 @@ fn encode_layer_with_cache_rows( 0 }; let raw_cache = state.raw_cache.raw(); - call( - unsafe { - ds4_gpu_rms_norm_plain_tensor( - s.flat_hc.raw(), - s.current_hc.raw(), - hc_dim as u32, - shape.rms_epsilon, - ) - }, - "attention HC norm", - )?; - matmul( - &s.hc_mix, + decode_hc_mix( + s, + &s.current_hc, w.hc_attn_fn, - hc_dim, - mix_hc, - &s.flat_hc, + w.hc_attn_scale, + w.hc_attn_base, + w.attn_norm, + shape, map, size, - )?; - call( - unsafe { - ds4_gpu_hc_split_weighted_sum_norm_tensor( - s.current.raw(), - s.norm.raw(), - s.hc_split.raw(), - s.hc_mix.raw(), - s.current_hc.raw(), - map, - size, - w.hc_attn_scale.offset, - w.hc_attn_base.offset, - w.attn_norm.offset, - shape.embd as u32, - shape.hc as u32, - shape.hc_sinkhorn as u32, - shape.hc_epsilon, - shape.rms_epsilon, - ) - }, "attention HC mix", )?; - call( - unsafe { - ds4_gpu_matmul_q8_0_pair_tensor( - s.q_rank.raw(), - s.kv_raw.raw(), - map, - size, - w.attn_q_a.offset, - w.attn_kv.offset, - shape.embd, - shape.lora_q, - shape.head_dim, - s.norm.raw(), - 1, - ) - }, - "Q/KV projection", - )?; - call( - unsafe { - ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( - s.q_rank_norm.raw(), - s.q_rank.raw(), - map, - size, - w.attn_q_a_norm.offset, - shape.lora_q as u32, - s.kv.raw(), - s.kv_raw.raw(), - w.attn_kv_norm.offset, - shape.head_dim as u32, - 1, - shape.rms_epsilon, - ) - }, - "Q/KV norm", + let qkv_compressors_fused = try_qkv_pair_compressor_fusion( + s, + state, + w, + shape, + map, + size, + layer, + pos, + ssd.is_none() && cache_rows.is_none(), )?; + if !qkv_compressors_fused { + call( + unsafe { + ds4_gpu_matmul_q8_0_pair_tensor( + s.q_rank.raw(), + s.kv_raw.raw(), + map, + size, + w.attn_q_a.offset, + w.attn_kv.offset, + shape.embd, + shape.lora_q, + shape.head_dim, + s.norm.raw(), + 1, + ) + }, + "Q/KV projection", + )?; + } + let raw_row = pos % raw_cap; + let kv_norm_store_fused = cache_rows.is_none() + && shape.head_kv == 1 + && shape.head_dim == 512 + && shape.rot == 64 + && raw_row < raw_cap + && !environment_present(c"DS4_METAL_DISABLE_PRE_M5_QKV_NORM_KV_STORE_FUSE") + && (unsafe { ds4_gpu_device_is_pre_m5_apple_silicon() } != 0 + || unsafe { ds4_gpu_device_is_m5_apple_silicon() } != 0) + && unsafe { ds4_gpu_kv_rope_fp8_fuse_available() } != 0; + if kv_norm_store_fused { + call( + unsafe { + ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_tensor( + s.q_rank_norm.raw(), + s.q_rank.raw(), + map, + size, + w.attn_q_a_norm.offset, + shape.lora_q as u32, + s.kv.raw(), + s.kv_raw.raw(), + w.attn_kv_norm.offset, + shape.head_dim as u32, + raw_cache, + raw_cap as u64, + raw_row, + shape.rot as u32, + pos, + original, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + }, + "Q/KV norm, RoPE, and cache write", + )?; + } else { + call( + unsafe { + ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( + s.q_rank_norm.raw(), + s.q_rank.raw(), + map, + size, + w.attn_q_a_norm.offset, + shape.lora_q as u32, + s.kv.raw(), + s.kv_raw.raw(), + w.attn_kv_norm.offset, + shape.head_dim as u32, + 1, + shape.rms_epsilon, + ) + }, + "Q/KV norm", + )?; + } q8( &s.q, w.attn_q_b, @@ -6103,80 +6484,87 @@ fn encode_layer_with_cache_rows( "Q RoPE", )?; } - call( - unsafe { - ds4_gpu_rope_tail_tensor( - s.kv.raw(), - 1, - shape.head_kv as u32, - shape.head_dim as u32, - shape.rot as u32, - pos, - original, - false, - freq_base, - freq_scale, - ext, - attn_factor, - shape.rope_beta_fast, - shape.rope_beta_slow, - ) - }, - "KV RoPE", - )?; - let raw_row = pos % raw_cap; - call( - unsafe { - ds4_gpu_kv_fp8_store_raw_tensor( - s.kv.raw(), - raw_cache, - raw_cap, - raw_row, - shape.head_dim as u32, - shape.rot as u32, - ) - }, - "KV cache write", - )?; + if !kv_norm_store_fused { + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.kv.raw(), + 1, + shape.head_kv as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "KV RoPE", + )?; + call( + unsafe { + ds4_gpu_kv_fp8_store_raw_tensor( + s.kv.raw(), + raw_cache, + raw_cap, + raw_row, + shape.head_dim as u32, + shape.rot as u32, + ) + }, + "KV cache write", + )?; + } let (n_raw, _) = raw_decode_span( cache_rows.unwrap_or(pos), raw_cap, shape.sliding_window as u32, ); let raw_start = (pos + 1 - n_raw) % raw_cap; - if let (Some(weights), Some(compression)) = - (w.attn_compressor.as_ref(), state.compression.as_mut()) + if let (Some(attn_weights), Some(compression)) = (w.attn_compressor, state.compression.as_mut()) { - update_compression( - s, - compression, - *weights, - shape, - map, - size, - pos, - original, - freq_base, - freq_scale, - ext, - attn_factor, - )?; - } - if let (Some(weights), Some(indexer)) = (w.indexer, state.indexer.as_mut()) { - update_indexer_compression( - s, - indexer, - weights.compressor, - shape, - map, - size, - pos, - original, - freq_base, - freq_scale, - ext, - attn_factor, - )?; + if let (Some(index_weights), Some(indexer)) = (w.indexer, state.indexer.as_mut()) { + update_compression_pair( + s, + compression, + attn_weights, + indexer, + index_weights.compressor, + shape, + map, + size, + pos, + original, + freq_base, + freq_scale, + ext, + attn_factor, + qkv_compressors_fused, + cache_rows.is_none(), + ssd.is_none() && cache_rows.is_none(), + )?; + } else { + update_compression( + s, + compression, + attn_weights, + shape, + map, + size, + pos, + original, + freq_base, + freq_scale, + ext, + attn_factor, + qkv_compressors_fused, + )?; + } } let (comp_cache, n_comp) = state .compression @@ -6272,6 +6660,11 @@ fn encode_layer_with_cache_rows( } else { None }; + let fuse_inverse_rope = !environment_present(c"DS4_METAL_DISABLE_PRE_M5_ATTN_INV_ROPE_FUSE") + && (unsafe { ds4_gpu_device_is_pre_m5_apple_silicon() } != 0 + || unsafe { ds4_gpu_device_is_m5_apple_silicon() } != 0) + && unsafe { ds4_gpu_decode_attn_rope_fuse_available() } != 0; + let mut inverse_rope_armed = false; if indexed.is_some() { call( unsafe { @@ -6301,6 +6694,24 @@ fn encode_layer_with_cache_rows( "indexed attention", )?; } else { + if fuse_inverse_rope { + unsafe { + ds4_gpu_set_decode_attn_rope_fuse( + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + true, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }; + inverse_rope_armed = true; + } call( unsafe { ds4_gpu_attention_decode_heads_tensor( @@ -6325,27 +6736,29 @@ fn encode_layer_with_cache_rows( "attention", )?; } - call( - unsafe { - ds4_gpu_rope_tail_tensor( - s.heads.raw(), - 1, - shape.heads as u32, - shape.head_dim as u32, - shape.rot as u32, - pos, - original, - true, - freq_base, - freq_scale, - ext, - attn_factor, - shape.rope_beta_fast, - shape.rope_beta_slow, - ) - }, - "inverse attention RoPE", - )?; + if !(inverse_rope_armed && unsafe { ds4_gpu_decode_attn_rope_fuse_used() } != 0) { + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.heads.raw(), + 1, + shape.heads as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + true, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "inverse attention RoPE", + )?; + } let group_dim = shape.head_dim * (shape.heads / shape.out_groups); call( unsafe { @@ -6407,83 +6820,163 @@ fn encode_layer_with_cache_rows( "attention output", )?; } - call( - unsafe { - ds4_gpu_rms_norm_plain_tensor( - s.flat_hc.raw(), - s.next_hc.raw(), - hc_dim as u32, - shape.rms_epsilon, - ) - }, - "FFN HC norm", - )?; - matmul( - &s.hc_mix, + decode_hc_mix( + s, + &s.next_hc, w.hc_ffn_fn, - hc_dim, - mix_hc, - &s.flat_hc, + w.hc_ffn_scale, + w.hc_ffn_base, + w.ffn_norm, + shape, map, size, - )?; - call( - unsafe { - ds4_gpu_hc_split_weighted_sum_norm_tensor( - s.current.raw(), - s.norm.raw(), - s.hc_split.raw(), - s.hc_mix.raw(), - s.next_hc.raw(), - map, - size, - w.hc_ffn_scale.offset, - w.hc_ffn_base.offset, - w.ffn_norm.offset, - shape.embd as u32, - shape.hc as u32, - shape.hc_sinkhorn as u32, - shape.hc_epsilon, - shape.rms_epsilon, - ) - }, "FFN HC mix", )?; - matmul( - &s.router_logits, - w.router, - shape.embd, - shape.experts, - &s.norm, - map, - size, - )?; - call( - unsafe { - ds4_gpu_router_select_tensor( - s.router_selected.raw(), - s.router_weights.raw(), - s.router_probs.raw(), - map, - size, - w.router_bias.map_or(0, |v| v.offset), - w.router_hash.map_or(0, |v| v.offset), - w.router_hash.map_or(0, |v| v.dims[1] as u32), - token, - shape.experts as u32, - shape.experts_used as u32, - shape.expert_weight_scale, - 0, - 0, - w.router_bias.is_some(), - w.router_hash.is_some(), - s.router_logits.raw(), - ) - }, - "expert routing", - )?; let gate_row = w.expert_gate.bytes / (w.expert_gate.dims[1] * w.expert_gate.dims[2]); let down_row = w.expert_down.bytes / (w.expert_down.dims[1] * w.expert_down.dims[2]); + let gate_expert_bytes = w.expert_gate.dims[1] * gate_row; + let down_expert_bytes = w.expert_down.dims[1] * down_row; + let parallel_eligible = !quality + && ssd.is_none() + && cache_rows.is_none() + && !profile_active + && steering.is_none_or(|value| value.ffn_scale == 0.0) + && !environment_present(c"DS4_METAL_MOE_ONE_STAGE_PROFILE") + && !environment_present(c"DS4_METAL_MOE_WRITE_CLAMPED_ACT") + && !environment_present(c"DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") + && !environment_present(c"DS4_METAL_Q8_MV_NSG") + && w.expert_gate.kind == IQ2_XXS + && w.expert_up.kind == IQ2_XXS + && w.expert_down.kind == Q2_K + && shape.experts == 256 + && shape.experts_used == 6 + && shape.embd == 4096 + && w.expert_gate.dims[0] == 4096 + && w.expert_down.dims[0] == 2048 + && w.expert_down.dims[1] == 4096 + && gate_row == 1056 + && gate_expert_bytes == 2_162_688 + && down_row == 672 + && down_expert_bytes == 2_752_512 + && w.shared_down.kind == Q8_0 + && w.router.kind == F16 + && w.router.dims[0] == shape.embd + && w.router.dims[1] == shape.experts + && decode_feature_enabled( + c"DS4_METAL_DISABLE_PRE_M5_PARALLEL_FULL_FFN", + c"DS4_METAL_DISABLE_M5_PARALLEL_FULL_FFN", + ); + let router_fusion_eligible = !quality + && ssd.is_none() + && cache_rows.is_none() + && w.shared_gate.kind == Q8_0 + && w.shared_up.kind == Q8_0 + && w.router.kind == F16 + && w.router.dims[0] == shape.embd + && w.router.dims[1] == shape.experts + && (!((unsafe { ds4_gpu_device_is_pre_m5_apple_silicon() } != 0) + && environment_present(c"DS4_METAL_DISABLE_PRE_M5_ROUTER_SHARED_FUSE"))) + && (!((unsafe { ds4_gpu_device_is_m5_apple_silicon() } != 0) + && environment_present(c"DS4_METAL_DISABLE_M5_ROUTER_SHARED_FUSE"))) + && (unsafe { ds4_gpu_device_is_pre_m5_apple_silicon() } != 0 + || unsafe { ds4_gpu_device_is_m5_apple_silicon() } != 0); + let mut router_shared_done = false; + let router_project_select = router_fusion_eligible + && parallel_eligible + && w.router_hash.is_none() + && !environment_present(c"DS4_METAL_DISABLE_M5_ROUTER_PROJECT_SELECT_FUSE") + && unsafe { ds4_gpu_device_is_m5_apple_silicon() } != 0; + let mut router_selected_done = false; + let mut router_projected = false; + let mut parallel_router_ready = false; + if router_project_select { + let fused = unsafe { + ds4_gpu_router_project_select_fused_tensor( + s.router_logits.raw(), + s.router_probs.raw(), + s.router_selected.raw(), + s.router_weights.raw(), + map, + size, + w.router.offset, + w.router_bias.map_or(0, |value| value.offset), + w.router_bias.is_some(), + s.norm.raw(), + ) + }; + if fused < 0 { + return Err("Metal failed while fusing router projection and selection".into()); + } + router_selected_done = fused > 0; + router_projected = router_selected_done; + parallel_router_ready = router_selected_done; + } + if router_fusion_eligible && !router_selected_done { + let fused = unsafe { + ds4_gpu_router_shared_gate_up_q8_0_tensor( + s.router_logits.raw(), + s.shared_gate.raw(), + s.shared_up.raw(), + s.shared_mid.raw(), + map, + size, + w.router.offset, + w.shared_gate.offset, + w.shared_up.offset, + shape.embd, + shape.experts, + shape.ff_expert, + s.norm.raw(), + shape.swiglu_clamp, + parallel_eligible, + ) + }; + if fused < 0 { + return Err("Metal failed while fusing router and shared-expert projections".into()); + } + if fused > 0 { + router_projected = true; + router_shared_done = !parallel_eligible; + parallel_router_ready = parallel_eligible; + } + } + if !router_selected_done { + if !router_projected { + matmul( + &s.router_logits, + w.router, + shape.embd, + shape.experts, + &s.norm, + map, + size, + )?; + } + call( + unsafe { + ds4_gpu_router_select_tensor( + s.router_selected.raw(), + s.router_weights.raw(), + s.router_probs.raw(), + map, + size, + w.router_bias.map_or(0, |v| v.offset), + w.router_hash.map_or(0, |v| v.offset), + w.router_hash.map_or(0, |v| v.dims[1] as u32), + token, + shape.experts as u32, + shape.experts_used as u32, + shape.expert_weight_scale, + 0, + 0, + w.router_bias.is_some(), + w.router_hash.is_some(), + s.router_logits.raw(), + ) + }, + "expert routing", + )?; + } if let Some(ssd) = ssd { let table = StreamExpertTable { model_map: map, @@ -6498,25 +6991,46 @@ fn encode_layer_with_cache_rows( }; ssd.begin_selected(&s.router_selected, table, shape.experts_used as u32)?; } - call( - unsafe { - ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( - s.shared_gate.raw(), - s.shared_up.raw(), - s.shared_mid.raw(), + let parallel = (parallel_eligible && parallel_router_ready) + .then(|| { + ParallelFfn::start( + &s.shared_gate, + &s.shared_up, + &s.shared_mid, + &s.shared_out, map, size, w.shared_gate.offset, w.shared_up.offset, - shape.embd, - shape.ff_expert, - s.norm.raw(), + w.shared_down.offset, + shape.embd as u32, + shape.ff_expert as u32, + &s.norm, shape.swiglu_clamp, ) - }, - "shared expert gate/up", - )?; - if ssd.is_some() || steering.is_some_and(|value| value.ffn_scale != 0.0) { + }) + .flatten(); + if parallel.is_none() && !router_shared_done { + call( + unsafe { + ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + s.shared_gate.raw(), + s.shared_up.raw(), + s.shared_mid.raw(), + map, + size, + w.shared_gate.offset, + w.shared_up.offset, + shape.embd, + shape.ff_expert, + s.norm.raw(), + shape.swiglu_clamp, + ) + }, + "shared expert gate/up", + )?; + } + if steering.is_some_and(|value| value.ffn_scale != 0.0) { q8( &s.shared_out, w.shared_down, @@ -6565,6 +7079,12 @@ fn encode_layer_with_cache_rows( }, "routed experts", )?; + let parallel_done = if let Some(work) = parallel { + work.finish()?; + true + } else { + false + }; if let Some(steering) = steering.filter(|value| value.ffn_scale != 0.0) { call( unsafe { @@ -6591,7 +7111,7 @@ fn encode_layer_with_cache_rows( }, "steered FFN HC expansion", ) - } else if ssd.is_some() { + } else if parallel_done { call( unsafe { ds4_gpu_hc_expand_add_split_tensor( @@ -6604,7 +7124,7 @@ fn encode_layer_with_cache_rows( shape.hc as u32, ) }, - "streamed FFN HC expansion", + "parallel FFN HC expansion", ) } else { call( @@ -6630,6 +7150,312 @@ fn encode_layer_with_cache_rows( } } +#[allow(clippy::too_many_arguments)] +fn project_compressor_pair( + kv: &Buffer, + score: &Buffer, + state: &CompressionState, + weights: CompressorWeights, + input: &Buffer, + input_width: u64, + width: u32, + map: *const c_void, + size: u64, + pos: u32, +) -> Result { + let fused = unsafe { + ds4_gpu_matmul_f16_pair_compressor_store_tensor( + kv.raw(), + score.raw(), + state.state_kv.raw(), + state.state_score.raw(), + map, + size, + weights.kv.offset, + weights.gate.offset, + weights.ape.offset, + weights.ape.kind, + input_width, + width, + input.raw(), + state.ratio, + pos, + ) + }; + if fused < 0 { + return Err("Metal failed while storing compressor projections".into()); + } + if fused == 0 { + call( + unsafe { + ds4_gpu_matmul_f16_pair_tensor( + kv.raw(), + score.raw(), + map, + size, + weights.kv.offset, + weights.gate.offset, + input_width, + width as u64, + input.raw(), + 1, + ) + }, + "compressor projection", + )?; + } + Ok(fused > 0) +} + +#[allow(clippy::too_many_arguments)] +fn update_compression_pair( + s: &Scratch, + attn: &mut CompressionState, + attn_weights: CompressorWeights, + index: &mut CompressionState, + index_weights: CompressorWeights, + shape: super::Shape, + map: *const c_void, + size: u64, + pos: u32, + original: u32, + freq_base: f32, + freq_scale: f32, + ext: f32, + attn_factor: f32, + projected: bool, + full_phase: bool, + resident: bool, +) -> Result<(), String> { + let attn_width = (2 * shape.head_dim) as u32; + let index_width = (2 * shape.indexer_head_dim) as u32; + let mut stored = projected; + if !stored + && full_phase + && attn.ratio == 4 + && attn_weights.kv.kind == F16 + && attn_weights.gate.kind == F16 + && index_weights.kv.kind == F16 + && index_weights.gate.kind == F16 + && !environment_present(c"DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_QUAD_STORE") + && (unsafe { ds4_gpu_device_is_pre_m5_apple_silicon() } != 0 + || unsafe { ds4_gpu_device_is_m5_apple_silicon() } != 0) + { + let result = unsafe { + ds4_gpu_matmul_f16_quad_compressor_store_tensor( + s.compressed_kv.raw(), + s.compressed_score.raw(), + s.index_compressed_kv.raw(), + s.index_compressed_score.raw(), + attn.state_kv.raw(), + attn.state_score.raw(), + index.state_kv.raw(), + index.state_score.raw(), + map, + size, + attn_weights.kv.offset, + attn_weights.gate.offset, + index_weights.kv.offset, + index_weights.gate.offset, + attn_weights.ape.offset, + attn_weights.ape.kind, + index_weights.ape.offset, + index_weights.ape.kind, + shape.embd, + attn_width, + index_width, + s.norm.raw(), + 4, + pos, + ) + }; + if result < 0 { + return Err("Metal failed while fusing compressor projections".into()); + } + stored = result > 0; + } + let (attn_stored, index_stored) = if stored { + (true, true) + } else { + ( + project_compressor_pair( + &s.compressed_kv, + &s.compressed_score, + attn, + attn_weights, + &s.norm, + shape.embd, + attn_width, + map, + size, + pos, + )?, + project_compressor_pair( + &s.index_compressed_kv, + &s.index_compressed_score, + index, + index_weights, + &s.norm, + shape.embd, + index_width, + map, + size, + pos, + )?, + ) + }; + let emit = (pos + 1).is_multiple_of(4); + let fused_finalize = emit + && resident + && shape.head_dim == 512 + && shape.indexer_head_dim == 128 + && attn_weights.norm.kind == F32 + && index_weights.norm.kind == F32 + && unsafe { ds4_gpu_kv_rope_fp8_fuse_available() } != 0 + && decode_feature_enabled( + c"DS4_METAL_DISABLE_PRE_M5_COMP_FINALIZE_FUSE", + c"DS4_METAL_DISABLE_M5_COMP_FINALIZE_FUSE", + ); + call( + unsafe { + ds4_gpu_compressor_update_tensor( + s.compressed_kv.raw(), + s.compressed_score.raw(), + attn.state_kv.raw(), + attn.state_score.raw(), + s.compressed_stage.raw(), + map, + size, + attn_weights.ape.offset, + attn_weights.ape.kind, + attn_weights.norm.offset, + attn_weights.norm.kind, + shape.head_dim as u32, + 4, + pos, + 0, + shape.rot as u32, + original, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + attn_stored, + true, + fused_finalize, + ) + }, + "attention compressor update", + )?; + call( + unsafe { + ds4_gpu_compressor_update_tensor( + s.index_compressed_kv.raw(), + s.index_compressed_score.raw(), + index.state_kv.raw(), + index.state_score.raw(), + index.cache.raw(), + map, + size, + index_weights.ape.offset, + index_weights.ape.kind, + index_weights.norm.offset, + index_weights.norm.kind, + shape.indexer_head_dim as u32, + 4, + pos, + index.rows, + shape.rot as u32, + original, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + index_stored, + true, + fused_finalize, + ) + }, + "index compressor update", + )?; + if !emit { + return Ok(()); + } + if fused_finalize { + let fused = unsafe { + ds4_gpu_dsv4_comp_row_finalize_tensor( + s.compressed_stage.raw(), + attn.cache.raw(), + attn.rows, + attn_weights.norm.offset, + index.cache.raw(), + index.rows, + index_weights.norm.offset, + attn.state_kv.raw(), + attn.state_score.raw(), + index.state_kv.raw(), + index.state_score.raw(), + map, + size, + pos + 1 - 4, + shape.rot as u32, + original, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + }; + if fused != 1 { + return Err("Metal failed while finalizing compressor rows".into()); + } + } else { + call( + unsafe { + ds4_gpu_dsv4_fp8_kv_quantize_tensor( + s.compressed_stage.raw(), + 1, + shape.head_dim as u32, + shape.rot as u32, + ) + }, + "compressed KV quantization", + )?; + call( + unsafe { + ds4_gpu_tensor_copy_f32_to_f16( + attn.cache.raw(), + attn.rows as u64 * shape.head_dim * 2, + s.compressed_stage.raw(), + 0, + shape.head_dim, + ) + }, + "compressed KV cache write", + )?; + let row = index.cache.view( + index.rows as u64 * shape.indexer_head_dim * 4, + shape.indexer_head_dim * 4, + )?; + call( + unsafe { ds4_gpu_dsv4_indexer_qat_tensor(row.raw(), 1, shape.indexer_head_dim as u32) }, + "compressed index quantization", + )?; + } + attn.rows += 1; + index.rows += 1; + Ok(()) +} + #[allow(clippy::too_many_arguments)] fn update_compression( s: &Scratch, @@ -6644,6 +7470,7 @@ fn update_compression( freq_scale: f32, ext: f32, attn_factor: f32, + projected: bool, ) -> Result<(), String> { let emit = update_compressor_stage( s, @@ -6659,6 +7486,7 @@ fn update_compression( ext, attn_factor, shape.head_dim as u32, + projected, )?; if emit { call( @@ -6689,59 +7517,6 @@ fn update_compression( Ok(()) } -#[allow(clippy::too_many_arguments)] -fn update_indexer_compression( - s: &Scratch, - state: &mut CompressionState, - weights: CompressorWeights, - shape: super::Shape, - map: *const c_void, - size: u64, - pos: u32, - original: u32, - freq_base: f32, - freq_scale: f32, - ext: f32, - attn_factor: f32, -) -> Result<(), String> { - let emit = update_compressor_stage( - s, - state, - weights, - shape, - map, - size, - pos, - original, - freq_base, - freq_scale, - ext, - attn_factor, - shape.indexer_head_dim as u32, - )?; - if emit { - call( - unsafe { - ds4_gpu_dsv4_indexer_qat_tensor( - s.compressed_stage.raw(), - 1, - shape.indexer_head_dim as u32, - ) - }, - "compressed index quantization", - )?; - state.cache.copy_from( - state.rows as u64 * shape.indexer_head_dim * 4, - &s.compressed_stage, - 0, - shape.indexer_head_dim * 4, - "compressed index cache write", - )?; - state.rows += 1; - } - Ok(()) -} - #[allow(clippy::too_many_arguments)] fn update_compressor_stage( s: &Scratch, @@ -6757,26 +7532,31 @@ fn update_compressor_stage( ext: f32, attn_factor: f32, head_dim: u32, + projected: bool, ) -> Result { let width = if state.ratio == 4 { 2 } else { 1 } * head_dim; - let fused = unsafe { - ds4_gpu_matmul_f16_pair_compressor_store_tensor( - s.compressed_kv.raw(), - s.compressed_score.raw(), - state.state_kv.raw(), - state.state_score.raw(), - map, - size, - weights.kv.offset, - weights.gate.offset, - weights.ape.offset, - weights.ape.kind, - shape.embd, - width, - s.norm.raw(), - state.ratio, - pos, - ) + let fused = if projected { + 1 + } else { + unsafe { + ds4_gpu_matmul_f16_pair_compressor_store_tensor( + s.compressed_kv.raw(), + s.compressed_score.raw(), + state.state_kv.raw(), + state.state_score.raw(), + map, + size, + weights.kv.offset, + weights.gate.offset, + weights.ape.offset, + weights.ape.kind, + shape.embd, + width, + s.norm.raw(), + state.ratio, + pos, + ) + } }; if fused < 0 { return Err("Metal failed while storing compressor projections".into()); @@ -6828,6 +7608,8 @@ fn update_compressor_stage( shape.rope_beta_slow, shape.rms_epsilon, fused > 0, + true, + false, ) }, "compressor update", @@ -7446,6 +8228,17 @@ mod tests { assert_eq!(super::half_to_f32(0x0001), 2.0_f32.powi(-24)); } + #[test] + fn decode_feature_policy_keeps_device_generations_and_rollbacks_separate() { + use super::decode_feature_allowed; + + assert!(decode_feature_allowed(true, false, false, false)); + assert!(!decode_feature_allowed(true, false, true, false)); + assert!(decode_feature_allowed(false, true, false, false)); + assert!(!decode_feature_allowed(false, true, false, true)); + assert!(!decode_feature_allowed(false, false, false, false)); + } + #[test] fn streaming_memory_plan_matches_ds4_graph_formulas() { assert_eq!(effective_prefill_cap(16_384, 0), 4_096); @@ -7487,6 +8280,226 @@ mod tests { } } + #[test] + #[ignore = "requires the installed 0731 Flash GGUF and an Apple M5 device"] + fn flash_0731_m5_decode_performance_gate() { + use super::{DeepSeekExecutor, Digest, Sha256, configure_sources}; + use crate::engine::{ChatTurn, Model}; + use crate::model::ModelChoice; + use crate::settings::{ + EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode, + }; + use std::time::Instant; + + configure_sources().unwrap(); + let path = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model; + let run = || { + let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash0731).unwrap(); + let prompt = model.render_conversation( + "", + &[ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "Count from one to two hundred, spelling out every number.".into(), + }], + ReasoningMode::Direct, + ); + let eos = model.eos_token(); + let streaming = std::env::var_os("DS4SERVER_BENCH_SSD").is_some(); + let mut executor = DeepSeekExecutor::open( + model, + 4096, + false, + 4096, + 100, + EngineSpeculativeSettings { + mtp_draft_tokens: 1, + mtp_margin: 3.0, + glm_mtp: false, + glm_mtp_timing: false, + dspark: false, + dspark_confidence_threshold: 0.9, + dspark_confidence_threshold_set: false, + dspark_strict: false, + dspark_exact_sampling: false, + }, + EngineSsdSettings { + enabled: streaming, + cold: false, + cache_experts: if streaming { 4_096 } else { 0 }, + cache_bytes: 0, + full_layers: 0, + full_layers_set: false, + preload_experts: if streaming { 4_096 } else { 0 }, + }, + EngineSteeringSettings { + file: None, + ffn_scale: 0.0, + attention_scale: 0.0, + }, + ) + .unwrap(); + let prefill_started = Instant::now(); + executor.prefill(&prompt, |_| true).unwrap(); + let prefill_seconds = prefill_started.elapsed().as_secs_f64(); + let mut generated = Vec::new(); + let mut latencies_ms = Vec::new(); + let measured = 128_u32; + let started = Instant::now(); + for _ in 0..measured { + let token = executor + .logits() + .iter() + .enumerate() + .filter(|(token, _)| *token as i32 != eos) + .max_by(|left, right| left.1.total_cmp(right.1)) + .map_or(-1, |(token, _)| token as i32); + generated.push(token); + let token_started = Instant::now(); + executor.eval(token).unwrap(); + latencies_ms.push(token_started.elapsed().as_secs_f64() * 1_000.0); + } + let seconds = started.elapsed().as_secs_f64(); + let tokens_per_second = f64::from(measured) / seconds; + let steady_seconds = latencies_ms[1..].iter().sum::() / 1_000.0; + let steady_tokens_per_second = f64::from(measured - 1) / steady_seconds; + let mut sorted = latencies_ms[1..].to_vec(); + sorted.sort_by(f64::total_cmp); + let p50 = sorted[sorted.len() / 2]; + let p95 = sorted[(sorted.len() * 95).div_ceil(100) - 1]; + let token_bytes = generated + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(); + let token_hash = Sha256::digest(token_bytes); + let token_hash = token_hash + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let stats = executor.execution_stats(); + eprintln!( + "DS4SERVER_METAL_PERF mode={} model=flash-0731 context=4096 prompt={} prefill_tps={:.6} measured={measured} seconds={seconds:.6} tps={tokens_per_second:.6} first_ms={:.6} steady_tps={steady_tokens_per_second:.6} p50_ms={p50:.6} p95_ms={p95:.6} cache_entries={} cache_hits={} cache_misses={} pread_bytes={}", + if streaming { "ssd" } else { "resident" }, + prompt.len(), + prompt.len() as f64 / prefill_seconds, + latencies_ms[0], + stats.ssd_cache_entries, + stats.ssd_cache_hits, + stats.ssd_cache_misses, + stats.ssd_pread_bytes, + ); + eprintln!("DS4SERVER_METAL_TOKEN_SHA256 {token_hash}"); + if std::env::var_os("DS4SERVER_BENCH_TOKENS").is_some() { + eprintln!("DS4SERVER_METAL_PROMPT_TOKENS {prompt:?}"); + eprintln!("DS4SERVER_METAL_GENERATED_TOKENS {generated:?}"); + } + assert!(steady_tokens_per_second.is_finite() && steady_tokens_per_second > 0.0); + steady_tokens_per_second + }; + let reference = std::env::var("DS4_REFERENCE_TPS") + .ok() + .map(|value| value.parse::().unwrap()); + let baseline = std::env::var("DS4SERVER_BASELINE_TPS") + .ok() + .map(|value| value.parse::().unwrap()); + let runs = if reference.is_some() || baseline.is_some() { + 3 + } else { + 1 + }; + let mut results = Vec::with_capacity(runs); + for run_index in 0..runs { + if run_index > 0 { + std::thread::sleep(std::time::Duration::from_secs(5)); + } + results.push(run()); + } + results.sort_by(f64::total_cmp); + let median = results[results.len() / 2]; + eprintln!("DS4SERVER_METAL_PERF_MEDIAN runs={runs} steady_tps={median:.6}"); + if let Some(reference) = reference { + assert!( + median >= reference * 0.95, + "DS4Server median {median:.3} tok/s is more than 5% below DS4 {reference:.3} tok/s" + ); + } + if let Some(baseline) = baseline { + assert!( + median > baseline, + "DS4Server median {median:.3} tok/s did not improve on {baseline:.3} tok/s" + ); + } + } + + #[test] + #[ignore = "requires the installed 0731 Flash GGUF and an Apple M5 device"] + fn flash_0731_long_context_crosses_indexed_prefill_boundary() { + use super::{DeepSeekExecutor, argmax, configure_sources}; + use crate::engine::{ChatTurn, Model}; + use crate::model::ModelChoice; + use crate::settings::{ + EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode, + }; + + configure_sources().unwrap(); + let path = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model; + let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash0731).unwrap(); + let prompt = model.render_conversation( + "", + &[ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "hi ".repeat(4_100), + }], + ReasoningMode::Direct, + ); + assert!(prompt.len() > 4_096 && prompt.len() < 8_192); + let mut executor = DeepSeekExecutor::open( + model, + 8_192, + false, + 4_096, + 100, + EngineSpeculativeSettings { + mtp_draft_tokens: 1, + mtp_margin: 3.0, + glm_mtp: false, + glm_mtp_timing: false, + dspark: false, + dspark_confidence_threshold: 0.9, + dspark_confidence_threshold_set: false, + dspark_strict: false, + dspark_exact_sampling: false, + }, + EngineSsdSettings { + enabled: false, + cold: false, + cache_experts: 0, + cache_bytes: 0, + full_layers: 0, + full_layers_set: false, + preload_experts: 0, + }, + EngineSteeringSettings { + file: None, + ffn_scale: 0.0, + attention_scale: 0.0, + }, + ) + .unwrap(); + assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len()); + executor.eval(argmax(executor.logits())).unwrap(); + assert_eq!(executor.position(), prompt.len() as u32 + 1); + } + #[test] #[ignore = "requires the installed 81 GiB Flash and legacy MTP GGUF fixtures"] fn legacy_mtp_runs_a_target_owned_greedy_cycle() { @@ -7626,7 +8639,7 @@ mod tests { } #[test] - #[ignore = "requires the installed 81 GiB Flash and DSpark GGUF fixtures"] + #[ignore = "requires the installed 0731 Flash and checkpoint-specific DSpark GGUF fixtures"] fn dspark_runs_a_target_owned_greedy_cycle() { use super::{DeepSeekExecutor, argmax, configure_sources}; use crate::engine::Model; @@ -7639,10 +8652,10 @@ mod tests { use std::sync::atomic::AtomicBool; configure_sources().unwrap(); - let artifacts = installed_artifacts(ModelChoice::DeepSeekV4Flash, false, true); + let artifacts = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true); let main_path = artifacts.model; let support_path = artifacts.mtp.unwrap(); - let mut model = Model::open_main(&main_path, ModelChoice::DeepSeekV4Flash).unwrap(); + let mut model = Model::open_main(&main_path, ModelChoice::DeepSeekV4Flash0731).unwrap(); let support = Gguf::open(&support_path).unwrap(); model.support_kind = Some(validate_support(&support, &model.shape).unwrap()); model.support = Some(support); @@ -7706,7 +8719,6 @@ mod tests { .unwrap(); generated.extend(cycle); } - assert_eq!(generated, [19_923, 3, 1_730, 588, 342, 1_694, 440, 4_316]); let dspark = executor.dspark.as_ref().unwrap(); assert!(dspark.drafted > 0); assert!(dspark.accepted > 0); @@ -7728,7 +8740,7 @@ mod tests { assert_eq!(cycle.len(), 1); target_only.extend(cycle); } - assert_eq!(target_only, [19_923, 3, 1_730, 588, 342, 1_694, 440, 4_316]); + assert_eq!(target_only, generated); let mut turns = Vec::new(); for _ in 0..4 { @@ -7951,7 +8963,7 @@ mod tests { #[test] #[ignore = "requires the installed 81 GiB Flash GGUF fixture and a Metal device"] - fn flash_resident_and_ssd_streaming_choose_the_same_tokens() { + fn flash_0731_resident_and_ssd_streaming_choose_the_same_tokens() { use super::{DeepSeekExecutor, argmax, configure_sources}; use crate::engine::Model; use crate::model::ModelChoice; @@ -7960,8 +8972,47 @@ mod tests { }; use std::path::Path; + struct FusionRollbacks; + + impl FusionRollbacks { + const NAMES: &'static [&'static str] = &[ + "DS4_METAL_DISABLE_PRE_M5_HC_NORM_MIX_FUSE", + "DS4_METAL_DISABLE_PRE_M5_HC_PRODUCER_PRE_NORM_FUSE", + "DS4_METAL_DISABLE_M5_HC_PRODUCER_PRE_NORM_FUSE", + "DS4_METAL_DISABLE_PRE_M5_QKV_PAIR_QUAD_FUSE", + "DS4_METAL_DISABLE_M5_QKV_PAIR_QUAD_FUSE", + "DS4_METAL_DISABLE_PRE_M5_QKV_PAIR_COMPRESSOR_FUSE", + "DS4_METAL_DISABLE_M5_QKV_PAIR_COMPRESSOR_FUSE", + "DS4_METAL_DISABLE_PRE_M5_QKV_NORM_KV_STORE_FUSE", + "DS4_METAL_DISABLE_PRE_M5_ATTN_INV_ROPE_FUSE", + "DS4_METAL_DISABLE_PRE_M5_PARALLEL_FULL_FFN", + "DS4_METAL_DISABLE_M5_PARALLEL_FULL_FFN", + "DS4_METAL_DISABLE_PRE_M5_ROUTER_SHARED_FUSE", + "DS4_METAL_DISABLE_M5_ROUTER_SHARED_FUSE", + "DS4_METAL_DISABLE_M5_ROUTER_PROJECT_SELECT_FUSE", + "DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_QUAD_STORE", + "DS4_METAL_DISABLE_PRE_M5_COMP_FINALIZE_FUSE", + "DS4_METAL_DISABLE_M5_COMP_FINALIZE_FUSE", + ]; + + fn activate() -> Self { + for name in Self::NAMES { + unsafe { std::env::set_var(name, "1") }; + } + Self + } + } + + impl Drop for FusionRollbacks { + fn drop(&mut self) { + for name in Self::NAMES { + unsafe { std::env::remove_var(name) }; + } + } + } + fn run(path: &Path, streaming: bool) -> Vec { - let model = Model::open_main(path, ModelChoice::DeepSeekV4Flash).unwrap(); + let model = Model::open_main(path, ModelChoice::DeepSeekV4Flash0731).unwrap(); let prompt = model.render_conversation( "", &[crate::engine::ChatTurn { @@ -8021,15 +9072,25 @@ mod tests { } configure_sources().unwrap(); - let path = installed_artifacts(ModelChoice::DeepSeekV4Flash, false, false).model; + let path = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model; let resident = run(&path, false); - assert_eq!(resident, [19_923, 3, 1_730, 588]); - assert_eq!(resident, run(&path, true)); + let resident_rolled_back = { + let _rollbacks = FusionRollbacks::activate(); + run(&path, false) + }; + assert_eq!(resident, resident_rolled_back); + let streaming = run(&path, true); + let streaming_rolled_back = { + let _rollbacks = FusionRollbacks::activate(); + run(&path, true) + }; + assert_eq!(streaming, streaming_rolled_back); + assert_eq!(resident, streaming); } #[test] #[ignore = "requires the installed 81 GiB Flash GGUF fixture and a Metal device"] - fn flash_ssd_streaming_maps_batched_prefill_layers() { + fn flash_0731_ssd_streaming_maps_and_seeds_batched_prefill_layers() { use super::{DeepSeekExecutor, argmax, configure_sources}; use crate::engine::Model; use crate::model::ModelChoice; @@ -8038,8 +9099,8 @@ mod tests { }; configure_sources().unwrap(); - let path = installed_artifacts(ModelChoice::DeepSeekV4Flash, false, false).model; - let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap(); + let path = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model; + let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash0731).unwrap(); let prompt = model.render_conversation( "", &[crate::engine::ChatTurn { @@ -8073,12 +9134,12 @@ mod tests { }, EngineSsdSettings { enabled: true, - cold: true, + cold: false, cache_experts: 16, cache_bytes: 0, full_layers: 0, full_layers_set: false, - preload_experts: 0, + preload_experts: 16, }, EngineSteeringSettings { file: None, @@ -8087,7 +9148,37 @@ mod tests { }, ) .unwrap(); + let fallback_layer = executor + .ssd + .as_ref() + .unwrap() + .preload_by_layer + .iter() + .position(|entries| !entries.is_empty()) + .unwrap(); + let ssd = executor.ssd.as_ref().unwrap(); + assert!( + !ssd.seed_mapped_layer( + &executor.model, + &executor.weights.layers[fallback_layer], + fallback_layer, + true, + ) + .unwrap() + ); + assert!( + ssd.seed_mapped_layer( + &executor.model, + &executor.weights.layers[fallback_layer], + fallback_layer, + false, + ) + .unwrap() + ); assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len()); + let stats = executor.execution_stats(); + assert_eq!(stats.ssd_preloaded_experts, 16); + assert!(stats.ssd_cache_entries > 0); executor.eval(argmax(executor.logits())).unwrap(); } diff --git a/src/engine/metal/glm.rs b/src/engine/metal/glm.rs index a5f83b2..b1fd27b 100644 --- a/src/engine/metal/glm.rs +++ b/src/engine/metal/glm.rs @@ -8,6 +8,11 @@ const CACHE_F16: bool = true; const DECODE_FLUSH_LAYERS: usize = 4; const AUTO_CACHE_BYTES: u64 = 12 * 1024 * 1024 * 1024; const STREAMING_TOKEN_PREFILL_MAX: u32 = 64; + +fn live_prefix_rewind_target(live: &[i32], incoming: &[i32]) -> Option { + (incoming.len() > 1 && incoming.len() < live.len() && live.starts_with(incoming)) + .then_some(incoming.len() - 1) +} const STREAMING_FULL_ATTN_CONTEXT: u32 = 8192; const LONG_CONTEXT_THRESHOLD: u32 = 65_536; const LONG_CONTEXT_FULL_ATTN_CONTEXT: u32 = 4096; @@ -2365,6 +2370,14 @@ impl GlmExecutor { } pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result { + if let Some(rewind) = live_prefix_rewind_target(&self.tokens, tokens) { + self.tokens.truncate(rewind); + if let Some(mtp) = &mut self.mtp { + mtp.pending = None; + mtp.min_pos = None; + } + return Ok(rewind); + } if !tokens.starts_with(&self.tokens) { self.reset()?; } @@ -3217,7 +3230,7 @@ fn f32_project_rows( mod tests { use super::{ GlmExecutor, argmax, dynamic_expert_budget, full_indexer_layer, indexed_prefill_rows, - streaming_token_prefill_eligible, + live_prefix_rewind_target, streaming_token_prefill_eligible, }; use crate::engine::{GLM, Model, ReasoningMode}; use crate::model::ModelChoice; @@ -3250,6 +3263,15 @@ mod tests { assert_eq!(indexed_prefill_rows(0, 17, 2048), 17); } + #[test] + fn repeated_glm_prompt_rewinds_one_token_for_logits() { + let live = [1, 2, 3, 4, 5, 6]; + assert_eq!(live_prefix_rewind_target(&live, &[1, 2, 3, 4]), Some(3)); + assert_eq!(live_prefix_rewind_target(&live, &[1]), None); + assert_eq!(live_prefix_rewind_target(&live, &live), None); + assert_eq!(live_prefix_rewind_target(&live, &[1, 2, 9]), None); + } + #[test] fn glm_byte_budget_reserves_prefill_before_dynamic_experts() { let mib = 1024 * 1024; diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index e0c8459..0b4da0f 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -67,6 +67,12 @@ unsafe extern "C" { expert_priorities: *const u32, experts: u32, ) -> i32; + pub(super) fn ds4_gpu_stream_expert_cache_seed_experts_gpu_copy( + table: *const StreamExpertTable, + expert_ids: *const i32, + expert_priorities: *const u32, + experts: u32, + ) -> i32; pub(super) fn ds4_gpu_stream_expert_cache_begin_selected_load( table: *const StreamExpertTable, selected_ids: *const i32, @@ -87,6 +93,26 @@ unsafe extern "C" { count: u32, ) -> i32; pub(super) fn ds4_gpu_flush_commands() -> i32; + pub(super) fn ds4_gpu_device_is_pre_m5_apple_silicon() -> i32; + pub(super) fn ds4_gpu_device_is_m5_apple_silicon() -> i32; + pub(super) fn ds4_gpu_set_decode_pipeline_fast_lookup(enabled: i32) -> i32; + pub(super) fn ds4_gpu_parallel_ffn_start( + gate: *mut GpuTensor, + up: *mut GpuTensor, + mid: *mut GpuTensor, + shared_out: *mut GpuTensor, + map: *const c_void, + size: u64, + gate_offset: u64, + up_offset: u64, + down_offset: u64, + model_dim: u32, + shared_dim: u32, + x: *const GpuTensor, + clamp: f32, + ) -> i32; + pub(super) fn ds4_gpu_parallel_ffn_finish() -> i32; + pub(super) fn ds4_gpu_parallel_ffn_abort(); pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor; pub(super) fn ds4_gpu_tensor_view( base: *const GpuTensor, @@ -664,6 +690,64 @@ unsafe extern "C" { ratio: u32, pos: u32, ) -> i32; + pub(super) fn ds4_gpu_matmul_f16_quad_compressor_store_tensor( + out0_kv: *mut GpuTensor, + out0_score: *mut GpuTensor, + out1_kv: *mut GpuTensor, + out1_score: *mut GpuTensor, + state0_kv: *mut GpuTensor, + state0_score: *mut GpuTensor, + state1_kv: *mut GpuTensor, + state1_score: *mut GpuTensor, + map: *const c_void, + size: u64, + weight0_kv: u64, + weight0_score: u64, + weight1_kv: u64, + weight1_score: u64, + ape0: u64, + ape0_type: u32, + ape1: u64, + ape1_type: u32, + input: u64, + width0: u32, + width1: u32, + x: *const GpuTensor, + ratio: u32, + pos: u32, + ) -> i32; + pub(super) fn ds4_gpu_qkv_pair_quad_compressor_store_tensor( + q_rank: *mut GpuTensor, + kv_raw: *mut GpuTensor, + out0_kv: *mut GpuTensor, + out0_score: *mut GpuTensor, + out1_kv: *mut GpuTensor, + out1_score: *mut GpuTensor, + state0_kv: *mut GpuTensor, + state0_score: *mut GpuTensor, + state1_kv: *mut GpuTensor, + state1_score: *mut GpuTensor, + map: *const c_void, + size: u64, + q_a: u64, + kv: u64, + weight0_kv: u64, + weight0_score: u64, + weight1_kv: u64, + weight1_score: u64, + ape0: u64, + ape0_type: u32, + ape1: u64, + ape1_type: u32, + input: u32, + q_rank_width: u32, + kv_width: u32, + width0: u32, + width1: u32, + x: *const GpuTensor, + ratio: u32, + pos: u32, + ) -> i32; pub(super) fn ds4_gpu_hc_split_weighted_sum_norm_tensor( out: *mut GpuTensor, norm: *mut GpuTensor, @@ -681,6 +765,38 @@ unsafe extern "C" { eps: f32, norm_eps: f32, ) -> i32; + pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_available() -> i32; + pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_tensor( + out: *mut GpuTensor, + x: *const GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + input: u32, + output: u32, + eps: f32, + ) -> i32; + pub(super) fn ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( + mix: *mut GpuTensor, + out: *mut GpuTensor, + norm: *mut GpuTensor, + split: *mut GpuTensor, + residual: *const GpuTensor, + map: *const c_void, + size: u64, + mix_weight: u64, + scale: u64, + base: u64, + norm_weight: u64, + input: u32, + mix_width: u32, + embd: u32, + hc: u32, + iterations: u32, + eps: f32, + hc_eps: f32, + norm_eps: f32, + ) -> i32; pub(super) fn ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( q_out: *mut GpuTensor, q: *const GpuTensor, @@ -695,6 +811,47 @@ unsafe extern "C" { rows: u32, eps: f32, ) -> i32; + pub(super) fn ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_tensor( + q_out: *mut GpuTensor, + q: *const GpuTensor, + map: *const c_void, + size: u64, + q_weight: u64, + q_width: u32, + kv_out: *mut GpuTensor, + kv: *const GpuTensor, + kv_weight: u64, + kv_width: u32, + raw_cache: *mut GpuTensor, + raw_cap: u64, + raw_row: u32, + rot: u32, + pos: u32, + original_context: u32, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + eps: f32, + ) -> i32; + pub(super) fn ds4_gpu_kv_rope_fp8_fuse_available() -> i32; + pub(super) fn ds4_gpu_decode_attn_rope_fuse_available() -> i32; + pub(super) fn ds4_gpu_decode_attn_rope_fuse_used() -> i32; + pub(super) fn ds4_gpu_set_decode_attn_rope_fuse( + head_dim: u32, + n_rot: u32, + pos0: u32, + n_ctx_orig: u32, + inverse: bool, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + ); pub(super) fn ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( out: *mut GpuTensor, half: *mut GpuTensor, @@ -898,6 +1055,33 @@ unsafe extern "C" { beta_slow: f32, rms_eps: f32, state_already_stored: bool, + decode_one_token: bool, + defer_finalize: bool, + ) -> i32; + pub(super) fn ds4_gpu_dsv4_comp_row_finalize_tensor( + attn_stage: *mut GpuTensor, + attn_cache: *mut GpuTensor, + attn_row: u32, + attn_norm: u64, + index_cache: *mut GpuTensor, + index_row: u32, + index_norm: u64, + attn_state_kv: *mut GpuTensor, + attn_state_score: *mut GpuTensor, + index_state_kv: *mut GpuTensor, + index_state_score: *mut GpuTensor, + map: *const c_void, + size: u64, + pos: u32, + rot: u32, + original_context: u32, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + rms_eps: f32, ) -> i32; pub(super) fn ds4_gpu_compressor_prefill_state_ratio4_tensor( state_kv: *mut GpuTensor, @@ -1258,6 +1442,35 @@ unsafe extern "C" { x: *const GpuTensor, clamp: f32, ) -> i32; + pub(super) fn ds4_gpu_router_shared_gate_up_q8_0_tensor( + router_logits: *mut GpuTensor, + gate: *mut GpuTensor, + up: *mut GpuTensor, + mid: *mut GpuTensor, + map: *const c_void, + size: u64, + router_weight: u64, + gate_weight: u64, + up_weight: u64, + input: u64, + experts: u64, + shared_width: u64, + x: *const GpuTensor, + clamp: f32, + router_only: bool, + ) -> i32; + pub(super) fn ds4_gpu_router_project_select_fused_tensor( + router_logits: *mut GpuTensor, + probs: *mut GpuTensor, + selected: *mut GpuTensor, + weights: *mut GpuTensor, + map: *const c_void, + size: u64, + router_weight: u64, + bias: u64, + has_bias: bool, + x: *const GpuTensor, + ) -> i32; pub(super) fn ds4_gpu_shared_down_hc_expand_q8_0_tensor( out_hc: *mut GpuTensor, shared_out: *mut GpuTensor, @@ -1308,6 +1521,10 @@ impl Context { unsafe { ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm); ds4_gpu_set_ssd_streaming(ssd_streaming); + // DS4 only enables this cache for the pre-M5 MXFP4 decode path. + // Rust does not accept MXFP4 weights yet, so keep the global + // native switch explicitly disabled until that path is admitted. + ds4_gpu_set_decode_pipeline_fast_lookup(0); } let recommended = unsafe { ds4_gpu_recommended_working_set_size() }; if admission_bytes != 0 && recommended != 0 && admission_bytes > recommended { @@ -1409,6 +1626,62 @@ impl Commands { } } +pub(super) struct ParallelFfn(bool); + +impl ParallelFfn { + #[allow(clippy::too_many_arguments)] + pub(super) fn start( + gate: &Buffer, + up: &Buffer, + mid: &Buffer, + shared_out: &Buffer, + map: *const c_void, + size: u64, + gate_weight: u64, + up_weight: u64, + down_weight: u64, + model_dim: u32, + shared_dim: u32, + x: &Buffer, + clamp: f32, + ) -> Option { + (unsafe { + ds4_gpu_parallel_ffn_start( + gate.raw(), + up.raw(), + mid.raw(), + shared_out.raw(), + map, + size, + gate_weight, + up_weight, + down_weight, + model_dim, + shared_dim, + x.raw(), + clamp, + ) + } != 0) + .then_some(Self(true)) + } + + pub(super) fn finish(mut self) -> Result<(), String> { + self.0 = false; + check( + unsafe { ds4_gpu_parallel_ffn_finish() }, + "joining parallel Metal FFN work", + ) + } +} + +impl Drop for ParallelFfn { + fn drop(&mut self) { + if self.0 { + unsafe { ds4_gpu_parallel_ffn_abort() }; + } + } +} + impl Drop for Commands { fn drop(&mut self) { if self.0 {