chore: import upstream snapshot with attribution

python/sglang/jit_kernel/dsv3_router_gemm.py

@@ -0,0 +1,92 @@
+"""
+JIT kernel for DeepSeek V3 router GEMM.
+
+Replaces the AOT sgl_kernel.dsv3_router_gemm for SM90+ (Hopper) GPUs.
+Supports num_experts in {256, 384}, hidden_dim a multiple of 1024, num_tokens 1-16.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_dsv3_router_gemm_module(
+    num_experts: int,
+    hidden_dim: int,
+    use_pdl: bool,
+    out_float: bool,
+) -> Module:
+    args = make_cpp_args(num_experts, hidden_dim, use_pdl, out_float)
+    return load_jit(
+        "dsv3_router_gemm",
+        *args,
+        cuda_files=["gemm/dsv3_router_gemm.cuh"],
+        cuda_wrappers=[
+            ("dsv3_router_gemm", f"DSV3RouterGemmKernel<{args}>::run"),
+        ],
+    )
+
+
+@register_custom_op(
+    op_name="dsv3_router_gemm",
+    mutates_args=["output"],
+)
+def _dsv3_router_gemm_custom_op(
+    hidden_states: torch.Tensor,
+    router_weights: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    num_experts = router_weights.shape[0]
+    hidden_dim = hidden_states.shape[1]
+    out_float = output.dtype == torch.float32
+    module = _jit_dsv3_router_gemm_module(
+        num_experts, hidden_dim, is_arch_support_pdl(), out_float
+    )
+    module.dsv3_router_gemm(hidden_states, router_weights, output)
+    return None
+
+
+@debug_kernel_api
+def dsv3_router_gemm(
+    hidden_states: torch.Tensor,
+    router_weights: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    DeepSeek V3 router GEMM kernel (JIT variant).
+
+    Args:
+        hidden_states: Input tensor of shape [num_tokens, hidden_dim], bfloat16.
+            hidden_dim must be a multiple of 1024 and num_tokens in [1, 16].
+        router_weights: Weight tensor of shape [num_experts, hidden_dim], bfloat16.
+        out_dtype: Output dtype, either torch.bfloat16 or torch.float32.
+        output: Optional pre-allocated output tensor.
+
+    Returns:
+        Output tensor of shape [num_tokens, num_experts].
+    """
+    if output is None:
+        output = torch.empty(
+            hidden_states.shape[0],
+            router_weights.shape[0],
+            device=hidden_states.device,
+            dtype=out_dtype,
+        )
+    _dsv3_router_gemm_custom_op(hidden_states, router_weights, output)
+    return output