[davzhuAMD]

Overview

On ROCm, DeepSeek-V4/V4.1's sqrtsoftplus routing never reaches vLLM's single-launch Triton top-k (dsv4_topk), because can_use_dsv4_topk() requires current_platform.is_cuda(); it runs the generic HIP topkGatingSoftplusSqrt kernel instead. This PR allows dsv4_topk on ROCm, and also calls it from the AITER fused-MoE branch of fused_topk_bias(), the branch the DSv4.1 ROCm recipe takes (VLLM_ROCM_USE_AITER_MOE=1). It adds no kernel and is Python only.

Credit: this change is @heachary's. He found the gate, measured it and wrote the patch; I ported it onto current main and extended the tests. The commit carries a Co-authored-by trailer for him.

Related:

• [ROCm][Kernel] Wave64 topk_softplus_sqrt fast path for DeepSeek-V4 routing #64208 (draft) speeds up the same call with a gfx950 Wave64 HIP path in the generic kernel. That is a different approach; see "Not duplicated".
• [ROCm] Heterogeneous fused shared experts for DeepSeek-V4.1 (native-FP8 shared slot, prefill-shaped batches) #63180 rewrites the same AITER sqrtsoftplus block for fused shared experts; see Details.
• [ROCm][DSv4.1][Perf] Fused router gate for gfx950 #57478 (closed) fused the gfx950 gate GEMM with the top-k.
• [Bugfix][Kernel] DSv4 router: compute sqrt(softplus) with log1p #64277 (draft) computes sqrt(softplus) with log1p in both dsv4_topk and the HIP kernel.

See Details.

Claims

• Scope: on ROCm, MoE layers that can_use_dsv4_topk() accepts now route through dsv4_topk, with or without AITER fused MoE. It accepts fp32 contiguous logits with 256 or 384 experts, top-6, renormalize, an fp32 correction bias, and int32/int64 ids, without a hash table. Hash layers, the 128-expert DSpark draft gate, other scoring functions and CUDA are unchanged.

• Kernel time: from @heachary's microbenchmark, not on this branch:

  • one idle MI355X, CUDA-graph replay, through GateLinear → fused_topk_bias();
  • the DSv4.1 router shape ([M, 5120] × 384 experts, top-6), vLLM main of Oct 21 (3f1bcfd8), one session.

  M	32	64	128	192	224	256	512
  top-k, µs before → after	9.24 → 3.27	9.24 → 3.34	9.27 → 3.33	9.32 → 3.38	9.27 → 3.40	9.22 → 3.43	9.43 → 3.50
  router chain (gate GEMM + top-k), µs	16.18 → 10.20	16.12 → 10.22	16.30 → 10.36	18.30 → 12.34	18.32 → 12.46	18.38 → 12.57	20.52 → 14.56

  The gate GEMM is untouched. The numbers on this branch are TODO (see Validation).

• Numerics: expert ids unchanged, weights not bitwise. @heachary checked against ops.topk_hash_softplus_sqrt (the kernel this replaces) on MI355X:

  • Expert ids were bit-identical on every row at M 1–8192, on:
    • real gate logits;
    • every softplus branch (above the 20.0 linear cut-over, underflow, all-zero, ±inf and NaN rows);
    • forced ties, and 0/50/100% padding;
    • bias_vl sentinel rows, and int32/int64 ids.
  • Routing weights differed by at most 9e-8 absolute (4e-7 relative). The cause is Triton's tl.log/tl.exp against the HIP kernel's fast-math intrinsics.
  • Both kernels break ties to the lowest expert id. Scores tied within ~1e-7 could still select differently in principle; none was found.

• E2E, supporting only (not on main; the A/B on main is a TODO in Validation):

  • MI355X: DeepSeek-V4.1-Flash, 2 × (TP4 + EP4) servers on 8 GPUs, AITER MoE, DSpark with 6 draft tokens, an agentic multi-turn workload at concurrency 64. Public nightly vllm/vllm-openai-rocm:nightly-rocm100-3966924d + [ROCm][DSv4.1][Perf] Enable prefix-cache SWA replay on gfx950 #64094, one machine, n = 1 per arm:

    arm	date	output tok/s	TPOT p50	workload score
    nightly + [ROCm][DSv4.1][Perf] Enable prefix-cache SWA replay on gfx950 #64094	Oct 30	8,027	4.671 ms	0.5323
    + an earlier revision of a split-K router GEMM (proposed separately)	Nov 4	8,088	4.624 ms	0.5298
    + that GEMM + this PR	Nov 4	8,191	4.545 ms	0.5274

    • This PR's effect is read by difference between the two Nov 4 runs: +1.3% output tok/s, −1.7% TPOT p50. It was added on top of the GEMM, not measured alone.
    • TPOT is client-side, and the client fell back to whitespace token counts on tool-call turns, so read it as relative only.
    • The runs used synthetic draft acceptance, so the score is not an accuracy check; GSM8K is a TODO.

  • MI350X (not in CI): DeepSeek-V4.1-Flash DP8 + EP8 with MoRI, agentic workload at concurrency 128, a vLLM nightly from Oct 21 plus other patches.

    • This change alone against the same stack without it, n = 2 runs per arm.
    • Output tok/s +1.30% (same-arm sd 0.54%); workload accuracy −0.0035 ± 0.0031.

• No new knobs: no new env var, config field or kernel.

Validation

• Unit tests: tests/kernels/moe/test_topk_softplus_sqrt.py.
  • CI coverage: the existing :amd: (MI355) MoE Kernels and :amd: (MI300) MoE Kernels steps run pytest kernels/moe and depend on vllm/model_executor/layers/fused_moe/. Both are optional, and no Buildkite change is needed.
  • Widened tests: test_dsv4_fast_topk (6 cases) and test_dsv4_fast_topk_bias_vl were CUDA-only and now run on ROCm too.
  • New test: test_dsv4_fast_topk_with_aiter_fused_moe (ROCm only; M 1, 48, 224 and 4096, with and without padding: 8 cases).
    • CI runs this file without the AITER env vars, so the test makes rocm_aiter_ops.is_fused_moe_enabled() return True.
    • It checks with a spy that fused_topk_bias() took dsv4_topk.
    • It compares against the file's torch reference: ids exactly, and weights within 2e-5, the tolerance test_dsv4_fast_topk already uses.

  # MI355X (gfx950); the whole file, so the HIP-kernel tests are re-checked too
  .venv/bin/python -m pytest -v -s tests/kernels/moe/test_topk_softplus_sqrt.py
  VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 .venv/bin/python -m pytest -v -s tests/kernels/moe/test_topk_softplus_sqrt.py
  TODO: paste both summaries
  

  • Expected pre-existing failures: test_hash_routing_backend_selection[triton_unfused-…] (2 cases) fails on ROCm on main too, as reported in [ROCm][Kernel] Wave64 topk_softplus_sqrt fast path for DeepSeek-V4 routing #64208.
  • TODO: the same file on MI300X (gfx942). is_cuda_alike() enables this path on every ROCm arch, and the MI300 step runs it.
  • TODO: the new test on main (test hunk only), to show it fails there at the dsv4_topk assertion.
  • CPU-only so far: the file collects (1,691 cases), and ruff check, ruff format and the SPDX hook pass. Full pre-commit: TODO: paste.
• Microbenchmark on this branch: TODO: MI355X, CUDA-graph replay, top-k and router-chain µs on main against this branch at M 1, 48, 128, 224 and 4096 (for example with [Benchmark][MoE] Add eager and graph top-k chain measurements #65298's top-k chain benchmark).
• Serving A/B on main: TODO:
  • MI355X, DeepSeek-V4.1-Flash, VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1;
  • DP8 + EP8 with --all2all-backend mori_high_throughput, and/or TP4 + EP4;
  • vllm bench serve on main against this branch, same image, n ≥ 2 interleaved runs per arm;
  • report output tok/s, TPOT p50 and the spread.

  TODO: server and bench commands; table per concurrency
  

• Accuracy: TODO: GSM8K (5-shot, greedy, 1,319 questions) before and after, on the same server.

Details

Why dsv4_topk was CUDA-only, and why it runs on ROCm.

• The kernel is plain Triton apart from the PDL intrinsics (gdc_wait, gdc_launch_dependents).
• Those sit behind the launch_pdl constexpr, which is current_platform.is_arch_support_pdl(). That is False on ROCm, so they compile out.
• One program (one warp) per row is the right shape on CDNA too. @heachary swept 1–16 rows per program × 1–4 warps on MI355X, and the shipped config was the fastest: 3.42 µs, against 14.97 µs at 16 rows per program.

Why it's faster. The generic kernel's time is flat in M (9.2–9.5 µs from 32 to 512 rows, above), which points at a fixed cost per call. #64208 traces it to 32 B per lane of scratch and the per-round argmax reductions.

Changes (+79/−26 in two source files, +63/−4 in one test file):

• router/dsv4_topk.py: can_use_dsv4_topk() and the guard around the kernel definition use is_cuda_alike() instead of is_cuda(), with a comment on PDL.
• router/fused_topk_bias_router.py:
  • A new _maybe_dsv4_topk() holds the existing eligibility test (sqrtsoftplus, no hash table, can_use_dsv4_topk()) and the dsv4_topk call; otherwise it returns None.
  • fused_topk_bias() calls it in the non-AITER branch, replacing the inline copy there; behaviour is unchanged.
  • It also calls it just before the AITER path's sqrtsoftplus fallback to vllm_topk_softplus_sqrt. That call site is new, and it is reached only when rocm_aiter_ops.is_fused_moe_enabled().
  • Both call sites pass get_moe_padding_mask() as before, so padded rows still get id −1 and weight 0.

Fused shared experts, and #63180.

• On main: FusedTopKBiasRouter._compute_routing() appends the fused shared-expert slots with torch.cat after fused_topk_bias() returns. So the append applies to dsv4_topk's output exactly as it does to the HIP kernel's.
• What [ROCm] Heterogeneous fused shared experts for DeepSeek-V4.1 (native-FP8 shared slot, prefill-shaped batches) #63180 changes: it adds out_topk_weights/out_topk_ids (persistent (M, top_k + n) buffers) and num_fused_shared_experts to fused_topk_bias(), so that the HIP kernel writes the shared columns itself in the AITER sqrtsoftplus block.
• What this PR does with that: _maybe_dsv4_topk() doesn't take those arguments, because they aren't on main, and it returns before that block.
• If both land as written:
  • the fast path ignores the buffers and returns (M, top_k);
  • [ROCm] Heterogeneous fused shared experts for DeepSeek-V4.1 (native-FP8 shared slot, prefill-shaped batches) #63180's width check in _compute_routing() then falls back to the cat append, so results stay correct;
  • but [ROCm] Heterogeneous fused shared experts for DeepSeek-V4.1 (native-FP8 shared slot, prefill-shaped batches) #63180's zero-kernel append (which it enables for prefill-shaped batches) is lost on the layers this path serves.
• TODO, whichever PR merges second: skip the fast path when out-buffers are passed, or teach dsv4_topk the append. To agree with @alexm-redhat.
• Textually, the two PRs conflict in context only: this PR's second call sits just above the block [ROCm] Heterogeneous fused shared experts for DeepSeek-V4.1 (native-FP8 shared slot, prefill-shaped batches) #63180 rewrites.

Limitations:

• Not bitwise: the weights differ by up to 9e-8 in the check above. Ids were identical in every case tested, but that isn't guaranteed for scores tied within ~1e-7.
• Coverage: hash layers, the 128-expert DSpark draft gate (top-3) and other expert counts keep the HIP kernel. [ROCm][Kernel] Wave64 topk_softplus_sqrt fast path for DeepSeek-V4 routing #64208 covers hash rows and 512 experts.
• Architectures: is_cuda_alike() turns this on for every ROCm GPU, but it has been measured only on gfx950 (MI355X, MI350X).
  • gfx942 is a TODO above, and RDNA (wave32) is untested.
  • If either is in doubt, I'll restrict the change to gfx9.
• Not yet run on this branch: every number above was measured on older builds; see the TODOs.

Not duplicated (searched 2026-11-04):

• Commands:

  • gh pr list --repo vllm-project/vllm --state all --search "<q> updated:>=2026-10-01", with q in: dsv4_topk, can_use_dsv4_topk, sqrtsoftplus, softplus_sqrt, topkGatingSoftplusSqrt, topk_hash_softplus_sqrt, "DSv4 top-k ROCm", "dsv4 topk rocm", "fused gate topk", "router top-k ROCm", "is_cuda_alike topk", "DeepSeek-V4 router ROCm", fused_topk_bias, FusedTopKBiasRouter;
  • gh pr view 64208 and gh pr view 63180 -R vllm-project/vllm.

• Findings: no open PR enables dsv4_topk on ROCm or calls it from the AITER branch.

• [ROCm][Kernel] Wave64 topk_softplus_sqrt fast path for DeepSeek-V4 routing #64208 (@amd-dlimpus, draft) targets the same call with a materially different approach. The µs figures below come from different benchmarks and builds, so they are not a controlled comparison:

  	[ROCm][Kernel] Wave64 topk_softplus_sqrt fast path for DeepSeek-V4 routing #64208	this PR
  approach	a new gfx950 Wave64 HIP path inside topkGatingSoftplusSqrtKernelLauncher (C++, needs a rebuild)	enables vLLM's existing Triton dsv4_topk on ROCm (Python only)
  time per call, MI355X, graph replay, 384 experts	5.38–5.80 µs at 1–256 rows (its description)	3.27–3.43 µs at 32–256 rows (above)
  weights vs the generic kernel	bit-identical	within 9e-8
  expert ids vs the generic kernel	bit-identical	bit-identical in every case tested
  coverage	256/384/512 experts, hash and non-hash rows; gfx950 only	256/384 experts, non-hash rows; all ROCm

  • They compose: if both land, this path takes the non-hash 256/384-expert layers, and [ROCm][Kernel] Wave64 topk_softplus_sqrt fast path for DeepSeek-V4 routing #64208 serves hash layers and 512 experts.
  • The tradeoff: this PR is faster per call and reuses a kernel already reviewed upstream, but gives up bitwise weights.
  • I'd like @amd-dlimpus's view on whether both should land.

• [ROCm][DSv4.1][Perf] Fused router gate for gfx950 #57478 (closed, "not significant e2e perf improvement") fused the gate GEMM with the top-k. This PR leaves the gate GEMM untouched, so the top-k consumes the same logits as before; it only swaps the top-k kernel.

• [Bugfix][Kernel] DSv4 router: compute sqrt(softplus) with log1p #64277 (draft) adds log1p to both kernels and keeps them consistent. If it lands first, I'll restate the numerics above.

• [ROCm][MoE] Accumulate the shared-expert output in AITER's stage 2; fold routed scaling into top-k #65207 (draft) edits the adjacent AITER-sigmoid lines in fused_topk_bias(); the conflict is context only. [MoE] Add a generic fused gate router and backend interface #65290 adds a generic fused-gate router interface (_compute_routing) and co-applies.

• [Bugfix][MoE] Fall back to torch routing for unsupported sqrtsoftplus expert counts #63687 and [Benchmark][MoE] Add eager and graph top-k chain measurements #65298: [Bugfix][MoE] Fall back to torch routing for unsupported sqrtsoftplus expert counts #63687 is a torch fallback for unsupported sqrtsoftplus expert counts and doesn't apply to main. [Benchmark][MoE] Add eager and graph top-k chain measurements #65298 is a top-k chain benchmark, usable for the TODO above.

AI assistance: AI tools (Claude, in Cursor) helped port @heachary's patch onto current main, write the added test, search for overlapping work and draft this description. I have reviewed every changed line; I will run the GPU tests, benchmarks and evals marked TODO above myself before this leaves draft. The commit carries a Co-authored-by: Claude trailer.

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀