spark_model/layers/moe_v41/
device_route.rs

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

spark_model/layers/moe_v41/
device_route.rs
