metrale_model_layers/layers/moe/
fp8_grouped_tc_w8a8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The opt-in W8A8 expert step of the grouped FP8 MoE decode
4//! (`moe_fp8_grouped_tc_w8a8.cu`): the layer input and the SiLU products are quantized to
5//! E4M3 per (row, 128 group), the checkpoint's declared `activation_scheme: dynamic`, and the
6//! FP8 weights go into E4M3 tensor-core MMAs undecoded.
7//!
8//! WHEN it runs is the weight-quantization policy's call, not this module's: the serve publishes
9//! whether the experts decode with FP8 activations ([`set_moe_expert_fp8_act`], the
10//! `--weight-quantization` policy's `fp8_decode_act` for the expert modules, published at serve
11//! setup); unpublished, the experts stay W8A16. This module says whether it CAN
12//! ([`MoeLayer::fp8_grouped_tc_w8a8_on`]: the tensor-core path is on, the three W8A8 kernels
13//! resolved and the shapes fit). Its output differs from the W8A16 kernels' by the activation
14//! rounding; a row's bits still do not depend on the other rows.
15//!
16//! Owner: model-layers (MoE).
17//! Invariants:
18//! - The first publication or read of the cell wins (`OnceLock`); the serve publishes before
19//!   the model runs, and anything that reads first fixes it at W8A16.
20//! - [`MoeLayer::run_fp8_grouped_w8a8`] launches nothing unless
21//!   [`MoeLayer::fp8_grouped_tc_w8a8_on`] holds and the buffers hold the layout.
22
23use super::*;
24
25static FP8_ACT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26
27/// 2026-09-28: Publish whether the MoE experts decode with FP8 activations (W8A8), as the
28/// weight-quantization policy decides for the expert modules. Returns the value in force; a
29/// caller that gets a different one should warn.
30pub fn set_moe_expert_fp8_act(fp8: bool) -> bool {
31    let _ = FP8_ACT.set(fp8);
32    *FP8_ACT.get().expect("just set")
33}
34
35/// 2026-09-28: Whether the MoE experts decode with FP8 activations: false (W8A16) unless the
36/// serve published otherwise.
37pub fn moe_expert_fp8_act() -> bool {
38    *FP8_ACT.get_or_init(|| false)
39}
40
41/// 2026-09-28: One grouped decode's inputs and outputs for the expert step.
42pub(super) struct GroupedExpertIo {
43    pub input: DevicePtr,
44    pub act: DevicePtr,
45    pub shared_act: DevicePtr,
46    pub expert_down_out: DevicePtr,
47    pub shared_out: DevicePtr,
48    pub rows: ops::Fp8GroupedW8a8Rows,
49}
50
51impl MoeLayer {
52    /// 2026-09-28: Whether this layer's grouped decode runs the W8A8 expert step: FP8 expert
53    /// activations published ([`moe_expert_fp8_act`]), the tensor-core W8A16 kernels'
54    /// conditions, the three W8A8 kernels, and gate+up and down in whole CTAs of their tiles.
55    pub(super) fn fp8_grouped_tc_w8a8_on(&self, hidden: usize, inter: usize) -> bool {
56        let (h, i) = (hidden as u32, inter as u32);
57        let k = &self.fp8_grouped_tc;
58        moe_expert_fp8_act()
59            && self.fp8_grouped_tc_on(hidden, inter)
60            && k.quant_w8a8.0 != 0
61            && k.gate_up_w8a8.0 != 0
62            && k.down_w8a8.0 != 0
63            && ops::fp8_grouped_tc_shape_ok(i, h, ops::FP8_GROUPED_GATE_UP_TC_W8A8)
64            && ops::fp8_grouped_tc_shape_ok(h, i, ops::FP8_GROUPED_DOWN_TC_W8A8)
65    }
66
67    /// 2026-09-28: Quantize the layer input, then W8A8 gate+up (quantized SiLU products) and
68    /// down, routed and shared, into `io.expert_down_out` / `io.shared_out`. The quantized
69    /// activations live in the SiLU buffers (`ops::Fp8GroupedW8a8Layout`).
70    pub(super) fn run_fp8_grouped_w8a8(
71        &self,
72        io: &GroupedExpertIo,
73        ctx: &ForwardContext,
74        stream: u64,
75    ) -> Result<()> {
76        let (Some(gp), Some(up), Some(dp), Some(sh)) = (
77            &self.fp8_gate_weight_ptrs,
78            &self.fp8_up_weight_ptrs,
79            &self.fp8_down_weight_ptrs,
80            &self.fp8_shared_expert,
81        ) else {
82            anyhow::bail!("W8A8 grouped decode: the FP8 expert tables are missing");
83        };
84        let cfg = ctx.config;
85        let (h, inter) = (cfg.hidden_size, cfg.moe_intermediate_size);
86        let m = io.rows.num_tokens as usize;
87        let lay = ops::Fp8GroupedW8a8Layout::new(
88            m,
89            cfg.num_experts_per_tok,
90            h,
91            inter,
92            ctx.buffers.expert_gate_out_bytes(),
93            ctx.buffers.logits_bytes(),
94        )?;
95        let k = &self.fp8_grouped_tc;
96        let (xq, xs) = (io.act.offset(lay.xq), io.act.offset(lay.xs));
97        let act = (io.act, io.act.offset(lay.act_s));
98        let sh_act = (io.shared_act, io.shared_act.offset(lay.sh_s));
99        let (h32, i32) = (h as u32, inter as u32);
100        ops::moe_act_quant_e4m3(
101            ctx.gpu,
102            k.quant_w8a8,
103            io.input,
104            xq,
105            xs,
106            m as u32,
107            h32,
108            stream,
109        )?;
110        ops::moe_expert_gate_up_act_fp8_grouped_tc_w8a8(
111            ctx.gpu,
112            k.gate_up_w8a8,
113            xq,
114            xs,
115            (gp.weight_ptrs, gp.scale_ptrs),
116            (up.weight_ptrs, up.scale_ptrs),
117            act,
118            &io.rows,
119            &sh.gate_proj,
120            &sh.up_proj,
121            sh_act,
122            i32,
123            h32,
124            stream,
125        )?;
126        ops::moe_expert_down_act_fp8_grouped_tc_w8a8(
127            ctx.gpu,
128            k.down_w8a8,
129            act,
130            (dp.weight_ptrs, dp.scale_ptrs),
131            io.expert_down_out,
132            &io.rows,
133            sh_act,
134            &sh.down_proj,
135            io.shared_out,
136            h32,
137            i32,
138            stream,
139        )
140    }
141}
metrale_model_layers/layers/moe/fp8_grouped_tc_w8a8.rs

metrale_model_layers/layers/moe/
fp8_grouped_tc_w8a8.rs
