rlx_cpu/
moe_topk_capture.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//! Capture MoE router [`Op::TopK`] outputs during CPU forward (TIDE refresh input).
6
7use std::sync::{Arc, Mutex};
8
9/// Shared capture buffer — one entry per MoE router TopK in schedule order.
10#[derive(Debug)]
11pub struct MoeTopkCapture {
12    pub num_experts: usize,
13    layers: Mutex<Vec<Vec<u32>>>,
14}
15
16impl MoeTopkCapture {
17    pub fn new(num_experts: usize) -> Arc<Self> {
18        Arc::new(Self {
19            num_experts,
20            layers: Mutex::new(Vec::new()),
21        })
22    }
23
24    pub fn clear(&self) {
25        self.layers.lock().unwrap().clear();
26    }
27
28    /// Record one router TopK output (`outer * k` f32-encoded expert ids).
29    pub fn push_topk_f32(&self, data: &[f32], axis_dim: usize) {
30        if axis_dim != self.num_experts {
31            return;
32        }
33        let flat: Vec<u32> = data.iter().map(|&v| v as u32).collect();
34        self.layers.lock().unwrap().push(flat);
35    }
36
37    pub fn take_layers(&self) -> Vec<Vec<u32>> {
38        std::mem::take(&mut *self.layers.lock().unwrap())
39    }
40}
rlx_cpu/moe_topk_capture.rs

rlx_cpu/
moe_topk_capture.rs
