Commit be9f88d
committed
gemm: add fp8 per-tensor grouped GEMM forward (M-grouped/MoE)
Adds a grouped (M-grouped / MoE) fp8 per-tensor GEMM forward kernel
(NT: out = a @ b^T) under kernels/gemm, plus a correctness+perf test.
The grid is over-launched to a host upper bound; each WG computes the
true tile count via an on-device O(G) scan (no host read of group lens)
and s_endpgm's when past the end. The same scan maps a tile id ->
(group, local tile). Per-group A/B addressing folds the group element
offset into the i64 SRD base (int64-safe > 4GB), and the C store passes
the absolute group-end row so partial-M tiles clamp cleanly at group
boundaries. Scales are scalar per-tensor.
Reuses the shared LDS coop-loaders / global swizzle / barriers from
fp8_gemm_utils; grouped-specific primitives are kept local. Ported from
Primus-Turbo, authored with FlyDSL.
Verified on gfx950 (MI355X) with flydsl 0.2.2: 5/5 cases pass
(aligned, ragged, empty-group, single-token, large-K).1 parent 8fcfee4 commit be9f88d
2 files changed
Lines changed: 930 additions & 0 deletions
0 commit comments