metrale_model_layers/layers/ops/
fp8_moe_grouped.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the cross-row grouped FP8 MoE decode kernels
4//! (`moe_shared_expert_fused_fp8_grouped.cu`, `moe_fp8_grouped_blend.cu`).
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8//!
9//! Rows are grouped by expert (`moe_sort_by_expert`). A gate+up or SiLU+down
10//! CTA owns one active expert and applies each weight row it streams to every
11//! row routed to that expert; the batch-2/3 kernels put one (token, slot) pair
12//! on each `blockIdx.y` instead. Intermediates are laid out by sorted
13//! position, and the blend maps each slot back through `token_to_perm`.
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
18
19use crate::weight_map::Fp8Weight;
20
21/// 2026-09-26: Output columns per gate+up CTA. Must equal `GU_COLS_PER_CTA` in the
22/// `.cu`.
23pub const FP8_GROUPED_GATE_UP_COLS_PER_CTA: u32 = 8;
24
25/// 2026-09-25: Output columns per down CTA. Must equal `DOWN_CTA_COLS` in the
26/// `.cu`: 2026-09-27, four column groups of 32.
27pub const FP8_GROUPED_DOWN_COLS_PER_CTA: u32 = 128;
28
29/// 2026-09-26: Rows per gate+up pass. Must equal `GU_GROUP_ROWS` in the `.cu`:
30/// the shared expert takes `ceil(num_tokens / this)` block rows.
31pub const FP8_GROUPED_GATE_UP_ROWS_PER_PASS: u32 = 4;
32
33/// 2026-09-26: Rows per down pass. Must equal `GROUP_ROWS` in the `.cu`.
34pub const FP8_GROUPED_DOWN_ROWS_PER_PASS: u32 = 4;
35
36/// 2026-09-28: Rows per pass of the tensor-core kernels (`moe_fp8_grouped_tc.cu`,
37/// `TC_ROWS`), gate+up and down alike.
38pub const FP8_GROUPED_TC_ROWS_PER_PASS: u32 = 8;
39
40/// 2026-09-28: Output columns per tensor-core gate+up CTA (`TC_GU_COLS`).
41pub const FP8_GROUPED_TC_GATE_UP_COLS_PER_CTA: u32 = 64;
42
43/// 2026-09-28: Output columns per tensor-core down CTA (`TC_DOWN_COLS`).
44pub const FP8_GROUPED_TC_DOWN_COLS_PER_CTA: u32 = 128;
45
46/// 2026-09-28: The launch shape of a grouped gate+up or down kernel: output columns per
47/// CTA (`grid.x = ceil(n / cols_per_cta)`), rows per pass (the shared expert takes
48/// `ceil(num_tokens / rows_per_pass)` block rows) and threads per CTA.
49#[derive(Clone, Copy, Debug, PartialEq, Eq)]
50pub struct Fp8GroupedGeometry {
51    pub cols_per_cta: u32,
52    pub rows_per_pass: u32,
53    pub threads: u32,
54}
55
56/// 2026-09-28: `moe_expert_gate_up_act_fp8_grouped` (FP32 SiLU products).
57pub const FP8_GROUPED_GATE_UP_SCALAR: Fp8GroupedGeometry = Fp8GroupedGeometry {
58    cols_per_cta: FP8_GROUPED_GATE_UP_COLS_PER_CTA,
59    rows_per_pass: FP8_GROUPED_GATE_UP_ROWS_PER_PASS,
60    threads: 128,
61};
62
63/// 2026-09-28: `moe_expert_down_act_fp8_grouped` (reads FP32 SiLU products).
64pub const FP8_GROUPED_DOWN_SCALAR: Fp8GroupedGeometry = Fp8GroupedGeometry {
65    cols_per_cta: FP8_GROUPED_DOWN_COLS_PER_CTA,
66    rows_per_pass: FP8_GROUPED_DOWN_ROWS_PER_PASS,
67    threads: 256,
68};
69
70/// 2026-09-28: `moe_expert_gate_up_act_fp8_grouped_tc` (FP32 SiLU products).
71pub const FP8_GROUPED_GATE_UP_TC: Fp8GroupedGeometry = Fp8GroupedGeometry {
72    cols_per_cta: FP8_GROUPED_TC_GATE_UP_COLS_PER_CTA,
73    rows_per_pass: FP8_GROUPED_TC_ROWS_PER_PASS,
74    threads: 128,
75};
76
77/// 2026-09-28: `moe_expert_down_act_fp8_grouped_tc` (reads FP32 SiLU products as BF16 hi + lo).
78pub const FP8_GROUPED_DOWN_TC: Fp8GroupedGeometry = Fp8GroupedGeometry {
79    cols_per_cta: FP8_GROUPED_TC_DOWN_COLS_PER_CTA,
80    rows_per_pass: FP8_GROUPED_TC_ROWS_PER_PASS,
81    threads: 128,
82};
83
84/// 2026-09-28: Whether the tensor-core kernels take an `n`-column, `k`-deep projection:
85/// both in whole 128 blocks (the FP8 scale blocks) and `n` in whole CTAs of `geometry`.
86pub fn fp8_grouped_tc_shape_ok(n: u32, k: u32, geometry: Fp8GroupedGeometry) -> bool {
87    n > 0
88        && k > 0
89        && k.is_multiple_of(128)
90        && n.is_multiple_of(128)
91        && n.is_multiple_of(geometry.cols_per_cta)
92}
93
94/// 2026-09-25: Cap on active experts, which sizes the grouped grids' Y extent:
95/// `num_tokens * top_k` rows can reach at most that many distinct experts. It
96/// does not depend on the routing, so a captured graph stays valid for every
97/// routing.
98pub fn fp8_grouped_active_cap(num_tokens: u32, top_k: u32, num_experts: u32) -> u32 {
99    (num_tokens * top_k).min(num_experts)
100}
101
102/// 2026-09-27: Experts `moe_fp8_grouped_sort` handles (`PMS_MAX_EXPERTS`).
103pub const FP8_GROUPED_SORT_MAX_EXPERTS: u32 = 1024;
104
105/// 2026-09-27: The outputs of [`moe_fp8_grouped_sort`]: `moe_sort_by_expert`'s
106/// four and the active-expert list with its length.
107pub struct Fp8GroupedSortOut {
108    pub sorted_token_ids: DevicePtr,
109    pub sorted_expert_ids: DevicePtr,
110    pub expert_offsets: DevicePtr,
111    pub token_to_perm: DevicePtr,
112    pub active_experts: DevicePtr,
113    pub active_count: DevicePtr,
114}
115
116/// 2026-09-27: Sorts the `total_expanded` slots of `topk_ids` by expert, as
117/// `moe_sort_by_expert` does, and writes `active_experts[0..count]` in
118/// ascending order with `active_count[0] = count`, in one launch on `stream`
119/// (the grouped kernels that read them must follow on it).
120#[allow(clippy::too_many_arguments)]
121pub fn moe_fp8_grouped_sort(
122    gpu: &dyn GpuBackend,
123    kernel: KernelHandle,
124    out: Fp8GroupedSortOut,
125    topk_ids: DevicePtr,
126    total_expanded: u32,
127    num_experts: u32,
128    top_k: u32,
129    stream: u64,
130) -> Result<()> {
131    anyhow::ensure!(
132        num_experts <= FP8_GROUPED_SORT_MAX_EXPERTS,
133        "moe_fp8_grouped_sort: {num_experts} experts exceed {FP8_GROUPED_SORT_MAX_EXPERTS}"
134    );
135    KernelLaunch::new(gpu, kernel)
136        .grid([1, 1, 1])
137        .block([256, 1, 1])
138        .arg_ptr(topk_ids)
139        .arg_ptr(out.sorted_token_ids)
140        .arg_ptr(out.sorted_expert_ids)
141        .arg_ptr(out.expert_offsets)
142        .arg_ptr(out.token_to_perm)
143        .arg_ptr(out.active_experts)
144        .arg_ptr(out.active_count)
145        .arg_u32(total_expanded)
146        .arg_u32(num_experts)
147        .arg_u32(top_k)
148        .launch(stream)
149}
150
151/// 2026-09-26: Grouped FP8 gate+up and SiLU. `cap` is [`fp8_grouped_active_cap`];
152/// the first `ceil(num_tokens / geometry.rows_per_pass)` block rows are the shared
153/// expert. Writes the product `silu(bf16(gate)) * bf16(up)`, `[positions, n]` for
154/// the routed experts into `act` and `[num_tokens, n]` for the shared expert into
155/// `sh_act`: FP32 from the scalar kernel ([`FP8_GROUPED_GATE_UP_SCALAR`]), BF16 from
156/// the tensor-core one ([`FP8_GROUPED_GATE_UP_TC`], 2026-09-28).
157#[allow(clippy::too_many_arguments)]
158pub fn moe_expert_gate_up_act_fp8_grouped(
159    gpu: &dyn GpuBackend,
160    kernel: KernelHandle,
161    geometry: Fp8GroupedGeometry,
162    input: DevicePtr,
163    gp_w: DevicePtr,
164    gp_s: DevicePtr,
165    up_w: DevicePtr,
166    up_s: DevicePtr,
167    act: DevicePtr,
168    expert_offsets: DevicePtr,
169    sorted_token_ids: DevicePtr,
170    active_experts: DevicePtr,
171    active_count: DevicePtr,
172    sh_gate: &Fp8Weight,
173    sh_up: &Fp8Weight,
174    sh_act: DevicePtr,
175    n: u32,
176    k: u32,
177    cap: u32,
178    num_tokens: u32,
179    stream: u64,
180) -> Result<()> {
181    KernelLaunch::new(gpu, kernel)
182        .grid([
183            div_ceil(n, geometry.cols_per_cta),
184            cap + div_ceil(num_tokens, geometry.rows_per_pass),
185            1,
186        ])
187        .block([geometry.threads, 1, 1])
188        .arg_ptr(input)
189        .arg_ptr(gp_w)
190        .arg_ptr(gp_s)
191        .arg_ptr(up_w)
192        .arg_ptr(up_s)
193        .arg_ptr(act)
194        .arg_ptr(expert_offsets)
195        .arg_ptr(sorted_token_ids)
196        .arg_ptr(active_experts)
197        .arg_ptr(active_count)
198        .arg_ptr(sh_gate.weight)
199        .arg_ptr(sh_gate.row_scale)
200        .arg_ptr(sh_up.weight)
201        .arg_ptr(sh_up.row_scale)
202        .arg_ptr(sh_act)
203        .arg_u32(n)
204        .arg_u32(k)
205        .arg_u32(cap)
206        .arg_u32(num_tokens)
207        .launch(stream)
208}
209
210/// 2026-09-26: Grouped FP8 down over the SiLU product, launched with `geometry`
211/// ([`FP8_GROUPED_DOWN_SCALAR`] reads FP32 products, [`FP8_GROUPED_DOWN_TC`] BF16).
212/// `k` is the intermediate width; rows of `act` and `output` are sorted
213/// positions, rows of `sh_act` and `sh_down_out` tokens.
214#[allow(clippy::too_many_arguments)]
215pub fn moe_expert_down_act_fp8_grouped(
216    gpu: &dyn GpuBackend,
217    kernel: KernelHandle,
218    geometry: Fp8GroupedGeometry,
219    act: DevicePtr,
220    down_w: DevicePtr,
221    down_s: DevicePtr,
222    output: DevicePtr,
223    expert_offsets: DevicePtr,
224    active_experts: DevicePtr,
225    active_count: DevicePtr,
226    sh_act: DevicePtr,
227    sh_down: &Fp8Weight,
228    sh_down_out: DevicePtr,
229    n: u32,
230    k: u32,
231    cap: u32,
232    num_tokens: u32,
233    stream: u64,
234) -> Result<()> {
235    KernelLaunch::new(gpu, kernel)
236        .grid([
237            div_ceil(n, geometry.cols_per_cta),
238            cap + div_ceil(num_tokens, geometry.rows_per_pass),
239            1,
240        ])
241        .block([geometry.threads, 1, 1])
242        .arg_ptr(act)
243        .arg_ptr(down_w)
244        .arg_ptr(down_s)
245        .arg_ptr(output)
246        .arg_ptr(expert_offsets)
247        .arg_ptr(active_experts)
248        .arg_ptr(active_count)
249        .arg_ptr(sh_act)
250        .arg_ptr(sh_down.weight)
251        .arg_ptr(sh_down.row_scale)
252        .arg_ptr(sh_down_out)
253        .arg_u32(n)
254        .arg_u32(k)
255        .arg_u32(cap)
256        .arg_u32(num_tokens)
257        .launch(stream)
258}
259
260/// 2026-09-25: Grouped blend, one `blockIdx.y` per token. `expert_out` rows
261/// are sorted positions; `token_to_perm[token * top_k + k]` maps a slot to its
262/// row. `gate_weight` may be NULL (ungated shared expert).
263#[allow(clippy::too_many_arguments)]
264pub fn moe_weighted_sum_blend_fp8_grouped(
265    gpu: &dyn GpuBackend,
266    kernel: KernelHandle,
267    output: DevicePtr,
268    expert_out: DevicePtr,
269    expert_weights: DevicePtr,
270    token_to_perm: DevicePtr,
271    shared_out: DevicePtr,
272    input: DevicePtr,
273    gate_weight: DevicePtr,
274    hidden: u32,
275    top_k: u32,
276    k: u32,
277    num_tokens: u32,
278    stream: u64,
279) -> Result<()> {
280    KernelLaunch::new(gpu, kernel)
281        .grid([div_ceil(hidden, 256), num_tokens, 1])
282        .block([256, 1, 1])
283        .arg_ptr(output)
284        .arg_ptr(expert_out)
285        .arg_ptr(expert_weights)
286        .arg_ptr(token_to_perm)
287        .arg_ptr(shared_out)
288        .arg_ptr(input)
289        .arg_ptr(gate_weight)
290        .arg_u32(hidden)
291        .arg_u32(top_k)
292        .arg_u32(k)
293        .launch(stream)
294}
295
296#[cfg(test)]
297#[path = "fp8_moe_grouped_tests.rs"]
298mod tests;
metrale_model_layers/layers/ops/fp8_moe_grouped.rs

metrale_model_layers/layers/ops/
fp8_moe_grouped.rs
