add topk_softplus kernel

[yzhou103]

Motivation

for deepseek v4

Technical Details

topk_softplus implements the DeepSeek V4-Pro sqrtsoftplus routing: sqrt(softplus(x)) + bias → topk selection → subtract bias → renormalize → scale. The original kernel uses shared memory to stage all expert scores, then performs topk sequential full scans of the shared-memory array (each scan finds one global maximum via warp reduce and marks the winner as -INF).

Test Plan

python op_tests/test_moe_topk_sigmoid.py --test softplus --num-experts 384 256 --num-tokens 1 1024 --topk 6

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests
ci:atom	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
ci:vllm	vLLM benchmark
ci:all	All of the above

Add labels via the sidebar or gh pr edit 2995 --add-label <label>

[Copilot]

Pull request overview

This PR introduces a new MoE routing primitive (topk_softplus, intended for DeepSeek v4-style sqrt(softplus) routing) and wires it through the C++/pybind interface, Python API, and op tests/benchmarks.

Changes:

• Add topk_softplus pybind exposure and C++ declaration.
• Add a Python aiter.topk_softplus entrypoint (JIT-compiled via module_moe_topk).
• Extend op_tests/test_moe_topk_sigmoid.py with topk_softplus correctness + benchmark paths and a --test selector.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
op_tests/test_moe_topk_sigmoid.py	Adds torch reference + fused invocation for topk_softplus, plus pytest correctness and CLI benchmarking updates.
csrc/include/rocm_ops.hpp	Extends MOE_TOPK_PYBIND to bind topk_softplus.
csrc/include/moe_op.h	Adds the C++ API declaration for aiter::topk_softplus.
aiter/ops/topk.py	Adds a compile_ops("module_moe_topk") Python stub for topk_softplus.
aiter/jit/optCompilerConfig.json	Adds a new source file entry to module_moe_topk build inputs for the topk_softplus implementation.

Comments suppressed due to low confidence (1)

csrc/include/rocm_ops.hpp:1181

• Minor: the binding default for routed_scaling_factor uses a double literal (1.0) while the C++ signature takes a float. Consider using 1.0f for consistency with other bindings in this header and to avoid an unnecessary implicit conversion.

          py::arg("local_expert_mask") = std::nullopt, \
          py::arg("num_local_tokens")  = std::nullopt, \
          py::arg("dispatch_policy")   = 0);

#define MOE_SORTING_OPUS_PYBIND                        \
    m.def("moe_sorting_opus_fwd",                      \
          &moe_sorting_opus_fwd,                       \
          py::arg("topk_ids"),                         \
          py::arg("topk_weights"),                     \

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yzhou103]