rlx_cpu/
moe_residency.rs

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

rlx_cpu/
moe_residency.rs
