Reach DS4 Metal performance parity

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.

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

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<usize> {
+    (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<usize, String> {
+        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;

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<Self> {
+        (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 {