metrale_model_arch/moe_v41/
device_route.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: The single-token routing on the device (`METRALE_DS41_DEVICE_ROUTE=1`).
5//!
6//! `moe_v41_route_select` computes the selection with ports of glibc's
7//! `expf` / `log1pf` for the softplus and ranks equal keys by the lower expert
8//! id, as the host's stable sort does. It writes `weight_dev` / `rows_dev` and
9//! the gate / up / down pointer table in plan order from a device slot table,
10//! and leaves one header for the host: the miss flag, the picks, the weight
11//! bits and the plan's slots. The host reads that header back once, fetches
12//! the picks through the cache, uploads the cache's slot changes to the device
13//! table, and re-uploads the pointer table from the fetched slots when a pick
14//! was missing or a slot moved.
15//!
16//! `METRALE_DS41_DEVICE_ROUTE_DIAG=1` also runs the host chain on the same
17//! logits and checks the picks, the weight bits, the plan weights and the
18//! pointer table on every token (a mismatch fails the step).
19//!
20//! Owner: model-arch (DeepSeek-V4.1 MoE).
21//! Invariants:
22//! - Outside tests, the selection kernel launches only for a layer that
23//!   `device_route_ok` accepted.
24
25use anyhow::{Result, bail, ensure};
26use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
27use metrale_gpu_runtime::kernel_args::KernelLaunch;
28use metrale_model_weights::weights::expert_stream::{ExpertLru, ExpertSource};
29
30use super::{MoeV41, MoeV41LayerWeights, SLOT_TABLE_LAYERS, route_from_logits};
31
32/// 2026-09-25: The two constants the kernel is written for.
33const ROUTE_NR: usize = 384;
34const ROUTE_TOPK: usize = 6;
35
36fn device_route() -> bool {
37    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
38    *V.get_or_init(|| std::env::var("METRALE_DS41_DEVICE_ROUTE").is_ok_and(|v| v == "1"))
39}
40
41fn device_route_diag() -> bool {
42    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
43    *V.get_or_init(|| std::env::var("METRALE_DS41_DEVICE_ROUTE_DIAG").is_ok_and(|v| v == "1"))
44}
45
46static DIAG_CHECKED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
47
48impl MoeV41 {
49    /// 2026-09-25: The device bias for a layer's router (`MoeV41LayerWeights::gate_bias_dev`).
50    pub fn upload_bias(gpu: &dyn GpuBackend, bias: &[f32]) -> Result<DevicePtr> {
51        let bytes: Vec<u8> = bias.iter().flat_map(|v| v.to_le_bytes()).collect();
52        let p = gpu.alloc(bytes.len().max(16))?;
53        gpu.copy_h2d(&bytes, p)?;
54        Ok(p)
55    }
56
57    /// 2026-09-25: Is the device selection on, and does this layer fit it?
58    /// A layer that does not fit logs one warning per process.
59    pub(super) fn device_route_ok(&self, w: &MoeV41LayerWeights) -> bool {
60        if !device_route() {
61            return false;
62        }
63        let c = &self.cfg;
64        let fits = c.n_routed == ROUTE_NR
65            && c.topk == ROUTE_TOPK
66            && (w.layer as usize) < SLOT_TABLE_LAYERS
67            && !w.gate_bias_dev.is_null();
68        if !fits {
69            static ONCE: std::sync::Once = std::sync::Once::new();
70            ONCE.call_once(|| {
71                tracing::warn!(
72                    "moe_v41: METRALE_DS41_DEVICE_ROUTE=1 ignored (n_routed {} topk {} layer {}): the kernel is written for {ROUTE_NR} x {ROUTE_TOPK}",
73                    c.n_routed,
74                    c.topk,
75                    w.layer
76                );
77            });
78        }
79        fits
80    }
81
82    /// 2026-09-25: Words in one layer's header: the miss flag, the picks, the
83    /// weight bits, the plan slots.
84    pub fn header_words(&self) -> usize {
85        1 + 3 * self.cfg.topk
86    }
87
88    /// 2026-09-25: Layer `layer`'s header inside `route_hdr`.
89    fn header_of(&self, layer: u32) -> DevicePtr {
90        DevicePtr(self.route_hdr.0 + (layer as usize * self.header_words() * 4) as u64)
91    }
92
93    /// 2026-09-25: The arena facts the selection kernel builds pointers from:
94    /// (base, slot bytes, gate / up / down offsets).
95    pub fn arena_of(lru: &ExpertLru) -> (u64, u64, u64, u64, u64) {
96        let lay = lru.layout();
97        (
98            lru.arena_dev().0,
99            lay.bytes as u64,
100            lay.gate_off as u64,
101            lay.up_off as u64,
102            lay.down_off as u64,
103        )
104    }
105
106    /// 2026-09-25: The device selection for a replayed step: the kernel
107    /// launched with its header in layer `w.layer`'s row, and no read-back
108    /// (the caller runs `read_headers` after the segment). The slot table must
109    /// already be up to date. Device work only; an error when
110    /// `device_route_ok` refuses the layer.
111    pub fn route_select_deferred(
112        &self,
113        gpu: &dyn GpuBackend,
114        w: &MoeV41LayerWeights,
115        arena: (u64, u64, u64, u64, u64),
116        stream: u64,
117    ) -> Result<()> {
118        ensure!(
119            self.device_route_ok(w),
120            "moe_v41: the step graph needs METRALE_DS41_DEVICE_ROUTE=1 (layer {})",
121            w.layer
122        );
123        self.route_select_launch(gpu, w.layer, w.gate_bias_dev, arena, stream)
124    }
125
126    /// 2026-09-25: Every layer's header in one read-back ordered on `stream`.
127    pub fn read_headers(&self, gpu: &dyn GpuBackend, stream: u64) -> Result<Vec<u8>> {
128        let mut hdr = vec![0u8; SLOT_TABLE_LAYERS * self.header_words() * 4];
129        gpu.copy_d2h_on_stream(self.route_hdr, &mut hdr, stream)?;
130        Ok(hdr)
131    }
132
133    /// 2026-09-25: `(miss flag, picks, weights)` of layer `layer` out of
134    /// `read_headers`; an error when a pick is outside `0..n_routed`.
135    pub fn parse_header(&self, hdr: &[u8], layer: u32) -> Result<(bool, Vec<usize>, Vec<f32>)> {
136        let k = self.cfg.topk;
137        let base = layer as usize * self.header_words();
138        let word = |i: usize| {
139            let o = 4 * (base + i);
140            i32::from_le_bytes([hdr[o], hdr[o + 1], hdr[o + 2], hdr[o + 3]])
141        };
142        let picks: Vec<usize> = (0..k).map(|i| word(1 + i) as usize).collect();
143        ensure!(
144            picks.iter().all(|&e| e < self.cfg.n_routed),
145            "device route L{layer}: pick outside 0..{}",
146            self.cfg.n_routed
147        );
148        let weights = (0..k)
149            .map(|i| f32::from_bits(word(1 + k + i) as u32))
150            .collect();
151        Ok((word(0) != 0, picks, weights))
152    }
153
154    /// 2026-09-25: The selection kernel on `self.logits` (row 0) for layer `layer`, the
155    /// pointer table from `arena` = (base, slot bytes, gate / up / down
156    /// offsets). Device work only; the header lands in layer `layer`'s row
157    /// of `self.route_hdr`.
158    pub(super) fn route_select_launch(
159        &self,
160        gpu: &dyn GpuBackend,
161        layer: u32,
162        bias_dev: DevicePtr,
163        arena: (u64, u64, u64, u64, u64),
164        stream: u64,
165    ) -> Result<()> {
166        let c = &self.cfg;
167        let row = DevicePtr(self.slot_table.0 + (layer as usize * c.n_routed * 4) as u64);
168        KernelLaunch::new(gpu, self.k.route_select)
169            .grid([1, 1, 1])
170            .block([ROUTE_NR as u32, 1, 1])
171            .arg_ptr(self.logits)
172            .arg_ptr(bias_dev)
173            .arg_f32(c.gate_temp)
174            .arg_f32(c.route_scale)
175            .arg_u32(c.norm_topk_prob as u32)
176            .arg_ptr(row)
177            .arg_u64(arena.0)
178            .arg_u64(arena.1)
179            .arg_u64(arena.2)
180            .arg_u64(arena.3)
181            .arg_u64(arena.4)
182            .arg_ptr(self.header_of(layer))
183            .arg_ptr(self.weight_dev)
184            .arg_ptr(self.rows_dev)
185            .arg_ptr(self.ptrs_dev)
186            .launch(stream)
187    }
188
189    /// 2026-09-25: Bring the device slot table up to date with the cache's
190    /// changes; an error for a triple outside the table.
191    pub(crate) fn slot_table_update(
192        &self,
193        gpu: &dyn GpuBackend,
194        changes: &[(u32, u32, i32)],
195        stream: u64,
196    ) -> Result<()> {
197        if changes.is_empty() {
198            return Ok(());
199        }
200        let c = &self.cfg;
201        // 2026-09-25: One triple per key, the last one: the kernel writes a
202        // chunk's triples in parallel, so a key assigned, evicted and assigned
203        // again inside one batch would race to a wrong final slot (-1 for a
204        // resident expert is a false miss flag).
205        let mut last: std::collections::HashMap<(u32, u32), i32> =
206            std::collections::HashMap::with_capacity(changes.len());
207        for &(l, e, s) in changes {
208            last.insert((l, e), s);
209        }
210        let changes: Vec<(u32, u32, i32)> = last.into_iter().map(|((l, e), s)| (l, e, s)).collect();
211        // 2026-09-25: The 12-byte triples are staged through `a_rows`, in
212        // chunks that fit it.
213        let per = (c.max_tokens * c.dim * 2 / 12).max(1);
214        for chunk in changes.chunks(per) {
215            let mut bytes: Vec<u8> = Vec::with_capacity(chunk.len() * 12);
216            for &(l, e, s) in chunk {
217                ensure!(
218                    (l as usize) < SLOT_TABLE_LAYERS && (e as usize) < c.n_routed,
219                    "slot table: layer {l} expert {e} outside the table"
220                );
221                bytes.extend_from_slice(&(l as i32).to_le_bytes());
222                bytes.extend_from_slice(&(e as i32).to_le_bytes());
223                bytes.extend_from_slice(&s.to_le_bytes());
224            }
225            gpu.copy_h2d_async(&bytes, self.a_rows, stream)?;
226            let n = chunk.len() as u32;
227            KernelLaunch::new(gpu, self.k.slot_table_set)
228                .grid([n.div_ceil(256), 1, 1])
229                .block([256, 1, 1])
230                .arg_ptr(self.slot_table)
231                .arg_u32(c.n_routed as u32)
232                .arg_ptr(self.a_rows)
233                .arg_u32(n)
234                .launch(stream)?;
235            // 2026-09-25: Synchronise before the next chunk overwrites `a_rows`.
236            gpu.synchronize(stream)?;
237        }
238        Ok(())
239    }
240
241    /// 2026-09-25: The device selection for one token: the header read back,
242    /// the picks fetched through the cache, the slot table and (on a miss or
243    /// a moved slot) the pointer table brought up to date. `route_launch` must
244    /// have run on `stream`. Returns the routing as `route_select` does.
245    pub(super) fn route_device_m1<S: ExpertSource + ?Sized>(
246        &self,
247        gpu: &dyn GpuBackend,
248        w: &MoeV41LayerWeights,
249        lru: &mut ExpertLru,
250        src: &S,
251        reader_threads: usize,
252        stream: u64,
253    ) -> Result<(Vec<f32>, Vec<usize>)> {
254        let c = &self.cfg;
255        let arena = Self::arena_of(lru);
256        // 2026-09-25: The table must know every slot the cache changed so far
257        // (prefill, the host-routed layers) before the kernel reads it.
258        let pending = lru.drain_slot_changes();
259        self.slot_table_update(gpu, &pending, stream)?;
260        self.route_select_launch(gpu, w.layer, w.gate_bias_dev, arena, stream)?;
261        let k = c.topk;
262        let mut hdr = vec![0u8; (1 + 3 * k) * 4];
263        gpu.copy_d2h_on_stream(self.header_of(w.layer), &mut hdr, stream)?;
264        let (miss_flag, indices, weights) = self.parse_header(&hdr, 0)?;
265        // 2026-09-25: The cache: hits touched, misses read (on the device
266        // arena, their copies enqueued on `stream` ahead of the experts).
267        lru.begin_token();
268        let keys: Vec<(u32, u32)> = indices.iter().map(|&e| (w.layer, e as u32)).collect();
269        let slots = lru.fetch_many_on(gpu, stream, src, &keys, &[], reader_threads)?;
270        let changes = lru.drain_slot_changes();
271        let changed = !changes.is_empty();
272        self.slot_table_update(gpu, &changes, stream)?;
273        if miss_flag || changed {
274            // 2026-09-25: A pick was not resident when the kernel looked, or
275            // the fetch moved a slot: rebuild the pointer table from the
276            // fetched slots, in plan order, as the host path builds it.
277            let (_, _, plan) = self.plan(&indices, &weights, 1);
278            ensure!(
279                plan.len() == k,
280                "device route: {} distinct picks, expected {k}",
281                plan.len()
282            );
283            self.upload_expert_table(gpu, &plan, &slots, stream)?;
284        }
285        if device_route_diag() {
286            self.device_route_check(gpu, w, &indices, &weights, &slots, stream)?;
287        }
288        Ok((weights, indices))
289    }
290
291    /// 2026-09-25: The host chain on the same logits against the device's
292    /// header, the plan weights and the pointer table; a mismatch fails the step.
293    fn device_route_check(
294        &self,
295        gpu: &dyn GpuBackend,
296        w: &MoeV41LayerWeights,
297        indices: &[usize],
298        weights: &[f32],
299        slots: &[metrale_model_weights::weights::expert_stream::ExpertSlot],
300        stream: u64,
301    ) -> Result<()> {
302        let c = &self.cfg;
303        let k = c.topk;
304        let mut bytes = vec![0u8; c.n_routed * 4];
305        gpu.copy_d2h_on_stream(self.logits, &mut bytes, stream)?;
306        let logits: Vec<f32> = bytes
307            .chunks_exact(4)
308            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
309            .collect();
310        let (hw, hi) = route_from_logits(&logits, 1, &w.gate_bias, c);
311        let l = w.layer;
312        if hi != indices {
313            bail!("device route L{l}: picks {indices:?} != host {hi:?}");
314        }
315        let (wb, hwb): (Vec<u32>, Vec<u32>) = (
316            weights.iter().map(|v| v.to_bits()).collect(),
317            hw.iter().map(|v| v.to_bits()).collect(),
318        );
319        if wb != hwb {
320            bail!("device route L{l}: weight bits {wb:08x?} != host {hwb:08x?}");
321        }
322        // 2026-09-25: The plan the device wrote: `weight_dev` in ascending
323        // expert id, then `ptrs_dev`.
324        let (_, w_host, plan) = self.plan(&hi, &hw, 1);
325        let mut w_dev = vec![0u8; k * 4];
326        gpu.copy_d2h_on_stream(self.weight_dev, &mut w_dev, stream)?;
327        if w_dev != w_host {
328            bail!("device route L{l}: plan weights differ from the host plan");
329        }
330        let mut ptrs = vec![0u8; 3 * k * 8];
331        gpu.copy_d2h_on_stream(self.ptrs_dev, &mut ptrs, stream)?;
332        let mut want: Vec<u8> = Vec::with_capacity(3 * k * 8);
333        for which in 0..3 {
334            for &(a0, _, _) in &plan {
335                let s = slots[a0];
336                want.extend_from_slice(&[s.gate, s.up, s.down][which].0.to_le_bytes());
337            }
338        }
339        if ptrs != want {
340            bail!("device route L{l}: pointer table differs from the fetched slots");
341        }
342        let n = DIAG_CHECKED.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
343        if n.is_multiple_of(4000) {
344            tracing::info!("device route diag: {n} layer-tokens checked, 0 mismatches");
345        }
346        Ok(())
347    }
348}
metrale_model_arch/moe_v41/device_route.rs

metrale_model_arch/moe_v41/
device_route.rs
