moe_fp8_grouped_gemm

Function moe_fp8_grouped_gemm 

Source
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)