metrale_model_layers/layers/moe/
forward_ep.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `MoeLayer::forward_ep_dispatch`, an EP token-dispatch entry point
4//! with no caller that serves through `forward`.
5//!
6//! Owner: model-layers (MoE).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11impl MoeLayer {
12    /// 2026-09-25: Route one token with the gate projection and plain
13    /// `moe_topk_softmax` (unlike `forward`, no router-LoRA fold, hash routing or
14    /// correction-bias scoring), copy its top-k to the host after a stream
15    /// synchronize, build an `EpRoutingTable` and log its local/remote counts at
16    /// debug level, then return `self.forward(...)`. The table does not affect
17    /// the output. Refuses a zero-expert router.
18    pub fn forward_ep_dispatch(
19        &self,
20        input: DevicePtr,
21        ctx: &ForwardContext,
22        stream: u64,
23    ) -> Result<DevicePtr> {
24        anyhow::ensure!(
25            self.router_logits_n as usize == ctx.config.num_experts,
26            "zero-expert MoE routing is not wired on this dispatch variant yet (forward_ep)"
27        );
28
29        use super::super::ep_dispatch::build_ep_routing_table;
30
31        let h = ctx.config.hidden_size as u32;
32        let num_experts = ctx.config.num_experts as u32;
33        let top_k = ctx.config.num_experts_per_tok as u32;
34        let (local_start, local_end) = ctx.config.local_expert_range();
35
36        let router_in = self.router_input(input, 1, h, ctx, stream)?;
37        let gate_logits = ctx.buffers.gate_logits();
38        if let Some(ref nvfp4) = self.gate_nvfp4 {
39            ops::w4a16_decode_gemv(
40                ctx.gpu,
41                self.w4a16_gemv,
42                self.w4a16_gemv_sw,
43                ctx.levers.gemv_sw,
44                router_in,
45                nvfp4,
46                gate_logits,
47                num_experts,
48                h,
49                stream,
50            )?;
51        } else {
52            ops::dense_gemv(
53                ctx.gpu,
54                self.dense_gemv,
55                router_in,
56                &self.weights.gate,
57                gate_logits,
58                num_experts,
59                h,
60                stream,
61            )?;
62        }
63
64        let scratch = ctx.buffers.scratch();
65        let indices_dev = scratch;
66        let weights_dev = scratch.offset(top_k as usize * 4);
67
68        ops::moe_topk_softmax(
69            ctx.gpu,
70            self.moe_topk,
71            gate_logits,
72            indices_dev,
73            weights_dev,
74            num_experts,
75            top_k,
76            ctx.config.norm_topk_prob,
77            stream,
78        )?;
79
80        ctx.gpu.synchronize(stream)?;
81        let k = top_k as usize;
82        let mut idx_buf = vec![0u8; k * 4];
83        let mut wt_buf = vec![0u8; k * 4];
84        ctx.gpu.copy_d2h(indices_dev, &mut idx_buf)?;
85        ctx.gpu.copy_d2h(weights_dev, &mut wt_buf)?;
86
87        let gate_indices: Vec<u32> = (0..k)
88            .map(|i| {
89                u32::from_le_bytes([
90                    idx_buf[i * 4],
91                    idx_buf[i * 4 + 1],
92                    idx_buf[i * 4 + 2],
93                    idx_buf[i * 4 + 3],
94                ])
95            })
96            .collect();
97        let gate_weights: Vec<f32> = (0..k)
98            .map(|i| {
99                f32::from_le_bytes([
100                    wt_buf[i * 4],
101                    wt_buf[i * 4 + 1],
102                    wt_buf[i * 4 + 2],
103                    wt_buf[i * 4 + 3],
104                ])
105            })
106            .collect();
107
108        let routing =
109            build_ep_routing_table(&gate_indices, &gate_weights, 1, k, local_start, local_end);
110
111        tracing::debug!(
112            "EP dispatch: local={} remote={} (rank {}, experts {}..{})",
113            routing.local_count(),
114            routing.remote_count(),
115            ctx.config.ep_rank,
116            local_start,
117            local_end,
118        );
119
120        let _ = routing;
121        self.forward(input, ctx, stream)
122    }
123}
metrale_model_layers/layers/moe/forward_ep.rs

metrale_model_layers/layers/moe/
forward_ep.rs
