feat(topk): optimized topk_gating kernel with sigmoid/softmax support (#3100)

* feat(topk): optimized topk_gating kernel with sigmoid/softmax support

Replace topk_softplus with a unified topk_gating kernel that supports
three scoring functions ("sqrtsoftplus" | "sigmoid" | "softmax") via a
templated SCORE_FUNC parameter, plus a Python-facing topk_gating()
wrapper.

Performance (vs reference kernels on MI355, bf16):
  sqrtsoftplus E=384 T=1024 K=6: 81.3 us -> 8.7 us  (9.3x)
  sigmoid     E=128 T=1   K=8 : CK   5.09 us -> 2.12 us (2.41x)
  sigmoid     E=128 T=1024 K=8 : CK   6.06 us -> 3.73 us (1.62x)
  softmax     E=256 T=1024 K=8 : vllm 7.95 us -> 6.32 us (1.26x)
  (vllm topk_softmax still wins for E<=32 / K<=2;
   we win for E>=64 + K>=4.)

Implementation highlights (register-only opt kernel):
- 1 warp = 1 token, no shared memory, no __syncthreads
- Each thread holds EPT = NUM_EXPERTS / WARP_SIZE elements
- Optimal sorting networks for N=2,3,4,6 (vs unrolled bubble sort)
- Fused DPP warp argmax (multithread_reduce_max_dpp + __ballot + readlane)
  -> ~9 instructions per warp argmax, NaN-safe via ballot fallback
- Cached unbiased scores so the merge phase reads no global memory
- fmaxf(sum, 1e-20f) renorm clamp to handle pathological all-zero rows
- Generic shared-memory fallback for non-power-of-WARP_SIZE expert counts
  (also handles softmax: max-reduce -> exp -> sum-reduce -> divide)

Stream management intentionally uses at::hip::getCurrentHIPStream()
for correctness with non-default-stream callers (CUDA graphs, EP/TP
comm streams).  Switching to aiter::getCurrentHIPStream() requires
migrating the whole pybind module to aiter_tensor_t + develop=True;
tracked as TODO in the kernel.

Test coverage:
- op_tests/test_moe_topk_gating.py (renamed from test_moe_topk_sigmoid.py)
  covers sqrtsoftplus / sigmoid / softmax across {fp16, bf16, fp32}
  inputs, {fp16, bf16, fp32} biases, num_experts in {64, 128, 256, 384},
  topk in {2, 4, 6, 8}, num_tokens in {1, 64, 256, 1024}.

Co-authored-by: Cursor <cursoragent@cursor.com>

* add more check

* fix test error

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: yzhou103

aiter/jit/optCompilerConfig.json

@@ -534,7 +534,7 @@
         "srcs": [
             "f'{AITER_CSRC_DIR}/pybind/moe_topk_pybind.cu'",
             "f'{AITER_CSRC_DIR}/py_itfs_ck/topk_sigmoid_kernels.cu'",
-            "f'{AITER_CSRC_DIR}/kernels/topk_softplus_kernels.cu'",
+            "f'{AITER_CSRC_DIR}/kernels/topk_gating_kernels.cu'",
             "f'{CK_DIR}/example/ck_tile/09_topk_softmax/topk_softmax_api.cpp'"
         ],
         "flags_extra_cc": [],

aiter/ops/topk.py

@@ -12,6 +12,12 @@
 from ..utility import dtypes
 
 
+# DEPRECATED: low-level binding kept for backward compatibility only.
+# Will be removed once all callers have migrated to topk_gating() below.
+# New code should use topk_gating(), which:
+#   - accepts an Optional[Tensor] correction_bias (None => no bias)
+#   - validates score_func string
+#   - exposes the same C++ kernel under a more accurate name
 @compile_ops("module_moe_topk")
 def topk_softplus(
     topk_weights: torch.Tensor,
@@ -20,9 +26,54 @@ def topk_softplus(
     correction_bias: torch.Tensor,
     need_renorm: bool,
     routed_scaling_factor: float = 1.0,
+    score_func: str = "sqrtsoftplus",
 ) -> None: ...
 
 
+_VALID_SCORE_FUNCS = {"sqrtsoftplus", "sigmoid", "softmax"}
+
+
+def topk_gating(
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    gating_output: torch.Tensor,
+    correction_bias: Optional[torch.Tensor] = None,
+    need_renorm: bool = True,
+    routed_scaling_factor: float = 1.0,
+    score_func: str = "sqrtsoftplus",
+) -> None:
+    """Unified fused topk gating for MoE routing.
+
+    Args:
+        score_func: one of {"sqrtsoftplus" (DeepSeek V4-Pro default),
+                            "sigmoid" (Llama4),
+                            "softmax" (DeepSeek V3 / classic MoE)}.
+        correction_bias: optional bias tensor, pass None for no bias.
+
+    Note: softmax is already normalized, so renorm is forced off.
+    """
+    assert (
+        score_func in _VALID_SCORE_FUNCS
+    ), f"Unknown score_func '{score_func}', expected one of {_VALID_SCORE_FUNCS}"
+    if correction_bias is None:
+        # Match gating dtype/device so dispatch picks DTYPE_B == DTYPE_I,
+        # avoiding extra kernel template instantiations.
+        correction_bias = torch.empty(
+            0, dtype=gating_output.dtype, device=gating_output.device
+        )
+    if score_func == "softmax":
+        need_renorm = False
+    topk_softplus(
+        topk_weights,
+        topk_indices,
+        gating_output,
+        correction_bias,
+        need_renorm,
+        routed_scaling_factor,
+        score_func,
+    )
+
+
 @compile_ops("module_moe_asm", fc_name="biased_grouped_topk")
 def biased_grouped_topk_hip(
     gating_output: torch.Tensor,

csrc/include/moe_op.h

@@ -55,7 +55,8 @@ void topk_softplus(torch::Tensor& topk_weights,
                     torch::Tensor& gating_output,
                     torch::Tensor& correction_bias,
                     bool need_renorm,
-                    float routed_scaling_factor = 1.0);
+                    float routed_scaling_factor = 1.0,
+                    const std::string& score_func = "sqrtsoftplus");
 
 void moe_align_block_size(torch::Tensor topk_ids,
                           int64_t num_experts,

csrc/include/rocm_ops.hpp

@@ -1180,7 +1180,8 @@ namespace py = pybind11;
           py::arg("correction_bias"),                    \
           py::arg("need_renorm"),                        \
           py::arg("routed_scaling_factor") = 1.0,        \
-          "Apply topk sqrtsoftplus to the gating outputs.");
+          py::arg("score_func") = "sqrtsoftplus",        \
+          "Fused topk gating: score_func='sqrtsoftplus'|'sigmoid'|'softmax'.");
 
 #define MOE_SORTING_PYBIND                             \
     m.def("moe_sorting_fwd",                           \