moe_bf16_grouped_gemm

Function moe_bf16_grouped_gemm 

Source
pub fn moe_bf16_grouped_gemm(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    input: DevicePtr,
    weight_ptrs: DevicePtr,
    output: DevicePtr,
    expert_offsets: DevicePtr,
    sorted_token_ids: DevicePtr,
    num_experts: u32,
    n: u32,
    k: u32,
    max_m_tiles: u32,
    stream: u64,
) -> Result<()>
Expand description

2026-09-25: BF16 grouped GEMM for sorted MoE prefill, without scales: input [total_tokens, K] BF16, weight_ptrs[e] → [N, K] BF16, output [total_expanded, N] BF16, sorted_token_ids [total_expanded] or NULL. The MoE layer’s long-prefill path uses it when its experts are BF16, which the loader sets up when METRALE_FP8_DEQUANT_MOE_TO_BF16=1 dequantizes FP8 experts at load (qwen35/load_layers.rs).