[AMD] Enable deepseek-v4 topk_transform v2 kernel

[1am9trash]

Motivation

Topk transform v2 kernel replaces v1's 8-bit coarse key and 128KB LDS candidate buffer with a 12-bit key, register-resident rows, so far fewer candidates reach the refinement pass and the GPU no longer runs one block per CU.

Modifications

• Enable topk v2 kernel.
• Enable amd code path to use topk v2 kernel in server_args.py gate.
• Enable unittest test_topk_v2.py.

Accuracy Tests

Unit test added in test/registered/kernels/ops/attention/test_topk_v2.py:

...
test/registered/kernels/ops/attention/test_topk_v2.py::test_topk_v2_ragged_no_row_starts[512] PASSED    [ 99%]
test/registered/kernels/ops/attention/test_topk_v2.py::test_topk_v2_ragged_no_row_starts[2048] PASSED   [100%]

 =================== 278 passed, 2 warnings in 14.81s  ===================

GSM8k, 1319:

100%|███████████████████| 1319/1319 [01:32<00:00, 14.19it/s]
Accuracy: 0.945
Invalid: 0.000
Latency: 92.956 s
Output throughput: 1251.910 token/s

Speed Tests and Profiling

Kernel side:

Batch	L (ctx_len / 4)	Baseline (v1)	v2	speedup
8192	2048 (ctx 8k)	256.0	87.3	2.93x
16384	2048 (ctx 8k)	504.0	169.8	2.97x
8192	17920 (ctx 70k)	687.0	257.9	2.66x
16384	17920 (ctx 70k)	1360.4	508.6	2.68x

E2E server run on long context:

• cc=2/4/8
• warmup=cc*2, num_prompt=cc*8
• ISL/OSL=70k/200

cc	v1 TTFT	v1 ITL	v1 TTT	v2 TTFT	v2 ITL	v2 TTT
2	3344.03	4.65	17903.61	3201.39	4.17	19226.93 (1.07x)
4	8960.77	17.40	16950.36	8463.05	17.08	17729.94 (1.05x)
8	15349.99	18.81	19274.98	14773.28	18.46	20056.60 (1.04x)

Server script to reproduce:

sglang serve \
    --model-path "${MODEL}" \
    --trust-remote-code \
    --tp 8 \
    --attention-backend dsv4 \
    --page-size 256 \
    --mem-fraction-static 0.89 \
    --swa-full-tokens-ratio 0.10 \
    --tool-call-parser deepseekv4 \
    --reasoning-parser deepseek-v4 \
    --chunked-prefill-size 16384 \
    --cuda-graph-max-bs 256 \
    --max-running-requests 256 \
    --kv-cache-dtype fp8_e4m3 \
    --enforce-shared-experts-fusion \
    --disable-radix-cache \
    --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \
    --port "${PORT}"

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 #33067512788
Latest PR Test (Extra): ✅ Run #33140291705
Latest PR Test (AMD ROCm 7.2): ❌ Run #33067512757

[HaiShaw]

HIP specific