[FlyDSL] Add FP4 LiteTopK prefill operator

[AMD-yanfeiwang]

Summary

Add a gfx950 FP4 LiteTopK path for DeepSeek-V4 paged-MQA prefill without materializing the dense [Q, S] score matrix.

The API is qualified for the released V4-Flash and V4-Pro contracts. Q=8192 is the shared benchmark workload; H/D/K come from each model config.

Model	Q	H	D	K	KV context
DeepSeek-V4-Flash	8192	64	128	512	C4 (raw / 4)
DeepSeek-V4-Pro	8192	64	128	1024	C4 (raw / 4)

Implementation

• Fuse FP4 paged-MQA scoring, calibrated candidate filtering, stable exact TopK, and page mapping.
• Export the explicit supported-K contract (512, 1024) and reject other values.
• Fail closed on overflow, underfill, invalid metadata/pages, nonfinite scores, and selector failures.
• Preserve caller device/stream ownership and 64-bit candidate/page addressing.

Validation

MI355X (gfx950), ROCm 7.2.4. These are operator benchmarks over deterministic tensors with the released model shapes, not full-model E2E runs. Each point is the median of three fresh processes; each process uses 5 warmup and 15 measured alternating dense/LiteTopK pairs. Exact TopK sets are checked outside timing.

Model	Raw context	C4 context	Dense + TopK	LiteTopK	Speedup
V4-Flash, K512	262,144	65,536	3.9780 ms	3.9221 ms	1.0142x
V4-Flash, K512	524,288	131,072	8.0164 ms	7.2235 ms	1.1098x
V4-Flash, K512	786,432	196,608	12.3227 ms	10.4344 ms	1.1810x
V4-Pro, K1024	262,144	65,536	4.0832 ms	4.1911 ms	0.9743x
V4-Pro, K1024	524,288	131,072	8.0133 ms	7.4099 ms	1.0814x
V4-Pro, K1024	786,432	196,608	12.5843 ms	10.6433 ms	1.1824x

• FP4 suite: 29 passed, 1 skipped; the two-GPU non-current-device case passes separately.
• Black, Ruff, and git diff --check: pass.

AITER_USE_SYSTEM_TRITON=1 pytest -q op_tests/test_flydsl_pa_mqa_litetopk_fp4.py
AITER_USE_SYSTEM_TRITON=1 python op_tests/op_benchmarks/bench_flydsl_pa_mqa_fp4_litetopk.py \
  --model deepseek-v4-flash --rows 8192 --raw-context 262144 --warmup 5 --iterations 15

References

• Original implementation: https://github.com/Heisenberg-Yin/LiteTopK
• Downstream integration: [ROCm][DSV4] Add FP4 LiteTopK indexer path sgl-project/sglang#38303

[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 5309 --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.