minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[JIT Kernel] Migrate dsv3_router_gemm from AOT sgl-kernel to JIT kernel (#21531)

Browse Source 
Co-authored-by: Guohao Shao <shao.gh.98@gmail.com>
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
meinie
2026-06-26 11:52:18 -07:00
committed by GitHub
co-authored by Guohao Shao Brayden Zhong
parent c98d31143d
commit 714011a40f
12 changed files with 395 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/dsv3_router_gemm.cuh
+184
View File
@@ -0,0 +1,184 @@
/*
* Adapted from
* https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
* https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
*
* Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <tvm/ffi/container/tensor.h>
#include <type_traits>
namespace {
using namespace device;
static constexpr int kDefaultNumExperts = 256;
static constexpr int kKimiK2NumExperts = 384;
static constexpr int kDefaultHiddenDim = 7168;
// kOutFloat: true = float32 output, false = bfloat16 output
template <
typename T,
typename OutT,
int kBlockSize,
int VPT,
int kNumTokens,
int kNumExperts,
int kHiddenDim,
bool kUsePDL>
__global__ __launch_bounds__(kBlockSize, 1) void router_gemm_kernel(OutT* out, T const* mat_a, T const* mat_b) {
constexpr int kWarpSize = 32;
constexpr int kNumWarps = kBlockSize / kWarpSize;
constexpr int kElemsPerKIter = VPT * kBlockSize;
static_assert(kHiddenDim % kElemsPerKIter == 0, "hidden_dim must be divisible by one K iteration");
constexpr int kIters = kHiddenDim / kElemsPerKIter;
// Padding to avoid shared memory bank conflicts when kNumTokens > 8
constexpr int kSmReductionPad = (kNumTokens > 8) ? 1 : 0;
static_assert(kSmReductionPad == 0 || kSmReductionPad == 1, "kSmReductionPad only supports 0 or 1");
int const n_idx = blockIdx.x;
int const tid = threadIdx.x;
int const warp_id = tid / kWarpSize;
int const lane_id = tid % kWarpSize;
float acc[kNumTokens] = {};
__shared__ float sm_reduction[kNumTokens][kNumWarps + kSmReductionPad];
T const* b_col = mat_b + n_idx * kHiddenDim;
PDLWaitPrimary<kUsePDL>();
int k_base = tid * VPT;
#pragma unroll
for (int ki = 0; ki < kIters; ++ki, k_base += kElemsPerKIter) {
AlignedVector<bf16_t, VPT> b_vec;
b_vec.load(b_col + k_base);
#pragma unroll
for (int m_idx = 0; m_idx < kNumTokens; ++m_idx) {
AlignedVector<bf16_t, VPT> a_vec;
a_vec.load(mat_a + m_idx * kHiddenDim + k_base);
#pragma unroll
for (int k = 0; k < VPT; ++k) {
acc[m_idx] += cast<float>(a_vec[k]) * cast<float>(b_vec[k]);
}
}
}
#pragma unroll
for (int m_idx = 0; m_idx < kNumTokens; ++m_idx) {
float sum = warp::reduce_sum(acc[m_idx]);
if (lane_id == 0) {
sm_reduction[m_idx][warp_id] = sum;
}
}
__syncthreads();
if (warp_id == 0 && lane_id < kNumTokens) {
float final_sum = 0.0f;
#pragma unroll
for (int w = 0; w < kNumWarps; ++w) {
final_sum += sm_reduction[lane_id][w];
}
out[lane_id * kNumExperts + n_idx] = cast<OutT>(final_sum);
}
PDLTriggerSecondary<kUsePDL>();
}
template <typename T, typename OutT, int kNumTokens, int kNumExperts, int kHiddenDim, bool kUsePDL>
void invokeRouterGemm(OutT* output, T const* mat_a, T const* mat_b, DLDevice device) {
constexpr int VPT = 16 / sizeof(T);
constexpr int kBlockSize = 128;
constexpr auto kernel = router_gemm_kernel<T, OutT, kBlockSize, VPT, kNumTokens, kNumExperts, kHiddenDim, kUsePDL>;
host::LaunchKernel(kNumExperts, kBlockSize, device).enable_pdl(kUsePDL)(kernel, output, mat_a, mat_b);
}
// Dispatch runtime num_tokens to compile-time template parameter [kBegin, kEnd]
template <int kBegin, int kEnd, typename OutT, int kNumExperts, int kHiddenDim, bool kUsePDL>
struct RouterGemmDispatcher {
static void run(int num_tokens, OutT* output, bf16_t const* mat_a, bf16_t const* mat_b, DLDevice device) {
if (num_tokens == kBegin) {
invokeRouterGemm<bf16_t, OutT, kBegin, kNumExperts, kHiddenDim, kUsePDL>(output, mat_a, mat_b, device);
} else {
RouterGemmDispatcher<kBegin + 1, kEnd, OutT, kNumExperts, kHiddenDim, kUsePDL>::run(
num_tokens, output, mat_a, mat_b, device);
}
}
};
// Base case: kBegin == kEnd
template <int kEnd, typename OutT, int kNumExperts, int kHiddenDim, bool kUsePDL>
struct RouterGemmDispatcher<kEnd, kEnd, OutT, kNumExperts, kHiddenDim, kUsePDL> {
static void run(int num_tokens, OutT* output, bf16_t const* mat_a, bf16_t const* mat_b, DLDevice device) {
if (num_tokens == kEnd) {
invokeRouterGemm<bf16_t, OutT, kEnd, kNumExperts, kHiddenDim, kUsePDL>(output, mat_a, mat_b, device);
} else {
host::panic({}, "dsv3_router_gemm: num_tokens must be between 1 and 16, got ", num_tokens);
}
}
};
// kNumExperts: compile-time 256 or 384
// kHiddenDim: compile-time hidden dim, any multiple of one K iteration (1024)
// kUsePDL: compile-time bool (true on SM90+)
// kOutFloat: compile-time bool (true = float32 output, false = bfloat16 output)
template <int kNumExperts, int kHiddenDim, bool kUsePDL, bool kOutFloat>
struct DSV3RouterGemmKernel {
static_assert(
kNumExperts == kDefaultNumExperts || kNumExperts == kKimiK2NumExperts,
"required num_experts == 256 or num_experts == 384");
using OutT = std::conditional_t<kOutFloat, fp32_t, bf16_t>;
static void
run(const tvm::ffi::TensorView mat_a, const tvm::ffi::TensorView mat_b, const tvm::ffi::TensorView output) {
using namespace host;
auto M = SymbolicSize{"num_tokens"};
auto K = SymbolicSize{"hidden_dim"};
auto N = SymbolicSize{"num_experts"};
auto device = SymbolicDevice{};
K.set_value(kHiddenDim);
N.set_value(kNumExperts);
device.set_options<kDLCUDA>();
TensorMatcher({M, K}).with_dtype<bf16_t>().with_device(device).verify(mat_a);
TensorMatcher({N, K}).with_dtype<bf16_t>().with_device(device).verify(mat_b);
TensorMatcher({M, N}).with_dtype<OutT>().with_device(device).verify(output);
const int num_tokens = static_cast<int>(M.unwrap());
RouterGemmDispatcher<1, 16, OutT, kNumExperts, kHiddenDim, kUsePDL>::run(
num_tokens,
static_cast<OutT*>(output.data_ptr()),
static_cast<bf16_t const*>(mat_a.data_ptr()),
static_cast<bf16_t const*>(mat_b.data_ptr()),
device.unwrap());
}
};
} // namespace
python/sglang/jit_kernel/dsv3_router_gemm.py
+92
View File
@@ -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
python/sglang/srt/models/deepseek_v2.py
+10 -37
View File
@@ -201,8 +201,11 @@ if _use_aiter:
pass
if _is_cuda:
from flashinfer.gemm import mm_M1_16_K7168_N256 as _raw_dsv3_router_gemm
from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm
from sgl_kernel import dsv3_fused_a_gemm
from sglang.jit_kernel.dsv3_router_gemm import (
dsv3_router_gemm as _jit_dsv3_router_gemm,
)
elif _is_npu:
from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
forward_dsa_core_npu,
@@ -213,7 +216,7 @@ elif _is_npu:
forward_mla_prepare_npu,
)
elif _is_musa:
from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm
from sgl_kernel import dsv3_fused_a_gemm
else:
pass
@@ -494,28 +497,16 @@ class MoEGate(nn.Module):
):
logits = F.linear(hidden_states, self.weight, None)
else:
# NOTE: For some unknown reason, router_gemm seems degrade accept length.
if (
_is_cuda
and hidden_states.shape[0] <= 16
and hidden_states.shape[1] == 7168
and hidden_states.shape[1] % 1024 == 0
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
and _device_sm >= 90
):
if _device_sm in [100, 103] and self.weight.shape[0] == 256:
# TODO: will check the dtype to be bf16
# router gemm output float32
logits = torch.empty(
hidden_states.shape[0],
self.weight.shape[0],
device=hidden_states.device,
dtype=torch.float32,
)
flashinfer_dsv3_router_gemm(logits, hidden_states, self.weight)
else:
logits = dsv3_router_gemm(
hidden_states, self.weight, out_dtype=torch.float32
)
logits = _jit_dsv3_router_gemm(
hidden_states, self.weight, out_dtype=torch.float32
)
elif _use_aiter:
logits = aiter_dsv3_router_gemm(hidden_states, self.weight)
@@ -2946,24 +2937,6 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
pass
@register_custom_op(
op_name="flashinfer_dsv3_router_gemm",
mutates_args=[],
fake_impl=lambda logits, hidden_states, weight: None,
)
def flashinfer_dsv3_router_gemm(
logits: torch.Tensor,
hidden_states: torch.Tensor,
weight: torch.Tensor,
) -> None:
_raw_dsv3_router_gemm(
hidden_states,
weight.t(),
logits,
launch_with_pdl=True,
)
@register_custom_op(out_shape="hidden_states")
def dsv2_flashinfer_moe_dual_stream_graph(
hidden_states: torch.Tensor,