rlx_cpu/
moe_residency.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Per-forward MoE expert residency mask (TIDE placement) for CPU dispatch.
6//!
7//! Set by `rlx_runtime::CompiledGraph::set_moe_resident_experts` before
8//! `run`. `crate::thunk::GroupedMatMul` reads the mask for accounting;
9//! numerics still use the full expert stack in the arena (lossless on CPU).
10
11use std::cell::{Cell, RefCell};
12use std::sync::{Arc, RwLock};
13
14/// Per-expert host pointers for one MoE layer (gate/up/down).
15#[derive(Debug, Clone)]
16pub struct LayerHostBind {
17    pub gate: Vec<*const f32>,
18    pub up: Vec<*const f32>,
19    pub down: Vec<*const f32>,
20    pub stride: usize,
21}
22
23/// Host weight lookup installed before forward (TIDE CPU/GPU fallback path).
24#[derive(Debug, Clone)]
25pub struct MoeHostBind {
26    pub layers: Vec<LayerHostBind>,
27}
28
29// Host weight pointers are installed on the calling thread; safe to share
30// across the RwLock used for TIDE residency bookkeeping.
31unsafe impl Send for LayerHostBind {}
32unsafe impl Sync for LayerHostBind {}
33unsafe impl Send for MoeHostBind {}
34unsafe impl Sync for MoeHostBind {}
35
36static HOST_BIND: RwLock<Option<MoeHostBind>> = RwLock::new(None);
37static LAST_STATS: RwLock<Option<MoeResidencyStats>> = RwLock::new(None);
38
39thread_local! {
40    /// Monotonic GroupedMatMul ordinal within one forward (layer = ord/3,
41    /// matrix = ord%3). **Thread-local** to match the residency [`CTX`]: forwards
42    /// run per-thread, so a process-global counter would let one thread's
43    /// `reset_gmm_counters` / `next_gmm_ord` clobber another's in-flight ordinal
44    /// sequence — corrupting the layer/matrix decode (wrong TIDE host-expert
45    /// weights + residency accounting).
46    static GMM_ORD: Cell<usize> = const { Cell::new(0) };
47}
48
49#[derive(Debug, Default, Clone)]
50pub struct MoeResidencyStats {
51    pub gpu_expert_calls: u64,
52    pub cpu_expert_calls: u64,
53    pub gpu_tokens: u64,
54    pub cpu_tokens: u64,
55}
56
57struct MoeResidencyCtx {
58    /// Union mask (legacy): expert on device if any layer has it resident.
59    merged: Option<Arc<[bool]>>,
60    /// TIDE per MoE layer (forward order); takes precedence over [`merged`].
61    per_layer: Option<Arc<Vec<Arc<[bool]>>>>,
62    stats: MoeResidencyStats,
63}
64
65thread_local! {
66    static CTX: RefCell<Option<MoeResidencyCtx>> = const { RefCell::new(None) };
67}
68
69/// Install merged (union) residency mask for the current thread until [`clear_mask`].
70pub fn set_mask(mask: Option<Arc<[bool]>>) {
71    CTX.with(|c| {
72        *c.borrow_mut() = Some(MoeResidencyCtx {
73            merged: mask,
74            per_layer: None,
75            stats: MoeResidencyStats::default(),
76        });
77    });
78}
79
80/// Install per-layer masks (one row per MoE FFN in forward order).
81pub fn set_per_layer_masks(layers: Option<Arc<Vec<Arc<[bool]>>>>) {
82    CTX.with(|c| {
83        *c.borrow_mut() = Some(MoeResidencyCtx {
84            merged: None,
85            per_layer: layers,
86            stats: MoeResidencyStats::default(),
87        });
88    });
89}
90
91pub fn clear_mask() {
92    CTX.with(|c| *c.borrow_mut() = None);
93}
94
95pub fn bind_host_weights(bind: Option<MoeHostBind>) {
96    *HOST_BIND.write().unwrap() = bind;
97}
98
99pub fn reset_gmm_counters() {
100    GMM_ORD.with(|c| c.set(0));
101}
102
103/// Next MoE GroupedMatMul ordinal for this forward (call once per kernel).
104pub fn next_gmm_ord() -> usize {
105    GMM_ORD.with(|c| {
106        let v = c.get();
107        c.set(v + 1);
108        v
109    })
110}
111
112/// Host expert weight pointer for `ord` (gate/up/down = ord%3, layer = ord/3).
113pub fn host_expert_weight_ptr(ord: usize, expert: usize) -> Option<*const f32> {
114    let bind = HOST_BIND.read().unwrap();
115    let bind = bind.as_ref()?;
116    let layer = bind.layers.get(ord / 3)?;
117    let ptrs = match ord % 3 {
118        0 => &layer.gate,
119        1 => &layer.up,
120        _ => &layer.down,
121    };
122    ptrs.get(expert).copied()
123}
124
125pub fn peek_stats() -> Option<MoeResidencyStats> {
126    CTX.with(|c| c.borrow().as_ref().map(|ctx| ctx.stats.clone()))
127}
128
129/// Stats from the most recent forward on this thread (set when the residency guard drops).
130pub fn take_last_forward_stats() -> Option<MoeResidencyStats> {
131    LAST_STATS.write().unwrap().take()
132}
133
134pub(crate) fn stash_last_forward_stats(stats: MoeResidencyStats) {
135    *LAST_STATS.write().unwrap() = Some(stats);
136}
137
138fn expert_on_device_inner(ctx: &MoeResidencyCtx, layer: Option<usize>, e: usize) -> bool {
139    if let Some(layers) = ctx.per_layer.as_ref() {
140        if let Some(li) = layer {
141            return layers
142                .get(li)
143                .and_then(|m| m.get(e).copied())
144                .unwrap_or(true);
145        }
146    }
147    ctx.merged
148        .as_ref()
149        .and_then(|m| m.get(e).copied())
150        .unwrap_or(true)
151}
152
153/// True when expert `e` is GPU-resident for MoE layer `layer` (GMM ord / 3).
154pub fn expert_on_device_for_layer(layer: usize, e: usize) -> bool {
155    CTX.with(|c| {
156        let borrow = c.borrow();
157        let Some(ctx) = borrow.as_ref() else {
158            return true;
159        };
160        expert_on_device_inner(ctx, Some(layer), e)
161    })
162}
163
164/// True when expert `e` is marked GPU-resident (merged mask if no per-layer table).
165pub fn expert_on_device(e: usize) -> bool {
166    expert_on_device_for_layer(0, e)
167}
168
169pub fn record_expert_tokens(layer: usize, e: usize, num_tokens: usize) {
170    if num_tokens == 0 {
171        return;
172    }
173    CTX.with(|c| {
174        let mut borrow = c.borrow_mut();
175        let Some(ctx) = borrow.as_mut() else {
176            return;
177        };
178        let on_device = expert_on_device_inner(ctx, Some(layer), e);
179        if on_device {
180            ctx.stats.gpu_expert_calls += 1;
181            ctx.stats.gpu_tokens += num_tokens as u64;
182        } else {
183            ctx.stats.cpu_expert_calls += 1;
184            ctx.stats.cpu_tokens += num_tokens as u64;
185        }
186    });
187}
188
189/// Take stats and clear the thread-local context.
190pub fn take_stats() -> Option<MoeResidencyStats> {
191    CTX.with(|c| c.borrow_mut().take().map(|ctx| ctx.stats))
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use std::sync::Arc;
198
199    #[test]
200    fn gmm_ordinal_is_thread_local() {
201        use std::sync::Barrier;
202        // Each forward runs on its own thread and expects its own ordinal
203        // sequence. Two threads resetting + advancing concurrently must not
204        // clobber each other (the old process-global atomic did).
205        let barrier = Arc::new(Barrier::new(2));
206        let handles: Vec<_> = (0..2)
207            .map(|_| {
208                let b = barrier.clone();
209                std::thread::spawn(move || {
210                    reset_gmm_counters();
211                    b.wait(); // both reset, then race the increments
212                    (0..5).map(|_| next_gmm_ord()).collect::<Vec<_>>()
213                })
214            })
215            .collect();
216        for h in handles {
217            assert_eq!(
218                h.join().unwrap(),
219                vec![0, 1, 2, 3, 4],
220                "GroupedMatMul ordinal must be per-thread monotonic"
221            );
222        }
223    }
224
225    #[test]
226    fn per_layer_masks_are_layer_local() {
227        let per = Arc::new(vec![
228            Arc::from([false, true, true, true]),
229            Arc::from([true, false, true, true]),
230        ]);
231        set_per_layer_masks(Some(per));
232        assert!(!expert_on_device_for_layer(0, 0));
233        assert!(expert_on_device_for_layer(0, 1));
234        assert!(expert_on_device_for_layer(1, 0));
235        assert!(!expert_on_device_for_layer(1, 1));
236        clear_mask();
237    }
238}
rlx_cpu/moe_residency.rs

rlx_cpu/
moe_residency.rs
