Repository navigation
Conversation
Integrate a persistent, work-stealing grouped GEMM around the existing FlyDSL MXFP8 kernel into GroupedLinear, dispatched inline at the forward GEMM point (like the Triton grouped path) so it rides off TE's existing MXFP8 quantization rather than quantizing itself. The grouped kernel is a compile-time specialization (group=True) of the same mxfp8_gemm source as the dense kernel; the hot MFMA pipeline is unchanged and runs in a process_tile closure. Grouping only shifts B/B-scale addressing to the packed [G*N, K] layout (expert*N column, G*N stride) under const_expr(group), so the dense binary is byte-identical. A fixed num-CU persistent grid claims tiles from a precomputed (pid_m, pid_n, expert) descriptor queue via an atomic counter with LDS broadcast; a static grid-stride schedule is also provided, selected by a host heuristic (NVTE_FLYDSL_GROUPED_SCHED override). Opt in with NVTE_USE_GROUPED_GEMM_FLYDSL=1 (HIP, gfx950, MXFP8 recipe, 128-aligned in/out features). Bitwise-equal to TE's default grouped MXFP8 GEMM and to a per-expert sequential FlyDSL dense reference, forward and backward. Co-Authored-By: Claude <noreply@anthropic.com>
Extend the persistent grouped MXFP8 GEMM to the NN layout so GroupedLinear's dgrad (dX = dY @ W) routes through FlyDSL, mirroring the TN forward. The core MFMA pipeline is untouched: NN only adds a const_expr(group) expert offset on the transpose-read B path. Unlike TN forward (experts pack on B's output-N axis), NN's contraction axis is N_out, so the packed weight [G*N_out, K_in] stacks experts on the CONTRACTION axis. The expert offset therefore folds into the k_base term (expert*N_out) and the B-scale k128 term (expert*N_out/128), not the output-N base. Shape-contract handling in mxfp8_matmul/do_gemm is made layout-aware: NN's B contraction dim is G*N_out (not equal to A's per-expert N_out) and the output width K_in is not packed. _run_mxfp8_grouped gains a layout param; its descriptor-queue and work-stealing logic are layout-independent and reused. Backward plumbing: forward records ctx.use_flydsl_mxfp8_grouped; the shared backward branches on it to pack dY (rowwise) + W (columnwise, already ensured) and call the grouped NN launch. wgrad still uses the default path. dgrad is bitwise-equal to the default TE grouped MXFP8 backward on both the worksteal and static schedules; dense (non-grouped) NN is unaffected. Co-Authored-By: Claude <noreply@anthropic.com>
Complete the grouped MXFP8 backward: wgrad (dW = dY^T @ X per expert) now routes through a dedicated FlyDSL NT kernel, so GroupedLinear forward + dgrad + wgrad all run on FlyDSL. wgrad contracts over the token axis m_e (variable per expert). Uniform padding: every expert's tokens are padded up to a common M_pad and one workgroup contracts the full padded token axis for each output tile (a single-WG contraction, NOT split-K -- no cross-WG reduction). Padded rows are zero and contribute nothing to the contraction, so the NT MFMA pipeline runs byte-identical with no atomics or partial-tile masking. dY and X pack on the contraction axis to [G*M_pad, *]; dW packs on the output-M axis to [G*N_out, K_in]. The kernel is a compile-time `wgrad` specialization of the same generator (distinct binary, no pipeline duplication): both operand contraction bases and their scale k128 terms shift by expert*M_pad, and the store row shifts by expert*N_out (store_m_expert_off) since the operands use local per-expert coords while the output is packed. Reuses the persistent descriptor-queue / work-stealing driver via `grouped = group or wgrad`. Plumbing: _run_mxfp8_grouped_wgrad (pad/pack + descriptor build), _mxfp8_matmul_grouped_wgrad (NT scale pack + do_gemm(wgrad=True)), _wgrad_flydsl_mxfp8_grouped module method, and a backward dispatch gated on ctx.use_flydsl_mxfp8_grouped. fuse_wgrad_accumulation, bias grad, and delayed wgrad still fall through to the default path. wgrad is bitwise-equal to the default TE grouped backward on both worksteal and static schedules; dense NT/NN/TN are unaffected. Co-Authored-By: Claude <noreply@anthropic.com>
e4125aa to
10f48f8
Compare
|
Re-reviewed the 5 commits added since the last pass ( Previous round: 9 of the 13 earlier findings verified genuinely fixed in the code, not just acknowledged — Verdict: approve with changes. The new work is well-structured and the launch keys are complete, but the in-place operand refactor loosened two guards that used to produce clean fallbacks:
Upstream compatibility is clean — all divergence is behind Copyright headers: OK (all 11 in-scope files). Details inline. |
|
Third review pass, scoped to the commits added since the last review ( Verdict: approve with minor changes. All three code fixes from the previous round check out under tracing:
Two non-blocking findings posted inline: a stale axis order in the new launch-guard comment (the bound itself is unaffected — same numel), plus a duplication note; and a residual wgrad-coverage gap in the new no-fallback test assertion, with a pointer to an existing helper that already does the same warning capture. Formatting is clean — the previously flagged Black violation is fixed and no added line exceeds 100 characters. Copyright headers: OK (all 11 in-scope files). |
| # With per-expert pointers, A (and wgrad's B) is one representative expert, but | ||
| # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit | ||
| # offsets, so the bound is on the packed extent. | ||
| a_rows, a_cols = a_shape if a_shape is not None else A.shape | ||
| b_rows, b_cols = b_shape if b_shape is not None else B.shape |
There was a problem hiding this comment.
The extent fix itself is right, but the replacement comment loses the rationale for b_shape=None and is inaccurate for grouped fwd/dgrad.
In fwd/dgrad, B is also one representative expert — _run_fp8_grouped always passes b_data_ptrs, and half_prec_gemm.py:1678 raises if group and b_data_ptrs is None. So "A (and wgrad's B) is one representative expert" undersells it, and "the bound is on the packed extent" is wrong for that B: b_shape=None there deliberately bounds the single expert, because the kernel addresses B from a per-expert base with expert-local offsets (half_prec_gemm.py:495-545). The clause deleted here — "Grouped fwd/dgrad weights use local offsets, so B's own size holds." — was the load-bearing reason, and without it a future reader is likely to "fix" b_shape=None on the fwd/dgrad wrappers.
mxfp8_gemm.py:1973-1976 already phrases this correctly ("b_shape likewise gives wgrad's packed B shape when B is read through b_data_ptrs"). Suggest matching that and keeping the local-offset sentence:
| # With per-expert pointers, A (and wgrad's B) is one representative expert, but | |
| # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit | |
| # offsets, so the bound is on the packed extent. | |
| a_rows, a_cols = a_shape if a_shape is not None else A.shape | |
| b_rows, b_cols = b_shape if b_shape is not None else B.shape | |
| # With per-expert pointers, A (and wgrad's B) is one representative expert, but | |
| # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit | |
| # offsets, so the bound is on the packed extent. Grouped fwd/dgrad weights are | |
| # addressed from a per-expert base with local offsets, so B's own size holds. | |
| a_rows, a_cols = a_shape if a_shape is not None else A.shape | |
| b_rows, b_cols = b_shape if b_shape is not None else B.shape |
Same text at half_prec_gemm.py:1701-1703. The new docstring sentences (fp8_gemm.py:1345-1346, half_prec_gemm.py:1672-1674) read the same way — "b_data_ptrs in wgrad" reads as if b_data_ptrs is wgrad-only, which contradicts the mandatory-for-all-grouped check four lines below in half_prec_gemm.py.
| expected.add("wgrad") | ||
| ran = {phase for phase, ok in flydsl_dispatches if ok} | ||
| fell_back = sorted({phase for phase, ok in flydsl_dispatches if not ok}) | ||
| assert not fell_back and expected <= ran, ( |
There was a problem hiding this comment.
Nit (low): the dispatch recording is a solid fix for the false pass, but expected <= ran is a subset check, so it only locks the lower bound. A regression that widens the production gate at grouped_linear.py:2804-2809 — e.g. FlyDSL wgrad running under fuse_wgrad_accumulation, where it would overwrite main_grad instead of accumulating — leaves ran == {"fwd","dgrad","wgrad"} against expected == {"fwd","dgrad"} and still passes this assertion.
expected == ran can't false-fail: phase is only ever "fwd"/"dgrad"/"wgrad", and with fuse_wgrad_accumulation or delay_wgrad_compute the gate short-circuits before _try_flydsl_grouped is called at all, so "wgrad" can never appear in ran for those cells. It makes the assertion verify the gate the comment above describes, for free.
|
Re-reviewed the increment since the last pass — commit Both prior findings are genuinely fixed, not just acknowledged:
Two new findings, both minor — see inline comments:
Copyright headers: OK (all 11 in-scope files). No correctness or upstream-compatibility issues in this increment — the CUDA path remains untouched and all new behavior stays behind |
|
I wonder if a PR like this should next time be issued as a PR-stack, breaking it into smaller easier-to-review pieces. As-is it is quite difficult for both AI review and human review. This isn't my area of expertise so if @matthiasdiener and others think this is an appropriate scope for a PR, then please disregard this comment. |
I agree it is a lot for a single PR. The scope was originally MXFP8 only, but from what I recall from discussion with @wenchenvincent the idea was to make it cover the dtypes in the same way that was done for PR676 #676, so this was primarily a grouped extension to that PR. |
| import torch | ||
| import transformer_engine_torch as tex | ||
|
|
||
| from . import grouped_gemm_utils as gd |
There was a problem hiding this comment.
Why is it called gd ?
There was a problem hiding this comment.
No particular reason. Gave it a more readable name in 87b343c
| """One cached 1-element buffer per (device, dtype) for dense kernel-signature | ||
| placeholders. | ||
|
|
||
| The FP32 binary requires the Bias/Aux epilogue tensor arguments to be present |
There was a problem hiding this comment.
The code uses "binary" in a few places, maybe "compiled kernel" or so could be clearer?
There was a problem hiding this comment.
Thanks for pointing this out. Made the change in 87b343c
|
|
||
|
|
||
| @functools.lru_cache(maxsize=None) | ||
| def _dense_placeholder(device: torch.device, dtype: torch.dtype, numel: int = 1): |
There was a problem hiding this comment.
This function exists multiple times in the code
|
|
||
|
|
||
| @functools.lru_cache(maxsize=None) | ||
| def _num_cus(device): |
There was a problem hiding this comment.
Duplicated (as get_num_cus)
| _COMPILED_LAUNCH_CACHE = {} | ||
|
|
||
|
|
||
| def _compiled_launch(key, jit_fn, args): |
There was a problem hiding this comment.
Duplicated function
| wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) | ||
|
|
||
| @staticmethod | ||
| def _fwd_flydsl_fp8_grouped(*, inputmats, weights, out, m_splits, m_splits_tensor): |
There was a problem hiding this comment.
These functions are all very similar, maybe they can be factored out a bit
| _select_grouped_sched = gd.select_grouped_sched | ||
|
|
||
|
|
||
| def _run_mxfp8_grouped( |
There was a problem hiding this comment.
These functions are all similar, maybe they can be factored out
| IS_HIP_EXTENSION | ||
| and get_device_compute_capability() == (9, 5) | ||
| and os.environ.get("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" | ||
| and (recipe.mxfp8() or recipe.float8_current_scaling() or recipe.delayed()) | ||
| ): |
There was a problem hiding this comment.
Could this call _flydsl_grouped_common_ok instead?
| # Select exactly one backend; leave the others' env vars unset. HipKittens and CK | ||
| # are gated behind NVTE_USE_CUTLASS_GROUPED_GEMM; FlyDSL has its own opt-in and | ||
| # must not see the cutlass/HK/CK env vars. | ||
| monkeypatch.delenv("NVTE_USE_CUTLASS_GROUPED_GEMM", raising=False) |
There was a problem hiding this comment.
nit: this delenv can be under else (flydsl) branch below
Description
Fixes https://github.com/AMD-ROCm-Internal/frameworks-internal/issues/338
Adds grouped GEMM support to the FlyDSL backend on gfx950. Every expert runs in
a single persistent kernel launch: workgroups pull output tiles from a per-tile work queue using either a static grid stride or atomic work-stealing.
Supported precisions: BF16/FP16, tensor-wise FP8 (current and delayed scaling), and MXFP8. Supported formats: TN/NN/NT.
Usage
Opt in with
NVTE_USE_GROUPED_GEMM_FLYDSL=1(gfx950,flydslinstalled). FlyDSL is thenselected in
GroupedLinearthe same way as the Triton grouped backend:general_grouped_gemm_flydslis a drop-in forgeneral_grouped_gemm.Anything the kernels can't serve falls back to the default grouped GEMM. That includes
bias, GELU, wgrad accumulation, fused/delayed wgrad, blockwise FP8, and features
not divisible by 128. Set
NVTE_FLYDSL_GEMM_WARN_FALLBACK=1to enable fallback warning messages.Optional tuning:
NVTE_FLYDSL_GROUPED_SCHED=auto|static|worksteal: override the schedule choice.NVTE_FLYDSL_WGRAD_MAX_BUCKETS: cap the number of wgrad launches.Design notes
place through per-expert base pointers.
expert doesn't force every expert to pad to its size. When the weight gradients are one
contiguous buffer, the kernel writes into it directly with no copy-out.
dtypes. A small per-dtype record (
_GroupedKernel) holds the parts that differ: howoperands and scales are read, and how the kernel is launched.
Other changes
flyc.compilepath (compiled_launch) skipsFlyDSL's per-call setup, which matters at small M.
fx.Index→fx.Int64:fx.Indexis deprecated per FlyDSL's kernel authoringguide;
fx.Int64is the recommended type for index and offset values.placeholder_arg,get_num_cus, andcompiled_launchnow live ingemm_common_utils.pyinstead of being copied into each kernel module.Testing
test_grouped_linear_accuracy_rocm_backendsgains aflydslbackend: 868 cases,420 passed / 0 failed. The 448 skips are configs FlyDSL doesn't support (blockwise FP8, FP32, unquantized fused wgrad). A recorder checks that each expected
phase actually ran on FlyDSL rather than silently falling back.
test_grouped_linear_flydsl_fp8_scaling: a multi-step accuracy test forcurrent and delayed scaling. Delayed scaling uses unit scales on its first step, so
the single-step test can't catch per-expert scale bugs; this test checks that the
compared step uses real, per-expert scales.
🤖 Generated with Claude Code
Performance
See: https://github.com/AMD-ROCm-Internal/frameworks-internal/issues/338#issuecomment-6327888906
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: