add flyDSL bwd FMHA kernel (#488)

[q10]

Summary:

Add FlyDSL FMHA backward kernel (gfx950 + gfx942 fallback)

Adds a FlyDSL FMHA backward implementation, registered as flydsl.BwOp (an
opt-in AttentionBwOpBase, following the same pattern as flash.BwOp /
flash3.BwOp). It is not wired into dispatch.py's live _dispatch_bw()
priority list (which still uses ck.BwOp on ROCm) — this PR adds the op and
its test coverage, not a change to production routing.

Two kernels ship together:

• gfx950 (CDNA4/MI350X): fused dQ+dV+dK kernel
  optimized for this architecture (hardware LDS transpose, XOR-swizzled
  layouts, register-resident K/V/KT, software-pipelined prefetch — see below).
• gfx942 (CDNA3) fallback: a 32x32 MFMA kernel (fmha_bwd_mfma.py) used
  when gfx950 isn't available.

Files changed

File	Description
mslk/attention/flydsl/fmha_bwd_mfma_gfx950.py	New gfx950 kernel (1710 lines)
mslk/attention/flydsl/fmha_bwd_mfma.py	New gfx942 fallback kernel (32x32 MFMA)
mslk/attention/flydsl/fmha_bwd_preprocess.py	New D_vec preprocess kernel
mslk/attention/flydsl/fmha_bwd_convert_dq.py	New dQ f32→output-dtype convert kernel
mslk/attention/fmha/flydsl.py	New flydsl.BwOp: gfx950 routes to the new kernel with ck_scope_dvdk=True; gfx942 falls back to fmha_bwd_mfma.py
mslk/attention/fmha/__init__.py	Register flydsl.BwOp in ALL_BW_OPS on ROCm (test enumeration only)
test/attention/fmha/test_backward.py	Add flydsl.BwOp to test_backward + test_backward_gqa; 3 new negative-path tests

• GQA (the production training config): End-to-end, FlyDSL is 1.04x–2.20x faster
  across all seqlens (1K–16K), winning at every single shape. Best result: 4K causal
  (2.20x).
• MHA: 1.02x–1.73x faster at seqlen ≥ 4K; near-parity at 2K (1.03x/0.96x);
  ~7–11% slower at 1K only (low grid occupancy with H=8, fewer blocks to fill 256 CUs).
• Both sides launch 3 GPU kernels: CK dispatches OGradDotO + DqDkDv + ConvertDq;
  FlyDSL dispatches its own FlyDSL preprocess (D_vec) + fused dQ+dV+dK + convert-dq
  kernels. The GQA .unflatten().sum() reduce for dK/dV is a PyTorch op on both sides
  (CK's C++ wrapper does the same reduce outside its kernels).
• CK baseline is the current MSLK production commit

Kernel breakdown (device-side, GQA H=64/Hkv=8)

Per-kernel device time via rocprofv3 --kernel-trace. Both sides launch 3 kernels:

Stage	CK kernel	FlyDSL kernel
Preprocess (D_vec)	FmhaBwdOGradDotOKernel	fmha_bwd_preprocess (d_vec_kernel)
Main (dQ+dV+dK)	FmhaBwdDQDKDVKernel	fmha_bwd_dqdkdv_mfma_gfx950
Convert (f32→bf16)	FmhaBwdConvertQGradKernel	fmha_bwd_convert_dq

GQA (H=64, Hkv=8)

Seqlen	Mask	Stage	CK (us)	FlyDSL (us)	Speedup
1024	none	pre	8.4	6.6	1.27x
		main	266.1	272.5	0.98x
		conv	9.7	9.1	1.06x
		total	284.2	288.2	0.99x
1024	causal	pre	8.1	6.6	1.22x
		main	242.4	215.4	1.13x
		conv	9.9	9.4	1.06x
		total	260.3	231.4	1.13x
2048	causal	pre	16.3	15.1	1.08x
		main	893.6	596.5	1.50x
		conv	19.4	16.6	1.17x
		total	929.3	628.2	1.48x
4096	causal	pre	40.3	25.8	1.57x
		main	3398.9	1872.1	1.82x
		conv	40.1	34.1	1.17x
		total	3479.3	1932.0	1.80x
8192	causal	pre	93.8	46.1	2.04x
		main	10486.8	6812.8	1.54x
		conv	87.7	65.7	1.33x
		total	10668.2	6924.6	1.54x
16384	none	pre	166.8	98.1	1.70x
		main	58309.1	42951.6	1.36x
		conv	195.7	132.2	1.48x
		total	58671.5	43182.0	1.36x
16384	causal	pre	167.5	94.6	1.77x
		main	34111.1	23885.8	1.43x
		conv	193.8	131.9	1.47x
		total	34472.4	24112.3	1.43x

Observations:

• Preprocess (D_vec): FlyDSL wins 1.06x–2.06x at every GQA shape. The
  multi-row-per-block kernel (16 rows/block for D=128) outperforms CK's OGradDotO.
• Main kernel (dQ+dV+dK): FlyDSL wins 1.13x–1.82x at every GQA shape
  (except 1K non-causal at 0.98x, within noise).
• Convert (f32→bf16): FlyDSL wins 1.00x–1.48x at every shape.

Test Plan:

• test_backward.py -k flydsl: 316 passed, 0 failed
• test_backward_gqa (flydsl.BwOp): 2 passed, 2 skipped (bf16 precision skip, matching CK's own)
• Kernel-level correctness suite (direct calls to compile_fmha_bwd_dqdkdv_mfma_gfx950,
  127 passed / 2 skipped across all supported D/dtype/causal/GQA/varlen/deterministic/
  packed-qkv combinations)
• Performance sweep (below): 20 shapes, D=128, bf16, seqlen 1K–16K, causal + non-causal, GQA + MHA

Performance: FlyDSL gfx950 FMHA Backward vs Production CK

Device-side kernel time measured via rocprofv3 --kernel-trace (true GPU dispatch
timestamps, no host overhead). CK baseline is the current MSLK production commit.
FlyDSL kernel uses ck_scope_dvdk=True for GQA shapes (dV/dK written per-query-head, reduced
outside the kernel — matches CK's own production C++ wrapper).

Config: B=1, D=128, bf16, atomic-add dQ, gfx950 (MI350X)

GQA (H=64, Hkv=8, heads_per_kv=8)

Seqlen	Mask	FlyDSL (us)	CK (us)	Speedup
1024	none	281.7	268.0	0.95x
1024	causal	218.1	240.7	1.10x
2048	none	839.5	948.5	1.13x
2048	causal	598.1	890.7	1.49x
4096	none	3028.7	3795.0	1.25x
4096	causal	1879.2	3398.4	1.81x
8192	none	10767.8	14777.0	1.37x
8192	causal	6785.1	10476.3	1.54x
16384	none	42967.7	58564.1	1.36x
16384	causal	23901.5	34050.3	1.42x

MHA (H=8, Hkv=8, heads_per_kv=1)

Seqlen	Mask	FlyDSL (us)	CK (us)	Speedup
1024	none	126.2	109.9	0.87x
1024	causal	126.7	114.5	0.90x
2048	none	242.0	214.5	0.89x
2048	causal	241.9	216.9	0.90x
4096	none	500.3	476.5	0.95x
4096	causal	473.6	428.8	0.91x
8192	none	1945.1	1766.5	0.91x
8192	causal	1170.3	1692.6	1.45x
16384	none	5845.2	7229.6	1.24x
16384	causal	3876.4	5046.1	1.30x

• CK (us) = CK's FmhaBwdDQDKDVKernel device time (the kernel directly comparable
  to FlyDSL's fused dQ+dV+dK kernel)
• Speedup = CK / FlyDSL. Values > 1.0x mean FlyDSL is faster; bold = FlyDSL wins
• This is a kernel-vs-kernel comparison of the main fused kernel only. For the full
  3-kernel-vs-3-kernel comparison (including preprocess and convert), see the
  "Kernel breakdown" section below.

End-to-end wall-clock: flydsl.BwOp vs ck.BwOp (backward only)

Wall-clock time calling BwOp.apply() directly — the real MSLK dispatch path,
including all host-side overhead. Both sides launch 3 GPU kernels for the main
GQA/MHA path: FlyDSL's preprocess (D_vec) + fused dQ+dV+dK + convert-dq kernels
vs CK's OGradDotO + FmhaBwdDQDKDVKernel + ConvertDQ. The GQA .unflatten().sum()
reduce for dK/dV remains a PyTorch op on both sides (CK's C++ wrapper does the
same reduce outside its kernels). Forward via ck.FwOp (shared, not timed).
JIT compilation cost is measured separately (first call only); steady-state
numbers reflect the cached kernel path.

Config: B=1, D=128, bf16, gfx950 (MI350X), N_WARMUP=3, N_ITER=20

GQA (H=64, Hkv=8, heads_per_kv=8)

Seqlen	Mask	FlyDSL (us)	CK (us)	Speedup
1024	none	348.9	362.0	1.04x
1024	causal	284.6	323.4	1.14x
2048	none	908.0	1260.0	1.39x
2048	causal	670.2	1169.8	1.75x
4096	none	3177.0	4952.4	1.56x
4096	causal	1995.4	4383.7	2.20x
8192	none	11059.8	19046.2	1.72x
8192	causal	7105.4	13382.5	1.88x
16384	none	43569.6	75153.7	1.72x
16384	causal	24580.4	43311.2	1.76x

MHA (H=8, Hkv=8, heads_per_kv=1)

Seqlen	Mask	FlyDSL (us)	CK (us)	Speedup
1024	none	169.6	158.5	0.93x
1024	causal	170.1	151.3	0.89x
2048	none	290.7	298.3	1.03x
2048	causal	294.0	283.3	0.96x
4096	none	540.1	587.5	1.09x
4096	causal	542.5	553.7	1.02x
8192	none	1996.9	2249.6	1.13x
8192	causal	1253.0	2163.8	1.73x
16384	none	5990.5	9415.2	1.57x
16384	causal	4018.6	6374.2	1.59x

• Speedup = CK wall-clock / FlyDSL wall-clock (>1.0 = FlyDSL faster; bold = FlyDSL wins)
• JIT first-call cost: ~3.0–3.7s per unique (causal, heads_per_kv) variant; amortized to zero after the first backward call per config

Reviewed By: bottler

Differential Revision: D114928768

Pulled By: q10

[meta-codesync]

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

[bottler]

Review automatically exported from Phabricator review in Meta.