metrale_model_layers/layers/ops/
fp8_moe_grouped_tc_w8a8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Launchers for the W8A8 tensor-core grouped FP8 MoE decode
4//! (`kernels/gb10/common/moe_fp8_grouped_tc_w8a8.cu`): the per-(row, 128) E4M3 activation
5//! quantization, gate+up with a re-quantized SiLU product, and down; plus where their
6//! quantized activations live inside the grouped decode's two SiLU buffers.
7//!
8//! Owner: model-layers ops.
9//! Invariants: [`Fp8GroupedW8a8Layout::new`] refuses buffers that cannot hold the layout.
10
11use anyhow::{Result, ensure};
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use super::{FP8_GROUPED_TC_ROWS_PER_PASS, Fp8GroupedGeometry};
16use crate::weight_map::Fp8Weight;
17
18/// 2026-09-28: `moe_expert_gate_up_act_fp8_grouped_tc_w8a8` (`TC8_GU_COLS`: one 128-column
19/// quantization group of the down projection's K per CTA).
20pub const FP8_GROUPED_GATE_UP_TC_W8A8: Fp8GroupedGeometry = Fp8GroupedGeometry {
21    cols_per_cta: 128,
22    rows_per_pass: FP8_GROUPED_TC_ROWS_PER_PASS,
23    threads: 128,
24};
25
26/// 2026-09-28: `moe_expert_down_act_fp8_grouped_tc_w8a8` (`TC8_DOWN_COLS`).
27pub const FP8_GROUPED_DOWN_TC_W8A8: Fp8GroupedGeometry = Fp8GroupedGeometry {
28    cols_per_cta: 256,
29    rows_per_pass: FP8_GROUPED_TC_ROWS_PER_PASS,
30    threads: 128,
31};
32
33fn align16(x: usize) -> usize {
34    x.div_ceil(16) * 16
35}
36
37/// 2026-09-28: Byte offsets of the W8A8 activations inside the grouped decode's buffers.
38/// The routed SiLU buffer (`expert_gate_out`, sized for `[te, inter]` FP32) holds the routed
39/// E4M3 products `[te, inter]`, their scales `[te, inter / 128]` FP32, the quantized layer
40/// input `[m, hidden]` E4M3 and its scales `[m, hidden / 128]`; the shared SiLU buffer
41/// (`logits`, `[m, inter]` FP32) holds the shared products and their scales.
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43pub struct Fp8GroupedW8a8Layout {
44    pub act_s: usize,
45    pub xq: usize,
46    pub xs: usize,
47    pub sh_s: usize,
48}
49
50impl Fp8GroupedW8a8Layout {
51    /// 2026-09-28: The layout for `m` rows, or an error when `hidden` or `inter` is not whole
52    /// 128 groups or either buffer is too small.
53    pub fn new(
54        m: usize,
55        top_k: usize,
56        hidden: usize,
57        inter: usize,
58        routed_bytes: usize,
59        shared_bytes: usize,
60    ) -> Result<Self> {
61        ensure!(
62            hidden.is_multiple_of(128) && inter.is_multiple_of(128),
63            "W8A8 grouped decode: hidden {hidden} and inter {inter} must be whole 128 groups"
64        );
65        let te = m * top_k;
66        let act_s = align16(te * inter);
67        let xq = align16(act_s + te * (inter / 128) * 4);
68        let xs = align16(xq + m * hidden);
69        let routed_end = xs + m * (hidden / 128) * 4;
70        let sh_s = align16(m * inter);
71        let shared_end = sh_s + m * (inter / 128) * 4;
72        ensure!(
73            routed_end <= routed_bytes && shared_end <= shared_bytes,
74            "W8A8 grouped decode: needs {routed_end} + {shared_end} bytes, buffers hold {routed_bytes} + {shared_bytes}"
75        );
76        Ok(Self {
77            act_s,
78            xq,
79            xs,
80            sh_s,
81        })
82    }
83}
84
85/// 2026-09-28: E4M3 per (row, 128-K group) of `x` `[rows, k]` BF16 into `q` / `s`.
86#[allow(clippy::too_many_arguments)]
87pub fn moe_act_quant_e4m3(
88    gpu: &dyn GpuBackend,
89    kernel: KernelHandle,
90    x: DevicePtr,
91    q: DevicePtr,
92    s: DevicePtr,
93    rows: u32,
94    k: u32,
95    stream: u64,
96) -> Result<()> {
97    KernelLaunch::new(gpu, kernel)
98        .grid([rows, 1, 1])
99        .block([128, 1, 1])
100        .arg_ptr(x)
101        .arg_ptr(q)
102        .arg_ptr(s)
103        .arg_u32(k)
104        .launch(stream)
105}
106
107/// 2026-09-28: The routed expert tables and the grouped decode's sort outputs the two W8A8
108/// kernels read.
109pub struct Fp8GroupedW8a8Rows {
110    pub expert_offsets: DevicePtr,
111    pub sorted_token_ids: DevicePtr,
112    pub active_experts: DevicePtr,
113    pub active_count: DevicePtr,
114    pub cap: u32,
115    pub num_tokens: u32,
116}
117
118/// 2026-09-28: W8A8 gate+up and SiLU: `(xq, xs)` the quantized layer input; writes the
119/// quantized products `(act_q, act_s)` by sorted position and `(sh_q, sh_s)` by token.
120#[allow(clippy::too_many_arguments)]
121pub fn moe_expert_gate_up_act_fp8_grouped_tc_w8a8(
122    gpu: &dyn GpuBackend,
123    kernel: KernelHandle,
124    xq: DevicePtr,
125    xs: DevicePtr,
126    gate: (DevicePtr, DevicePtr),
127    up: (DevicePtr, DevicePtr),
128    act: (DevicePtr, DevicePtr),
129    rows: &Fp8GroupedW8a8Rows,
130    sh_gate: &Fp8Weight,
131    sh_up: &Fp8Weight,
132    sh: (DevicePtr, DevicePtr),
133    n: u32,
134    k: u32,
135    stream: u64,
136) -> Result<()> {
137    let g = FP8_GROUPED_GATE_UP_TC_W8A8;
138    KernelLaunch::new(gpu, kernel)
139        .grid([
140            div_ceil(n, g.cols_per_cta),
141            rows.cap + div_ceil(rows.num_tokens, g.rows_per_pass),
142            1,
143        ])
144        .block([g.threads, 1, 1])
145        .arg_ptr(xq)
146        .arg_ptr(xs)
147        .arg_ptr(gate.0)
148        .arg_ptr(gate.1)
149        .arg_ptr(up.0)
150        .arg_ptr(up.1)
151        .arg_ptr(act.0)
152        .arg_ptr(act.1)
153        .arg_ptr(rows.expert_offsets)
154        .arg_ptr(rows.sorted_token_ids)
155        .arg_ptr(rows.active_experts)
156        .arg_ptr(rows.active_count)
157        .arg_ptr(sh_gate.weight)
158        .arg_ptr(sh_gate.row_scale)
159        .arg_ptr(sh_up.weight)
160        .arg_ptr(sh_up.row_scale)
161        .arg_ptr(sh.0)
162        .arg_ptr(sh.1)
163        .arg_u32(n)
164        .arg_u32(k)
165        .arg_u32(rows.cap)
166        .arg_u32(rows.num_tokens)
167        .launch(stream)
168}
169
170/// 2026-09-28: W8A8 down over the quantized SiLU products: routed rows into `output` by
171/// sorted position, the shared rows into `sh_down_out` by token, both BF16.
172#[allow(clippy::too_many_arguments)]
173pub fn moe_expert_down_act_fp8_grouped_tc_w8a8(
174    gpu: &dyn GpuBackend,
175    kernel: KernelHandle,
176    act: (DevicePtr, DevicePtr),
177    down: (DevicePtr, DevicePtr),
178    output: DevicePtr,
179    rows: &Fp8GroupedW8a8Rows,
180    sh: (DevicePtr, DevicePtr),
181    sh_down: &Fp8Weight,
182    sh_down_out: DevicePtr,
183    n: u32,
184    k: u32,
185    stream: u64,
186) -> Result<()> {
187    let g = FP8_GROUPED_DOWN_TC_W8A8;
188    KernelLaunch::new(gpu, kernel)
189        .grid([
190            div_ceil(n, g.cols_per_cta),
191            rows.cap + div_ceil(rows.num_tokens, g.rows_per_pass),
192            1,
193        ])
194        .block([g.threads, 1, 1])
195        .arg_ptr(act.0)
196        .arg_ptr(act.1)
197        .arg_ptr(down.0)
198        .arg_ptr(down.1)
199        .arg_ptr(output)
200        .arg_ptr(rows.expert_offsets)
201        .arg_ptr(rows.active_experts)
202        .arg_ptr(rows.active_count)
203        .arg_ptr(sh.0)
204        .arg_ptr(sh.1)
205        .arg_ptr(sh_down.weight)
206        .arg_ptr(sh_down.row_scale)
207        .arg_ptr(sh_down_out)
208        .arg_u32(n)
209        .arg_u32(k)
210        .arg_u32(rows.cap)
211        .arg_u32(rows.num_tokens)
212        .launch(stream)
213}
metrale_model_layers/layers/ops/fp8_moe_grouped_tc_w8a8.rs

metrale_model_layers/layers/ops/
fp8_moe_grouped_tc_w8a8.rs
