Skip to content

FlyDSL persistent work-stealing grouped GEMM enablement (gfx950) - #746

Open
aris134 wants to merge 28 commits into
devfrom
amartin/flydsl-mxfp8-group-gemm
Open

aris134 wants to merge 28 commits into
devfrom
amartin/flydsl-mxfp8-group-gemm

Conversation

@aris134

@aris134 aris134 commented Sep 29, 2026 •

Copy link
Copy Markdown
Contributor

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, flydsl installed). FlyDSL is then
selected in GroupedLinear the same way as the Triton grouped backend:
general_grouped_gemm_flydsl is a drop-in for general_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=1 to 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

  • No host-side concatenation: activations, weights, and scales are read per expert in
    place through per-expert base pointers.
  • wgrad is K-bucketed: experts are grouped by their padded token count, so one large
    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.
  • Shared implementation: one forward/dgrad runner and one wgrad runner cover all
    dtypes. A small per-dtype record (_GroupedKernel) holds the parts that differ: how
    operands and scales are read, and how the kernel is launched.

Other changes

  • Faster kernel launch: a cached flyc.compile path (compiled_launch) skips
    FlyDSL's per-call setup, which matters at small M.
  • fx.Index → fx.Int64: fx.Index is deprecated per FlyDSL's kernel authoring
    guide; fx.Int64 is the recommended type for index and offset values.
  • Shared helpers: placeholder_arg, get_num_cus, and compiled_launch now live in
    gemm_common_utils.py instead of being copied into each kernel module.

Testing

  • test_grouped_linear_accuracy_rocm_backends gains a flydsl backend: 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.
  • New test_grouped_linear_flydsl_fp8_scaling: a multi-step accuracy test for
    current 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

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

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>
@aris134 aris134 self-assigned this Sep 29, 2026
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>
@aris134 aris134 added the ci-level 1 CI test level 1 label Sep 30, 2026
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>
@aris134
aris134 force-pushed the amartin/flydsl-mxfp8-group-gemm branch from e4125aa to 10f48f8 Compare September 30, 2026 16:02
@aris134 aris134 changed the title FlyDSL persistent work-stealing MXFP8 grouped GEMM (gfx950) FlyDSL persistent work-stealing grouped GEMM enablement (gfx950) Oct 9, 2026
@Micky774 Micky774 added the P2 Priority 2: normal label Oct 16, 2026
@aris134
aris134 marked this pull request as ready for review October 18, 2026 02:52
Comment thread transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py Outdated
Comment thread transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py Outdated
Comment thread transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py
Comment thread tests/pytorch/test_grouped_linear.py
Comment thread transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py Outdated
@github-actions

Copy link
Copy Markdown

Re-reviewed the 5 commits added since the last pass (a710a9aa..ff139a49, 8 files, +1147/-670): per-expert in-place operand reads via pointer arrays, the shared act_ptrs path across all three dtypes, the non-blocking host dispatch, and the "remove redundant checks" cleanup.

Previous round: 9 of the 13 earlier findings verified genuinely fixed in the code, not just acknowledged — m_splits alignment now raises, num_tiles_bound is gone and num_work is threaded exactly from the host, _dummy_arg is cached, the alignment hook is recipe-restricted, the pointer cache is bounded, _transpose_invalid is checked, the device index is in all four launch keys, use_grouped_gemm_triton is in the gate, and _packed_wgrad_view is reused by all three wgrad paths.

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:

  • The 2 GiB require_launch_size check now measures a single-expert representative rather than the packed span the kernel addresses, while offsets are still narrowed to i32 — a previously-rejected shape can now wrap into OOB reads instead of falling back. (medium-high)
  • act_ptrs is derived from a_data_ptrs but gates the b_data_ptrs read in half_prec_gemm.py and mxfp8_gemm.py; fp8_gemm.py already has the correct independent-flag shape. Latent today, but it's the same class as the deref just fixed. (medium)
  • The new flydsl test parametrization never asserts FlyDSL actually ran, so a silent fallback would compare the default backend against itself. NVTE_FLYDSL_GEMM_WARN_FALLBACK + the test_numerics.py warning-capture pattern already solve this. (medium)
  • Plus the dropped wgrad C-shape validation and an uncached get_num_cus(). (low)

Upstream compatibility is clean — all divergence is behind IS_HIP_EXTENSION and flydsl_kernels/ is a ROCm-only tree, so the CUDA path is byte-identical.

Copyright headers: OK (all 11 in-scope files).

Details inline.

Comment thread transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py Outdated
Comment thread tests/pytorch/test_grouped_linear.py Outdated
@github-actions

Copy link
Copy Markdown

Third review pass, scoped to the commits added since the last review (ff139a49..6635016c, 5 files, +96/-23). The origin/dev merge in the head commit was excluded from scope.

Verdict: approve with minor changes. All three code fixes from the previous round check out under tracing:

  • The 2 GiB launch-size guard now measures the packed operand extent rather than the per-expert slice, which is the arithmetically correct bound given how stacked_operand rebases pointers to a shared base.
  • The act_ptrs / b_data_ptrs dereference hole is closed by the new host-side guards; the compile-flag conjunction is now sound.
  • The C-shape validation restoration in half_prec_gemm.py is a byte-identical revert, so the dense path is unaffected.

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).

Comment on lines +1374 to +1378
# 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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Suggested change
# 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, (

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@github-actions

Copy link
Copy Markdown

Re-reviewed the increment since the last pass — commit e169938d4 ("address bot comments"), 4 files, +77/-53.

Both prior findings are genuinely fixed, not just acknowledged:

  • The launch-size guard now takes explicit a_shape/b_shape from the wrappers instead of re-deriving the packed extent in-kernel. I checked the arithmetic at all 6 call sites: MXFP8/BF16 wgrad operands are [m_e, N_out]/[m_e, K_in] padded on dim 0, FP8 wgrad operands are the transposes [N_out, m_e]/[K_in, m_e] padded on the trailing axis — the previously-flagged axis order is now correct, and b_shape=None on fwd/dgrad is arithmetically right since the kernel uses expert-local B offsets.
  • The accuracy test now records every _try_flydsl_grouped dispatch and asserts on the phases that actually ran, which closes the false pass on wgrad. Verified it can't false-fail: fwd/dgrad/wgrad are each reachable in every non-skipped cell, the 256-byte alignment hook in get_align_size_for_quantization keeps the quantized configs from falling back, and monkeypatch.setattr restores the original classmethod descriptor rather than a bound method.

Two new findings, both minor — see inline comments:

  • The replacement comment and docstring on the launch guard drop the rationale for b_shape=None and read as inaccurate for grouped fwd/dgrad, where B is also a single representative expert. Documentation only, no behavior impact.
  • expected <= ran in the new test assertion is a subset check, so it doesn't catch a regression that widens the production wgrad gate. One-character tightening.

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 NVTE_USE_GROUPED_GEMM_FLYDSL plus the gfx950 check.

Comment thread transformer_engine/pytorch/quantization.py Outdated
@Micky774

Copy link
Copy Markdown
Contributor

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.

@aris134

aris134 commented Oct 22, 2026

Copy link
Copy Markdown
Contributor Author

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is it called gd ?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The code uses "binary" in a few places, maybe "compiled kernel" or so could be clearer?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function exists multiple times in the code

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the duplicate in 87b343c



@functools.lru_cache(maxsize=None)
def _num_cus(device):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Duplicated (as get_num_cus)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the duplicates in 87b343c

_COMPILED_LAUNCH_CACHE = {}


def _compiled_launch(key, jit_fn, args):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Duplicated function

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 87b343c

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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These functions are all very similar, maybe they can be factored out a bit

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Refactored in 87b343c

_select_grouped_sched = gd.select_grouped_sched


def _run_mxfp8_grouped(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These functions are all similar, maybe they can be factored out

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Refactored in 87b343c

Comment on lines +350 to +354
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())
):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this call _flydsl_grouped_common_ok instead?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Refactored in 87b343c

Comment thread tests/pytorch/test_grouped_linear.py Outdated
# 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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this delenv can be under else (flydsl) branch below

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 87b343c

Comment thread transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci-level 1 CI test level 1 P2 Priority 2: normal

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants