Add FlyDSL grouped groupwise FP8 GEMM for ROCm by aryaman-gupta · Pull Request #456 · meta-pytorch/MSLK

aryaman-gupta · 2026-07-27T11:35:42Z

Adds a FlyDSL implementation of the grouped groupwise FP8 GEMM for ROCm and makes it the ROCm backend for mslk::f8f8bf16_groupwise_grouped.

Two ops share one kernel, selected by a compile-time b_preshuffled flag:

mslk::f8f8bf16_groupwise_grouped — plain row-major [G, N, K] weights.
mslk::f8f8bf16_groupwise_grouped_preshuffle — new ROCm-only sibling taking weights already swizzled into the MFMA B layout by mslk.quantize.shuffle.preshuffle_b_mfma. Callers shuffle once at load time and the op does no shuffling.

Both require N and K to be multiples of 128, matching the weight scale-block granularity. Activation scales are consumed in the per-group block layout produced by quantize_fp8_group(m_sizes=...).

Performance

CUDA-graph replay latency in µs on gfx950 (MI350X), G=8. Both implementations are tuned: Triton through its own @triton.autotune, FlyDSL through MSLK autotune over its full candidate tile set.

The last two columns are the latency reduction relative to Triton, (t_triton − t_flydsl) / t_triton × 100, so higher is better and a negative value is a slowdown.

Representative rows (two per shape) from a 50-row sweep. Preshuffle is faster on 45 of 50 shapes, median +28%, range −25% to +66%.

totM	N	K	triton µs	plain µs	preshuffle µs	plain	preshuffle
32	1280	8192	42.9	65.2	53.5	−52%	−25%
128	1280	8192	82.0	66.0	54.5	+20%	+34%
1024	2048	7168	99.4	62.2	53.0	+37%	+47%
16384	2048	7168	699.5	470.9	405.4	+33%	+42%
1024	7168	2304	74.1	42.8	43.5	+42%	+41%
16384	7168	2304	623.6	547.0	482.6	+12%	+23%
32	7168	8192	111.8	119.3	97.0	−7%	+13%
128	7168	8192	156.8	116.2	102.0	+26%	+35%
1	8192	1024	20.7	13.4	12.0	+35%	+42%
128	8192	1024	29.6	19.8	19.9	+33%	+33%
32	8192	3584	46.2	41.9	39.9	+9%	+14%
128	8192	3584	62.8	43.4	41.0	+31%	+35%
32	13312	6656	158.8	154.0	146.1	+3%	+8%
128	13312	6656	231.8	157.2	148.6	+32%	+36%
32	13312	16384	373.2	453.5	367.2	−22%	+2%
128	13312	16384	549.0	463.6	380.7	+16%	+31%
32	16384	6656	266.7	167.1	162.6	+37%	+39%
128	16384	6656	507.5	186.2	174.2	+63%	+66%

FlyDSL is slower on 1280 × 8192 up to totM 96, where the output is narrow enough that it cannot produce as many tiles as Triton; it leads on every other shape measured.

Every shape is checked against a per-group bf16 reference before timing, and additionally against Triton: the two agree to within one bf16 ulp on every shape, and the residual against the reference is FP8 quantization error, identical for both (~3.5% of max|ref|).

Tile selection

With MSLK_AUTOTUNE_ENABLE set, FlyDSL autotune benchmarks the candidate tiles on a cache miss and persists the winner to disk, keyed on (nextPow2(total_M), N, K, b_preshuffled). This follows the CUTLASS precedent in f8f8bf16_groupwise_grouped.cu, including the env gate: unset — the default, and what CI uses — a fixed tile is used and nothing is benchmarked.

Benchmarking cannot run inside a graph capture, so a shape must be warmed before it is captured. Everything else about dispatch is capture-safe: the grid extent is derived from static shapes alone, nothing on the host reads tensor contents, and no helper kernels are launched, so eager and captured dispatch issue the same single kernel.

Architecture support

gfx950 uses the wide mfma_scale_f32_16x16x128_f8f6f4 with a neutral E8M0 scale and applies the FP32 scales in software. gfx942 lacks that instruction and takes a narrow 16x16x32 path. LDS budgets are checked per architecture (64 KiB on gfx942, 160 KiB on gfx950), so a tile that fits on MI350 is rejected at compile time on MI300 rather than failing in the backend.

Correctness has been verified on gfx950 (MI350) and gfx942 (MI300).

Testing

test/gemm/gemm_test.py::test_f8f8bf16_groupwise_grouped covers both variants across four group splits, including uneven and non-tile-aligned ones, and requires FlyDSL on ROCm.

The tolerance there is tightened from atol=8e-2 to 1e-2. Outputs on those shapes are O(1e-2), so the previous bound sat several times above the signal and would pass regardless of which scales the kernel applied.

Notes for reviewers

ROCm CI already triggers on this change; mslk/gemm/flydsl/** is added to its path filter so later changes confined to the op wrappers do not skip it.

The FlyDSL modules vendored here are also vendored, at the same path, by #434 and #444. If either lands first, this PR will be refactored to drop the repeated files.

mslk/flydsl/kernels/{common,mma}/ and gemm/fp8_gemm_utils.py are vendored unmodified from FlyDSL and are best reviewed as an import. The kernel itself is mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_{contiguous,common}.py, and the op wrapper is mslk/gemm/flydsl/fp8_groupwise_grouped_gemm.py.

Add the cross-cutting FlyDSL kernel building blocks that MSLK's FlyDSL-backed GEMM/MoE kernels depend on, under mslk/flydsl/kernels/ mirroring FlyDSL's own common/mma/gemm domain layout. These modules are not part of the installable flydsl package and must be vendored to be importable by MSLK. Co-Authored-By: Claude <noreply@anthropic.com>

Adds the ROCm preshuffle sibling of f8f8bf16_groupwise_grouped, backed by a FlyDSL contiguous grouped GEMM kernel that dispatches over a linearized per-tile map (compact unpadded output, no tile spanning a group boundary) and consumes weights pre-swizzled into the MFMA B layout. Registers the op when FlyDSL is available, adds preshuffle_b_mfma for the one-time weight swizzle, and covers correctness in the grouped GEMM test. Co-Authored-By: Claude <noreply@anthropic.com>

Extends the FlyDSL contiguous grouped groupwise FP8 GEMM kernel two ways, both unified in grouped_gemm_blockscale_contiguous.py: - Wide-MFMA software-scaling path: on gfx950 the FP32-scale path now uses the wide mfma_scale_f32_16x16x128 with a neutral E8M0 scale + per-scale-block FP32 apply (gfx942 keeps the narrow 16x16x32 path), matching the working blockscale_preshuffle_gemm kernel. Large speedup on K-heavy shapes. - Plain (non-preshuffled) B via a b_preshuffled flag (default True). When False, B is plain row-major [G,N,K] and is staged HBM->LDS->registers (mirroring the A pipeline: coalesced dwordx4 loads, xor16-swizzled bank-conflict-free LDS) instead of the preshuffle HBM->registers path. This is the apples-to-apples input for the Triton f8f8bf16_groupwise_grouped op. The two variants share the whole kernel body (tile-map group dispatch, scale-apply, wide-MFMA, CShuffle epilogue); only the B load stage + its LDS allocation differ. grouped_gemm_blockscale_plain.py is a thin delegator. Correctness verified on gfx950 for both variants (plain grouped test + existing preshuffle acceptance cases). Co-Authored-By: Claude <noreply@anthropic.com>

Registers the FlyDSL plain-B kernel as the ROCm impl of the plain op mslk::f8f8bf16_groupwise_grouped (C++ schema in gemm_ops.cpp), consuming plain row-major [G,N,K] weights via the unified kernel with b_preshuffled=False. The Triton impl now registers only as a fallback when FlyDSL is unavailable (it is slated for removal). Drops the redundant grouped_gemm_blockscale_plain delegator in favor of calling compile_grouped_gemm_blockscale_contiguous(..., b_preshuffled=False) directly. Verified: both grouped ops pass the acceptance suite end-to-end via torch.ops.mslk (plain cases 0-3 + preshuffle cases 4-7), 8/8 on gfx950. Co-Authored-By: Claude <noreply@anthropic.com>

The FlyDSL wrapper rebuilt a per-tile dispatch map every call via ~8 tiny tensor-op kernels (arange/cumsum/searchsorted/where). Under CUDA-graph capture these replay as separate kernels with fixed per-call overhead that dominated small-M shapes. Replace the host tile-map with in-kernel group resolution: the kernel takes M_sizes directly and, for its flat M-tile id, runs a compile-time unrolled loop over the groups (num_groups is constexpr) to find the owning group, row start, and row limit. The grid keeps its compact upper-bound extent. Also wire per-shape tile selection through get_tile_config (static tables) and unify the plain and preshuffle ops behind _dispatch_grouped_gemm. Co-Authored-By: Claude <noreply@anthropic.com>

@autotune

Replace the static tile table with FlyDSL's native @autotune, gated on MSLK_AUTOTUNE_ENABLE (CUTLASS precedent): when set, autotune benchmarks the candidate tiles on a cache-miss and persists the winner to disk; when unset, a fixed default tile is used with no benchmarking (the CI / graph-capture path). The autotuner wraps a thin target that forwards to the existing compile factory, so the kernel is unchanged. The tuning key is (nextPow2(TotalM), N, K, b_preshuffled): TotalM is bucketed to match graph-capture buckets and bound the pre-warm set, N/K separate problem shapes, and b_preshuffled keeps the two kernels' entries distinct in one shared disk cache. Invalid tiles are pruned per shape before benchmarking. FlyDSL's Autotuner discards the target return value, so the wrapper returns the in-place output buffer itself. Co-Authored-By: Claude <noreply@anthropic.com>

mslk/utils/flydsl.py was removed when FlyDSL became a required ROCm dependency (meta-pytorch#447); is_flydsl_available/require_flydsl now live in mslk.flydsl.common and run_compiled in mslk.flydsl.jit. Point the grouped-GEMM op, its Triton fallback gate, the gemm package init, and the test at the new modules.

The preshuffle path had no host-side LDS check, so an oversized tile reached the compiler backend, which reports an LDS overflow as a hard error that kills the process -- autotune could not catch and skip it. Add validate_lds_budget_preshuffle (A ping-pong aliased with the CShuffle epilogue) and run it alongside the plain check before tracing. The capacity was also hardcoded to 64 KiB. It is arch-dependent: 64 KiB on gfx942 but 160 KiB on gfx950. Source it from FlyDSL's SMEM_CAPACITY_MAP so the limit matches the compiler that enforces it, falling back to 64 KiB for unknown architectures. A tile_m=256 tile that fits on MI350 is now correctly rejected on MI300, and the wider gfx950 tiles are no longer needlessly excluded.

The plain f8f8bf16_groupwise_grouped op is FlyDSL-backed on ROCm, so require FlyDSL for both variants there rather than exercising the Triton fallback. Drop test/gemm/flydsl_plain_grouped_test.py: it covered the same four group shapes as the parameterized test in gemm_test.py.

quantize_fp8_group(m_sizes=...) stores scale_a as per-group blocks: group g starts at M_start*scale_k and holds element (local_m, k_g) at local_m + k_g*M_g. The kernel instead indexed it as a global [scale_k, TotalM] transpose (k_g*TotalM + row_global), which only coincides with the real layout when there is a single group or a single K-block. With G > 1 and K > 128 every group past the first applied another group's scales, silently producing wrong results. Carry the owning group's M_start and M_g out of the group-resolution loop and use them to address scale_a. Output now matches the Triton implementation bit-exactly on the shapes that previously diverged. Tighten the grouped test tolerance to match: outputs are O(1e-2) and FP8 block quantization accounts for ~4e-3, so the old atol of 8e-2 was several times the signal and passed regardless of the scales used.

The kernel resolves group ownership from m_sizes, but the module docstring and the launcher comment still described the tile_group/tile_row_start/ tile_row_limit arrays that preceded it, and the kernel still declared i32_num_m_tiles without using it (the launcher keeps it for the grid extent). Both docstrings also described scale_a as a global [scale_k, M_total] transpose; document the per-group block layout the kernel actually reads. Note in the wrapper docstring and the gemm package init that FlyDSL registers the plain op as well as the preshuffle sibling. Rename the local i32-constant helper to _i32 so it no longer shadows the _c compile-constants namedtuple from the enclosing scope.

The consuming kernels build their B layout with NLane=16, KLane=4 and KPack=16 hardcoded (make_preshuffle_b_layout), so the NLane parameter could only ever produce a layout they cannot read. Drop it and fix the extents. Add the divisibility precondition as well: an N or K that does not divide evenly previously surfaced as a bare reshape error naming internal dimensions. Accept **kwargs in _prune_tiles so it works with either the FlyDSL (configs, sig_args) or Triton (configs, named_args, **meta) calling convention.

The wrapper hands the FP8 operands to the kernel as raw bytes, so weights quantized in the other FP8 format were applied with the wrong exponent bias instead of being rejected. Check both operands against the format the device uses (fnuz on gfx942, OCP elsewhere), and restore the M_sizes rank and length checks the Triton implementation performed. Read m_sizes directly as int64 by loading the low dword of each element, which removes the per-call narrowing kernel the wrapper previously enqueued.

Describe what the code does rather than how it came to look that way, and correct the plain-B K-loop docstring: the loop is software-pipelined and uses two barriers per K-tile, not one, and single-buffered LDS is a budget tradeoff rather than unfinished work.

The schema comment still pointed readers at the Triton implementation.

Every other op in mslk/gemm declares its schema in csrc/gemm/gemm_ops.cpp, so move this one there too, next to the other ROCm-only preshuffle GEMMs. The implementation stays in Python, matching ops such as bf16i4bf16_rowwise. Register both ops through one guarded helper: the schema is absent in a python-only build, and a repeated import must not fail on rebinding.

Argument marshalling packs each memref extent as int32, so a flattened operand overflowed once it reached 2**31 elements: at G=8, N=K=16384 the B operand is exactly 2**31 bytes and no configuration could launch, though the Triton implementation handles that shape. Pass the operands with their natural shape; the kernel addresses them as flat byte buffers through explicit descriptors either way.

meta-codesync · 2026-07-28T19:48:00Z

@q10 has imported this pull request. If you are a Meta employee, you can view this in D113947727.

mslk/gemm/flydsl/ holds the op wrappers for the FlyDSL GEMM kernels but was not in the ROCm trigger paths, so a change confined to it would skip ROCm CI.

q10 · 2026-08-03T23:51:02Z

@aryaman-gupta could you rebase this PR on latest main?

Summary: Pull Request resolved: meta-pytorch#470 Adds a FlyDSL implementation of the grouped groupwise FP8 GEMM for ROCm and makes it the ROCm backend for `mslk::f8f8bf16_groupwise_grouped`. Two ops share one kernel, selected by a compile-time `b_preshuffled` flag: - **`mslk::f8f8bf16_groupwise_grouped`** — plain row-major `[G, N, K]` weights. - **`mslk::f8f8bf16_groupwise_grouped_preshuffle`** — new ROCm-only sibling taking weights already swizzled into the MFMA B layout by `mslk.quantize.shuffle.preshuffle_b_mfma`. Callers shuffle once at load time and the op does no shuffling. Both require `N` and `K` to be multiples of 128, matching the weight scale-block granularity. Activation scales are consumed in the per-group block layout produced by `quantize_fp8_group(m_sizes=...)`. ## Performance CUDA-graph replay latency in µs on gfx950 (MI350X), G=8. **Both implementations are tuned**: Triton through its own `triton.autotune`, FlyDSL through MSLK autotune over its full candidate tile set. The last two columns are speedup over Triton, `(t_triton / t_flydsl − 1) × 100`, so **higher is better** and a negative value is a slowdown. Representative rows (two per shape) from a 50-row sweep. Preshuffle is faster on 45 of 50 shapes, median **+38%**, range −20% to +191%. | totM | N | K | triton µs | plain µs | preshuffle µs | plain | preshuffle | |-----:|----:|-----:|----------:|---------:|--------------:|------:|-----------:| | 32 | 1280 | 8192 | 42.9 | 65.2 | 53.5 | −34% | −20% | | 128 | 1280 | 8192 | 82.0 | 66.0 | 54.5 | +24% | +50% | | 1024 | 2048 | 7168 | 99.4 | 62.2 | 53.0 | +60% | +88% | | 16384 | 2048 | 7168 | 699.5 | 470.9 | 405.4 | +49% | +73% | | 1024 | 7168 | 2304 | 74.1 | 42.8 | 43.5 | +73% | +70% | | 16384 | 7168 | 2304 | 623.6 | 547.0 | 482.6 | +14% | +29% | | 32 | 7168 | 8192 | 111.8 | 119.3 | 97.0 | −6% | +15% | | 128 | 7168 | 8192 | 156.8 | 116.2 | 102.0 | +35% | +54% | | 1 | 8192 | 1024 | 20.7 | 13.4 | 12.0 | +54% | +72% | | 128 | 8192 | 1024 | 29.6 | 19.8 | 19.9 | +49% | +49% | | 32 | 8192 | 3584 | 46.2 | 41.9 | 39.9 | +10% | +16% | | 128 | 8192 | 3584 | 62.8 | 43.4 | 41.0 | +45% | +53% | | 32 | 13312 | 6656 | 158.8 | 154.0 | 146.1 | +3% | +9% | | 128 | 13312 | 6656 | 231.8 | 157.2 | 148.6 | +47% | +56% | | 32 | 13312 | 16384 | 373.2 | 453.5 | 367.2 | −18% | +2% | | 128 | 13312 | 16384 | 549.0 | 463.6 | 380.7 | +18% | +44% | | 32 | 16384 | 6656 | 266.7 | 167.1 | 162.6 | +60% | +64% | | 128 | 16384 | 6656 | 507.5 | 186.2 | 174.2 | +173% | +191% | FlyDSL is slower on `1280 × 8192` up to totM 96, where the output is narrow enough that it cannot produce as many tiles as Triton; it leads on every other shape measured. Every shape is checked against a per-group bf16 reference before timing, and additionally against Triton: the two agree to within one bf16 ulp on every shape, and the residual against the reference is FP8 quantization error, identical for both (~3.5% of max|ref|). ## Tile selection With `MSLK_AUTOTUNE_ENABLE` set, FlyDSL autotune benchmarks the candidate tiles on a cache miss and persists the winner to disk, keyed on `(nextPow2(total_M), N, K, b_preshuffled)`. This follows the CUTLASS precedent in `f8f8bf16_groupwise_grouped.cu`, including the env gate: unset — the default, and what CI uses — a fixed tile is used and nothing is benchmarked. Benchmarking cannot run inside a graph capture, so a shape must be warmed before it is captured. Everything else about dispatch is capture-safe: the grid extent is derived from static shapes alone, nothing on the host reads tensor contents, and no helper kernels are launched, so eager and captured dispatch issue the same single kernel. ## Architecture support gfx950 uses the wide `mfma_scale_f32_16x16x128_f8f6f4` with a neutral E8M0 scale and applies the FP32 scales in software. gfx942 lacks that instruction and takes a narrow `16x16x32` path. LDS budgets are checked per architecture (64 KiB on gfx942, 160 KiB on gfx950), so a tile that fits on MI350 is rejected at compile time on MI300 rather than failing in the backend. Correctness has been verified on gfx950 (MI350) and gfx942 (MI300). ## Testing `test/gemm/gemm_test.py::test_f8f8bf16_groupwise_grouped` covers both variants across four group splits, including uneven and non-tile-aligned ones, and requires FlyDSL on ROCm. The tolerance there is tightened from `atol=8e-2` to `1e-2`. Outputs on those shapes are O(1e-2), so the previous bound sat several times above the signal and would pass regardless of which scales the kernel applied. ## Notes for reviewers `mslk/flydsl/kernels/{common,mma}/` and `gemm/fp8_gemm_utils.py` are vendored unmodified from FlyDSL and are best reviewed as an import. The kernel itself is `mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_{contiguous,common}.py`, and the op wrapper is `mslk/gemm/flydsl/fp8_groupwise_grouped_gemm.py`. Pull Request resolved: meta-pytorch#456 Reviewed By: jwfromm Differential Revision: D113947727 Pulled By: q10

meta-codesync · 2026-08-06T03:31:05Z

@q10 merged this pull request in d40c04f.

aryaman-gupta and others added 8 commits July 27, 2026 09:49

meta-cla Bot added the cla signed label Jul 27, 2026

aryaman-gupta added 9 commits July 27, 2026 13:55

Tidy grouped GEMM comments

bbcffbd

Describe what the code does rather than how it came to look that way, and correct the plain-B K-loop docstring: the loop is software-pipelined and uses two barriers per K-tile, not one, and single-buffered LDS is a budget tradeoff rather than unfinished work.

Name FlyDSL as the ROCm impl of f8f8bf16_groupwise_grouped

1eac28f

The schema comment still pointed readers at the Triton implementation.

Run ROCm CI when the FlyDSL GEMM ops change

8a0e6a2

mslk/gemm/flydsl/ holds the op wrappers for the FlyDSL GEMM kernels but was not in the ROCm trigger paths, so a change confined to it would skip ROCm CI.

aryaman-gupta changed the title ~~Flydsl groupwise grouped~~ Add FlyDSL grouped groupwise FP8 GEMM for ROCm Jul 31, 2026

aryaman-gupta marked this pull request as ready for review July 31, 2026 10:38

meta-codesync Bot closed this in d40c04f Aug 6, 2026

meta-codesync Bot added the Merged label Aug 6, 2026

Navigation Menu

Uh oh!

Add FlyDSL grouped groupwise FP8 GEMM for ROCm - #456

Add FlyDSL grouped groupwise FP8 GEMM for ROCm#456
aryaman-gupta wants to merge 18 commits into
meta-pytorch:mainfrom
aryaman-gupta:flydsl-groupwise-grouped

aryaman-gupta commented Jul 27, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented Jul 28, 2026

Uh oh!

q10 commented Aug 3, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

aryaman-gupta commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance

Tile selection

Architecture support

Testing

Notes for reviewers

Uh oh!

meta-codesync Bot commented Jul 28, 2026

Uh oh!

q10 commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

aryaman-gupta commented Jul 27, 2026 •

edited

Loading

q10 commented Aug 3, 2026 •

edited

Loading