rlx_runtime/
moe_expert_store.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-expert F32 weight slabs for MoE offload (TIDE-style migration source).
6
7use crate::ExpertPool;
8use std::sync::Arc;
9
10/// One expert projection stack in GroupedMatMul layout `[num_experts, k, n]`.
11#[derive(Debug, Clone)]
12pub struct ExpertStackF32 {
13    pub num_experts: usize,
14    pub k: usize,
15    pub n: usize,
16    pub data: Arc<[f32]>,
17}
18
19impl ExpertStackF32 {
20    pub fn new(data: Vec<f32>, num_experts: usize, k: usize, n: usize) -> Self {
21        assert_eq!(data.len(), num_experts * k * n);
22        Self {
23            num_experts,
24            k,
25            n,
26            data: Arc::from(data),
27        }
28    }
29
30    pub fn expert_stride(&self) -> usize {
31        self.k * self.n
32    }
33
34    pub fn expert_slice(&self, e: usize) -> &[f32] {
35        let stride = self.expert_stride();
36        let start = e * stride;
37        &self.data[start..start + stride]
38    }
39
40    pub fn as_slice(&self) -> &[f32] {
41        &self.data
42    }
43}
44
45/// Gate / up / down expert stacks for one decoder layer.
46#[derive(Debug, Clone)]
47pub struct LayerMoeWeights {
48    pub layer_index: usize,
49    pub gate: ExpertStackF32,
50    pub up: ExpertStackF32,
51    pub down: ExpertStackF32,
52}
53
54/// Host-side expert weights for all MoE layers (migration source of truth).
55#[derive(Debug, Clone)]
56pub struct MoeExpertStore {
57    pub layers: Vec<LayerMoeWeights>,
58}
59
60impl MoeExpertStore {
61    pub fn num_layers(&self) -> usize {
62        self.layers.len()
63    }
64
65    /// Apply captured TopK indices to per-layer pools (TIDE refresh).
66    pub fn refresh_pools(
67        &self,
68        pools: &mut [ExpertPool],
69        captured: &[Vec<u32>],
70        decode_step: usize,
71        is_prefill_block: bool,
72    ) -> bool {
73        let n = self.layers.len().min(pools.len()).min(captured.len());
74        if n == 0 {
75            return false;
76        }
77        let refresh =
78            pools[0].should_refresh(crate::MoEExecMode::Reuse, decode_step, is_prefill_block);
79        if !refresh {
80            return false;
81        }
82        for i in 0..n {
83            pools[i].refresh_from_indices(&captured[i]);
84        }
85        true
86    }
87
88    /// Push full host stacks into compiled params (lossless; refreshes arena bytes).
89    pub fn apply_to_compiled(&self, compiled: &mut crate::CompiledGraph) {
90        for layer in &self.layers {
91            let il = layer.layer_index;
92            compiled.set_param(
93                &format!("blk.{il}.ffn_gate_exps.weight"),
94                layer.gate.as_slice(),
95            );
96            compiled.set_param(&format!("blk.{il}.ffn_up_exps.weight"), layer.up.as_slice());
97            compiled.set_param(
98                &format!("blk.{il}.ffn_down_exps.weight"),
99                layer.down.as_slice(),
100            );
101        }
102    }
103}
rlx_runtime/moe_expert_store.rs

rlx_runtime/
moe_expert_store.rs
