feat: add FlyDSL batched preshuffle GEMM for FP8 rowwise scaling (WP-G2) by kudomcho · Pull Request #444 · meta-pytorch/MSLK

kudomcho · 2026-07-22T17:58:31Z

Motivation

Add batched FP8 preshuffle GEMM via FlyDSL for gfx950 (MI350). Derivative of WP-G1 — builds on the non-batched preshuffle kernel from PR #434.

Depends on WP-G1 PR #434. Rebased on main after PR #447 merge.

Technical Details

New API: mslk.gemm.flydsl.preshuffle_gemm

flydsl_preshuffle_batched_gemm(XQ, WQ, x_scale, w_scale, ...) — batched GEMM with Grid-Z batching (single kernel launch for all B batches)
Preshuffled weights cached by data_ptr — preshuffle once, reuse on subsequent calls
Registered as torch.library.impl("mslk::f8f8bf16_rowwise_batched", "CUDA") on gfx950

Performance optimizations:

Grid-Z batching — uses gpu.block_id("z") to index into the batch dimension, launching all B batches in a single kernel with grid=(gx, gy, B). Each batch's buffer resource addresses are offset by bz * batch_stride_bytes. Eliminates Python dispatch overhead entirely (was ~258 us for 16 run_compiled calls).
XCD swizzle + waves_per_eu tuning — full parameter sweep across tile configs × xcd_swizzle (0,1,2,4) × waves_per_eu (0,1,2) identified optimal settings per config. xcd_swizzle=1 improves L2 cache reuse across chiplets; waves_per_eu=2 improves scheduling on large shapes.
Profile-guided shape overrides — _SHAPE_OVERRIDES_GFX950 lookup table maps (m_range, N, K) to sweep-optimal configs for N=1280 and N=8192 shapes where the heuristic picks suboptimal tiles.
Batch-aware occupancy heuristic — select_default_config(batch=B) factors Grid-Z parallelism into the occupancy threshold (m_tiles * n_tiles * B >= 64), enabling larger tile configs that were previously rejected.

Benchmark: FP8RowwiseBatchedPreshuffleFlyDSL — new benchmark class in bench/gemm/gemm_ops.py targeting AMD_GFX950, gated on is_flydsl_available().

Test Plan

# Correctness test
pytest test/gemm/gemm_test.py -k FlyDSLPreshuffleBatchedGemmTest

Reproducer: Host-side end-to-end (CUDA events, separate processes)

# CK (no FlyDSL op override):
cat > /tmp/bench_ck_host.py << 'PYEOF'
import torch
from mslk.quantize.triton.fp8_quantize import quantize_fp8_row
SHAPES = [(16,1,1280,8192),(16,128,1280,8192),(16,1024,1280,8192),(16,4096,8192,1024)]
for B,M,N,K in SHAPES:
    xq,xs = quantize_fp8_row(torch.randn(B,M,K,dtype=torch.bfloat16,device='cuda')*0.1)
    wq,ws = quantize_fp8_row(torch.randn(B,N,K,dtype=torch.bfloat16,device='cuda')*0.01)
    for _ in range(10): torch.ops.mslk.f8f8bf16_rowwise_batched(xq,wq,xs,ws)
    torch.cuda.synchronize()
    s,e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    s.record()
    for _ in range(200): torch.ops.mslk.f8f8bf16_rowwise_batched(xq,wq,xs,ws)
    e.record(); torch.cuda.synchronize()
    print(f"CK  ({B},{M},{N},{K}): {s.elapsed_time(e)*1000/200:.1f} us")
PYEOF
MSLK_FLYDSL_DISABLE=1 python /tmp/bench_ck_host.py

# FlyDSL:
cat > /tmp/bench_fly_host.py << 'PYEOF'
import torch
from mslk.quantize.triton.fp8_quantize import quantize_fp8_row
from mslk.gemm.flydsl.preshuffle_gemm import flydsl_preshuffle, flydsl_preshuffle_batched_gemm
SHAPES = [(16,1,1280,8192),(16,128,1280,8192),(16,1024,1280,8192),(16,4096,8192,1024)]
for B,M,N,K in SHAPES:
    xq,xs = quantize_fp8_row(torch.randn(B,M,K,dtype=torch.bfloat16,device='cuda')*0.1)
    wq,ws = quantize_fp8_row(torch.randn(B,N,K,dtype=torch.bfloat16,device='cuda')*0.01)
    wq_shuf = torch.stack([flydsl_preshuffle(wq[i]) for i in range(B)])
    out = torch.empty(B,M,N,dtype=torch.bfloat16,device='cuda')
    for _ in range(10): flydsl_preshuffle_batched_gemm(xq,wq_shuf,xs,ws,out=out)
    torch.cuda.synchronize()
    s,e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    s.record()
    for _ in range(200): flydsl_preshuffle_batched_gemm(xq,wq_shuf,xs,ws,out=out)
    e.record(); torch.cuda.synchronize()
    print(f"FLY ({B},{M},{N},{K}): {s.elapsed_time(e)*1000/200:.1f} us")
PYEOF
python /tmp/bench_fly_host.py

Test Result

Correctness: 8/8 shapes pass

test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_0 PASSED [ 12%]
test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_1 PASSED [ 25%]
test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_2 PASSED [ 37%]
test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_3 PASSED [ 50%]
test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_4 PASSED [ 62%]
test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_5 PASSED [ 75%]
test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_6 PASSED [ 87%]
test/gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_7 PASSED [100%]
================ 13 passed, 210 deselected, 1 warning in 8.73s ================

Host-side end-to-end: FlyDSL vs CK (B=16, gfx950 MI350, CUDA events, 200 iters)

FlyDSL uses Grid-Z batching (single kernel launch). CK uses a single C++ batched kernel. Speedup > 1.0 means FlyDSL is faster.

Shape (B=16)	CK (us)	FlyDSL (us)	Speedup (CK/FlyDSL)
(1, 1280, 8192)	30.3	27.6	1.10x
(32, 1280, 8192)	31.1	30.2	1.03x
(128, 1280, 8192)	61.2	57.2	1.07x
(512, 1280, 8192)	120.6	132.9	0.91x
(1024, 1280, 8192)	267.5	243.2	1.10x
(4096, 1280, 8192)	729.7	763.4	0.96x
(1, 8192, 1024)	21.0	22.1	0.95x
(128, 8192, 1024)	44.1	35.9	1.23x
(512, 8192, 1024)	132.8	115.1	1.15x
(1024, 8192, 1024)	257.3	220.6	1.17x
(4096, 8192, 1024)	958.6	818.8	1.17x

FlyDSL beats CK on 8 of 11 shapes (up to 1.23x faster). Remaining 3 shapes are within 4–9% of CK.

GPU kernel time vs host time breakdown (rocprof + CUDA events)

Confirms the speedups are from genuine GPU kernel improvements, not host-side artifacts. GPU time measured via rocprof --stats, host time via CUDA events (separate runs, no rocprof overhead).

Shape (B=16)	CK GPU (us)	FLY GPU (us)	GPU speedup	CK Host (us)	FLY Host (us)	Host speedup
(1, 1280, 8192)	30.8	28.3	1.09x	30.4	27.6	1.10x
(128, 1280, 8192)	56.2	49.5	1.14x	60.6	57.5	1.05x
(512, 1280, 8192)	112.1	124.0	0.90x	120.0	131.7	0.91x
(1024, 1280, 8192)	270.9	236.7	1.14x	267.1	241.6	1.11x
(4096, 1280, 8192)	830.2	882.6	0.94x	731.9	764.4	0.96x
(1, 8192, 1024)	21.2	21.3	1.00x	21.2	22.2	0.95x
(128, 8192, 1024)	36.4	33.3	1.09x	46.5	34.7	1.34x
(1024, 8192, 1024)	284.6	201.0	1.42x	259.2	218.7	1.19x
(4096, 8192, 1024)	1031.4	941.6	1.10x	962.1	818.1	1.18x

FlyDSL GPU kernel is faster on 7 of 9 shapes (up to 1.42x at M=1024, K=1024). The xcd_swizzle + waves_per_eu tuning drives the kernel-level gains. CK's (128,8192,1024) host overhead (22%) is notably higher than FlyDSL's (4%), explaining the larger host-level speedup on that shape.

Optimization impact summary

Optimization	Impact
Grid-Z batching	Eliminated 16-launch Python dispatch overhead (~258 us → single launch). Small-M speedup from 8–12x slower to parity or faster.
xcd_swizzle=1	+5–20% on most shapes via improved L2 cache reuse across XCDs
waves_per_eu=2	+5–15% on large-M shapes via better wave scheduling
Shape overrides	+10–60% on N=1280 shapes vs heuristic-only config selection
Occupancy heuristic	Up to 4x speedup on N=1280 shapes by avoiding low-tile-count configs

Submission Checklist

cthi

Thanks, please also check the CI lint as it is failing.

Left some comments, also would be good to nail down the code structure for the new FlyDSL gemms. For exmaple do you plan to have flydsl/_kernels/[...] as building blocks for FlyDSL? if so why do we put the kernel in mslk/gemm/flydsl/_kernels/preshuffle_gemm.py ?

cthi · 2026-07-24T15:50:18Z

Largely looks good. Please resolve the remaining module imports.

cthi · 2026-07-24T15:47:43Z

+    def supported(self) -> bool:
+        if not super().supported:
+            return False
+        from mslk.utils.flydsl import is_flydsl_available


Use top-level import here.

Moved to top-level, guarded by if is_flydsl_available().

cthi · 2026-07-24T15:48:16Z

+    def supported(self) -> bool:
+        if get_current_accelerator() not in self.supported_accelerators:
+            return False
+        from mslk.utils.flydsl import is_flydsl_available


Use top-level import here.

Same, top-level now.

cthi · 2026-07-24T15:49:15Z

+    @classmethod
+    def setUpClass(cls):
+        cls.device = torch.accelerator.current_accelerator()
+        from mslk.utils.flydsl import is_flydsl_available


Use top-level import here.

Done, moved to the top-level import block.

cthi · 2026-07-27T13:09:39Z

Seeing illegal memory writes in unit test from both runs:

gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_0 PASSED
gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_1 PASSED
gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_2 PASSED
gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_3 PASSED
gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_4 PASSED
gemm/gemm_test.py::FlyDSLPreshuffleBatchedGemmTest::test_gemm_5 PASSED
Memory access fault by GPU node-2 (Agent handle: 0x583c88c7ee80) on address 0x70992b801000. Reason: Write access to a read-only page.
Fatal Python error: Aborted

kudomcho · 2026-07-27T19:57:46Z

memory access fault resolved. Test logs are now provided on PR info

Add FlyDSL JIT backend for FP8 rowwise preshuffle GEMM on gfx950 (MI350). Provides single and batched FP8 preshuffle GEMM via FlyDSL, with optional HIP graph acceleration for the batched path. Also registers as the ROCm implementation of the mslk rowwise FP8 ops on gfx950. New module: mslk.gemm.flydsl - flydsl_preshuffle(src) — weight shuffle into FlyDSL layout - flydsl_preshuffle_gemm() — JIT-compile and run the kernel - _configs.py — default tile configs for gfx950 heuristic selection - _kernels/ — ported FlyDSL kernel compiler Shared kernel infra: mslk.flydsl.kernels.mma - mfma_epilogues.py, mfma_preshuffle_pipeline.py Op wiring (gfx950, gated on is_flydsl_available()): - torch.ops.mslk.f8f8bf16_rowwise → FlyDSL (bf16) - torch.ops.mslk.f8f8f16_rowwise → FlyDSL (fp16) - torch.ops.mslk.f8f8bf16_rowwise_out → FlyDSL (out-tensor) - torch.ops.mslk.f8f8bf16_rowwise_batched → FlyDSL batched Benchmark: FP8RowwisePreshuffleFlyDSL in bench/gemm/gemm_ops.py Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

… B batches Use gpu.block_id("z") to index into batch dimension, offsetting all buffer resource addresses by batch stride. Eliminates 16 separate kernel launches in favor of one launch with grid=(gx, gy, B).

meta-codesync · 2026-08-04T16:16:46Z

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

Remove the old FP8RowwisePreshuffleFlyDSL definition that used is_flydsl_version_at_least() with module-level imports, keeping the newer version with lazy imports. Remove unused is_flydsl_available import from gemm_test.py.

cthi · 2026-08-07T16:49:57Z

+            _batched_preshuffle_cache: dict = {}
+
+            def _get_batched_preshuffled(WQ: Tensor) -> Tensor:
+                key = WQ.data_ptr()


I think I missed this in a prior PR, but basically caching stuff based on the data_ptr is not really valid PyTorch, CCA could re-use these pointer for something else. Can we avoid doing this? Why do we need to cache the preshuffled W? User can simply pass it pre-shuffled properly.

Removed the data_ptr cache entirely. The batched dispatch now expects pre-shuffled weights from the caller.

cthi · 2026-08-07T16:54:10Z

+                WQ: Tensor,
+                x_scale: Tensor,
+                w_scale: Tensor,
+                bias: Optional[Tensor] = None,


As bias is not supported, raise an exception if it is present.

Added, raises NotImplementedError if bias is not None.

cthi · 2026-08-07T16:54:38Z

+                use_fast_accum: bool = True,
+                output: Optional[Tensor] = None,
+            ) -> Tensor:
+                WQ_shuf = _get_batched_preshuffled(WQ)


I commented about this above, but we should let user explicitly preshuffle as needed instead of trying to cache it.

Removed the cache. Caller is responsible for preshuffling now.

- Empty __init__.py re-exports, follow triton convention (direct imports) - Move all method-level FlyDSL imports to top-level in bench and test - Remove data_ptr-based preshuffle cache (invalid under CCA pointer reuse) - Batched dispatch now expects caller to pass pre-shuffled weights - Raise NotImplementedError if bias is passed to batched op

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

cthi requested changes Jul 23, 2026

View reviewed changes

kudomcho force-pushed the wp-g2/flydsl-batched-gemm branch 2 times, most recently from ae3d13e to 81e35b4 Compare July 24, 2026 15:21

kudomcho force-pushed the wp-g2/flydsl-batched-gemm branch 2 times, most recently from 7c59584 to 3b6632d Compare July 24, 2026 16:29

cthi reviewed Jul 24, 2026

View reviewed changes

kudomcho force-pushed the wp-g2/flydsl-batched-gemm branch 3 times, most recently from 3f248d0 to 3d52b6b Compare July 27, 2026 19:56

kudomcho force-pushed the wp-g2/flydsl-batched-gemm branch 4 times, most recently from 9e95d95 to 7d4bc9d Compare July 29, 2026 18:10

aryaman-gupta mentioned this pull request Jul 31, 2026

Add FlyDSL grouped groupwise FP8 GEMM for ROCm #456

Closed

kudomcho and others added 4 commits August 4, 2026 16:03

style: fix ufmt import ordering in gemm_ops.py

affd021

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

perf: Grid-Z batching for batched GEMM — single kernel launch for all…

b739f75

… B batches Use gpu.block_id("z") to index into batch dimension, offsetting all buffer resource addresses by batch stride. Eliminates 16 separate kernel launches in favor of one launch with grid=(gx, gy, B).

style: fix ufmt lint issues in batched launcher

cf78e5b

kudomcho force-pushed the wp-g2/flydsl-batched-gemm branch from 3ec250c to cf78e5b Compare August 4, 2026 16:06

cthi reviewed Aug 7, 2026

View reviewed changes

Uh oh!

Conversation

kudomcho commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Reproducer: Host-side end-to-end (CUDA events, separate processes)

Test Result

Correctness: 8/8 shapes pass

Host-side end-to-end: FlyDSL vs CK (B=16, gfx950 MI350, CUDA events, 200 iters)

GPU kernel time vs host time breakdown (rocprof + CUDA events)

Optimization impact summary

Submission Checklist

Uh oh!

cthi left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cthi commented Jul 24, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

kudomcho commented Jul 27, 2026

Uh oh!

meta-codesync Bot commented Aug 4, 2026

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kudomcho commented Jul 22, 2026 •

edited

Loading

cthi commented Jul 27, 2026 •

edited

Loading

cthi Aug 7, 2026 •

edited

Loading

kudomcho Aug 7, 2026 •

edited

Loading

kudomcho Aug 7, 2026 •

edited

Loading