feat: add FlyDSL batched preshuffle GEMM for FP8 rowwise scaling (WP-G2) (#444)

[cthi]

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

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

• Correctness: 8/8 shapes pass (both clang and gcc CI), no crashes
• Standalone API: flydsl_preshuffle_batched_gemm()
• Benchmark class: FP8RowwiseBatchedPreshuffleFlyDSL
• Op registration: f8f8bf16_rowwise_batched on gfx950
• Rebased on main
• Gated on is_flydsl_available() — graceful fallback
• Grid-Z batching: single kernel launch for all B batches
• XCD swizzle + waves_per_eu tuning: sweep-optimized per config
• Profile-guided shape overrides for N=1280 and N=8192
• Batch-aware occupancy heuristic
• Host-side benchmark with reproducer

Reviewed By: q10

Differential Revision: D114744187

Pulled By: cthi

[meta-codesync]

@cthi has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114744187.