metrale_model_layers/layers/moe/
helpers_b.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The lazy down_proj scratch transpose, the layout predicates
4//! (`use_t_layout_for_prefill`, `use_t_layout_for_decode`, `grouped_decode_ok`)
5//! and the routed grouped-GEMM launcher.
6//!
7//! Owner: model-layers (MoE).
8//! Invariants: none beyond the types.
9
10use super::*;
11
12impl MoeLayer {
13    /// 2026-09-25: Wire the shared down_proj scratch and its transposed pointer
14    /// tables. The factory (`m2_setup`) calls this for every layer after the
15    /// gate+up-only transpose; one scratch serves all MoE layers, refilled by
16    /// each layer's prefill. `scale2_vals` is the untransposed `down_ptrs` table.
17    pub fn set_down_transpose_scratch(
18        &mut self,
19        scratch_packed: DevicePtr,
20        scratch_scale: DevicePtr,
21        packed_ptrs_t: DevicePtr,
22        scale_ptrs_t: DevicePtr,
23    ) {
24        self.down_t_scratch_packed = Some(scratch_packed);
25        self.down_t_scratch_scale = Some(scratch_scale);
26        self.down_ptrs_t = Some(ExpertPtrTable {
27            packed_ptrs: packed_ptrs_t,
28            scale_ptrs: scale_ptrs_t,
29            scale2_vals: self.down_ptrs.scale2_vals,
30        });
31    }
32
33    /// 2026-09-25: Fill the down scratch from this layer's untransposed
34    /// `down_ptrs` (`[hidden, inter/2]` packed and `[hidden, inter/16]` scales,
35    /// transposed per expert). `forward_prefill` calls it before the routed
36    /// GEMMs. A no-op unless the scratch is wired.
37    pub(crate) fn transpose_down_into_scratch(
38        &self,
39        ctx: &crate::layer::ForwardContext,
40        stream: u64,
41    ) -> Result<()> {
42        let Some(dpt) = self.down_ptrs_t.as_ref() else {
43            return Ok(());
44        };
45        if self.down_t_scratch_packed.is_none() {
46            return Ok(());
47        }
48        let num_experts = ctx.config.num_experts as u32;
49        let h = ctx.config.hidden_size as u32;
50        let inter = ctx.config.moe_intermediate_size as u32;
51        crate::layers::ops::moe_transpose_u8_batched(
52            ctx.gpu,
53            self.moe_transpose_u8_batched_k,
54            self.down_ptrs.packed_ptrs,
55            dpt.packed_ptrs,
56            h,
57            inter / 2,
58            num_experts,
59            stream,
60        )?;
61        crate::layers::ops::moe_transpose_u8_batched(
62            ctx.gpu,
63            self.moe_transpose_u8_batched_k,
64            self.down_ptrs.scale_ptrs,
65            dpt.scale_ptrs,
66            h,
67            inter / 16,
68            num_experts,
69            stream,
70        )?;
71        Ok(())
72    }
73
74    /// 2026-09-25: Unused. The down scratch transpose on `prefill_stream`, after
75    /// `compute_stream` reaches `event_a`, recording `event_b` when done.
76    #[allow(dead_code)]
77    pub(crate) fn kick_off_lazy_transpose(
78        &self,
79        ctx: &crate::layer::ForwardContext,
80        compute_stream: u64,
81    ) -> Result<()> {
82        let Some(dpt) = self.down_ptrs_t.as_ref() else {
83            return Ok(());
84        };
85        if self.down_t_scratch_packed.is_none() {
86            return Ok(());
87        }
88        ctx.gpu.record_event(self.event_a, compute_stream)?;
89        ctx.gpu
90            .stream_wait_event(self.prefill_stream, self.event_a)?;
91
92        let num_experts = ctx.config.num_experts as u32;
93        let h = ctx.config.hidden_size as u32;
94        let inter = ctx.config.moe_intermediate_size as u32;
95        crate::layers::ops::moe_transpose_u8_batched(
96            ctx.gpu,
97            self.moe_transpose_u8_batched_k,
98            self.down_ptrs.packed_ptrs,
99            dpt.packed_ptrs,
100            h,
101            inter / 2,
102            num_experts,
103            self.prefill_stream,
104        )?;
105        crate::layers::ops::moe_transpose_u8_batched(
106            ctx.gpu,
107            self.moe_transpose_u8_batched_k,
108            self.down_ptrs.scale_ptrs,
109            dpt.scale_ptrs,
110            h,
111            inter / 16,
112            num_experts,
113            self.prefill_stream,
114        )?;
115        ctx.gpu.record_event(self.event_b, self.prefill_stream)?;
116        Ok(())
117    }
118
119    /// 2026-09-25: Unused. Whether the down scratch is wired.
120    #[allow(dead_code)]
121    pub(crate) fn has_overlapped_transpose(&self) -> bool {
122        self.down_t_scratch_packed.is_some()
123    }
124
125    /// 2026-09-25: Unused. The event `kick_off_lazy_transpose` records.
126    #[allow(dead_code)]
127    pub(crate) fn lazy_transpose_done_event(&self) -> u64 {
128        self.event_b
129    }
130
131    /// 2026-09-25: True when `forward_batched` uses the transposed `_t` kernels:
132    /// unified or hybrid layout (METRALE_UNIFIED_MOE_LAYOUT / METRALE_HYBRID_MOE_LAYOUT,
133    /// read at construction), transposed tables for all three projections, and
134    /// no lazy down scratch (it holds one layer's down weights and only
135    /// `forward_prefill` fills it).
136    #[inline]
137    pub(crate) fn use_t_layout_for_prefill(&self) -> bool {
138        (self.unified_layout || self.hybrid_layout)
139            && self.gate_ptrs_t.is_some()
140            && self.up_ptrs_t.is_some()
141            && self.down_ptrs_t.is_some()
142            && self.down_t_scratch_packed.is_none()
143    }
144
145    #[inline]
146    /// 2026-09-25: Whether decode may take the grouped `forward_prefill` path: no
147    /// BF16 or FP8 expert tables and no hash routing (`tid2eid_dev`). A mixed
148    /// BF16 shared expert is allowed (`forward_prefill` runs it separately). EP
149    /// is not checked here.
150    pub(crate) fn grouped_decode_ok(&self) -> bool {
151        self.bf16_gate_weight_ptrs.is_none()
152            && self.fp8_gate_weight_ptrs.is_none()
153            && self.tid2eid_dev.is_none()
154    }
155
156    /// 2026-09-25: True when decode (`forward`, `forward_k2`, `forward_k3`) uses the
157    /// transposed `_t` kernels: unified layout without hybrid, transposed tables
158    /// for all three projections, and no lazy down scratch.
159    pub(crate) fn use_t_layout_for_decode(&self) -> bool {
160        self.unified_layout
161            && !self.hybrid_layout
162            && self.gate_ptrs_t.is_some()
163            && self.up_ptrs_t.is_some()
164            && self.down_ptrs_t.is_some()
165            && self.down_t_scratch_packed.is_none()
166    }
167}
168
169impl super::MoeLayer {
170    /// 2026-09-25: The routed grouped-GEMM kernel: the wider-K twin when it
171    /// resolved (METRALE_MOE_GROUPED_K32=1 and the target ships it), else the
172    /// base kernel. Both launch through `moe_w4a16_grouped_gemm_ptrtable`.
173    pub(super) fn grouped_gemm_kernel(&self) -> metrale_gpu_runtime::gpu::KernelHandle {
174        if self.moe_grouped_gemm_k32.0 != 0 {
175            self.moe_grouped_gemm_k32
176        } else {
177            self.moe_grouped_gemm
178        }
179    }
180}
181
182impl super::MoeLayer {
183    /// 2026-09-25: Launch the routed grouped GEMM. With the M_TILE=256 kernel
184    /// (METRALE_MOE_GROUPED_M256=1 and resolved) the M=64 tile count
185    /// `max_m_tiles` is divided by 4, rounded up; otherwise `grouped_gemm_kernel()`.
186    ///
187    /// Both wide kernels are opt-in. Measured 2026-08-27 on GB10 (qwen4_exp, 28K
188    /// prefill, back to back): k32 +1.0% and m256 +1.4% over the arm without them,
189    /// inside run-to-run noise, and m256 averaged 31.30 ms per call against 26.17 ms for the base kernel.
190    #[allow(clippy::too_many_arguments)]
191    pub(super) fn launch_grouped_gemm(
192        &self,
193        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
194        a: metrale_gpu_runtime::gpu::DevicePtr,
195        packed_ptrs: metrale_gpu_runtime::gpu::DevicePtr,
196        scale_ptrs: metrale_gpu_runtime::gpu::DevicePtr,
197        scale2_vals: metrale_gpu_runtime::gpu::DevicePtr,
198        c: metrale_gpu_runtime::gpu::DevicePtr,
199        expert_offsets: metrale_gpu_runtime::gpu::DevicePtr,
200        sorted_token_ids: metrale_gpu_runtime::gpu::DevicePtr,
201        num_experts: u32,
202        n_out: u32,
203        k: u32,
204        max_m_tiles: u32,
205        stream: u64,
206    ) -> anyhow::Result<()> {
207        if self.moe_grouped_gemm_m256.0 != 0 {
208            return crate::layers::ops::moe_w4a16_grouped_gemm_ptrtable_m256(
209                gpu,
210                self.moe_grouped_gemm_m256,
211                a,
212                packed_ptrs,
213                scale_ptrs,
214                scale2_vals,
215                c,
216                expert_offsets,
217                sorted_token_ids,
218                num_experts,
219                n_out,
220                k,
221                max_m_tiles.div_ceil(4).max(1),
222                stream,
223            );
224        }
225        crate::layers::ops::moe_w4a16_grouped_gemm_ptrtable(
226            gpu,
227            self.grouped_gemm_kernel(),
228            a,
229            packed_ptrs,
230            scale_ptrs,
231            scale2_vals,
232            c,
233            expert_offsets,
234            sorted_token_ids,
235            num_experts,
236            n_out,
237            k,
238            max_m_tiles,
239            stream,
240        )
241    }
242}
metrale_model_layers/layers/moe/helpers_b.rs

metrale_model_layers/layers/moe/
helpers_b.rs
