[AMD] Enable GLM DSA prefill top-k to the v2 kernel

[EricKing626]

Motivation

Extend the DSA top-k v2 kernel to packed score rows and route GLM-5.x prefill through it. The prefill top-k kernel drops ~73%; at ISL 70000 / OSL 300 that is +4.9% token throughput per GPU, −3.5% median TPOT and −4.8% median TTFT (geomean over concurrency 4-64); GSM8k 0.927.

PR #36684 and PR #36851 turned the v2 fused top-k on for GLM-5.x on ROCm, but only decode ever reached it. Prefill still ran the legacy paged transform, which gathers the wide page_size=1 table.

The blocker was addressing, not math. The v2 kernel assumed each row starts at column 0 of its own score row and owns one page-table row. DSA extend breaks both: all requests share one packed score buffer, so a row starts at some column in the middle, and every row of a request reads the same page-table row.

Two small lookup tables fix that -- one for where a row starts, one for which page-table row it uses. No copy of the score buffer, and no per-row page table, which at 327K context would be hundreds of MB per forward.

Modifications

topk_v2.cuh, dsv4/topk.py -- add two optional (rows,) int32 indirections to topk_transform_512_v2:

	omitted (decode layout)	supplied (packed extend)
score window of row i	scores[i, :seq_lens[i]]	scores[i, row_starts[i] : row_starts[i] + seq_lens[i]]
page-table row	page_tables[i]	page_tables[row_to_batch[i]]

Both default to null and give exactly the old addressing, so nothing changes for existing callers. Shifting the score pointer keeps the selected index row-local, which is what the transform already expects, so the output path is untouched. The page-table row count has to match the score row count only when row_to_batch is absent; with it, the caller owns that bound.

dsa/dsa_topk_backend.py -- route packed PAGED extend to _topk_transform_v2_paged, passing ks as row_starts and token_to_batch_idx as row_to_batch. GLM-5.x already qualifies as PAGED for EXTEND, and the plan needs no new work: it is already built once per forward over dsa_seqlens_expanded, whose row count is what v2 sees.

Two shapes still take the legacy path. The dispatch checks for them itself, so they fall back instead of tripping the helper's asserts:

• Chunked-prefix-cache extend, where the plan covers the whole forward but each call sees only one chunk. Not to be confused with scheduler-level chunked prefill: under --chunked-prefill-size each chunk is its own forward with its own plan, and does take the v2 path.
• Row stride not a multiple of 4, which the vectorized load needs. An extend row stride is the batch's total KV length, so it is aligned only by luck.

The decode condition is left byte-identical, because dsa_drop_wide_page_table drops the page_size=1 table on exactly that condition and the two must stay in sync.

dsa_backend.py -- comment only, recording that packed PAGED extend is now a plan consumer.

Accuracy Tests

GSM8k, 8x MI355X (gfx950), GLM-5.2-MXFP4, TP=4, non-MTP, with this PR:

$ python3 benchmark/gsm8k/bench_sglang.py --port 8000 --num-questions 1200 --parallel 1200
Accuracy: 0.927
Invalid: 0.000
Latency: 39.462 s
Output throughput: 3095.724 token/s

test/registered/kernels/ops/attention/test_topk_v2.py on gfx950: 250 passed, 3 warnings in 19.99s, including the new test_topk_v2_packed_rows cases that check the row_starts / row_to_batch layout against a row-local reference.

Speed Tests and Profiling

Baseline is #29xxx as merged (top-k v2 on, prefill still legacy); this PR is the prefill routing on top of it. Same build, same flags, non-MTP, 8x MI355X, TP4.

Kernel-level -- topk_main_kernel now replaces topk_transform_prefill_kernel on every prefill launch; the old kernel is gone from the traces entirely. ISL 70000 / OSL 300, per TP rank, matching launch grids:

conc	grid.x	baseline	this PR	Δ
4	16384	2710.2 µs/launch	721.9 µs/launch	−73.4%
4	4615 (tail)	1390.3 µs/launch	359.1 µs/launch	−74.2%
64	16384	2694.2 µs/launch	726.4 µs/launch	−73.0%

Per rank at conc 4: 882.9 ms → 234.9 ms of top-k over the profiled window.

End-to-end -- ISL 70000 / OSL 300, --max-running-requests 8. Geomeans are over the per-config ratios.

conc	Interactivity (tok/s/user)	Token TPUT per GPU (tok/s)	Δ TPUT	Median TTFT (ms)	Median TPOT (ms)	Median ITL (ms)
4	29.34 → 30.35	4573.4 → 4761.6	+4.1%	3276 → 3128	34.08 → 32.95	10.63 → 10.52
8	14.13 → 15.08	4938.7 → 5139.9	+4.1%	5774 → 5646	70.77 → 66.33	12.28 → 12.24
16	12.44 → 13.08	5064.3 → 5296.6	+4.6%	27576 → 26563	80.38 → 76.43	12.21 → 11.99
32	12.94 → 13.06	5004.1 → 5301.0	+5.9%	79535 → 74147	77.25 → 76.57	12.20 → 12.30
64	12.72 → 13.00	4996.4 → 5297.1	+6.0%	181367 → 169622	78.60 → 76.91	12.11 → 12.01
geomean	+3.7%		+4.9%	−4.8%	−3.5%	not claimed

Interactivity is 1000 / median TPOT; Token TPUT per GPU is (input + output) tok/s / 4 for one TP4 server. Concurrency above 8 exceeds the admission cap, so TTFT there is mostly queueing time.

Reproduce:

export SGLANG_ROCM_FUSED_DECODE_MLA=0
export ROCM_QUICK_REDUCE_QUANTIZATION=INT4   # FP4 model

python3 -m sglang.launch_server --model GLM-5.2-MXFP4 \
  --tp 4 --trust-remote-code --tool-call-parser glm47 --reasoning-parser glm45 \
  --mem-fraction-static 0.85 --kv-cache-dtype fp8_e4m3 \
  --disable-radix-cache --chunked-prefill-size 16384 \
  --dsa-prefill-backend triton --dsa-decode-backend triton \
  --enable-aiter-allreduce-fusion --tokenizer-worker-num 8 \
  --max-running-requests 8

python3 -m sglang.bench_serving --backend sglang --dataset-name random \
  --random-input-len 70000 --random-output-len 300 --random-range-ratio 0.8 \
  --max-concurrency $CONC --num-prompts $((5 * CONC))

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): 🚫 Run #34846434360
Latest PR Test (Extra): ❌ Run #34846433791
Latest PR Test (AMD ROCm 10): ❌ Run #34846434075

[KingRei]

/rerun-failed-ci

[EricKing626]

/rerun-failed-ci

[DarkSharpness]

Hi. Could you please temporarily hold on until #38798 ? We have some critical update on v2 algorithm, which introduce many new mechanisms for that (like padding).

Also, some comments on this PR: please try to avoid adding too many comments in the code (especially for parts that remain unchanged). Please keep the comments brief within 1 ~ 2 lines

[EricKing626]

/rerun-failed-ci

[EricKing626]

/rerun-failed-ci