chore: import upstream snapshot with attribution

python/sglang/jit_kernel/moe_topk_sigmoid.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_moe_topk_sigmoid_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "moe_topk_sigmoid",
+        *args,
+        cuda_files=["moe/moe_topk_sigmoid.cuh"],
+        cuda_wrappers=[("topk_sigmoid", f"topk_sigmoid<{args}>")],
+        extra_cuda_cflags=["--use_fast_math"],
+    )
+
+
+@register_custom_op(
+    op_name="moe_topk_sigmoid_out",
+    mutates_args=["topk_weights", "topk_ids"],
+)
+def moe_topk_sigmoid_out(
+    gating_output: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    workspace: torch.Tensor,
+    renormalize: bool,
+    correction_bias: Optional[torch.Tensor],
+    routed_scaling_factor: float,
+    num_fused_shared_experts: int,
+) -> None:
+    """
+    Fused sigmoid top-k MoE gate (destination-passing style).
+
+    Args:
+        gating_output: [num_tokens, num_experts], fp32/fp16/bf16
+        topk_weights:  [num_tokens, topk], float32, pre-allocated output
+        topk_ids:      [num_tokens, topk], int32,   pre-allocated output
+        workspace:     [num_tokens * num_experts] float32 scratch (may be size 1
+                       when num_experts is a supported power-of-2 ≤ 256)
+        renormalize:   whether to renormalize weights to sum to 1 per row
+        correction_bias: [num_experts] float32 per-expert bias, or None
+        routed_scaling_factor: [num_tokens, num_experts] float32, or None
+    """
+    module = _jit_moe_topk_sigmoid_module(gating_output.dtype)
+    module.topk_sigmoid(
+        gating_output,
+        topk_weights,
+        topk_ids,
+        workspace,
+        renormalize,
+        correction_bias,
+        routed_scaling_factor,
+        num_fused_shared_experts,
+    )
+
+
+def topk_sigmoid(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    correction_bias: Optional[torch.Tensor] = None,
+    routed_scaling_factor: float = 1.0,
+    num_fused_shared_experts: int = 0,
+) -> None:
+    """
+    Fused sigmoid top-k MoE gate with the same call signature as
+    ``sgl_kernel.topk_sigmoid`` (destination-passing, in-place).
+
+    Args:
+        topk_weights:  [num_tokens, topk] float32, written in-place
+        topk_ids:      [num_tokens, topk] int32,   written in-place
+        gating_output: [num_tokens, num_experts] fp32/fp16/bf16
+        renormalize:   whether to renormalize weights to sum to 1 per row
+        correction_bias: [num_experts] float32 per-expert bias, or None
+    """
+    num_tokens = gating_output.shape[0]
+    num_experts = gating_output.shape[1]
+
+    is_pow2 = num_experts != 0 and (num_experts & (num_experts - 1)) == 0
+    needs_workspace = not is_pow2 or num_experts > 256
+    workspace_size = num_tokens * num_experts if needs_workspace else 1
+    workspace = torch.empty(
+        workspace_size, dtype=torch.float32, device=gating_output.device
+    )
+
+    moe_topk_sigmoid_out(
+        gating_output,
+        topk_weights,
+        topk_ids,
+        workspace,
+        renormalize,
+        correction_bias,
+        routed_scaling_factor,
+        num_fused_shared_experts,
+    )