add flyDSL bwd FMHA kernel (#467)

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:
- [x] `test_backward.py -k flydsl`: 316 passed, 0 failed
- [x] `test_backward_gqa` (flydsl.BwOp): 2 passed, 2 skipped (bf16 precision skip, matching CK's own)
- [x] 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) 
- [x] 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

Author: amd-weisun

mslk/attention/flydsl/fmha_bwd_convert_dq.py

@@ -0,0 +1,71 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""FMHA backward dQ convert: f32 accumulator -> output dtype (bf16/fp16).
+
+A simple elementwise cast over the flat [B*M*H*D] dQ accumulator buffer.
+
+Target: gfx950 (CDNA4, wave64).
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from mslk.attention.flydsl.fmha_bwd_mfma import dtype_to_elem_type
+
+BLOCK_THREADS = 256
+VEC_WIDTH = 4
+
+
+def compile_fmha_bwd_convert_dq(*, dtype_str: str = "bf16"):
+    """Compile the dQ f32->output-dtype convert kernel.
+
+    Returns:
+        launch_fn(dq_f32, dq_out, n_elems, stream)
+          dq_f32  : [n_elems, 1] float32 input
+          dq_out  : [n_elems, 1] output dtype (bf16/fp16)
+          n_elems : total element count (B*M*H*D)
+    """
+    elem_dtype = dtype_to_elem_type(dtype_str)
+
+    @flyc.kernel(known_block_size=[BLOCK_THREADS, 1, 1])
+    def convert_dq_kernel(
+        dq_f32: fx.Tensor,
+        dq_out: fx.Tensor,
+        n_elems: fx.Int32,
+    ):
+        from flydsl.expr import buffer_ops as _bops
+        from flydsl.expr.typing import Vector as Vec
+
+        bid = fx.block_idx.x
+        tid = fx.thread_idx.x
+        n_elems_idx = fx.Index(n_elems)
+
+        src_rsrc = _bops.create_buffer_resource(dq_f32)
+        dst_rsrc = _bops.create_buffer_resource(dq_out)
+
+        base = (fx.Index(bid) * BLOCK_THREADS + fx.Index(tid)) * VEC_WIDTH
+        if base < n_elems_idx:
+            v = _bops.buffer_load(src_rsrc, base, vec_width=VEC_WIDTH, dtype=fx.Float32)
+            v_out = Vec(v).to(elem_dtype)
+            _bops.buffer_store(v_out.ir_value(), dst_rsrc, base)
+
+    @flyc.jit
+    def launch_fn(
+        dq_f32: fx.Tensor,
+        dq_out: fx.Tensor,
+        n_elems: fx.Int32,
+        stream: fx.Stream,
+    ):
+        n_blocks = (fx.Index(n_elems) + (BLOCK_THREADS * VEC_WIDTH) - 1) // (
+            BLOCK_THREADS * VEC_WIDTH
+        )
+        convert_dq_kernel(dq_f32, dq_out, n_elems).launch(
+            grid=(fx.Int32(n_blocks), 1, 1),
+            block=(BLOCK_THREADS, 1, 1),
+            stream=stream,
+        )
+
+    return launch_fn