spark_model/layers/moe_v41/
route.rs

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

spark_model/layers/moe_v41/
route.rs
