metrale_model_arch/deepseek_v41_ref/
moe.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: DeepSeek-V4.1 MoE (the reference's `Gate`, `Expert`, `MoE`), CPU
5//! reference against the golden: sqrt-softplus routing with a correction
6//! bias that picks experts but does not scale them, top-k renormalisation,
7//! `route_scale`, the SwiGLU clamp, and one shared expert every token goes
8//! through.
9//!
10//! Precision: the gate runs in f32 from a bf16 `x`. Each expert projection is
11//! `linear_bf16` (f32 accumulate, bf16 result), the SwiGLU is f32, the
12//! routing weight multiplies in f32, the product is rounded to bf16 before
13//! `w2`, and the per-expert outputs accumulate into an f32 `y` rounded to
14//! bf16 once at the end.
15//!
16//! Owner: model-arch, DeepSeek-V4.1 reference.
17//! Invariants: none beyond the types.
18
19use super::to_bf16_rne;
20
21/// 2026-09-25: softplus with `beta = 1` and linear above `threshold = 20`.
22fn softplus(x: f32) -> f32 {
23    if x > 20.0 { x } else { x.exp().ln_1p() }
24}
25
26fn silu(x: f32) -> f32 {
27    x / (1.0 + (-x).exp())
28}
29
30/// 2026-09-25: Linear on bf16 operands: `y[i][o] = bf16(sum_d x[i][d] * w[o][d])`,
31/// f32 accumulate. `w` is `[out, in]` row-major.
32pub fn linear_bf16(x: &[f32], w: &[f32], rows: usize, in_dim: usize, out_dim: usize) -> Vec<f32> {
33    let mut y = vec![0f32; rows * out_dim];
34    for i in 0..rows {
35        let xi = &x[i * in_dim..(i + 1) * in_dim];
36        for o in 0..out_dim {
37            let wo = &w[o * in_dim..(o + 1) * in_dim];
38            let acc: f32 = xi.iter().zip(wo).map(|(a, b)| a * b).sum();
39            y[i * out_dim + o] = to_bf16_rne(acc);
40        }
41    }
42    y
43}
44
45/// 2026-09-25: The gate for `score_func = "sqrtsoftplus"`. `x`: `[tokens, dim]`
46/// bf16 values; `w`: `[n_routed, dim]` bf16 values; `bias`: `[n_routed]`
47/// f32. Returns (`weights[tokens, topk]` f32, `indices[tokens, topk]`),
48/// indices descending by `score + bias` (a stable sort).
49pub fn gate(
50    x: &[f32],
51    w: &[f32],
52    bias: &[f32],
53    tokens: usize,
54    dim: usize,
55    n_routed: usize,
56    topk: usize,
57    gate_temp: f32,
58    norm_topk_prob: bool,
59    route_scale: f32,
60) -> (Vec<f32>, Vec<usize>) {
61    let mut weights = Vec::with_capacity(tokens * topk);
62    let mut indices = Vec::with_capacity(tokens * topk);
63    for t in 0..tokens {
64        let xt = &x[t * dim..(t + 1) * dim];
65        let scores: Vec<f32> = (0..n_routed)
66            .map(|e| {
67                let s: f32 = xt
68                    .iter()
69                    .zip(&w[e * dim..(e + 1) * dim])
70                    .map(|(a, b)| a * b)
71                    .sum();
72                softplus(s / gate_temp).sqrt()
73            })
74            .collect();
75        // 2026-09-25: The bias picks experts but does not scale them.
76        let mut order: Vec<usize> = (0..n_routed).collect();
77        order.sort_by(|&a, &b| {
78            (scores[b] + bias[b])
79                .partial_cmp(&(scores[a] + bias[a]))
80                .expect("finite scores")
81        });
82        let picked = &order[..topk];
83        let mut wt: Vec<f32> = picked.iter().map(|&e| scores[e]).collect();
84        if norm_topk_prob && topk > 1 {
85            let sum: f32 = wt.iter().sum::<f32>() + 1e-20;
86            for v in &mut wt {
87                *v /= sum;
88            }
89        }
90        for v in &mut wt {
91            *v *= route_scale;
92        }
93        weights.extend(wt);
94        indices.extend_from_slice(picked);
95    }
96    (weights, indices)
97}
98
99/// 2026-09-25: One SwiGLU expert. `w1`, `w3`: `[inter, dim]`; `w2`: `[dim,
100/// inter]`; all bf16 values. `x`: `[rows, dim]` bf16 values; `weights`:
101/// per-row routing weight (f32) or `None` for the shared expert. Returns
102/// `[rows, dim]` bf16 values.
103pub fn expert(
104    x: &[f32],
105    w1: &[f32],
106    w2: &[f32],
107    w3: &[f32],
108    rows: usize,
109    dim: usize,
110    inter: usize,
111    swiglu_limit: f32,
112    weights: Option<&[f32]>,
113) -> Vec<f32> {
114    let gate_b = linear_bf16(x, w1, rows, dim, inter);
115    let up_b = linear_bf16(x, w3, rows, dim, inter);
116    let mut h = vec![0f32; rows * inter];
117    for r in 0..rows {
118        for j in 0..inter {
119            let mut g = gate_b[r * inter + j];
120            let mut u = up_b[r * inter + j];
121            if swiglu_limit > 0.0 {
122                u = u.clamp(-swiglu_limit, swiglu_limit);
123                g = g.min(swiglu_limit);
124            }
125            let mut v = silu(g) * u;
126            if let Some(wr) = weights {
127                v *= wr[r];
128            }
129            h[r * inter + j] = to_bf16_rne(v);
130        }
131    }
132    linear_bf16(&h, w2, rows, inter, dim)
133}
134
135/// 2026-09-25: The MoE parameters; `experts[i]` and `shared` are `(w1, w2, w3)`.
136pub struct MoeWeights<'a> {
137    pub gate_w: &'a [f32],
138    pub gate_bias: &'a [f32],
139    pub experts: Vec<(&'a [f32], &'a [f32], &'a [f32])>,
140    pub shared: (&'a [f32], &'a [f32], &'a [f32]),
141}
142
143pub struct MoeCfg {
144    pub dim: usize,
145    pub inter: usize,
146    pub n_routed: usize,
147    pub topk: usize,
148    pub gate_temp: f32,
149    pub norm_topk_prob: bool,
150    pub route_scale: f32,
151    pub swiglu_limit: f32,
152}
153
154/// 2026-09-25: The MoE forward on one rank. Returns (`y[tokens, dim]` bf16
155/// values, routing `weights`, routing `indices`).
156pub fn moe(
157    x: &[f32],
158    tokens: usize,
159    w: &MoeWeights,
160    c: &MoeCfg,
161) -> (Vec<f32>, Vec<f32>, Vec<usize>) {
162    let (weights, indices) = gate(
163        x,
164        w.gate_w,
165        w.gate_bias,
166        tokens,
167        c.dim,
168        c.n_routed,
169        c.topk,
170        c.gate_temp,
171        c.norm_topk_prob,
172        c.route_scale,
173    );
174    let mut y = vec![0f32; tokens * c.dim];
175    for (e, (w1, w2, w3)) in w.experts.iter().enumerate() {
176        // 2026-09-25: Rows routed to expert e, in token order, with the matching
177        // routing weight.
178        let mut rows = Vec::new();
179        let mut rw = Vec::new();
180        for t in 0..tokens {
181            for k in 0..c.topk {
182                if indices[t * c.topk + k] == e {
183                    rows.push(t);
184                    rw.push(weights[t * c.topk + k]);
185                }
186            }
187        }
188        if rows.is_empty() {
189            continue;
190        }
191        let xs: Vec<f32> = rows
192            .iter()
193            .flat_map(|&t| x[t * c.dim..(t + 1) * c.dim].iter().copied())
194            .collect();
195        let out = expert(
196            &xs,
197            w1,
198            w2,
199            w3,
200            rows.len(),
201            c.dim,
202            c.inter,
203            c.swiglu_limit,
204            Some(&rw),
205        );
206        for (i, &t) in rows.iter().enumerate() {
207            for d in 0..c.dim {
208                y[t * c.dim + d] += out[i * c.dim + d];
209            }
210        }
211    }
212    let (s1, s2, s3) = w.shared;
213    let shared = expert(x, s1, s2, s3, tokens, c.dim, c.inter, c.swiglu_limit, None);
214    for (yv, sv) in y.iter_mut().zip(&shared) {
215        *yv += sv;
216    }
217    (y.into_iter().map(to_bf16_rne).collect(), weights, indices)
218}
219
220#[cfg(test)]
221#[path = "moe_tests.rs"]
222mod tests;
metrale_model_arch/deepseek_v41_ref/moe.rs

metrale_model_arch/deepseek_v41_ref/
moe.rs
