metrale_model_arch/moe_v41/
route.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: The MoE router: the logits on the GPU, the top-k selection on
5//! the CPU (`route`, `route_launch`, `route_select`), and the next layer's
6//! router on this layer's input for prefetching (`predict_launch`,
7//! `predict_select`, `route_predict`).
8//!
9//! Owner: model-arch (DeepSeek-V4.1 MoE).
10//! Invariants:
11//! - Every selection in this module is `route_from_logits` on logits read
12//!   back from the device.
13
14use anyhow::{Result, ensure};
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
16use metrale_gpu_runtime::kernel_args::KernelLaunch;
17
18use super::{MoeV41, MoeV41LayerWeights, RouterWeights, route_from_logits};
19
20impl MoeV41 {
21    /// 2026-09-25: Router logits on the GPU (f32 out), the selection on the
22    /// CPU. Drains `stream`.
23    pub fn route(
24        &self,
25        gpu: &dyn GpuBackend,
26        w: &MoeV41LayerWeights,
27        x: DevicePtr,
28        m: usize,
29        stream: u64,
30    ) -> Result<(Vec<f32>, Vec<usize>)> {
31        self.route_launch(gpu, w, x, m, stream)?;
32        self.route_select(gpu, w, m, stream)
33    }
34
35    /// 2026-09-25: The router GEMM/GEMV into `self.logits`: device work only,
36    /// so a CUDA graph can hold it (segment A of the step graph ends with it).
37    /// Errors when `m` is outside `1..=max_tokens`.
38    pub fn route_launch(
39        &self,
40        gpu: &dyn GpuBackend,
41        w: &MoeV41LayerWeights,
42        x: DevicePtr,
43        m: usize,
44        stream: u64,
45    ) -> Result<()> {
46        let c = &self.cfg;
47        ensure!(
48            m >= 1 && m <= c.max_tokens,
49            "moe_v41: {m} tokens outside 1..={}",
50            c.max_tokens
51        );
52        // 2026-09-25: At m <= 8 a GEMV. By default the products form
53        // (`moe_v41_router_gemv_f32out_products`, one gate row a block: the
54        // 256 threads write the bf16 x bf16 products into shared memory and
55        // thread 0 adds them in k order); `METRALE_DS41_ROUTER_STAGED=0`
56        // selects the direct form, one thread per logit. Above 8 rows the
57        // tiled `dense_gemm_bf16_f32out`.
58        let (kernel, grid, block, smem) = if m <= 8 {
59            if router_staged() {
60                (
61                    self.k.router_gemv_staged,
62                    [c.n_routed as u32, m as u32, 1],
63                    [256, 1, 1],
64                    c.dim as u32 * 4,
65                )
66            } else {
67                (
68                    self.k.router_gemv,
69                    [(c.n_routed as u32).div_ceil(64), m as u32, 1],
70                    [64, 1, 1],
71                    0,
72                )
73            }
74        } else {
75            (
76                self.k.gemm_f32out,
77                [(c.n_routed as u32).div_ceil(16), (m as u32).div_ceil(16), 1],
78                [16, 16, 1],
79                0,
80            )
81        };
82        KernelLaunch::new(gpu, kernel)
83            .grid(grid)
84            .block(block)
85            .shared_mem(smem)
86            .arg_ptr(x)
87            .arg_ptr(w.gate_w)
88            .arg_ptr(self.logits)
89            .arg_u32(m as u32)
90            .arg_u32(c.n_routed as u32)
91            .arg_u32(c.dim as u32)
92            .launch(stream)
93    }
94
95    /// 2026-09-25: The selection: drain the stream, download the logits
96    /// `route_launch` wrote, pick the top-k on the CPU with
97    /// `route_from_logits`. Errors when the host bias length is not
98    /// `n_routed`.
99    pub fn route_select(
100        &self,
101        gpu: &dyn GpuBackend,
102        w: &MoeV41LayerWeights,
103        m: usize,
104        stream: u64,
105    ) -> Result<(Vec<f32>, Vec<usize>)> {
106        let c = &self.cfg;
107        let mut bytes = vec![0u8; m * c.n_routed * 4];
108        gpu.copy_d2h_on_stream(self.logits, &mut bytes, stream)?;
109        let logits: Vec<f32> = bytes
110            .chunks_exact(4)
111            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
112            .collect();
113        ensure!(
114            w.gate_bias.len() == c.n_routed,
115            "gate bias has {} entries for {} experts",
116            w.gate_bias.len(),
117            c.n_routed
118        );
119        Ok(route_from_logits(&logits, m, &w.gate_bias, c))
120    }
121}
122
123impl MoeV41 {
124    /// 2026-09-25: The next layer's router GEMV on this layer's single-token
125    /// input, into `pred_logits`. Device work only; `predict_select` reads it
126    /// back.
127    pub fn predict_launch(
128        &self,
129        gpu: &dyn GpuBackend,
130        next: &RouterWeights,
131        x: DevicePtr,
132        stream: u64,
133    ) -> Result<()> {
134        let c = &self.cfg;
135        let (kernel, grid, block, smem) = if router_staged() {
136            (
137                self.k.router_gemv_staged,
138                [c.n_routed as u32, 1, 1],
139                [256, 1, 1],
140                c.dim as u32 * 4,
141            )
142        } else {
143            (
144                self.k.router_gemv,
145                [(c.n_routed as u32).div_ceil(64), 1, 1],
146                [64, 1, 1],
147                0,
148            )
149        };
150        KernelLaunch::new(gpu, kernel)
151            .grid(grid)
152            .block(block)
153            .shared_mem(smem)
154            .arg_ptr(x)
155            .arg_ptr(next.gate_w)
156            .arg_ptr(self.pred_logits)
157            .arg_u32(1)
158            .arg_u32(c.n_routed as u32)
159            .arg_u32(c.dim as u32)
160            .launch(stream)
161    }
162
163    /// 2026-09-25: The `k` best experts of the predicted layer by
164    /// `score + bias`, as `(layer, expert)` cache keys, from the logits
165    /// `predict_launch` wrote. Drains `stream`.
166    pub fn predict_select(
167        &self,
168        gpu: &dyn GpuBackend,
169        next: &RouterWeights,
170        k: usize,
171        stream: u64,
172    ) -> Result<Vec<(u32, u32)>> {
173        let c = &self.cfg;
174        let mut bytes = vec![0u8; c.n_routed * 4];
175        gpu.copy_d2h_on_stream(self.pred_logits, &mut bytes, stream)?;
176        let logits: Vec<f32> = bytes
177            .chunks_exact(4)
178            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
179            .collect();
180        let mut wide = c.clone();
181        wide.topk = k.min(c.n_routed);
182        let (_, ids) = route_from_logits(&logits, 1, &next.gate_bias, &wide);
183        Ok(ids.into_iter().map(|e| (next.layer, e as u32)).collect())
184    }
185
186    /// 2026-09-25: The router of this layer on `x` (one token), the `k` best
187    /// experts by `score + bias` in `route_from_logits` order. Drains the
188    /// stream; overwrites `self.logits`.
189    pub fn route_predict(
190        &self,
191        gpu: &dyn GpuBackend,
192        w: &MoeV41LayerWeights,
193        x: DevicePtr,
194        k: usize,
195        stream: u64,
196    ) -> Result<Vec<usize>> {
197        let c = &self.cfg;
198        self.route_launch(gpu, w, x, 1, stream)?;
199        let mut bytes = vec![0u8; c.n_routed * 4];
200        gpu.copy_d2h_on_stream(self.logits, &mut bytes, stream)?;
201        let logits: Vec<f32> = bytes
202            .chunks_exact(4)
203            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
204            .collect();
205        let mut wide = c.clone();
206        wide.topk = k.min(c.n_routed);
207        Ok(route_from_logits(&logits, 1, &w.gate_bias, &wide).1)
208    }
209}
210
211fn router_staged() -> bool {
212    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
213    *ON.get_or_init(|| !std::env::var("METRALE_DS41_ROUTER_STAGED").is_ok_and(|v| v == "0"))
214}
metrale_model_arch/moe_v41/route.rs

metrale_model_arch/moe_v41/
route.rs
