metrale_model_layers/layers/
ep_dispatch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Partition of top-k (token, expert) assignments into this rank's
4//! experts and the other ranks' experts, for EP token dispatch.
5//!
6//! Its callers are `MoeLayer::forward_ep_dispatch` (`moe/forward_ep.rs`),
7//! which has no caller and serves through `forward`, and the
8//! `glm53_ep_semantics` test. The served EP MoE is masked-local plus
9//! all-reduce (`moe/forward.rs`).
10//!
11//! Owner: model-layers (MoE).
12//! Invariants: none beyond the types.
13
14/// 2026-09-25: Top-k assignments split into a local bucket (expert in this
15/// rank's range) and a remote bucket (any other expert). Entry `i` of a bucket
16/// is the triple (`*_token_indices[i]`, `*_expert_ids[i]`, `*_weights[i]`), in
17/// the order the assignments were given.
18#[derive(Debug)]
19pub struct EpRoutingTable {
20    pub local_token_indices: Vec<u32>,
21    /// 2026-09-25: Global expert ids, not rank-relative.
22    pub local_expert_ids: Vec<u32>,
23    pub local_weights: Vec<f32>,
24    pub remote_token_indices: Vec<u32>,
25    /// 2026-09-25: Global expert ids, not rank-relative.
26    pub remote_expert_ids: Vec<u32>,
27    pub remote_weights: Vec<f32>,
28}
29
30impl EpRoutingTable {
31    pub fn local_count(&self) -> usize {
32        self.local_token_indices.len()
33    }
34
35    pub fn remote_count(&self) -> usize {
36        self.remote_token_indices.len()
37    }
38
39    /// 2026-09-25: Equals `num_tokens * top_k` of the build call.
40    pub fn total_count(&self) -> usize {
41        self.local_count() + self.remote_count()
42    }
43}
44
45/// 2026-09-25: Build the table from row-major `[num_tokens, top_k]` expert ids
46/// and weights. An expert is local when it lies in
47/// `local_expert_start..local_expert_end`.
48///
49/// # Panics
50/// If `gate_indices.len()` or `gate_weights.len()` is not
51/// `num_tokens * top_k`.
52pub fn build_ep_routing_table(
53    gate_indices: &[u32],
54    gate_weights: &[f32],
55    num_tokens: usize,
56    top_k: usize,
57    local_expert_start: usize,
58    local_expert_end: usize,
59) -> EpRoutingTable {
60    let total = num_tokens * top_k;
61    assert_eq!(gate_indices.len(), total, "gate_indices length mismatch");
62    assert_eq!(gate_weights.len(), total, "gate_weights length mismatch");
63
64    let mut local_token_indices = Vec::with_capacity(total);
65    let mut local_expert_ids = Vec::with_capacity(total);
66    let mut local_weights = Vec::with_capacity(total);
67    let mut remote_token_indices = Vec::with_capacity(total);
68    let mut remote_expert_ids = Vec::with_capacity(total);
69    let mut remote_weights = Vec::with_capacity(total);
70
71    for token_idx in 0..num_tokens {
72        for k in 0..top_k {
73            let flat_idx = token_idx * top_k + k;
74            let expert_id = gate_indices[flat_idx];
75            let weight = gate_weights[flat_idx];
76            let eid = expert_id as usize;
77
78            if eid >= local_expert_start && eid < local_expert_end {
79                local_token_indices.push(token_idx as u32);
80                local_expert_ids.push(expert_id);
81                local_weights.push(weight);
82            } else {
83                remote_token_indices.push(token_idx as u32);
84                remote_expert_ids.push(expert_id);
85                remote_weights.push(weight);
86            }
87        }
88    }
89
90    EpRoutingTable {
91        local_token_indices,
92        local_expert_ids,
93        local_weights,
94        remote_token_indices,
95        remote_expert_ids,
96        remote_weights,
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_all_local() {
106        let indices = vec![3u32, 7, 100, 200];
107        let weights = vec![0.6f32, 0.4, 0.55, 0.45];
108        let table = build_ep_routing_table(&indices, &weights, 2, 2, 0, 256);
109
110        assert_eq!(table.local_count(), 4);
111        assert_eq!(table.remote_count(), 0);
112        assert_eq!(table.total_count(), 4);
113        assert_eq!(table.local_token_indices, vec![0, 0, 1, 1]);
114        assert_eq!(table.local_expert_ids, vec![3, 7, 100, 200]);
115    }
116
117    #[test]
118    fn test_all_remote() {
119        let indices = vec![300u32, 400, 256, 511];
120        let weights = vec![0.6f32, 0.4, 0.55, 0.45];
121        let table = build_ep_routing_table(&indices, &weights, 2, 2, 0, 256);
122
123        assert_eq!(table.local_count(), 0);
124        assert_eq!(table.remote_count(), 4);
125        assert_eq!(table.remote_token_indices, vec![0, 0, 1, 1]);
126        assert_eq!(table.remote_expert_ids, vec![300, 400, 256, 511]);
127    }
128
129    #[test]
130    fn test_mixed_routing() {
131        let indices = vec![
132            10u32, 300, // 2026-09-25: token 0: expert 10 local, 300 remote.
133            255, 256, // 2026-09-25: token 1: expert 255 local, 256 remote.
134            400, 500, // 2026-09-25: token 2: both remote.
135        ];
136        let weights = vec![0.7f32, 0.3, 0.5, 0.5, 0.6, 0.4];
137        let table = build_ep_routing_table(&indices, &weights, 3, 2, 0, 256);
138
139        assert_eq!(table.local_count(), 2);
140        assert_eq!(table.remote_count(), 4);
141        assert_eq!(table.local_token_indices, vec![0, 1]);
142        assert_eq!(table.local_expert_ids, vec![10, 255]);
143        assert_eq!(table.local_weights, vec![0.7, 0.5]);
144        assert_eq!(table.remote_token_indices, vec![0, 1, 2, 2]);
145        assert_eq!(table.remote_expert_ids, vec![300, 256, 400, 500]);
146        assert_eq!(table.remote_weights, vec![0.3, 0.5, 0.6, 0.4]);
147    }
148
149    #[test]
150    fn test_rank1_perspective() {
151        let indices = vec![
152            10u32, 300, // 2026-09-25: token 0: expert 10 remote, 300 local.
153            255, 256, // 2026-09-25: token 1: expert 255 remote, 256 local.
154        ];
155        let weights = vec![0.7f32, 0.3, 0.5, 0.5];
156        let table = build_ep_routing_table(&indices, &weights, 2, 2, 256, 512);
157
158        assert_eq!(table.local_count(), 2);
159        assert_eq!(table.remote_count(), 2);
160        assert_eq!(table.local_token_indices, vec![0, 1]);
161        assert_eq!(table.local_expert_ids, vec![300, 256]);
162        assert_eq!(table.remote_token_indices, vec![0, 1]);
163        assert_eq!(table.remote_expert_ids, vec![10, 255]);
164    }
165
166    #[test]
167    fn test_single_token() {
168        let indices = vec![5u32, 260, 100];
169        let weights = vec![0.5f32, 0.3, 0.2];
170        let table = build_ep_routing_table(&indices, &weights, 1, 3, 0, 256);
171
172        assert_eq!(table.local_count(), 2);
173        assert_eq!(table.remote_count(), 1);
174        assert_eq!(table.total_count(), 3);
175    }
176
177    #[test]
178    #[should_panic(expected = "gate_indices length mismatch")]
179    fn index_length_mismatch_panics() {
180        let indices = vec![1u32, 2, 3];
181        let weights = vec![0.5f32; 4];
182        build_ep_routing_table(&indices, &weights, 2, 2, 0, 256);
183    }
184
185    #[test]
186    #[should_panic(expected = "gate_weights length mismatch")]
187    fn weight_length_mismatch_panics() {
188        let indices = vec![1u32, 2, 3, 4];
189        let weights = vec![0.5f32; 3];
190        build_ep_routing_table(&indices, &weights, 2, 2, 0, 256);
191    }
192}
metrale_model_layers/layers/ep_dispatch.rs

metrale_model_layers/layers/
ep_dispatch.rs
