[FlyDSL][DSv4] Prototype FP4 MQA streaming TopK

[AMD-yanfeiwang]

Motivation

The DeepSeek-V4 FP4 indexer currently materializes an FP32 [query_rows, c4_context] score matrix before TopK. At 768K full-token context this is about 0.75 MiB per query row; a 1024-row local prefill is about 768 MiB per C4 layer invocation. SGLang currently avoids allocator fragmentation with a persistent 2 GiB slab (sgl-project/sglang#37660), but that reserves context-sized storage instead of removing the intermediate.

This draft adds an AITER operator that computes exact TopK without materializing the full logits matrix.

Changes

• Add flydsl_pa_mqa_topk_fp4_prefill for gfx950 FP4 paged MQA.
• Fuse score production with CTA-local exact TopK:
  • scores stream through a bounded LDS pool;
  • four-byte score radix plus four-byte logical-index radix implements total order (score descending, logical index ascending);
  • NaNs sort below numeric values and +0 > -0 is preserved;
  • each CTA emits only (candidate_value, raw_index, valid_count).
• Merge split candidates exactly, then map raw C4 positions through the page table.
• Preserve the existing short-row contract: when length <= topk, return every logical position sequentially and pad with -1.
• Add a bounded context-tiled implementation as a correctness/reference fallback.
• Keep the existing full-logits API source-compatible.

For the common SGLang prefill shape rows=1024, parallel_unit_num=1024, topk=1024, internal candidate + merge scratch is about 20 MiB, independent of context length, versus roughly 768 MiB of logits at 768K context. No device-global context-sized allocation is introduced.

Scope

Initial production specialization:

• gfx950 / wave64
• H=64, D=128
• FP4 E2M1 payload with shuffled UE8M0 scales
• KV page size 64, score block 256
• TopK 512 or 1024
• ragged prefill windows

Graph capture is rejected explicitly in this first version; SGLang will retain its existing graph path as fallback. A dependent SGLang draft will capability-probe this API and keep old AITER compatibility. sgl-project/sglang#38086 remains an independent workspace-based alternative.

Validation

Completed:

• Black / Ruff / Python syntax checks
• CPU reference/property tests for exact candidate union and special-value ordering
• offline gfx950 FlyDSL compilation for TopK 512/1024 and the legacy scorer
• safety review for empty windows, nonzero starts, speculative page lookahead, stream affinity, and launcher lifetime

Queued on MI355X:

• runtime fused vs full-logits correctness for TopK 512/1024
• shuffled page tables, empty/short/ragged windows, and non-current streams
• fused/tiled/full-logits latency and peak-memory benchmarks
• dependent SGLang adapter tests

This stays draft until those GPU and end-to-end long-context results are attached.

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

[AMD-yanfeiwang]

MI355X runtime update (gfx950, ROCm 7.2):

• focused AITER suite: 18 passed (K=512/1024, exact candidate merge, empty/ragged windows, streams)
• dependent SGLang adapter suite: 5 passed
• bounded and true-fused outputs match the full-logits selected set for K=512/1024

Production-grid microbenchmark (rows=1024, parallel_unit_num=1024, score_batch_chunks=16):

• C4 width 196,608 (~768K full context), K=512: legacy 1.908 ms / 868 MiB path storage; fused 4.522 ms / 16 MiB
• C4 width 196,608, K=1024: legacy 1.935 ms / 872 MiB; fused 4.810 ms / 32 MiB
• C4 width 262,144 (~1M full context), K=512: legacy 2.505 ms / 1,156 MiB; fused 6.059 ms / 16 MiB
• C4 width 262,144, K=1024: legacy 2.537 ms / 1,160 MiB; fused 6.438 ms / 32 MiB

So the prototype removes 96–99% of score-path storage, but is currently about 2.4–2.5x slower at the core operator. That does not meet the default-path performance gate yet; this PR remains Draft while the 8-GPU end-to-end run finishes. The result supports keeping sgl-project/sglang#38086 as the near-term alternative unless further local-selection optimization closes the gap.