metrale_model_layers/layers/moe/
helpers_a.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Load-time MoE setup: the pre-expert norm and GeGLU setters, the
4//! routed/shared expert transposes for the prefill kernels (full, gate+up only,
5//! unified, hybrid), and the CUTLASS grouped host tables.
6//!
7//! Owner: model-layers (MoE).
8//! Invariants: none beyond the types.
9
10use super::*;
11
12impl MoeLayer {
13    /// 2026-09-25: Install the pre-expert norm: the router reads the un-normed
14    /// input and the experts read it normed with this weight (`forward`,
15    /// `forward_prefill`). The Gemma-4 loader passes pre_feedforward_layernorm_2.
16    pub fn set_pre_expert_norm(&mut self, norm: crate::weight_map::DenseWeight) {
17        self.pre_expert_norm = Some(norm);
18    }
19
20    /// 2026-09-25: Switch the expert activation kernel `moe_act_mul` to GeGLU
21    /// (`gelu_mul`) and set `gelu_activation`, which only disables the fused
22    /// SiLU+FP8-quant kernel in `forward_prefill_fp8`. Fused decode kernels are
23    /// not steered away by it.
24    pub fn set_gelu_activation(&mut self, gpu: &dyn GpuBackend) -> Result<()> {
25        self.moe_act_mul = gpu.kernel("gelu", "gelu_mul")?;
26        self.gelu_activation = true;
27        Ok(())
28    }
29
30    /// 2026-09-25: Transpose every routed expert's gate, up and down `[N, K/2]`
31    /// weights and scales to `[K/2, N]` (the `*_ptrs_t` tables), and the shared
32    /// expert's when it has one. The originals are kept, so the expert weights
33    /// are held twice.
34    pub fn transpose_for_prefill(
35        &mut self,
36        gpu: &dyn GpuBackend,
37        config: &metrale_config::ModelConfig,
38    ) -> Result<()> {
39        self.transpose_for_prefill_impl(gpu, config, true)
40    }
41
42    /// 2026-09-25: As `transpose_for_prefill`, for gate and up only:
43    /// `down_ptrs_t` stays `None`, so prefill runs the fused transposed gate/up
44    /// kernel and the untransposed grouped down GEMM.
45    pub fn transpose_gate_up_for_prefill(
46        &mut self,
47        gpu: &dyn GpuBackend,
48        config: &metrale_config::ModelConfig,
49    ) -> Result<()> {
50        self.transpose_for_prefill_impl(gpu, config, false)
51    }
52
53    pub(super) fn transpose_for_prefill_impl(
54        &mut self,
55        gpu: &dyn GpuBackend,
56        config: &metrale_config::ModelConfig,
57        include_down: bool,
58    ) -> Result<()> {
59        let h = config.hidden_size;
60        let inter = config.moe_intermediate_size;
61        let shared_inter = config.shared_expert_intermediate_size;
62
63        let num_experts = self.weights.experts.len();
64        let mut gate_t = Vec::with_capacity(num_experts);
65        let mut up_t = Vec::with_capacity(num_experts);
66        let mut down_t = Vec::with_capacity(num_experts);
67
68        // 2026-09-25: The scale block is 32 for E8M0 (native MXFP4) routed experts
69        // and 16 for NVFP4; the scale transpose must use the matching size.
70        let routed_gs =
71            if self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Mxfp4E8m0 {
72                32
73            } else {
74                16
75            };
76        for expert in &self.weights.experts {
77            if expert.gate_proj.is_null() {
78                gate_t.push(QuantizedWeight::null());
79                up_t.push(QuantizedWeight::null());
80                if include_down {
81                    down_t.push(QuantizedWeight::null());
82                }
83            } else {
84                gate_t.push(
85                    expert
86                        .gate_proj
87                        .transpose_for_gemm_gs(gpu, inter, h, routed_gs)?,
88                );
89                up_t.push(
90                    expert
91                        .up_proj
92                        .transpose_for_gemm_gs(gpu, inter, h, routed_gs)?,
93                );
94                if include_down {
95                    down_t.push(
96                        expert
97                            .down_proj
98                            .transpose_for_gemm_gs(gpu, h, inter, routed_gs)?,
99                    );
100                }
101            }
102        }
103
104        self.gate_ptrs_t = Some(build_ptr_table_from_qw(&gate_t, gpu)?);
105        self.up_ptrs_t = Some(build_ptr_table_from_qw(&up_t, gpu)?);
106        if include_down {
107            self.down_ptrs_t = Some(build_ptr_table_from_qw(&down_t, gpu)?);
108        }
109
110        if !self.weights.shared_expert.gate_proj.is_null() && shared_inter > 0 {
111            self.shared_gate_t = Some(self.weights.shared_expert.gate_proj.transpose_for_gemm(
112                gpu,
113                shared_inter,
114                h,
115            )?);
116            self.shared_up_t = Some(self.weights.shared_expert.up_proj.transpose_for_gemm(
117                gpu,
118                shared_inter,
119                h,
120            )?);
121            if include_down {
122                self.shared_down_t =
123                    Some(self.weights.shared_expert.down_proj.transpose_for_gemm(
124                        gpu,
125                        h,
126                        shared_inter,
127                    )?);
128            }
129        }
130
131        Ok(())
132    }
133
134    /// 2026-09-25: Unified layout: transpose gate and up for every expert, free
135    /// their untransposed copies, then the same for down, so only one half is
136    /// ever held twice. Afterwards the `[N, K/2]` kernels cannot run.
137    ///
138    /// The caller sets `METRALE_UNIFIED_MOE_LAYOUT=1` (without the hybrid
139    /// variable) so `use_t_layout_for_decode()` sends decode to the `_t`
140    /// kernels, calls this instead of `transpose_for_prefill` /
141    /// `transpose_gate_up_for_prefill`, and does not build the CUTLASS grouped
142    /// tables afterwards (they point at the freed weights).
143    pub fn transpose_for_prefill_unified(
144        &mut self,
145        gpu: &dyn GpuBackend,
146        config: &metrale_config::ModelConfig,
147    ) -> Result<()> {
148        self.transpose_for_prefill_unified_inner(gpu, config, false)
149    }
150
151    /// 2026-09-25: Hybrid layout: the same build as `transpose_for_prefill_unified`
152    /// but the untransposed originals are kept, so decode and verify stay on the
153    /// `[N, K/2]` kernels (`use_t_layout_for_decode()` is false in hybrid mode).
154    /// The expert weights are held twice; the caller must check that they fit.
155    pub fn transpose_for_prefill_hybrid(
156        &mut self,
157        gpu: &dyn GpuBackend,
158        config: &metrale_config::ModelConfig,
159    ) -> Result<()> {
160        self.transpose_for_prefill_unified_inner(gpu, config, true)
161    }
162
163    /// 2026-09-25: Transpose gate and up (routed and shared), then down. With
164    /// `keep_originals` false, each half's untransposed weights are freed right
165    /// after it is transposed; with true, nothing is freed. An error leaves the
166    /// layer partly transposed.
167    pub(super) fn transpose_for_prefill_unified_inner(
168        &mut self,
169        gpu: &dyn GpuBackend,
170        config: &metrale_config::ModelConfig,
171        keep_originals: bool,
172    ) -> Result<()> {
173        let h = config.hidden_size;
174        let inter = config.moe_intermediate_size;
175        let shared_inter = config.shared_expert_intermediate_size;
176        let _num_experts = self.weights.experts.len();
177
178        let routed_gs =
179            if self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Mxfp4E8m0 {
180                32
181            } else {
182                16
183            };
184        let gate_src: Vec<QuantizedWeight> = self
185            .weights
186            .experts
187            .iter()
188            .map(|e| {
189                if e.gate_proj.is_null() {
190                    QuantizedWeight::null()
191                } else {
192                    e.gate_proj
193                }
194            })
195            .collect();
196        let up_src: Vec<QuantizedWeight> = self
197            .weights
198            .experts
199            .iter()
200            .map(|e| {
201                if e.gate_proj.is_null() {
202                    QuantizedWeight::null()
203                } else {
204                    e.up_proj
205                }
206            })
207            .collect();
208        let gate_t = self.transpose_experts_gpu(gpu, &gate_src, inter, h, routed_gs)?;
209        let up_t = self.transpose_experts_gpu(gpu, &up_src, inter, h, routed_gs)?;
210        self.gate_ptrs_t = Some(build_ptr_table_from_qw(&gate_t, gpu)?);
211        self.up_ptrs_t = Some(build_ptr_table_from_qw(&up_t, gpu)?);
212        if !self.weights.shared_expert.gate_proj.is_null() && shared_inter > 0 {
213            self.shared_gate_t = Some(self.weights.shared_expert.gate_proj.transpose_for_gemm(
214                gpu,
215                shared_inter,
216                h,
217            )?);
218            self.shared_up_t = Some(self.weights.shared_expert.up_proj.transpose_for_gemm(
219                gpu,
220                shared_inter,
221                h,
222            )?);
223        }
224
225        if !keep_originals {
226            // 2026-09-25: The `gate_ptrs`/`up_ptrs` tables keep the freed
227            // addresses. Decode avoids them under `use_t_layout_for_decode()`;
228            // the CUTLASS grouped tables would still read them.
229            for expert in &mut self.weights.experts {
230                if !expert.gate_proj.weight.is_null() {
231                    gpu.free(expert.gate_proj.weight)?;
232                    gpu.free(expert.gate_proj.weight_scale)?;
233                    expert.gate_proj.weight = DevicePtr::NULL;
234                    expert.gate_proj.weight_scale = DevicePtr::NULL;
235                }
236                if !expert.up_proj.weight.is_null() {
237                    gpu.free(expert.up_proj.weight)?;
238                    gpu.free(expert.up_proj.weight_scale)?;
239                    expert.up_proj.weight = DevicePtr::NULL;
240                    expert.up_proj.weight_scale = DevicePtr::NULL;
241                }
242            }
243            if !self.weights.shared_expert.gate_proj.weight.is_null() && shared_inter > 0 {
244                gpu.free(self.weights.shared_expert.gate_proj.weight)?;
245                gpu.free(self.weights.shared_expert.gate_proj.weight_scale)?;
246                self.weights.shared_expert.gate_proj.weight = DevicePtr::NULL;
247                self.weights.shared_expert.gate_proj.weight_scale = DevicePtr::NULL;
248                gpu.free(self.weights.shared_expert.up_proj.weight)?;
249                gpu.free(self.weights.shared_expert.up_proj.weight_scale)?;
250                self.weights.shared_expert.up_proj.weight = DevicePtr::NULL;
251                self.weights.shared_expert.up_proj.weight_scale = DevicePtr::NULL;
252            }
253        }
254
255        let down_src: Vec<QuantizedWeight> = self
256            .weights
257            .experts
258            .iter()
259            .map(|e| {
260                if e.down_proj.is_null() {
261                    QuantizedWeight::null()
262                } else {
263                    e.down_proj
264                }
265            })
266            .collect();
267        let down_t = self.transpose_experts_gpu(gpu, &down_src, h, inter, routed_gs)?;
268        self.down_ptrs_t = Some(build_ptr_table_from_qw(&down_t, gpu)?);
269        if !self.weights.shared_expert.down_proj.is_null() && shared_inter > 0 {
270            self.shared_down_t = Some(self.weights.shared_expert.down_proj.transpose_for_gemm(
271                gpu,
272                h,
273                shared_inter,
274            )?);
275        }
276
277        if !keep_originals {
278            for expert in &mut self.weights.experts {
279                if !expert.down_proj.weight.is_null() {
280                    gpu.free(expert.down_proj.weight)?;
281                    gpu.free(expert.down_proj.weight_scale)?;
282                    expert.down_proj.weight = DevicePtr::NULL;
283                    expert.down_proj.weight_scale = DevicePtr::NULL;
284                }
285            }
286            if !self.weights.shared_expert.down_proj.weight.is_null() && shared_inter > 0 {
287                gpu.free(self.weights.shared_expert.down_proj.weight)?;
288                gpu.free(self.weights.shared_expert.down_proj.weight_scale)?;
289                self.weights.shared_expert.down_proj.weight = DevicePtr::NULL;
290                self.weights.shared_expert.down_proj.weight_scale = DevicePtr::NULL;
291            }
292        }
293
294        Ok(())
295    }
296
297    /// 2026-09-25: Transpose one projection of every routed expert on the GPU with
298    /// `moe_transpose_u8_batched`, into one packed slab and one scale slab.
299    ///
300    /// `src` gives each expert's untransposed `[n, k/2]` packed bytes and
301    /// `[n, k/group_size]` scales. The returned weights point into the slabs and
302    /// copy `weight_scale_2`, `input_scale` and `weight_scale_2_vec` from the
303    /// source. A NULL source gets a NULL slot, which the kernel skips.
304    #[allow(clippy::too_many_arguments)]
305    fn transpose_experts_gpu(
306        &self,
307        gpu: &dyn GpuBackend,
308        src: &[QuantizedWeight],
309        n: usize,
310        k: usize,
311        group_size: usize,
312    ) -> Result<Vec<QuantizedWeight>> {
313        let num_experts = src.len();
314        let packed_each = n * (k / 2);
315        let scale_each = n * (k / group_size);
316        anyhow::ensure!(
317            packed_each > 0 && scale_each > 0,
318            "transpose_experts_gpu: zero-sized projection (n={n} k={k} gs={group_size})"
319        );
320
321        let packed_slab = gpu.alloc(num_experts * packed_each)?;
322        let scale_slab = gpu.alloc(num_experts * scale_each)?;
323
324        let mut out = Vec::with_capacity(num_experts);
325        for (e, w) in src.iter().enumerate() {
326            if w.is_null() {
327                out.push(QuantizedWeight::null());
328            } else {
329                out.push(QuantizedWeight {
330                    weight: packed_slab.offset(e * packed_each),
331                    weight_scale: scale_slab.offset(e * scale_each),
332                    weight_scale_2: w.weight_scale_2,
333                    input_scale: w.input_scale,
334                    weight_scale_2_vec: w.weight_scale_2_vec,
335                    act: w.act,
336                });
337            }
338        }
339
340        let src_tbl = build_ptr_table_from_qw(src, gpu)?;
341        let dst_tbl = build_ptr_table_from_qw(&out, gpu)?;
342        let stream = gpu.default_stream();
343        crate::layers::ops::moe_transpose_u8_batched(
344            gpu,
345            self.moe_transpose_u8_batched_k,
346            src_tbl.packed_ptrs,
347            dst_tbl.packed_ptrs,
348            n as u32,
349            (k / 2) as u32,
350            num_experts as u32,
351            stream,
352        )?;
353        crate::layers::ops::moe_transpose_u8_batched(
354            gpu,
355            self.moe_transpose_u8_batched_k,
356            src_tbl.scale_ptrs,
357            dst_tbl.scale_ptrs,
358            n as u32,
359            (k / group_size) as u32,
360            num_experts as u32,
361            stream,
362        )?;
363        gpu.synchronize(stream)?;
364        gpu.free(src_tbl.packed_ptrs)?;
365        gpu.free(src_tbl.scale_ptrs)?;
366        gpu.free(src_tbl.scale2_vals)?;
367        gpu.free(dst_tbl.packed_ptrs)?;
368        gpu.free(dst_tbl.scale_ptrs)?;
369        gpu.free(dst_tbl.scale2_vals)?;
370        Ok(out)
371    }
372
373    /// 2026-09-25: Build the host tables of the CUTLASS grouped NVFP4 path
374    /// (METRALE_HOLO_MOE_GROUPED_CUTLASS). Per expert, `pack_weight_sfb` swizzles
375    /// the gate/up scales (and down's, when a down scale table exists) into the
376    /// CUTLASS SFB layout, from the transposed `[K/16, N]` scales when
377    /// `gate_ptrs_t`/`up_ptrs_t` exist and from the untransposed `[N, K/16]` ones
378    /// otherwise. The snapshot pairs them with the untransposed packed pointers
379    /// and scale2 values. A no-op when a gate or up scale table is NULL.
380    pub fn build_cutlass_grouped_sfb(
381        &mut self,
382        gpu: &dyn GpuBackend,
383        config: &metrale_config::ModelConfig,
384        stream: u64,
385    ) -> Result<()> {
386        let h = config.hidden_size;
387        let inter = config.moe_intermediate_size;
388        let num = self.weights.experts.len();
389        // 2026-09-25: Swizzled SFB size in bytes: round_up(N, 128) * round_up(K/16, 4).
390        let sfb_len = |n: usize, k: usize| n.div_ceil(128) * 128 * (k / 16).div_ceil(4) * 4;
391        // 2026-09-25: Without transposed scales, read the untransposed [N, K/16]
392        // ones with `src_n_major` set, rather than transposing only for this.
393        let (gate_scale_dev, up_scale_dev, src_n_major) =
394            match (self.gate_ptrs_t.as_ref(), self.up_ptrs_t.as_ref()) {
395                (Some(g), Some(u)) => (g.scale_ptrs, u.scale_ptrs, false),
396                _ => (self.gate_ptrs.scale_ptrs, self.up_ptrs.scale_ptrs, true),
397            };
398        if gate_scale_dev.is_null() || up_scale_dev.is_null() {
399            return Ok(());
400        }
401        let down_scale_dev = match self.down_ptrs_t.as_ref() {
402            Some(d) => Some(d.scale_ptrs),
403            None if !self.down_ptrs.scale_ptrs.is_null() => Some(self.down_ptrs.scale_ptrs),
404            None => None,
405        };
406        let mut owned: Vec<DevicePtr> = Vec::new();
407        // 2026-09-25: `n`/`k` are the projection's GEMM dims (gate/up: inter, hidden;
408        // down: hidden, inter). Returns the host vector of per-expert SFB
409        // pointers, which goes into the snapshot below; a NULL scale pointer
410        // leaves 0.
411        let mut build_one = |scale_ptrs_dev: DevicePtr, n: usize, k: usize| -> Result<Vec<u64>> {
412            let len = sfb_len(n, k);
413            let sp = crate::layers::ops::read_expert_ptrs_u64(gpu, scale_ptrs_dev, num)?;
414            let mut sfb_ptrs = vec![0u64; num];
415            for (e, &sptr) in sp.iter().enumerate() {
416                if sptr == 0 {
417                    continue;
418                }
419                let sfb = gpu.alloc(len)?;
420                metrale_gpu_runtime::cutlass::pack_weight_sfb(
421                    sptr,
422                    sfb.0,
423                    n as u32,
424                    k as u32,
425                    src_n_major,
426                    stream,
427                )?;
428                sfb_ptrs[e] = sfb.0;
429                owned.push(sfb);
430            }
431            gpu.synchronize(stream)?;
432            Ok(sfb_ptrs)
433        };
434        let gate_sfb = build_one(gate_scale_dev, inter, h)?;
435        let up_sfb = build_one(up_scale_dev, inter, h)?;
436        let down = match down_scale_dev {
437            Some(ds) => Some((
438                self.down_ptrs.packed_ptrs,
439                build_one(ds, h, inter)?,
440                self.down_ptrs.scale2_vals,
441            )),
442            None => None,
443        };
444        self.cutlass_grouped_host = Some(crate::layers::ops::MoeCutlassHostTables::snapshot(
445            gpu,
446            num,
447            self.gate_ptrs.packed_ptrs,
448            gate_sfb,
449            self.gate_ptrs.scale2_vals,
450            self.up_ptrs.packed_ptrs,
451            up_sfb,
452            self.up_ptrs.scale2_vals,
453            down,
454        )?);
455        self._cutlass_sfb_owned = owned;
456        tracing::info!(
457            "CUTLASS grouped SFB: built {num} experts gate/up (N={inter} K={h}) + down (N={h} K={inter})"
458        );
459        Ok(())
460    }
461}
metrale_model_layers/layers/moe/helpers_a.rs

metrale_model_layers/layers/moe/
helpers_a.rs
