spark_model/layers/ops/
moe_grouped_fp4.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Native-FP4 grouped MoE launcher. Split from `moe_grouped_a.rs` (500-LoC cap).
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use super::*;
12
13/// Grouped W4A4 expert UP GEMM with fused relu^2 (native FP4 tensor cores).
14/// A is the pre-quantized NVFP4 latent (packed E2M1 + per-16 E4M3 scales);
15/// B comes from the per-expert NVFP4 pointer tables unchanged.
16/// Grid: (ceil(n_out/128), max_m_tiles, num_experts)  Block: (128, 1, 1)
17#[allow(clippy::too_many_arguments)]
18pub fn moe_w4a4_grouped_gemm_relu2(
19    gpu: &dyn GpuBackend,
20    kernel: KernelHandle,
21    a_packed: DevicePtr,
22    a_sf: DevicePtr,
23    b_packed_ptrs: DevicePtr,
24    b_scale_ptrs: DevicePtr,
25    scale2_vals: DevicePtr,
26    output: DevicePtr,
27    expert_offsets: DevicePtr,
28    sorted_token_ids: DevicePtr,
29    num_experts: u32,
30    n_out: u32,
31    k: u32,
32    max_m_tiles: u32,
33    stream: u64,
34) -> Result<()> {
35    KernelLaunch::new(gpu, kernel)
36        .grid([div_ceil(n_out, 128), max_m_tiles, num_experts])
37        .block([128, 1, 1])
38        .arg_ptr(a_packed)
39        .arg_ptr(a_sf)
40        .arg_ptr(b_packed_ptrs)
41        .arg_ptr(b_scale_ptrs)
42        .arg_ptr(scale2_vals)
43        .arg_ptr(output)
44        .arg_ptr(expert_offsets)
45        .arg_ptr(sorted_token_ids)
46        .arg_u32(num_experts)
47        .arg_u32(n_out)
48        .arg_u32(k)
49        .launch(stream)
50}
spark_model/layers/ops/moe_grouped_fp4.rs

spark_model/layers/ops/
moe_grouped_fp4.rs
