gemm: add fp8 per-tensor grouped GEMM forward (M-grouped/MoE)#887
Open
kyle-256 wants to merge 1 commit into
Open
gemm: add fp8 per-tensor grouped GEMM forward (M-grouped/MoE)#887kyle-256 wants to merge 1 commit into
kyle-256 wants to merge 1 commit into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new fp8 per-tensor M-grouped (MoE-style) GEMM forward kernel to kernels/gemm, plus a GPU correctness/perf test harness to exercise grouped edge cases (ragged groups, empty group, single-token group).
Changes:
- Introduces
kernels/gemm/fp8_grouped_gemm.py: a grouped NT fp8 GEMM forward kernel that over-launches the grid and maps tiles to groups on-device, with i64 SRD rebasing and row-band bounded stores. - Adds a public Python entrypoint
grouped_gemm_fp8_forward(...)with a compile cache keyed by(N, K, G, out_fp16, cbsz, blgp). - Adds
tests/kernels/test_fp8_grouped_gemm.pycovering correctness and reporting TFLOPS across representative grouped shapes.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
kernels/gemm/fp8_grouped_gemm.py |
New grouped fp8 GEMM forward kernel + Python entrypoint/compile cache. |
tests/kernels/test_fp8_grouped_gemm.py |
New grouped fp8 GEMM correctness + perf harness and benchmark CLI. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+753
to
+758
| assert a.ndim == 2 and b_T.ndim == 3, f"a{tuple(a.shape)} b_T{tuple(b_T.shape)}" | ||
| M_total, K = a.shape | ||
| G, N, K_b = b_T.shape | ||
| assert K == K_b, f"K mismatch a={K} b_T={K_b}" | ||
|
|
||
| out = torch.empty((M_total, N), device=a.device, dtype=out_dtype) |
Comment on lines
+759
to
+763
| # The kernel reads group_offs as int64 low-words via a free int32-view (no | ||
| # .to(int32) cast); int32 callers are upcast to int64 once. | ||
| _go64 = group_offs if group_offs.dtype == torch.int64 else group_offs.to(torch.int64) | ||
| go32 = _go64.view(torch.int32) | ||
| out_fp16 = out_dtype == torch.float16 |
Comment on lines
+767
to
+772
| key = (N, K, G, out_fp16, cbsz, blgp) | ||
| launch = _GROUPED_FWD_CACHE.get(key) | ||
| if launch is None: | ||
| launch = compile_fp8_grouped_gemm( | ||
| K=K, G=G, out_fp16=out_fp16, cbsz=cbsz, blgp=blgp | ||
| ) |
Comment on lines
+208
to
+220
| gSA = fx.rocdl.make_buffer_tensor(A_scale, max_size=False, num_records_bytes=4) # 1 fp32 | ||
| gSB = fx.rocdl.make_buffer_tensor(B_scale, max_size=False, num_records_bytes=4) # 1 fp32 | ||
| self.sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) | ||
| self.sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) | ||
| self.scale_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) | ||
| self.reg_f32_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) | ||
|
|
||
| def _load_scalar(self, div): | ||
| fx.copy(self.scale_atom_1, fx.slice(div, (None, fx.Int32(0))), self.reg_f32_1) | ||
| return Vec(fx.memref_load_vec(self.reg_f32_1))[0] | ||
|
|
||
| def store(self, c_frag, base_row, base_col): | ||
| scale = self._load_scalar(self.sa_div) * self._load_scalar(self.sb_div) |
| ok = verify_output(out.to(torch.float32), c_ref, rtol=0.1, atol=0.1) | ||
| assert ok, f"correctness failed for G={G} M_total={M_total} N={N} K={K}" | ||
|
|
||
| _, us = run_perftest(lambda _o: _launch(), out, num_iters=num_iters, num_warmup=num_warmups) |
kyle-256
force-pushed
the
dev/kyle/grouped_gemm
branch
from
July 23, 2026 07:20
51ca972 to
fa9c760
Compare
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).
kyle-256
force-pushed
the
dev/kyle/grouped_gemm
branch
from
July 23, 2026 08:14
fa9c760 to
be9f88d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
FlyDSL currently ships dense fp8 GEMMs (
kernels/gemm/fp8_gemm_4wave.py/fp8_gemm_8wave.py) but has no grouped (M-grouped / MoE) GEMM. MoE modelsneed a grouped GEMM where several per-expert matmuls with a shared K but
different row counts run as one launch. This PR adds the fp8 per-tensor
grouped GEMM forward (NT:
out = a @ b^T), reusing the existing dense fp8GEMM primitives so it stays consistent with the rest of
kernels/gemm.Technical Details
New kernel
kernels/gemm/fp8_grouped_gemm.py:a=[M_total, K]fp8 (the groups concatenated along M),b_T=[G, N, K]fp8 (per-group weights),out=[M_total, N],group_offs=[G+1]int64 splittingM_totalintoGgroups.group_offsis read as an int32 view (a free reinterpret; low word at
i32[2*idx]).(ceil(M_total/BLOCK_M) + G) * n_blocks(no host read of the group lengths).Each WG computes the true
total_tileson-device via an O(G) scan ands_endpgms when its tile id is past the end (one tile / WG, full-devicedefault). The same scan maps a tile id →
(group_idx, local tile).(
m_start*K/group_idx*N*K) into the i64 SRD base, so in-tile offsets stayint32 even for A/B beyond 2^31 elems / 4 GB; the SRD
num_recordsbound clampsthe last over-read to 0. The C store passes
c_rows = group_offs[group_idx+1](the absolute group-end row) so a partial M-tile's extra rows (which belong to
the next group) are dropped by the row-band SRD bound — no spill across groups.
G2SLoader/S2RLoader), globalswizzle, and barriers come from
fp8_gemm_utils. The grouped-specificprimitives — i64-rebased SRD, row-band scalar store, in-place AGPR MFMA,
K-tail mask, L2-reuse tile swizzle, and XCD remap — are local to the new file.
a_scale,b_scale).Public entry:
grouped_gemm_fp8_forward(a, b_T, a_scale, b_scale, group_offs, out_dtype=torch.bfloat16), with a compile cache keyed by the static dims(N, K, G, out_fp16, cbsz, blgp).Also adds
tests/kernels/test_fp8_grouped_gemm.py(correctness + perf; a torchper-group reference and a CLI benchmark entry).
Test Plan
Run on gfx950 (MI355X) with
flydsl 0.2.2:The parametrized cases deliberately cover the group-handling edge paths:
aligned groups, ragged (non-BLOCK_M-multiple) groups, an empty group
(
M_g = 0), a single-token group (M_g = 1), and a large-K shape — so theon-device group scan, K-tail masking, and partial-M-tile / group-boundary clamps
are all exercised. Correctness is checked against a torch per-group reference
with
verify_output(rtol=0.1, atol=0.1).Test Result
5/5 pass:
Submission Checklist