[feat] FP8 (DeepSeek-V4 layout) sparse paged prefill attention

[carlushuang]

Summary

Adds FP8 KV-cache support to the OPUS sparse paged prefill attention (pa_sparse_prefill_opus) using the DeepSeek-V4 / asm-v4 mixed-precision head-dim layout:

D = 512 = 448 NOPE (fp8 e4m3, per-64-element-tile e8m0 scale) + 64 ROPE (bf16), for both Q and KV (prefix unified_kv + extend kv).

The existing bf16 kernel and all current behavior are unchanged.

What's added (two ops, both validated on gfx950)

Op	Approach	Perf	Notes
pa_sparse_prefill_opus_fp8	dequant-prepass: a standalone fp8→bf16 device kernel writes transient bf16 scratch, then the existing fast bf16 attention runs	~1.04–1.15× bf16	The performant path. Persistent KV-cache stays fp8 (capacity preserved); only a per-call transient bf16 scratch.
pa_sparse_prefill_opus_fp8_fused	fully fused single __global__: fp8 QK-nope MFMA + software per-64-tile e8m0 scale, bf16 QK-rope, bf16 PV with on-chip V dequant, 4-way KV split (prefix/extend × nope/rope), per-head sink, no bf16 scratch	~26× bf16 (H=128)	Correctness-first; occupancy/latency-bound. Perf-rewrite path is characterized (see below).

Design notes

• The mxfp8 QK uses the proven non-scaled fp8 16×16×32 MFMA + software per-64-tile scaling (the opus_gemm a8w8_scale pattern), validated against the canonical op_tests/opus/device/test_mxfp.cu. This sidesteps the hardware MX scale-operand encoding and naturally handles the 64-element quant tiles (no dup×2 needed).
• Scales are emitted as a separate fp32 [*,7] tensor (e8m0-rounded powers of two) rather than asm-v4's in-buffer packed layout — identical math, much cleaner for a fresh HIP kernel.

Validation (gfx950, rocm/atom 7.2.3 image)

• 48/48 tests pass: 24 bf16 (no regression) + 12 fp8 dequant-prepass + 12 fp8 fused — sparse/dense/empty × H∈{16,32,64,128}, vs an fp8-dequant reference within fp8 tolerance.
• Standalone MFMA probes validate the fp8 QK math to ~5e-5 (op_tests/opus/qk_{nope,full}_probe.cc).

Perf state (fused kernel)

The fused kernel is correct but ~26× off the hand-tuned bf16 baseline. Root cause (measured): VGPR-occupancy + memory-latency bound — incremental knobs (LDS reduction +2.6×, multi-warp +5–20%, larger tiles / forced occupancy = regressions). The dequant-prepass is the recommended performant path today; closing the fused kernel to parity requires porting the bf16 kernel's async double-buffered pipeline onto the validated fp8 QK path (future work).

Files

• csrc/include/pa_sparse_prefill_opus.h — dequant + fused __global__ kernels, fp8 kargs, public API
• csrc/py_itfs_cu/pa_sparse_prefill_opus_kernels.cu — fp8 host launchers
• csrc/include/rocm_ops.hpp, csrc/pybind/pa_sparse_prefill_opus_pybind.cu — pybind registration
• aiter/ops/pa_sparse_prefill_opus.py — python ops + wrappers
• op_tests/test_pa_sparse_prefill_opus_fp8.py, op_tests/pa_sparse_prefill_opus_fp8_quant.py — oracle test + v4 fp8 quant helpers
• op_tests/bench_pa_sparse_prefill_opus_fp8.py — benchmark
• op_tests/opus/qk_{nope,full}_probe.cc — standalone MFMA validation probes

Test plan

• pytest op_tests/test_pa_sparse_prefill_opus_fp8.py (gfx950) — 24/24 fp8 pass
• pytest op_tests/test_pa_sparse_prefill_opus.py (gfx950) — 24/24 bf16, no regression
• CI on MI350 runner

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3583 --add-label <label>