gemm: add fp8 per-tensor grouped GEMM forward (M-grouped/MoE) - #887
Open
kyle-256 wants to merge 3 commits into
Open
gemm: add fp8 per-tensor grouped GEMM forward (M-grouped/MoE)#887kyle-256 wants to merge 3 commits into
kyle-256 wants to merge 3 commits 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
2 times, most recently
from
July 23, 2026 08:14
fa9c760 to
be9f88d
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).
The rebase onto main was textually clean and semantically broken. main moved buffer_ops out of the DSL -- python/flydsl/expr/buffer_ops.py is gone and the module now lives in the repo's own kernels/common/ -- so `from flydsl.expr import buffer_ops` fails at import, and the whole kernel with it. The four entry points this kernel uses (extract_base_index, _get_buffer_flags, create_buffer_resource_from_addr, buffer_store) came across with identical signatures, including buffer_store's mask= and offset_is_bytes=, so only the import moves. Checked every other `from flydsl...` in the file against the tree: all still resolve; this was the only casualty.
Rebasing onto main moved this kernel from flydsl 0.2.4 to 0.3.x, and it started returning numbers uncorrelated with the reference -- logits diff 0.98, i.e. not drift, wrong data -- while compiling and running clean. The cause is the inline-asm MMA's accumulator constraint. asm mode "2" emits v_mfma_f32_16x16x128_f8f6f4 with "=a,v,v,0": an AGPR-class output tied to srcC. On 0.2.4 that was correct; on 0.3.x it silently miscompiles, and the shape of the damage is distinctive -- with B = I, so the output must equal A, every fourth output row comes back entirely zero and the rest are ~25% populated. Four accumulator slots per lane map to four consecutive rows, so that is one whole accumulator slot lost. Isolated to the constraint: the operand types entering the asm are right (vector<8xi32>, vector<8xi32>, vector<4xf32> against a vector<4xf32> result), dropping the amdgpu-agpr-alloc="128,128" passthrough does not help, and the identical asm with "=v" (mode "3") is correct. So acc_mode defaults to "vgpr" -- still the in-place asm accumulate this kernel is built around, just in the VGPR file -- and "agpr" now asserts instead of returning wrong numbers, since nothing about the failure is visible at compile time. _raw -> as_ir_value on the asm operands is not the fix; it is the unwrap main moved to in #913 when it ported the parent fp8_gemm_4wave, and it is what the passing configuration runs. Measured on gfx950 with flydsl 0.3.2: 5 passed, and 2751 TFLOPS on 2x2048/n4096/k7168 against 2708 for the intrinsic MMA (agpr_inplace=False) -- the same within single-run noise, so the asm path keeps its reason to exist. Worth recording for whoever hits this next: the parent kernel does not use inline-asm MFMA at all, so this path had no coverage in main's 0.2.4 -> 0.3.x ports. Everything else about the kernel survived the jump.
kyle-256
force-pushed
the
dev/kyle/grouped_gemm
branch
from
September 5, 2026 08:44
be9f88d to
afd0f6d
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