spark_model/layers/
moe_v41.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! DeepSeek-V4.1 Flash MoE on streamed, still-quantized experts.
5//!
6//! The routed experts never leave their Q2_K / Q3_K blocks: per token the
7//! router picks `topk` of 384, the [`ExpertLru`](spark_runtime::weights::expert_stream::ExpertLru) gathers those experts' raw
8//! slices into its device-visible slots (misses read from the SSD), and each
9//! expert runs as three K-quant GEMVs on the raw blocks (`kquant_mmvq_q2_k_w` for
10//! gate and up, `kquant_mmvq_q3_k_w` for down) with the activation quantised to
11//! q8_1 in between. The shared expert every token goes through is resident bf16
12//! and runs on the dense GEMM.
13//!
14//! Routing follows the CPU reference (`deepseek_v41_ref::moe::gate`) exactly:
15//! the logits are an f32-accumulated GEMM of the bf16 input against the bf16
16//! gate weight, then `sqrt(softplus(logit / temp))`, and the top-k is chosen by
17//! `score + correction_bias` while the weights are the unbiased scores,
18//! renormalised and scaled by `route_scale`. The selection runs on the CPU from
19//! the downloaded logits (`[tokens, 384]` f32), which keeps the tie-breaking the
20//! reference's and costs nothing at decode.
21//!
22//! Oracle: `moe_v41_tests.rs`, synthetic Q2_K / Q3_K experts on a pinned arena
23//! against the CPU decoders + q8_1 emulation + the reference's expert math.
24
25use spark_runtime::gpu::{DevicePtr, KernelHandle};
26
27use crate::layers::ops::ResidentMat;
28
29mod device_route;
30mod forward;
31
32/// Layers the device slot table covers (V4.1 has 40; the loader checks).
33pub(crate) const SLOT_TABLE_LAYERS: usize = 64;
34mod init;
35mod route;
36mod shared;
37mod single;
38
39const MODULE: &str = "moe_v41";
40const GEMM_MODULE: &str = "gemm";
41
42#[derive(Clone, Debug)]
43pub struct MoeV41Cfg {
44    pub dim: usize,
45    pub inter: usize,
46    pub n_routed: usize,
47    pub topk: usize,
48    pub gate_temp: f32,
49    pub norm_topk_prob: bool,
50    pub route_scale: f32,
51    pub swiglu_limit: f32,
52    pub max_tokens: usize,
53}
54
55/// One layer's resident MoE weights: the router and the shared expert.
56pub struct MoeV41LayerWeights {
57    pub layer: u32,
58    /// bf16 `[n_routed, dim]`
59    pub gate_w: DevicePtr,
60    /// f32 `[n_routed]`, host: the selection runs on the CPU
61    pub gate_bias: Vec<f32>,
62    /// The same bias on the device, for the device-side selection.
63    pub gate_bias_dev: DevicePtr,
64    /// `[inter, dim]`, `[dim, inter]`, `[inter, dim]`: bf16 or the GGUF's
65    /// Q2_K (w1, w3) / Q3_K (w2) blocks on the routed experts' kernels
66    pub shared_w1: ResidentMat,
67    pub shared_w2: ResidentMat,
68    pub shared_w3: ResidentMat,
69}
70
71/// The router of one layer alone: what predicting a layer's selection from
72/// an earlier layer's input needs (see `MoeV41::predict_launch`).
73pub struct RouterWeights {
74    pub layer: u32,
75    /// bf16 `[n_routed, dim]`
76    pub gate_w: DevicePtr,
77    pub gate_bias: Vec<f32>,
78}
79
80/// `ATLAS_DS41_PREFETCH_K` (default 0 = off), read once: how many of the
81/// next layer's predicted experts to start reading while this layer computes
82/// (needs `ATLAS_DS41_READER_POOL=1`). Off by default: on the 09-19 standard
83/// the predictor's router launch (2.7 ms a token) and the wasted reads on
84/// the shared disk cost more than the caught misses saved (K=4: 12.50/16.67,
85/// K=6: 13.63/16.39, against 14.20/17.73 without).
86pub fn prefetch_k() -> usize {
87    static K: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
88    *K.get_or_init(|| {
89        std::env::var("ATLAS_DS41_PREFETCH_K")
90            .ok()
91            .and_then(|v| v.parse().ok())
92            .unwrap_or(0)
93    })
94}
95
96struct Kernels {
97    gemm: KernelHandle,
98    /// bf16 shared expert at m = 1 (the tiled GEMM idles 15 of 16 rows)
99    gemv: KernelHandle,
100    gemm_f32out: KernelHandle,
101    /// router logits at m <= 8: strict-order GEMV, bit-identical to gemm_f32out
102    router_gemv: KernelHandle,
103    /// the same logits from the products staged in shared memory, N blocks (decode)
104    router_gemv_staged: KernelHandle,
105    q8_rows: KernelHandle,
106    mmvq_q2k: KernelHandle,
107    mmvq_q3k: KernelHandle,
108    /// the single-token arm: every routed expert in one launch a projection
109    mmvq_q2k_experts: KernelHandle,
110    mmvq_q3k_experts: KernelHandle,
111    /// rows (warps) a block of the expert batch, 2 / 4 / 8 (ATLAS_DS41_EXPERT_WARPS)
112    experts_warps: u32,
113    swiglu: KernelHandle,
114    /// the swiglu and the q8_1 rows of its output in one launch (decode)
115    swiglu_q8: KernelHandle,
116    accumulate: KernelHandle,
117    finish: KernelHandle,
118    gather: KernelHandle,
119    scatter_add: KernelHandle,
120    sum_rows: KernelHandle,
121    quant_d2s6: KernelHandle,
122    quant_d4: KernelHandle,
123    mmq_q2k_nc: KernelHandle,
124    mmq_q2k_wc: KernelHandle,
125    mmq_q3k_nc: KernelHandle,
126    mmq_q3k_wc: KernelHandle,
127    route_select: KernelHandle,
128    slot_table_set: KernelHandle,
129}
130
131/// Where one call's time went (wall clock, host side).
132#[derive(Clone, Copy, Debug, Default)]
133pub struct MoeV41Timing {
134    pub route_ms: f64,
135    pub fetch_ms: f64,
136    pub compute_ms: f64,
137    pub hits: u64,
138    pub misses: u64,
139    pub bytes_read: u64,
140}
141
142impl MoeV41Timing {
143    pub fn add(&mut self, o: &MoeV41Timing) {
144        self.route_ms += o.route_ms;
145        self.fetch_ms += o.fetch_ms;
146        self.compute_ms += o.compute_ms;
147        self.hits += o.hits;
148        self.misses += o.misses;
149        self.bytes_read += o.bytes_read;
150    }
151}
152
153/// What `stage_m1` leaves for `compute_m1`: the distinct expert count and
154/// the routing (for callers and diagnostics), with the host-span timing.
155pub struct MoeV41Stage {
156    pub ne: usize,
157    pub weights: Vec<f32>,
158    pub indices: Vec<usize>,
159    pub timing: MoeV41Timing,
160}
161
162pub struct MoeV41 {
163    /// The last forward's timing.
164    pub last: std::cell::Cell<MoeV41Timing>,
165    /// sync at the end of the forward so the step timer reads GPU time (diag only)
166    timing_sync: bool,
167    pub cfg: MoeV41Cfg,
168    k: Kernels,
169    logits: DevicePtr,
170    /// the next layer's router on this layer's input (`[n_routed]` f32)
171    pred_logits: DevicePtr,
172    /// gathered rows of one expert group, `[m, dim]` bf16
173    a_rows: DevicePtr,
174    /// the group's q8_1 activations (plain rows or the MMQ layout)
175    a_q8: DevicePtr,
176    gate_out: DevicePtr,
177    up_out: DevicePtr,
178    h: DevicePtr,
179    h_q8: DevicePtr,
180    down_out: DevicePtr,
181    /// `[m * topk]` i32 token rows and f32 routing weights, group-major
182    rows_dev: DevicePtr,
183    weight_dev: DevicePtr,
184    /// The device selection's header (miss flag, picks, weight bits, plan
185    /// slots) and the `[64 layers][n_routed]` slot table (-1 = not resident).
186    route_hdr: DevicePtr,
187    slot_table: DevicePtr,
188    /// gate / up / down block pointers of the token's experts, `3 * topk`
189    ptrs_dev: DevicePtr,
190    sg: DevicePtr,
191    su: DevicePtr,
192    sh: DevicePtr,
193    sd: DevicePtr,
194    acc: DevicePtr,
195    out: DevicePtr,
196    /// The shared expert's own q8_1 scratch for its SwiGLU output (`[1,
197    /// inter]`) on the side stream, so it never shares `h_q8` with the routed
198    /// experts running at the same time; the token's input rows (`a_q8`) are
199    /// written once on the forking stream and read by both.
200    sh_q8: DevicePtr,
201    /// The side stream the single-token shared expert runs on, and the two
202    /// events that fence it: `ev_in` (the input is ready, recorded on the
203    /// main stream) and `ev_out` (`sd` is ready, recorded on the side).
204    side: u64,
205    ev_in: u64,
206    ev_out: u64,
207}
208
209/// `ATLAS_DS41_SHARED_SIDE` (default on; `0` = off), read once: at decode the
210/// shared expert runs on a side stream while the main stream does the router
211/// and the host picks the experts (the GPU is otherwise idle for that span);
212/// its output is added into `acc` after the routed experts, as before, so the
213/// accumulation order and the bits are unchanged.
214pub fn shared_side() -> bool {
215    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
216    *V.get_or_init(|| !std::env::var("ATLAS_DS41_SHARED_SIDE").is_ok_and(|v| v == "0"))
217}
218
219fn softplus(x: f32) -> f32 {
220    if x > 20.0 { x } else { x.exp().ln_1p() }
221}
222
223/// The reference's `Gate.forward` on f32 logits: returns
224/// (`weights[tokens, topk]`, `indices[tokens, topk]`) in torch's top-k order.
225pub fn route_from_logits(
226    logits: &[f32],
227    tokens: usize,
228    bias: &[f32],
229    c: &MoeV41Cfg,
230) -> (Vec<f32>, Vec<usize>) {
231    let n = c.n_routed;
232    let mut weights = Vec::with_capacity(tokens * c.topk);
233    let mut indices = Vec::with_capacity(tokens * c.topk);
234    for t in 0..tokens {
235        let scores: Vec<f32> = (0..n)
236            .map(|e| softplus(logits[t * n + e] / c.gate_temp).sqrt())
237            .collect();
238        let mut order: Vec<usize> = (0..n).collect();
239        order.sort_by(|&a, &b| {
240            (scores[b] + bias[b])
241                .partial_cmp(&(scores[a] + bias[a]))
242                .expect("finite scores")
243        });
244        let picked = &order[..c.topk];
245        let mut wt: Vec<f32> = picked.iter().map(|&e| scores[e]).collect();
246        if c.norm_topk_prob && c.topk > 1 {
247            let sum: f32 = wt.iter().sum::<f32>() + 1e-20;
248            for v in &mut wt {
249                *v /= sum;
250            }
251        }
252        for v in &mut wt {
253            *v *= c.route_scale;
254        }
255        weights.extend(wt);
256        indices.extend_from_slice(picked);
257    }
258    (weights, indices)
259}
260
261/// A `[tokens, dim]` bf16 device buffer's bytes, for callers staging inputs.
262pub fn bf16_bytes(v: &[f32]) -> Vec<u8> {
263    v.iter()
264        .flat_map(|&x| {
265            let b = x.to_bits();
266            let lsb = (b >> 16) & 1;
267            ((b.wrapping_add(0x7FFF + lsb) >> 16) as u16).to_le_bytes()
268        })
269        .collect()
270}
271
272#[cfg(test)]
273#[path = "moe_v41_route_tests.rs"]
274mod route_tests;
275#[cfg(test)]
276#[path = "moe_v41_tests.rs"]
277mod tests;
spark_model/layers/moe_v41.rs

spark_model/layers/
moe_v41.rs
