pub fn moe_fp8_grouped_gemm(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
input: DevicePtr,
weight_ptrs: DevicePtr,
scale_ptrs: DevicePtr,
output: DevicePtr,
expert_offsets: DevicePtr,
sorted_token_ids: DevicePtr,
num_experts: u32,
n: u32,
k: u32,
worklist: DevicePtr,
total_tiles: DevicePtr,
max_tiles: u32,
stream: u64,
) -> Result<()>Expand description
FP8 grouped GEMM for sorted MoE prefill — grid-compaction over the COMPACTED
work-list built by moe_build_tile_worklist. THE routed-expert FP8 prefill
kernel.
The kernel grid-strides by gridDim.x, so the launch is sized to
max_tiles — the caller’s exact upper bound on the work-item (tile) count
(wl_cap_items). This covers the whole work-list in ~one pass instead of
serializing dozens of tiles per CTA behind sync barriers (the old fixed
96-CTA persistent grid left the GPU >90% idle: ~0.2% occupancy / ~16%
MemUnitBusy, measured on gfx1151). Oversubscription is safe (extra CTAs
exit the loop immediately); undersizing is merely slower, never wrong.
max_tiles is clamped to MAX_GRID_CTAS so a pathological worklist bound
cannot request an unbounded grid.
SAME-STREAM INVARIANT: MUST be launched on the SAME stream as the
preceding moe_build_tile_worklist (read-after-write of total_tiles).
Grid: (max_tiles.clamp(1, MAX_GRID_CTAS), 1, 1) Block: (256, 1, 1)