spark_model/layers/moe/
forward_ep.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! MoeLayer::forward_ep_dispatch.
4
5use super::*;
6
7impl MoeLayer {
8    /// 4. Computes local experts on local + received tokens
9    /// 5. Sends results back (combine)
10    /// 6. Weighted sum into output
11    ///
12    /// Currently scaffolding only — builds routing table and logs statistics.
13    /// Expert compute and actual dispatch use the existing per-token path.
14    /// The all-reduce fallback is used for the actual output until dispatch
15    /// kernels are implemented.
16    pub fn forward_ep_dispatch(
17        &self,
18        input: DevicePtr,
19        ctx: &ForwardContext,
20        stream: u64,
21    ) -> Result<DevicePtr> {
22        // LongCat zero-experts are wired only on the single-token decode
23        // + prefill paths (v1); this variant would silently mis-route the
24        // 384-wide router. Named refusal, not silent wrongness.
25        anyhow::ensure!(
26            self.router_logits_n as usize == ctx.config.num_experts,
27            "zero-expert MoE routing is not wired on this dispatch variant yet (forward_ep)"
28        );
29
30        use super::super::ep_dispatch::build_ep_routing_table;
31
32        let h = ctx.config.hidden_size as u32;
33        let num_experts = ctx.config.num_experts as u32;
34        let top_k = ctx.config.num_experts_per_tok as u32;
35        let (local_start, local_end) = ctx.config.local_expert_range();
36
37        // Gemma-4 router pre-norm (no-op for other models). EP dispatch is
38        // per-token so uses num_tokens=1.
39        let router_in = self.router_input(input, 1, h, ctx, stream)?;
40        // Step 1: Gate projection (same as forward())
41        let gate_logits = ctx.buffers.gate_logits();
42        if let Some(ref nvfp4) = self.gate_nvfp4 {
43            ops::w4a16_decode_gemv(
44                ctx.gpu,
45                self.w4a16_gemv,
46                self.w4a16_gemv_sw,
47                ctx.levers.gemv_sw,
48                router_in,
49                nvfp4,
50                gate_logits,
51                num_experts,
52                h,
53                stream,
54            )?;
55        } else {
56            ops::dense_gemv(
57                ctx.gpu,
58                self.dense_gemv,
59                router_in,
60                &self.weights.gate,
61                gate_logits,
62                num_experts,
63                h,
64                stream,
65            )?;
66        }
67
68        // Step 2: Top-K routing
69        let scratch = ctx.buffers.scratch();
70        let indices_dev = scratch;
71        let weights_dev = scratch.offset(top_k as usize * 4);
72
73        ops::moe_topk_softmax(
74            ctx.gpu,
75            self.moe_topk,
76            gate_logits,
77            indices_dev,
78            weights_dev,
79            num_experts,
80            top_k,
81            ctx.config.norm_topk_prob,
82            stream,
83        )?;
84
85        // Step 3: Build routing table (requires D2H copy of indices/weights)
86        // This is CPU-side work — acceptable for scaffolding, will move to
87        // GPU-side routing table construction in the optimized path.
88        ctx.gpu.synchronize(stream)?;
89        let k = top_k as usize;
90        let mut idx_buf = vec![0u8; k * 4];
91        let mut wt_buf = vec![0u8; k * 4];
92        ctx.gpu.copy_d2h(indices_dev, &mut idx_buf)?;
93        ctx.gpu.copy_d2h(weights_dev, &mut wt_buf)?;
94
95        let gate_indices: Vec<u32> = (0..k)
96            .map(|i| {
97                u32::from_le_bytes([
98                    idx_buf[i * 4],
99                    idx_buf[i * 4 + 1],
100                    idx_buf[i * 4 + 2],
101                    idx_buf[i * 4 + 3],
102                ])
103            })
104            .collect();
105        let gate_weights: Vec<f32> = (0..k)
106            .map(|i| {
107                f32::from_le_bytes([
108                    wt_buf[i * 4],
109                    wt_buf[i * 4 + 1],
110                    wt_buf[i * 4 + 2],
111                    wt_buf[i * 4 + 3],
112                ])
113            })
114            .collect();
115
116        let routing =
117            build_ep_routing_table(&gate_indices, &gate_weights, 1, k, local_start, local_end);
118
119        tracing::debug!(
120            "EP dispatch: local={} remote={} (rank {}, experts {}..{})",
121            routing.local_count(),
122            routing.remote_count(),
123            ctx.config.ep_rank,
124            local_start,
125            local_end,
126        );
127
128        // Steps 4-6: For now, fall back to the existing forward() path
129        // which uses all-reduce. The routing table is built but not yet
130        // used for actual dispatch. This will be replaced with:
131        //   - comm.group_start()
132        //   - comm.send_to() for remote tokens
133        //   - comm.recv_from() for incoming tokens
134        //   - comm.group_end()
135        //   - Local expert compute on local + received tokens
136        //   - comm.group_start()
137        //   - comm.send_to() results back
138        //   - comm.recv_from() results from partner
139        //   - comm.group_end()
140        //   - Weighted sum into output
141        let _ = routing; // suppress unused warning until dispatch is wired
142        self.forward(input, ctx, stream)
143    }
144}
spark_model/layers/moe/forward_ep.rs

spark_model/layers/moe/
forward_ep.rs
