metrale_model_arch/
moe_v41.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: DeepSeek-V4.1 Flash MoE on streamed, still-quantized experts.
5//!
6//! The routed experts stay in their GGUF blocks: gate and up run on Q2_K
7//! kernels, down on Q3_K kernels. Per token the router picks `topk` of
8//! `n_routed` experts, the [`ExpertLru`](metrale_model_weights::weights::expert_stream::ExpertLru) gathers those experts' raw
9//! slices into its device-visible slots (misses are read from the
10//! `ExpertSource`), and the kernels read the blocks directly with the
11//! activation quantised to q8_1: GEMVs for groups of up to eight rows, and by
12//! default MMQ above. The shared expert is resident, bf16 or K-quant (`ResidentMat`).
13//!
14//! Routing matches `deepseek_v41_ref::moe::gate`: f32-accumulated router
15//! logits, scores `sqrt(softplus(logit / gate_temp))`, the top-k chosen by
16//! `score + gate_bias` with a stable sort, and the weights the unbiased scores,
17//! renormalised when `norm_topk_prob` and scaled by `route_scale`. By default
18//! the selection runs on the CPU from the downloaded logits; `device_route`
19//! holds the device version.
20//!
21//! Tested in `moe_v41_tests.rs` on synthetic Q2_K / Q3_K experts against the
22//! reference `gate` and a CPU emulation of the numerics.
23//!
24//! Owner: model-arch (DeepSeek-V4.1 MoE).
25//! Invariants: none beyond the types.
26
27use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
28
29use metrale_model_layers::layers::ops::ResidentMat;
30
31mod device_route;
32mod forward;
33
34/// 2026-09-25: Layers the device slot table and the route headers cover;
35/// `device_route_ok` refuses the device selection for a layer at or past it.
36pub(crate) const SLOT_TABLE_LAYERS: usize = 64;
37mod init;
38mod route;
39mod shared;
40mod single;
41
42const MODULE: &str = "moe_v41";
43const GEMM_MODULE: &str = "gemm";
44
45#[derive(Clone, Debug)]
46pub struct MoeV41Cfg {
47    pub dim: usize,
48    pub inter: usize,
49    pub n_routed: usize,
50    pub topk: usize,
51    pub gate_temp: f32,
52    pub norm_topk_prob: bool,
53    pub route_scale: f32,
54    pub swiglu_limit: f32,
55    pub max_tokens: usize,
56}
57
58/// 2026-09-25: One layer's resident MoE weights: the router and the shared expert.
59pub struct MoeV41LayerWeights {
60    pub layer: u32,
61    /// 2026-09-25: bf16 `[n_routed, dim]`.
62    pub gate_w: DevicePtr,
63    /// 2026-09-25: f32 `[n_routed]` on the host, for the CPU selection.
64    pub gate_bias: Vec<f32>,
65    /// 2026-09-25: The same bias on the device, for the device-side selection.
66    pub gate_bias_dev: DevicePtr,
67    /// 2026-09-25: w1, w2, w3 as `[inter, dim]`, `[dim, inter]`, `[inter, dim]`:
68    /// bf16, or the GGUF's K-quant blocks run on the routed experts' kernels.
69    pub shared_w1: ResidentMat,
70    pub shared_w2: ResidentMat,
71    pub shared_w3: ResidentMat,
72}
73
74/// 2026-09-25: The router of one layer alone: what predicting a layer's
75/// selection from an earlier layer's input needs (see `MoeV41::predict_launch`).
76pub struct RouterWeights {
77    pub layer: u32,
78    /// 2026-09-25: bf16 `[n_routed, dim]`.
79    pub gate_w: DevicePtr,
80    pub gate_bias: Vec<f32>,
81}
82
83/// 2026-09-25: `METRALE_DS41_PREFETCH_K`, read once (default 0, off): how
84/// many of the next layer's predicted experts to start reading while this
85/// layer computes. It acts only at one token, with the host selection, and
86/// with the reader pool (`METRALE_DS41_READER_POOL=1`).
87pub fn prefetch_k() -> usize {
88    static K: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
89    *K.get_or_init(|| {
90        std::env::var("METRALE_DS41_PREFETCH_K")
91            .ok()
92            .and_then(|v| v.parse().ok())
93            .unwrap_or(0)
94    })
95}
96
97struct Kernels {
98    gemm: KernelHandle,
99    /// 2026-09-25: The bf16 shared expert at one token.
100    gemv: KernelHandle,
101    gemm_f32out: KernelHandle,
102    /// 2026-09-25: Router logits at up to eight tokens with
103    /// `METRALE_DS41_ROUTER_STAGED=0`, and for the prediction.
104    router_gemv: KernelHandle,
105    /// 2026-09-25: Router logits at up to eight tokens by default: one block
106    /// per expert row, the products staged in shared memory.
107    router_gemv_staged: KernelHandle,
108    q8_rows: KernelHandle,
109    mmvq_q2k: KernelHandle,
110    mmvq_q3k: KernelHandle,
111    /// 2026-09-25: The single-token arm: all selected experts in one launch
112    /// per projection.
113    mmvq_q2k_experts: KernelHandle,
114    mmvq_q3k_experts: KernelHandle,
115    /// 2026-09-25: Warps per block of the expert batch: 2, 4 or 8
116    /// (`METRALE_DS41_EXPERT_WARPS`, default 8).
117    experts_warps: u32,
118    swiglu: KernelHandle,
119    /// 2026-09-25: The SwiGLU and the q8_1 rows of its output in one launch
120    /// (up to eight rows).
121    swiglu_q8: KernelHandle,
122    accumulate: KernelHandle,
123    finish: KernelHandle,
124    gather: KernelHandle,
125    scatter_add: KernelHandle,
126    sum_rows: KernelHandle,
127    quant_d2s6: KernelHandle,
128    quant_d4: KernelHandle,
129    mmq_q2k_nc: KernelHandle,
130    mmq_q2k_wc: KernelHandle,
131    mmq_q3k_nc: KernelHandle,
132    mmq_q3k_wc: KernelHandle,
133    route_select: KernelHandle,
134    slot_table_set: KernelHandle,
135}
136
137/// 2026-09-25: Where one call's time went (wall clock, host side).
138#[derive(Clone, Copy, Debug, Default)]
139pub struct MoeV41Timing {
140    pub route_ms: f64,
141    pub fetch_ms: f64,
142    pub compute_ms: f64,
143    pub hits: u64,
144    pub misses: u64,
145    pub bytes_read: u64,
146}
147
148impl MoeV41Timing {
149    pub fn add(&mut self, o: &MoeV41Timing) {
150        self.route_ms += o.route_ms;
151        self.fetch_ms += o.fetch_ms;
152        self.compute_ms += o.compute_ms;
153        self.hits += o.hits;
154        self.misses += o.misses;
155        self.bytes_read += o.bytes_read;
156    }
157}
158
159/// 2026-09-25: What `stage_m1` leaves for `compute_m1`: the distinct expert
160/// count and the routing (for callers and diagnostics), with the host-span
161/// timing.
162pub struct MoeV41Stage {
163    pub ne: usize,
164    pub weights: Vec<f32>,
165    pub indices: Vec<usize>,
166    pub timing: MoeV41Timing,
167}
168
169pub struct MoeV41 {
170    /// 2026-09-25: The last forward's timing.
171    pub last: std::cell::Cell<MoeV41Timing>,
172    /// 2026-09-25: `METRALE_DS41_DIAG=1`: synchronise at the end of `forward`
173    /// so `compute_ms` is GPU time.
174    timing_sync: bool,
175    pub cfg: MoeV41Cfg,
176    k: Kernels,
177    logits: DevicePtr,
178    /// 2026-09-25: The next layer's router on this layer's input
179    /// (`[n_routed]` f32).
180    pred_logits: DevicePtr,
181    /// 2026-09-25: Gathered rows of one expert group, `[m, dim]` bf16; also
182    /// the staging buffer of `slot_table_update`.
183    a_rows: DevicePtr,
184    /// 2026-09-25: The group's q8_1 activations (plain rows or the MMQ layout).
185    a_q8: DevicePtr,
186    gate_out: DevicePtr,
187    up_out: DevicePtr,
188    h: DevicePtr,
189    h_q8: DevicePtr,
190    down_out: DevicePtr,
191    /// 2026-09-25: `[m * topk]` i32 token rows and f32 routing weights,
192    /// grouped by ascending expert id.
193    rows_dev: DevicePtr,
194    weight_dev: DevicePtr,
195    /// 2026-09-25: The device selection's headers (`header_words` i32 words
196    /// per layer: miss flag, picks, weight bits, plan slots) and the
197    /// `[SLOT_TABLE_LAYERS][n_routed]` i32 slot table (-1 = not resident).
198    route_hdr: DevicePtr,
199    slot_table: DevicePtr,
200    /// 2026-09-25: Gate, up and down block pointers of the token's experts,
201    /// `3 * topk` u64.
202    ptrs_dev: DevicePtr,
203    sg: DevicePtr,
204    su: DevicePtr,
205    sh: DevicePtr,
206    sd: DevicePtr,
207    acc: DevicePtr,
208    out: DevicePtr,
209    /// 2026-09-25: The shared expert's own q8_1 scratch for its SwiGLU output
210    /// (`[1, inter]`) on the side stream, so it never shares `h_q8` with the routed
211    /// experts running at the same time; the token's input rows (`a_q8`) are
212    /// written once on the forking stream and read by both.
213    sh_q8: DevicePtr,
214    /// 2026-09-25: The side stream the single-token shared expert runs on, and the two
215    /// events that fence it: `ev_in` (the input is ready, recorded on the
216    /// main stream) and `ev_out` (`sd` is ready, recorded on the side).
217    side: u64,
218    ev_in: u64,
219    ev_out: u64,
220}
221
222/// 2026-09-25: `METRALE_DS41_SHARED_SIDE`, read once (default on; `0` turns it
223/// off): at one token the shared expert runs on a side stream while the main
224/// stream runs the router and the host picks the experts. Its output is added
225/// into `acc` after the routed experts, the same point as on the main stream,
226/// so the accumulation order does not change.
227pub fn shared_side() -> bool {
228    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
229    *V.get_or_init(|| !std::env::var("METRALE_DS41_SHARED_SIDE").is_ok_and(|v| v == "0"))
230}
231
232fn softplus(x: f32) -> f32 {
233    if x > 20.0 { x } else { x.exp().ln_1p() }
234}
235
236/// 2026-09-25: `deepseek_v41_ref::moe::gate` on f32 logits: returns
237/// (`weights[tokens, topk]`, `indices[tokens, topk]`), each token's picks in
238/// descending `score + bias` order, ties to the lower expert id. Panics on a
239/// NaN key.
240pub fn route_from_logits(
241    logits: &[f32],
242    tokens: usize,
243    bias: &[f32],
244    c: &MoeV41Cfg,
245) -> (Vec<f32>, Vec<usize>) {
246    let n = c.n_routed;
247    let mut weights = Vec::with_capacity(tokens * c.topk);
248    let mut indices = Vec::with_capacity(tokens * c.topk);
249    for t in 0..tokens {
250        let scores: Vec<f32> = (0..n)
251            .map(|e| softplus(logits[t * n + e] / c.gate_temp).sqrt())
252            .collect();
253        let mut order: Vec<usize> = (0..n).collect();
254        order.sort_by(|&a, &b| {
255            (scores[b] + bias[b])
256                .partial_cmp(&(scores[a] + bias[a]))
257                .expect("finite scores")
258        });
259        let picked = &order[..c.topk];
260        let mut wt: Vec<f32> = picked.iter().map(|&e| scores[e]).collect();
261        if c.norm_topk_prob && c.topk > 1 {
262            let sum: f32 = wt.iter().sum::<f32>() + 1e-20;
263            for v in &mut wt {
264                *v /= sum;
265            }
266        }
267        for v in &mut wt {
268            *v *= c.route_scale;
269        }
270        weights.extend(wt);
271        indices.extend_from_slice(picked);
272    }
273    (weights, indices)
274}
275
276/// 2026-09-25: `v` rounded to bf16 (nearest, ties to even) as little-endian
277/// bytes, for callers staging inputs.
278pub fn bf16_bytes(v: &[f32]) -> Vec<u8> {
279    v.iter()
280        .flat_map(|&x| {
281            let b = x.to_bits();
282            let lsb = (b >> 16) & 1;
283            ((b.wrapping_add(0x7FFF + lsb) >> 16) as u16).to_le_bytes()
284        })
285        .collect()
286}
287
288#[cfg(test)]
289#[path = "moe_v41_route_tests.rs"]
290mod route_tests;
291#[cfg(test)]
292#[path = "moe_v41_tests.rs"]
293mod tests;
metrale_model_arch/moe_v41.rs

metrale_model_arch/
moe_v41.rs
