metrale_model_layers/layers/ops/
bf16_moe_grouped.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: Launchers for the grouped BF16 MoE decode kernels
4//! (`kernels/gb10/common/moe_bf16_grouped_tc.cu`, the BF16 point of the tensor-core grouped
5//! family). The layout is that of the grouped FP8 and NVFP4 decodes: rows grouped by expert by
6//! `moe_fp8_grouped_sort`, intermediates by sorted position, the blend
7//! `moe_weighted_sum_blend_fp8_grouped`. Weights are row-major BF16 per expert (pointer tables).
8//!
9//! Owner: model-layers ops.
10//! Invariants: none beyond the types.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use super::Fp8GroupedGeometry;
17
18/// 2026-10-02: `moe_expert_gate_up_act_bf16_grouped_tc` (`BTC_GU_COLS` columns per CTA; writes the
19/// SiLU product as BF16 hi + lo).
20pub const BF16_GROUPED_GATE_UP_TC: Fp8GroupedGeometry = Fp8GroupedGeometry {
21    cols_per_cta: 64,
22    rows_per_pass: 8,
23    threads: 128,
24};
25
26/// 2026-10-02: `moe_expert_down_act_bf16_grouped_tc` (`BTC_DOWN_COLS` columns per CTA).
27pub const BF16_GROUPED_DOWN_TC: Fp8GroupedGeometry = Fp8GroupedGeometry {
28    cols_per_cta: 128,
29    rows_per_pass: 8,
30    threads: 128,
31};
32
33/// 2026-10-02: Whether the BF16 tensor-core kernels take an `n`-column, `k`-deep projection: `k`
34/// in whole load groups (two 32-K chunks; 128 for margin), `n` in whole CTAs of `geometry`.
35pub fn bf16_grouped_tc_shape_ok(n: u32, k: u32, geometry: Fp8GroupedGeometry) -> bool {
36    n > 0 && k > 0 && k.is_multiple_of(128) && n.is_multiple_of(geometry.cols_per_cta)
37}
38
39/// 2026-10-02: Grouped BF16 gate+up and SiLU. `gate_ptrs` / `up_ptrs` are device tables of each
40/// routed expert's `[n, k]` BF16 weight; `sh_gate` / `sh_up` the shared expert's. Writes the SiLU
41/// product as BF16 hi + lo rows (`[positions, 2n]` into `act`, `[num_tokens, 2n]` into `sh_act`).
42#[allow(clippy::too_many_arguments)]
43pub fn moe_expert_gate_up_act_bf16_grouped(
44    gpu: &dyn GpuBackend,
45    kernel: KernelHandle,
46    input: DevicePtr,
47    gate_ptrs: DevicePtr,
48    up_ptrs: DevicePtr,
49    act: DevicePtr,
50    expert_offsets: DevicePtr,
51    sorted_token_ids: DevicePtr,
52    active_experts: DevicePtr,
53    active_count: DevicePtr,
54    sh_gate: DevicePtr,
55    sh_up: DevicePtr,
56    sh_act: DevicePtr,
57    n: u32,
58    k: u32,
59    cap: u32,
60    num_tokens: u32,
61    stream: u64,
62) -> Result<()> {
63    let g = BF16_GROUPED_GATE_UP_TC;
64    KernelLaunch::new(gpu, kernel)
65        .grid([
66            div_ceil(n, g.cols_per_cta),
67            cap + div_ceil(num_tokens, g.rows_per_pass),
68            1,
69        ])
70        .block([g.threads, 1, 1])
71        .arg_ptr(input)
72        .arg_ptr(gate_ptrs)
73        .arg_ptr(up_ptrs)
74        .arg_ptr(act)
75        .arg_ptr(expert_offsets)
76        .arg_ptr(sorted_token_ids)
77        .arg_ptr(active_experts)
78        .arg_ptr(active_count)
79        .arg_ptr(sh_gate)
80        .arg_ptr(sh_up)
81        .arg_ptr(sh_act)
82        .arg_u32(n)
83        .arg_u32(k)
84        .arg_u32(cap)
85        .arg_u32(num_tokens)
86        .launch(stream)
87}
88
89/// 2026-10-02: Grouped BF16 down over the hi + lo SiLU rows. `k` is the intermediate width; rows
90/// of `act` and `output` are sorted positions, rows of `sh_act` and `sh_down_out` tokens.
91#[allow(clippy::too_many_arguments)]
92pub fn moe_expert_down_act_bf16_grouped(
93    gpu: &dyn GpuBackend,
94    kernel: KernelHandle,
95    act: DevicePtr,
96    down_ptrs: DevicePtr,
97    output: DevicePtr,
98    expert_offsets: DevicePtr,
99    active_experts: DevicePtr,
100    active_count: DevicePtr,
101    sh_act: DevicePtr,
102    sh_down: DevicePtr,
103    sh_down_out: DevicePtr,
104    n: u32,
105    k: u32,
106    cap: u32,
107    num_tokens: u32,
108    stream: u64,
109) -> Result<()> {
110    let g = BF16_GROUPED_DOWN_TC;
111    KernelLaunch::new(gpu, kernel)
112        .grid([
113            div_ceil(n, g.cols_per_cta),
114            cap + div_ceil(num_tokens, g.rows_per_pass),
115            1,
116        ])
117        .block([g.threads, 1, 1])
118        .arg_ptr(act)
119        .arg_ptr(down_ptrs)
120        .arg_ptr(output)
121        .arg_ptr(expert_offsets)
122        .arg_ptr(active_experts)
123        .arg_ptr(active_count)
124        .arg_ptr(sh_act)
125        .arg_ptr(sh_down)
126        .arg_ptr(sh_down_out)
127        .arg_u32(n)
128        .arg_u32(k)
129        .arg_u32(cap)
130        .arg_u32(num_tokens)
131        .launch(stream)
132}
metrale_model_layers/layers/ops/bf16_moe_grouped.rs

metrale_model_layers/layers/ops/
bf16_moe_grouped.rs
