spark_model/layers/deepseek_v41_ref/
moe.rs

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

spark_model/layers/deepseek_v41_ref/
moe.rs
