Conversation
Collaborator
|
Hi @kudomcho Could you share the data you mentioned CK wins remaining 20 shapes (small-M dominated by FlyDSL ~22µs dispatch overhead).. Note that we need to on-pair with CK. Please specify the dispatch overhead with concreate data to support |
Collaborator
Author
|
@amd-weisun Re-Validated with kernel and host dispatch time using rocprof and reproducer provided on the PR. Thank you. |
|
@q10 has imported this pull request. If you are a Meta employee, you can view this in D112057807. |
11 tasks
kudomcho
force-pushed
the
wp-g1/flydsl-preshuffle-gemm
branch
4 times, most recently
from
July 24, 2026 19:27
9fae6cb to
8193af2
Compare
Port FlyDSL preshuffle GEMM kernel from ROCm/aiter to MSLK, enabling FP8 rowwise GEMM with FlyDSL JIT backend on gfx950 (MI350). - Add mslk.gemm.flydsl module with flydsl_preshuffle() and flydsl_preshuffle_gemm() public API - Port FlyDSL kernel compiler files (preshuffle_gemm, mfma_epilogues, mfma_preshuffle_pipeline, tensor_shim) to _kernels subpackage - Add FP8RowwisePreshuffleFlyDSL benchmark class targeting gfx950 - Add FlyDSLPreshuffleGemmTest correctness test Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
- E226: add whitespace around arithmetic operator in mfma_epilogues.py - F841: suppress unused variable warnings (_c64_elem, _K1) with noqa - F811: suppress intentional redefinition of load_fp4_scale_chunk with noqa Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Fix import ordering and formatting to pass MSLK CI lint checks. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
…ops on gfx950 Register torch.library.impl for CUDA dispatch key on ROCm so that: - torch.ops.mslk.f8f8bf16_rowwise() -> FlyDSL (bf16 output) - torch.ops.mslk.f8f8f16_rowwise() -> FlyDSL (fp16 output) - torch.ops.mslk.f8f8bf16_rowwise_out() -> FlyDSL (out-tensor variant) All three ops preshuffle weights on-the-fly and dispatch to flydsl_preshuffle_gemm. Replaces CK backend on gfx950. Tested on gfx950 MI350: 18/18 shapes pass (bf16 + fp16 + out-tensor), max_diff=0.0 vs dequantized reference on all shapes. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
The op impl was calling flydsl_preshuffle(WQ) on every inference call, adding ~1us overhead that made FlyDSL appear slower than CK on small shapes. In production, weights are static — preshuffle once and cache by data pointer. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Move mfma_epilogues.py and mfma_preshuffle_pipeline.py from the preshuffle-GEMM-private _kernels/ directory to mslk/flydsl/kernels/mma/, matching the shared FlyDSL infra layout convention. Preshuffle-GEMM- specific files (preshuffle_gemm.py, tensor_shim.py) stay in _kernels/. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
… naming - Move implementation from __init__.py to preshuffle_gemm.py (follows triton convention); __init__.py now just re-exports - Merge _configs.py into preshuffle_gemm.py (config is kernel-specific) - Use top-level imports in bench/gemm_ops.py instead of method-level - Rename test method: test_flydsl_preshuffle_gemm -> test_gemm - Move test imports to setUpClass instead of per-test Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
kudomcho
force-pushed
the
wp-g1/flydsl-preshuffle-gemm
branch
4 times, most recently
from
July 28, 2026 00:32
1996423 to
61f95f3
Compare
Update mslk.utils.flydsl → mslk.flydsl.common and fix ufmt sort order. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
kudomcho
force-pushed
the
wp-g1/flydsl-preshuffle-gemm
branch
from
July 28, 2026 00:36
61f95f3 to
ddf3bfe
Compare
kudomcho
added a commit
to kudomcho/MSLK
that referenced
this pull request
Jul 28, 2026
The _cleanup_gpu_caches atexit handler triggers a memory access fault on ROCm 7.0 when deallocating GPU tensors after HIP runtime shutdown. PR meta-pytorch#434 (single GEMM) has no atexit handler and passes ROCm 7.0 CI.
kudomcho
added a commit
to kudomcho/MSLK
that referenced
this pull request
Jul 29, 2026
The _cleanup_gpu_caches atexit handler triggers a memory access fault on ROCm 7.0 when deallocating GPU tensors after HIP runtime shutdown. PR meta-pytorch#434 (single GEMM) has no atexit handler and passes ROCm 7.0 CI.
Contributor
|
Thanks, this looks good. Internally we need to do some small stuff to make it build/run properly. |
11 tasks
cthi
pushed a commit
to cthi/MSLK-1
that referenced
this pull request
Aug 13, 2026
…G2) (meta-pytorch#482) Summary: ## Motivation Add batched FP8 preshuffle GEMM via FlyDSL for gfx950 (MI350). Derivative of WP-G1 — builds on the non-batched preshuffle kernel from PR meta-pytorch#434. Depends on WP-G1 PR meta-pytorch#434. Rebased on main after PR meta-pytorch#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:** 1. **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). 2. **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. 3. **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. 4. **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: ```bash # Correctness test pytest test/gemm/gemm_test.py -k FlyDSLPreshuffleBatchedGemmTest ``` ### Reproducer: Host-side end-to-end (CUDA events, separate processes) ```bash # 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 - [x] Correctness: 8/8 shapes pass (both clang and gcc CI), no crashes - [x] Standalone API: `flydsl_preshuffle_batched_gemm()` - [x] Benchmark class: `FP8RowwiseBatchedPreshuffleFlyDSL` - [x] Op registration: `f8f8bf16_rowwise_batched` on gfx950 - [x] Rebased on main - [x] Gated on `is_flydsl_available()` — graceful fallback - [x] Grid-Z batching: single kernel launch for all B batches - [x] XCD swizzle + waves_per_eu tuning: sweep-optimized per config - [x] Profile-guided shape overrides for N=1280 and N=8192 - [x] Batch-aware occupancy heuristic - [x] Host-side benchmark with reproducer Reviewed By: q10 Differential Revision: D114744187 Pulled By: cthi
cthi
pushed a commit
to cthi/MSLK-1
that referenced
this pull request
Aug 14, 2026
…G2) (meta-pytorch#482) Summary: ## Motivation Add batched FP8 preshuffle GEMM via FlyDSL for gfx950 (MI350). Derivative of WP-G1 — builds on the non-batched preshuffle kernel from PR meta-pytorch#434. Depends on WP-G1 PR meta-pytorch#434. Rebased on main after PR meta-pytorch#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:** 1. **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). 2. **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. 3. **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. 4. **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: ```bash # Correctness test pytest test/gemm/gemm_test.py -k FlyDSLPreshuffleBatchedGemmTest ``` ### Reproducer: Host-side end-to-end (CUDA events, separate processes) ```bash # 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 - [x] Correctness: 8/8 shapes pass (both clang and gcc CI), no crashes - [x] Standalone API: `flydsl_preshuffle_batched_gemm()` - [x] Benchmark class: `FP8RowwiseBatchedPreshuffleFlyDSL` - [x] Op registration: `f8f8bf16_rowwise_batched` on gfx950 - [x] Rebased on main - [x] Gated on `is_flydsl_available()` — graceful fallback - [x] Grid-Z batching: single kernel launch for all B batches - [x] XCD swizzle + waves_per_eu tuning: sweep-optimized per config - [x] Profile-guided shape overrides for N=1280 and N=8192 - [x] Batch-aware occupancy heuristic - [x] Host-side benchmark with reproducer Reviewed By: q10 Differential Revision: D114744187 Pulled By: cthi
cthi
pushed a commit
to cthi/MSLK-1
that referenced
this pull request
Aug 14, 2026
…G2) (meta-pytorch#482) Summary: ## Motivation Add batched FP8 preshuffle GEMM via FlyDSL for gfx950 (MI350). Derivative of WP-G1 — builds on the non-batched preshuffle kernel from PR meta-pytorch#434. Depends on WP-G1 PR meta-pytorch#434. Rebased on main after PR meta-pytorch#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:** 1. **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). 2. **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. 3. **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. 4. **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: ```bash # Correctness test pytest test/gemm/gemm_test.py -k FlyDSLPreshuffleBatchedGemmTest ``` ### Reproducer: Host-side end-to-end (CUDA events, separate processes) ```bash # 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 - [x] Correctness: 8/8 shapes pass (both clang and gcc CI), no crashes - [x] Standalone API: `flydsl_preshuffle_batched_gemm()` - [x] Benchmark class: `FP8RowwiseBatchedPreshuffleFlyDSL` - [x] Op registration: `f8f8bf16_rowwise_batched` on gfx950 - [x] Rebased on main - [x] Gated on `is_flydsl_available()` — graceful fallback - [x] Grid-Z batching: single kernel launch for all B batches - [x] XCD swizzle + waves_per_eu tuning: sweep-optimized per config - [x] Profile-guided shape overrides for N=1280 and N=8192 - [x] Batch-aware occupancy heuristic - [x] Host-side benchmark with reproducer Reviewed By: q10 Differential Revision: D114744187 Pulled By: cthi
meta-codesync Bot
pushed a commit
that referenced
this pull request
Aug 14, 2026
…G2) (#482) Summary: Pull Request resolved: #482 ## 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:** 1. **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). 2. **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. 3. **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. 4. **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()`. Pull Request resolved: #444 Test Plan: ```bash # Correctness test pytest test/gemm/gemm_test.py -k FlyDSLPreshuffleBatchedGemmTest ``` ### Reproducer: Host-side end-to-end (CUDA events, separate processes) ```bash # 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 - [x] Correctness: 8/8 shapes pass (both clang and gcc CI), no crashes - [x] Standalone API: `flydsl_preshuffle_batched_gemm()` - [x] Benchmark class: `FP8RowwiseBatchedPreshuffleFlyDSL` - [x] Op registration: `f8f8bf16_rowwise_batched` on gfx950 - [x] Rebased on main - [x] Gated on `is_flydsl_available()` — graceful fallback - [x] Grid-Z batching: single kernel launch for all B batches - [x] XCD swizzle + waves_per_eu tuning: sweep-optimized per config - [x] Profile-guided shape overrides for N=1280 and N=8192 - [x] Batch-aware occupancy heuristic - [x] Host-side benchmark with reproducer Reviewed By: q10 Differential Revision: D114744187 Pulled By: cthi fbshipit-source-id: 4cec3916103da18a602f62178ad20f6ca6b4235b
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
3 participants
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
Add FlyDSL JIT backend for FP8 rowwise preshuffle GEMM on gfx950 (MI350). MSLK's existing
f8f8bf16_rowwise_preshuffleop only has a CK backend targeting gfx942. FlyDSL provides a Python-only JIT-compiled alternative with faster GPU kernels on most shapes, though Python dispatch overhead reduces net gains on small-M shapes.Technical Details
Port the FlyDSL preshuffle GEMM kernel from ROCm/aiter into MSLK as a self-contained module with no aiter dependency.
New module:
mslk.gemm.flydslflydsl_preshuffle(src)— weight shuffle into FlyDSL layout (different fromck_preshuffle)flydsl_preshuffle_gemm(XQ, WQ, x_scale, w_scale, ...)— JIT-compile and run the kernel_configs.py— default tile configs for gfx950 heuristic selection_kernels/— ported FlyDSL kernel compiler (preshuffle_gemm, mfma_epilogues, mfma_preshuffle_pipeline, tensor_shim)Op wiring (gfx950, gated on
is_flydsl_available()):torch.ops.mslk.f8f8bf16_rowwise→ FlyDSL (bf16 output)torch.ops.mslk.f8f8f16_rowwise→ FlyDSL (fp16 output)torch.ops.mslk.f8f8bf16_rowwise_out→ FlyDSL (out-tensor variant)Benchmark:
FP8RowwisePreshuffleFlyDSL— new benchmark class inbench/gemm/gemm_ops.pytargetingAMD_GFX950, gated onis_flydsl_available().Test Plan
Test Result
Correctness: 18/18 shapes pass (bf16 + fp16 + out-tensor, max_diff=0.0)
FlyDSL vs CK bpreshuffle benchmark (gfx950 MI350)
Both backends use the same bpreshuffle weight layout. CK forced to CK-only kernel path. All numbers are mean ± stddev over 3 independent runs.
GPU kernel time (rocprof --stats, 10 iters per run)
GPU kernel: FlyDSL on-par or faster on 7/11 shapes, up to 1.63x. CK faster on 4 shapes (0.81–0.97x).
Host-side end-to-end (CUDA events, 500 iters per run)
Host-side: FlyDSL wins 6/11 shapes after ~10–15 µs Python dispatch overhead. Crossover at ~M=1024.
Dispatch overhead breakdown
FlyDSL host floor is ~18.5 µs; CK host floor is ~8 µs. The ~10 µs delta is Python dispatch cost (
ptr_argwrapping,flyccache lookup,fx.Stream). On shapes where GPU kernel time > 20 µs, this overhead is amortized and FlyDSL net wins.Op dispatch benchmark (30 shapes, via torch.ops.mslk.f8f8bf16_rowwise)
CK wins 5/30, FlyDSL wins 1/30, Ties 24/30. Geo mean 0.978x. CK wins on 5 small attn_out shapes (M=1-128, K=1024) due to ~1us on-the-fly preshuffle overhead in the op impl. Pre-shuffled weights (direct API) would tie.
Submission Checklist
is_flydsl_available()— graceful fallback