Skip to content

Commit be9f88d

Browse files
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

File tree

0 commit comments

Comments
 (0)