[FlyDSL] Enable FP8 FMHA split-K for short cached-prefix queries

[gbyu-amd]

Adds split-K support for short queries (Q < 384) in noncausal cross-sequence FP8 FMHA, with automatic and explicit split selection.

Measured on gfx950: B=1, Q=70, H=12, Dq=Dk=192, Dv=128, return_lse=True. Both versions use automatic split selection.

KV tokens	Base (246efe910)	PR (38f239361)	Speedup
16,384	284.868 µs	38.464 µs	7.41×
11,008	194.872 µs	36.414 µs	5.35×

Median of 7 rounds × 100 CUDA graph replays per mode, after 5 warmups. Includes attention and split combine; excludes quantization and JIT compilation. PyTorch 2.10.0 / ROCm 7.2.4, FlyDSL 0.3.2.

Accuracy validation: 27 targeted GPU tests passed in op_tests/test_flydsl_fmha.py, covering short-query splits, output/LSE references, variable-length inputs, custom softmax scales, graph replay, and rejection of causal short-query split-K. Black and Ruff passed.

[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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
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 5508 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[gbyu-amd]

@coderfeli @binding7012 hi, could you help review again?

[zufayu]

Advisory review (static + hand-run; not a merge gate). Validation/Perf ran no GPU stage — reasons are on their lines below. Findings tagged [verified] are traced to code/repro.

Lets short (Sq<384) noncausal cross-sequence FP8 queries on gfx950 use split-KV attention — explicit num_kv_splits and the auto-split heuristic — instead of forcing split=1, and adds short-query split shapes to the FMHA tests.

Review (advisory): ⚠️ NEEDS WORK
Validation (deterministic): NOT RUN — REVIEW_AUTO_VALIDATE=0 and no gfx950 GPU on this box (environment gap, not a PR defect); triage required runtime+test, target op_tests/test_flydsl_fmha.py, runtime path aiter/ops/flydsl/kernels/flash_attn_func_fp8_gfx950.py
Perf (advisory): NOT RUN — the triaged target op_tests/test_flydsl_fmha.py exposes no benchmark entry point and auto-validation is disabled (REVIEW_AUTO_VALIDATE=0); the PR's own 7.44x/5.30x (B=1, Q=70, H=12, D=192/Dv=128, KV=16384/11008, split=1 vs auto split-K, head-only) are unverified here

⚠️ [verified] aiter/ops/flydsl/kernels/flash_attn_func_fp8_gfx950.py sits on the default fp8 attention dispatch chain (aiter/ops/mha.py:2842/3798 route fp8 e4m3 gfx950 through aiter/ops/flydsl/fmha_kernels.py:358/487), and those callers never pass num_kv_splits — so once this merges, every downstream fp8 FMHA caller on gfx950 with a short noncausal cross query (Sq<384, Skv≳1985) silently flips from split=1 to auto split-K (up to 16 splits plus a per-call fp32 workspace), with only aiter's own gfx950 op tests as evidence and no ci: label on the PR. Reviewer should ask which downstream model (ATOM/SGLang/vLLM) exercises fp8 e4m3 FMHA on gfx950 and add the matching ci: label (ci:all if untraceable), green before merge.

⚠️ [verified] The only timing evidence is the PR's own head-only A/B (split=1 vs auto split-K: 278.459→37.413 µs at KV=16384, 190.633→35.997 µs at KV=11008, B=1/Q=70/H=12), and op_tests/test_flydsl_fmha.py ships no benchmark entry point, so the claimed 5.3–7.4x cannot be re-run from the repo and no base-vs-head number exists. Author must add a timing harness (or attach base-vs-head numbers with sample counts) for the short-query split-K shapes so the win is reproducible.

📝 [verified] The re-scoped rejection (causal with Sq<384 and explicit splits must raise ValueError) is pinned by no test — op_tests/test_flydsl_fmha.py has no pytest.raises case for the split-K gate, so a future edit could silently admit causal short-Q split-K against the docstring's stated contract. Reviewer should ask for one pytest.raises(ValueError) case, e.g. causal=True, Sq=70, Skv=16384, num_kv_splits=4.

[gbyu-amd]

@zufayu Thanks for the review. Addressed the missing rejection test and the base/head timing evidence; I do not think the suggested CI label would validate the intended consumer yet.

• Downstream coverage: the concrete consumer is Kimi-K3 cached-prefix prefill in ATOM #2191. The current ci:atom workflow checks out ATOM main and does not enable ATOM_USE_FLYDSL_FP8_PREFILL_ATTN. I have left that label off because it would not exercise this integration. Prior full-model runs with 1,024-token cached chunks select split=1, so they are not evidence for short-query split-K. Direct validation here was run explicitly with pytest on gfx950.
• Performance: attached fresh base-vs-head measurements below, using the review's alternative to adding a benchmark entry point. This keeps the test-file changes small.
• Causal rejection: added test_fp8_short_causal_split_kv_rejected in 1fc14083e, reusing _run_fp8_shape: Q=70, KV=16,384, causal=True, splits=4 must raise the split-K ValueError. No kernel changes or new test files.

B=1, Q=70, H=12, Dq=Dk=192, Dv=128, noncausal, return_lse=True; same FP8 inputs and GPU for both versions:

KV tokens	Base auto (246efe910)	PR auto (38f239361)	Speedup	PR explicit split=1
16,384	284.868 µs	38.464 µs	7.41×	284.650 µs
11,008	194.872 µs	36.414 µs	5.35×	194.645 µs

Each cell is the median of 7 rounds × 100 CUDA graph replays (700 timed calls per mode/shape/version), after 5 eager warmups and 5 graph warmups. Auto and explicit split=1 alternate order between rounds within each version. Includes attention + combine; quantization, reference checks and JIT compilation are outside timing. Base and head used separate JIT caches. Environment: gfx950, PyTorch 2.10.0 / ROCm 7.2.4, FlyDSL 0.3.2. The new commit only adds the rejection test; timed kernel code is unchanged.

Validation: 27 passed, including output/LSE reference checks, variable-length inputs, custom scales and graph replay. Black/Ruff passed.

python -m pytest -q op_tests/test_flydsl_fmha.py \
  -k 'split_kv_cross_length or short_causal_split_kv_rejected or softmax_scale_varlen or softmax_scale_graph_replay'

[binding7012]

LGTM