feat: RMSNorm + RoPE fusion for WAN: flashinfer.diffusion_ops.fused_qk_rmsnorm_rope (#3148)

<!-- .github/pull_request_template.md -->

## 📌 Description
https://github.com/flashinfer-ai/flashinfer/issues/2971

Add a fused CUDA kernel for across-heads QK RMSNorm + 3D Rotary Position
Embeddings (RoPE) + V copy, targeting video generation DIT (Diffusion
Transformer) self-attention workloads such as WAN 2.1/2.2.

This kernel fuses three operations into a single launch:
1. **Across-heads RMSNorm** on Q and K (normalizes over the full
`hidden_dim = num_heads * head_dim`, not per-head)
2. **3D RoPE** with frame/height/width spatial decomposition (each head
dimension is split into temporal, height, and width frequency channels)
3. **V passthrough copy** to a contiguous output buffer

Optional **FP8 E4M3 quantized output** with SM89+ vectorized PTX
conversion and SM100+ Blackwell FFMA2 intrinsics (all with scalar
fallbacks for SM80+).

### API

User-facing API follows the `dsv3_ops` pattern — implementation lives in
`flashinfer/norm/` (alongside `rmsnorm`, `gemma_rmsnorm`, etc.) with a
re-export facade at `flashinfer/diffusion_ops/`:

```python
from flashinfer.diffusion_ops import fused_qk_rmsnorm_rope

q, k, v = fused_qk_rmsnorm_rope(
    qkv,               # [batch, seq_len, (nq+nk+nv)*head_dim] or [num_tokens, ...]
    q_weight, k_weight, # RMSNorm weights [num_heads_x * head_dim]
    ppf=5, pph=12, ppw=32,
    num_frame_channels=44, num_height_channels=42, num_width_channels=42,
    num_heads_q=24, num_heads_k=24, num_heads_v=24,
    head_dim=128,
)
```

### Benchmark Results (B200 (sm100), CUPTI)

```
python /workspace/flashinfer/benchmarks/bench_fused_qk_rmsnorm_rope.py
GPU: NVIDIA B200
Config: WAN 2.2 5B (num_heads=24, head_dim=128)

Shape                                                Eager (ms)   Fused (ms)    Speedup
------------------------------------------------------------------------------------------
B=1 5x12x32= 1920 (480p production (1920 tokens))        0.2420       0.0389      6.22x
B=1 5x12x8=  480 (480p small (480 tokens))               0.2363       0.0114     20.68x
B=1 5x48x32= 7680 (720p large (7680 tokens))             0.4948       0.1462      3.38x
B=2 5x12x32= 1920 (batch=2 (3840 tokens))                0.2864       0.0750      3.82x
B=1 5x6x4=  120 (tiny (120 tokens))                      0.2467       0.0045     54.67x
B=4 5x12x32= 1920 (batch=4 (7680 tokens))                0.4770       0.1462      3.26x
B=1 5x12x16=  960 (half seq (960 tokens))                0.2376       0.0207     11.49x
B=1 10x12x32= 3840 (double frames (3840 tokens))         0.2823       0.0749      3.77x
------------------------------------------------------------------------------------------
```

### Bug Fix

Found and fixed a bug in the NeoX (non-interleaved) RoPE path: `pos_id`
was computed from `dim_idx_x`'s mapped value but reused for `dim_idx_y`
without recomputation. Since the `(dim_idx * 2) & mask` mapping can
place adjacent elements in a float2 pair into different spatial slices
(e.g., height vs width), the y component received an incorrect position
ID. Fix: one line to recompute `pos_id` from `dim_idx_y`. The
interleaved path (used in production) was unaffected.

### Architecture Support

| SM | Architecture | Support Level |
|----|-------------|--------------|
| SM80 | Ampere (A100) | Full — BF16 path, FP8 via software emulation |
| SM86 | Ampere (RTX 3090) | Same as SM80 |
| SM89 | Ada (L40, RTX 4090) | + native FP8 conversion |
| SM90 | Hopper (H100) | Primary target |
| SM100 | Blackwell (B200, GB200, RTX 5090) | Primary target + FFMA2 |
| SM103 | Blackwell (B300, GB300) | Primary target + FFMA2 |
| SM110–SM121 | Blackwell variants | Expected to work (FFMA2 + FP8) |

### Known Limitations

- **`num_heads ≤ 32`**: The kernel uses one warp per head, so `max_heads
× 32 = 1024` threads per block (CUDA maximum). WAN 14B (40 heads) is
unsupported. Supporting it would require a kernel redesign (multi-head
per warp).
- **BF16 input only**: FP16/FP32 input would need new template
instantiations.
- **3D RoPE is specialized**: The frame/height/width decomposition
targets video-gen DIT models. This is not a general-purpose RoPE.

### Files Changed

```
include/flashinfer/fused_qk_rmsnorm_rope.cuh     # CUDA kernel + utilities (754 lines)
csrc/norm.cu                                    # TVM-FFI launcher (added to norm module)
csrc/flashinfer_norm_binding.cu                 # TVM-FFI export (added to norm module)
flashinfer/norm/fused_qk_rmsnorm_rope.py           # Python API with validation
flashinfer/norm/__init__.py                     # Re-export
flashinfer/diffusion_ops/__init__.py            # User-facing facade (like dsv3_ops/)
tests/norm/test_fused_qk_rmsnorm_rope.py           # 26 tests
benchmarks/bench_fused_qk_rmsnorm_rope.py          # Benchmark script
```

## 🔍 Related Issues

https://github.com/flashinfer-ai/flashinfer/issues/2971

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

**Test summary: 26 tests (25 pass, 1 xfail)**
- 8 interleaved correctness shapes (WAN 2.2 5B config)
- 3 NeoX (non-interleaved) correctness shapes
- V passthrough (exact BF16 copy)
- Destination-passing style
- 2D `[num_tokens, hidden]` input
- 3 FP8 output scales (1.0, 0.5, 2.0)
- RoPE-only mode (`is_qk_norm=False`)
- 3 multi-config: WAN 1.3B (12 heads), WAN 5B (24 heads), WAN 14B (40
heads — xfail, exceeds 32-head limit)
- 5 error-case validation tests

**Validated on 3 GPU architectures:**
- NVIDIA A100 (SM80, Ampere)
- NVIDIA L40S (SM89, Ada)
- NVIDIA H100 NVL (SM90, Hopper)
- NVIDIA B200 (sm100, Blackwell)

## Reviewer Notes

- The kernel compiles as part of the existing `norm` JIT module
(`gen_norm_module()`), so no new JIT spec or AOT registration is needed.
- The NeoX RoPE bugfix is a one-line change in
`fused_qk_rmsnorm_rope.cuh` — reviewers may want to verify the `pos_id`
recomputation logic for `dim_idx_y`.
- The `diffusion_ops/` facade follows the exact same pattern as
`dsv3_ops/` — pure re-export, no implementation.
- Future video-gen kernels (e.g., fused cross-attention, `rmsnorm_silu`
for WAN) can be added to `diffusion_ops/`.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Fused Q/K RMSNorm + 3D rotary embeddings for video self-attention with
optional scaling and FP8 E4M3 output modes.
* Public Python API and package re-exports for easy invocation; outputs
can be preallocated or auto-allocated and V is passed through unchanged.
* GPU benchmark script to compare fused implementation vs. reference
timings across representative shapes.

* **Tests**
* Comprehensive CUDA test suite validating BF16/FP8 modes, RoPE
variants, correctness, output semantics, and many error cases.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Cursor <cursoragent@cursor.com>