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

[Kernel] Replace dsv3_router_gemm with the unified tiny GEMM (#34693)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-09-01 23:00:21 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent ee462b5899
commit cb6dd58fbe
15 changed files with 445 additions and 637 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh
-184
View File
@@ -1,184 +0,0 @@
/*
* 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 sglang {
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 sglang
python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh
+146 -101
View File
@@ -9,41 +9,115 @@
#include <tvm/ffi/container/tensor.h>
#include <array>
#include <bit>
#include <cstddef>
#include <cstdint>
#include <utility>
namespace sglang {
using namespace device;
constexpr uint32_t kMaxBlockThreads = 1024;
constexpr uint32_t kTinyNGemmVecSize = kMaxVecBytes / sizeof(bf16_t);
template <uint32_t N_, uint32_t K_, uint32_t N_SPLIT_, uint32_t kBytes_>
struct GEMMTraitN {
static constexpr uint32_t N = N_;
static constexpr uint32_t K = K_;
static constexpr uint32_t N_SPLIT = N_SPLIT_; // per block n size
static constexpr uint32_t kBytes = kBytes_;
template <uint32_t M, uint32_t N, uint32_t K, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
__global__ __launch_bounds__(K / kTinyNGemmVecSize, 1) // 1 block per SM
void tiny_n_gemm_kernel(OutT* __restrict__ out, const bf16_t* __restrict__ x, const bf16_t* __restrict__ w) {
constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize;
static_assert(N % N_SPLIT == 0, "N must be divisible by n_split");
static_assert((K * sizeof(bf16_t)) % kBytes == 0, "K must be divisible by kBytes");
static_assert(kBytes % device::kMaxVecBytes == 0);
static constexpr uint32_t kNumBlocks = N / N_SPLIT;
static constexpr uint32_t kBlockSize = (K * sizeof(bf16_t)) / kBytes;
// always use the largest possible vector
static constexpr uint32_t kVecSize = device::kMaxVecBytes / sizeof(bf16_t);
// may reduce block size for less thread usage
static constexpr uint32_t kUnroll = kBytes / device::kMaxVecBytes;
static_assert(kBlockSize % device::kWarpThreads == 0, "block size must be divisible by warp size");
static_assert(kBlockSize <= kMaxBlockThreads, "block size exceeds the maximum block size");
};
template <uint32_t N_, uint32_t K_, uint32_t N_SPLIT_, uint32_t kBytes_>
struct GEMMTraitK {
static constexpr uint32_t N = N_;
static constexpr uint32_t K = K_;
static constexpr uint32_t N_SPLIT = N_SPLIT_; // per block n size
static constexpr uint32_t kBytes = kBytes_;
static_assert(N % N_SPLIT == 0, "N must be divisible by n_split");
static_assert((K * sizeof(bf16_t)) % kBytes == 0, "K must be divisible by kBytes");
static_assert(device::kMaxVecBytes % kBytes == 0);
static constexpr uint32_t kNumBlocks = N / N_SPLIT;
static constexpr uint32_t kNumKLanes = (K * sizeof(bf16_t)) / kBytes;
static constexpr uint32_t kVecSize = kBytes / sizeof(bf16_t);
static constexpr uint32_t kBlockSize = N_SPLIT * K / kVecSize;
static_assert(device::kWarpThreads % kNumKLanes == 0, "K reduction must fit in a warp");
// A partial warp would leave reduce_sum's kFullMask naming absent lanes.
static_assert(kBlockSize % device::kWarpThreads == 0, "block size must be divisible by warp size");
static_assert(kBlockSize <= kMaxBlockThreads, "block size exceeds the maximum block size");
};
#define TINY_GEMM_KERNEL __global__ __launch_bounds__(Trait::kBlockSize, 1) // grid: 1 block per SM
struct TinyGEMMParams {
void* __restrict__ out;
const bf16_t* __restrict__ x;
const bf16_t* __restrict__ w;
int64_t stride_x;
};
template <std::size_t N>
SGL_DEVICE void dot_product(device::AlignedVector<bf16x2_t, N> a, device::AlignedVector<bf16x2_t, N> b, float& acc) {
using namespace device;
#pragma unroll
for (uint32_t i = 0; i < N; ++i) {
#if SGL_ARCH_BLACKWELL_OR_GREATER
acc = device::math::fma_f32_bf16(a[i].x, b[i].x, acc);
acc = device::math::fma_f32_bf16(a[i].y, b[i].y, acc);
#else
const auto [a0, a1] = cast<fp32x2_t>(a[i]);
const auto [b0, b1] = cast<fp32x2_t>(b[i]);
acc += a0 * b0;
acc += a1 * b1;
#endif
}
}
template <typename Trait, uint32_t M, typename Out, bool kUsePDL>
TINY_GEMM_KERNEL void tiny_n_gemm_kernel(const TinyGEMMParams params) {
using namespace device;
constexpr uint32_t N = Trait::N;
constexpr uint32_t K = Trait::K;
constexpr uint32_t N_SPLIT = Trait::N_SPLIT;
constexpr uint32_t kVecSize = Trait::kVecSize;
constexpr uint32_t kUnroll = Trait::kUnroll;
constexpr uint32_t kBlockSize = Trait::kBlockSize;
constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
static_assert(M * N_SPLIT <= kBlockSize, "output tile must fit one thread each for the final reduce");
using vec_t = AlignedVector<bf16_t, kTinyNGemmVecSize>;
using vec_t = AlignedVector<bf16x2_t, kVecSize / 2>;
const uint32_t bx = blockIdx.x;
const uint32_t tx = threadIdx.x;
const bf16_t* w_tile = w + bx * (N_SPLIT * K);
const bf16_t* w_tile = params.w + bx * (N_SPLIT * K);
// Weight prefetch: address is input-independent, load before the PDL wait.
vec_t wv[N_SPLIT];
// prefetch weight before PDL
vec_t wv[N_SPLIT][kUnroll];
#pragma unroll
for (uint32_t n = 0; n < N_SPLIT; ++n) {
wv[n].load(w_tile + n * K, tx);
#pragma unroll
for (uint32_t u = 0; u < kUnroll; ++u) {
wv[n][u].load(w_tile + n * K, tx + u * kBlockSize);
}
}
PDLWaitPrimary<kUsePDL>();
vec_t xv[M];
vec_t xv[M][kUnroll];
#pragma unroll
for (uint32_t m = 0; m < M; ++m) {
xv[m].load(x + m * K, tx);
#pragma unroll
for (uint32_t u = 0; u < kUnroll; ++u) {
xv[m][u].load(params.x + m * params.stride_x, tx + u * kBlockSize);
}
}
__shared__ float s_acc[kNumWarps][M * N_SPLIT];
@@ -54,26 +128,18 @@ __global__ __launch_bounds__(K / kTinyNGemmVecSize, 1) // 1 block per SM
#pragma unroll
for (uint32_t n = 0; n < N_SPLIT; ++n) {
float acc = 0.0f;
#if SGL_ARCH_BLACKWELL_OR_GREATER
#pragma unroll
for (uint32_t i = 0; i < kTinyNGemmVecSize; ++i) {
acc = device::math::fma_f32_bf16(xv[m][i], wv[n][i], acc);
for (uint32_t u = 0; u < kUnroll; ++u) {
dot_product(xv[m][u], wv[n][u], acc);
}
#else
for (uint32_t i = 0; i < kTinyNGemmVecSize / 2; ++i) {
const auto [x0, x1] = cast<fp32x2_t>(bf16x2_t{xv[m][2 * i], xv[m][2 * i + 1]});
const auto [w0, w1] = cast<fp32x2_t>(bf16x2_t{wv[n][2 * i], wv[n][2 * i + 1]});
acc = fmaf(x0, w0, acc);
acc = fmaf(x1, w1, acc);
}
#endif
// NOTE: broadcast write (all lanes hold the reduced value), safe here.
s_acc[warp_id][m * N_SPLIT + n] = warp::reduce_sum(acc);
}
}
PDLTriggerSecondary<kUsePDL>();
__syncthreads();
static_assert(M * N_SPLIT <= kBlockSize);
if (tx < M * N_SPLIT) {
float acc[kNumWarps];
#pragma unroll
@@ -86,38 +152,25 @@ __global__ __launch_bounds__(K / kTinyNGemmVecSize, 1) // 1 block per SM
}
const uint32_t m = tx / N_SPLIT;
const uint32_t n = tx % N_SPLIT;
out[m * N + bx * N_SPLIT + n] = cast<OutT>(acc[0]);
static_cast<Out*>(params.out)[m * N + bx * N_SPLIT + n] = cast<Out>(acc[0]);
}
}
SGL_DEVICE void cp_async_cg_16(void* smem_dst, const void* gmem_src, int32_t vec_offset) {
const uint32_t offset = static_cast<uint32_t>(vec_offset * 16);
#if defined(USE_ROCM)
*reinterpret_cast<uint4*>(static_cast<char*>(smem_dst) + offset) =
*reinterpret_cast<const uint4*>(static_cast<const char*>(gmem_src) + offset);
#else
const uint32_t smem_addr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst)) + offset;
const uint64_t gmem_addr = static_cast<uint64_t>(__cvta_generic_to_global(gmem_src)) + offset;
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" : : "r"(smem_addr), "l"(gmem_addr) : "memory");
#endif
}
template <typename Trait, uint32_t M, typename Out, bool kUsePDL>
TINY_GEMM_KERNEL void tiny_k_gemm_kernel(const TinyGEMMParams params) {
using namespace device;
constexpr uint32_t N = Trait::N;
constexpr uint32_t K = Trait::K;
constexpr uint32_t N_SPLIT = Trait::N_SPLIT;
constexpr uint32_t kVecSize = Trait::kVecSize;
constexpr uint32_t kNumKLanes = Trait::kNumKLanes;
using vec_t = AlignedVector<bf16x2_t, kVecSize / 2>;
constexpr uint32_t kTinyKGemmVecSize = 16 / sizeof(bf16_t); // NOTE: no need to be large
template <uint32_t M, uint32_t N, uint32_t K, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
__global__ __launch_bounds__(N_SPLIT* K / kTinyKGemmVecSize, 1) // control the block size
void tiny_k_gemm_kernel(
OutT* __restrict__ out, const bf16_t* __restrict__ x, const bf16_t* __restrict__ w, const int64_t dx) {
using vec_t = AlignedVector<bf16_t, kTinyKGemmVecSize>;
constexpr uint32_t kNumKLanes = K / kTinyKGemmVecSize;
static_assert(std::has_single_bit(kNumKLanes), "K / vec_size must be a power of 2");
static_assert(kNumKLanes <= kWarpThreads, "require in-warp reduction");
static_assert((N_SPLIT * K / kTinyKGemmVecSize) % kWarpThreads == 0);
const uint32_t bx = blockIdx.x;
const uint32_t tx = threadIdx.x;
const uint32_t n_idx = bx * N_SPLIT + tx / kNumKLanes;
const uint32_t work_id = tx % kNumKLanes;
const bf16_t* w_tile = w + n_idx * K;
const bf16_t* w_tile = params.w + n_idx * K;
// Weight prefetch: address is input-independent, load before the PDL wait.
vec_t wv;
@@ -127,38 +180,29 @@ __global__ __launch_bounds__(N_SPLIT* K / kTinyKGemmVecSize, 1) // control the
vec_t xv[M];
#pragma unroll
for (uint32_t m = 0; m < M; ++m) {
xv[m].load(x + m * dx, work_id);
xv[m].load(params.x + m * params.stride_x, work_id);
}
#pragma unroll
for (uint32_t m = 0; m < M; ++m) {
float acc = 0.0f;
#pragma unroll
for (uint32_t i = 0; i < kTinyKGemmVecSize; ++i) {
acc = device::math::fma_f32_bf16(xv[m][i], wv[i], acc);
}
dot_product(xv[m], wv, acc);
// Broadcast store: every lane of the group holds the reduced sum.
out[m * N + n_idx] = cast<OutT>(warp::reduce_sum<kNumKLanes>(acc));
const auto sum = warp::reduce_sum<kNumKLanes>(acc);
static_cast<Out*>(params.out)[m * N + n_idx] = cast<Out>(sum);
}
PDLTriggerSecondary<kUsePDL>();
}
template <uint32_t N, uint32_t K, uint32_t kMaxM, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
struct TinyNGemmKernel {
static constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize;
static constexpr uint32_t kNumBlocks = N / N_SPLIT;
static_assert(K % kTinyNGemmVecSize == 0, "K must be divisible by the vector width");
static_assert(kBlockSize % kWarpThreads == 0, "K / vec_size must be a multiple of the warp size");
static_assert(kBlockSize <= 1024, "K / vec_size exceeds the maximum block size");
static_assert(N % N_SPLIT == 0, "N must be divisible by split_n");
static_assert(kMaxM * N_SPLIT <= kBlockSize, "max_m * split_n must fit one thread each for the final reduce");
using KernelFn = void (*)(OutT*, const bf16_t*, const bf16_t*);
using Trait = GEMMTraitN<N, K, N_SPLIT, 32>;
using KernelFn = void (*)(TinyGEMMParams);
template <std::size_t... I>
static constexpr auto make_table(std::index_sequence<I...>) {
return std::array<KernelFn, kMaxM + 1>{nullptr, tiny_n_gemm_kernel<I + 1, N, K, N_SPLIT, OutT, kUsePDL>...};
return std::array<KernelFn, kMaxM + 1>{nullptr, tiny_n_gemm_kernel<Trait, I + 1, OutT, kUsePDL>...};
}
static constexpr auto kTable = make_table(std::make_index_sequence<kMaxM>{});
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView w, const tvm::ffi::TensorView out) {
@@ -167,36 +211,37 @@ struct TinyNGemmKernel {
auto M = SymbolicSize{"num_tokens"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({M, K}).with_dtype<bf16_t>().with_device(device).verify(x);
TensorMatcher({M, K}).with_strides({-1, 1}).with_dtype<bf16_t>().with_device(device).verify(x);
TensorMatcher({N, K}).with_dtype<bf16_t>().with_device(device).verify(w);
TensorMatcher({M, N}).with_dtype<OutT>().with_device(device).verify(out);
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM);
LaunchKernel(kNumBlocks, kBlockSize, device.unwrap())
.enable_pdl(kUsePDL)(
kTable[num_tokens],
static_cast<OutT*>(out.data_ptr()),
static_cast<const bf16_t*>(x.data_ptr()),
static_cast<const bf16_t*>(w.data_ptr()));
if (num_tokens == 0) return;
CHECK_HOST(num_tokens >= 1 && num_tokens <= kMaxM);
// x may be a row-sliced view of a wider buffer, but the rows are loaded as
// whole vectors, so every row start must stay vector-aligned.
CHECK_HOST(x.stride(0) % Trait::kVecSize == 0)
<< "x rows must stay aligned to the vector width, got stride " << x.stride(0);
const auto params = TinyGEMMParams{
.out = out.data_ptr(),
.x = static_cast<const bf16_t*>(x.data_ptr()),
.w = static_cast<const bf16_t*>(w.data_ptr()),
.stride_x = static_cast<int64_t>(x.stride(0)),
};
LaunchKernel(Trait::kNumBlocks, Trait::kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kTable[num_tokens], params);
}
};
template <uint32_t N, uint32_t K, uint32_t kMaxM, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
struct TinyKGemmKernel {
static constexpr uint32_t kNumKLanes = K / kTinyKGemmVecSize;
static constexpr uint32_t kBlockSize = N_SPLIT * kNumKLanes;
static constexpr uint32_t kNumBlocks = N / N_SPLIT;
static_assert(K % kTinyKGemmVecSize == 0, "K must be divisible by the vector width");
static_assert(N % N_SPLIT == 0, "N must be divisible by split_n");
static_assert(kBlockSize % kWarpThreads == 0, "split_n * K-lanes must fill whole warps");
static_assert(kBlockSize <= 1024, "split_n * K-lanes exceeds the maximum block size");
using KernelFn = void (*)(OutT*, const bf16_t*, const bf16_t*, int64_t);
using Trait = GEMMTraitK<N, K, N_SPLIT, 16>;
using KernelFn = void (*)(TinyGEMMParams);
template <std::size_t... I>
static constexpr auto make_table(std::index_sequence<I...>) {
return std::array<KernelFn, kMaxM + 1>{nullptr, tiny_k_gemm_kernel<I + 1, N, K, N_SPLIT, OutT, kUsePDL>...};
return std::array<KernelFn, kMaxM + 1>{nullptr, tiny_k_gemm_kernel<Trait, I + 1, OutT, kUsePDL>...};
}
static constexpr auto kTable = make_table(std::make_index_sequence<kMaxM>{});
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView w, const tvm::ffi::TensorView out) {
@@ -205,24 +250,24 @@ struct TinyKGemmKernel {
auto M = SymbolicSize{"num_tokens"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
// x may be a row-sliced view of a wider fused buffer: allow stride != K.
TensorMatcher({M, K}).with_dtype<bf16_t>().with_strides({-1, 1}).with_device(device).verify(x);
TensorMatcher({M, K}).with_strides({-1, 1}).with_dtype<bf16_t>().with_device(device).verify(x);
TensorMatcher({N, K}).with_dtype<bf16_t>().with_device(device).verify(w);
TensorMatcher({M, N}).with_dtype<OutT>().with_device(device).verify(out);
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
const auto x_stride = static_cast<int64_t>(x.stride(0));
RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM);
RuntimeCheck(
x_stride * sizeof(bf16_t) % (kTinyKGemmVecSize * sizeof(bf16_t)) == 0,
"x rows must stay aligned to the vector width, got stride ",
x_stride);
LaunchKernel(kNumBlocks, kBlockSize, device.unwrap())
.enable_pdl(kUsePDL)(
kTable[num_tokens],
static_cast<OutT*>(out.data_ptr()),
static_cast<const bf16_t*>(x.data_ptr()),
static_cast<const bf16_t*>(w.data_ptr()),
x_stride);
if (num_tokens == 0) return;
CHECK_HOST(num_tokens >= 1 && num_tokens <= kMaxM);
// x may be a row-sliced view of a wider buffer, but the rows are loaded as
// whole vectors, so every row start must stay vector-aligned.
CHECK_HOST(x.stride(0) % Trait::kVecSize == 0)
<< "x rows must stay aligned to the vector width, got stride " << x.stride(0);
const auto params = TinyGEMMParams{
.out = out.data_ptr(),
.x = static_cast<const bf16_t*>(x.data_ptr()),
.w = static_cast<const bf16_t*>(w.data_ptr()),
.stride_x = static_cast<int64_t>(x.stride(0)),
};
LaunchKernel(Trait::kNumBlocks, Trait::kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kTable[num_tokens], params);
}
};
python/sglang/kernels/ops/gemm/__init__.py
+16 -14
View File
@@ -201,15 +201,15 @@ register_kernel(
)
register_kernel(
KernelSpec(
op="gemm.dsv3_router_gemm",
op="gemm.tiny_gemm",
backend=KernelBackend.JIT,
target="sglang.kernels.ops.gemm.dsv3_router_gemm:dsv3_router_gemm",
target="sglang.kernels.ops.gemm.tiny_gemm:tiny_gemm_bf16",
capabilities=_CUDA,
format_signature=FormatSignature(
supported_dtypes=("bfloat16",),
description="DeepSeek-V3 router GEMM; num_tokens in [1, 16]",
description="skinny [m, k] @ [n, k].T with m in [1, max_m]; serves MoE routers and gate projections",
),
description="DeepSeek-V3 router GEMM (sglang.kernels.jit, JIT-only).",
description="Tiny bf16 GEMM (sglang.kernels.jit, JIT-only).",
)
)
@@ -249,17 +249,19 @@ def dsv3_fused_a_gemm(
return get_kernel("gemm.dsv3_fused_a_gemm", KernelBackend.AOT)(mat_a, mat_b, output)
def dsv3_router_gemm(
hidden_states: torch.Tensor,
router_weights: torch.Tensor,
# NOTE: named after the kernel entry, not the op -- a bare `tiny_gemm` here would
# be shadowed by the submodule of the same name as soon as it is imported.
def tiny_gemm_bf16(
x: torch.Tensor,
w: torch.Tensor,
out: Optional[torch.Tensor] = None,
*,
out_dtype: Optional[torch.dtype] = None,
output: Optional[torch.Tensor] = None,
max_m: int = 16,
) -> torch.Tensor:
"""DeepSeek-V3 router GEMM (JIT-backed). ``out_dtype`` defaults to bfloat16."""
impl = get_kernel("gemm.dsv3_router_gemm", KernelBackend.JIT)
if out_dtype is None:
return impl(hidden_states, router_weights, output=output)
return impl(hidden_states, router_weights, out_dtype, output)
"""Tiny bf16 GEMM ``x[m, k] @ w[n, k].T``, support bf16/fp32 out"""
impl = get_kernel("gemm.tiny_gemm", KernelBackend.JIT)
return impl(x, w, out, out_dtype=out_dtype, max_m=max_m)
__all__ = [
@@ -267,7 +269,7 @@ __all__ = [
"fp8_scaled_mm",
"bmm_fp8",
"dsv3_fused_a_gemm",
"dsv3_router_gemm",
"tiny_gemm_bf16",
]
python/sglang/kernels/ops/gemm/dsv3_router_gemm.py
-92
View File
@@ -1,92 +0,0 @@
"""
JIT kernel for DeepSeek V3 router GEMM.
Runtime-compiled CUDA C++ kernel 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.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernels.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 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 = 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/kernels/ops/gemm/tiny_gemm.py
+123 -82
View File
@@ -1,3 +1,11 @@
"""Tiny bf16 GEMM for skinny ``x[m, k] @ w[n, k].T`` with a handful of rows.
Two kernels behind one entry point, picked by whichever of N / K is the tiny
dimension: the N variant spreads K across a block and walks N, the K variant
reduces K inside a warp and spreads N across the grid. See
:func:`tiny_gemm_bf16`.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
@@ -10,135 +18,168 @@ from sglang.kernels.jit.utils import (
load_jit,
make_cpp_args,
)
from sglang.kernels.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
_MAX_M_DEFAULT: int = 16
# NOTE: CUDA constraints
_WARP_THREADS: int = 32
_MAX_BLOCK_THREADS: int = 1024
# Mirrors GEMMTraitN in tiny_gemm.cuh: one thread owns kBytes = 32 bytes of K,
# so the block is K / 16 threads on every arch (k_unroll absorbs the vector-width
# difference between Hopper and Blackwell).
_N_GEMM_ELEMS_PER_THREAD: int = 16
# Mirrors GEMMTraitK: the K variant always uses 16-byte vectors.
_K_GEMM_ELEMS_PER_THREAD: int = 8
def _prefer_k_variant(n: int, k: int) -> bool:
"""Find the **tiny** dimension"""
return k <= n
@cache_once
def _jit_tiny_gemm_module(
n: int, k: int, max_m: int, split_n: int, out_dtype: torch.dtype
n: int, k: int, max_m: int, n_split: int, out_dtype: torch.dtype
) -> Module:
args = make_cpp_args(n, k, max_m, split_n, out_dtype, is_arch_support_pdl())
use_k_variant = _prefer_k_variant(n, k)
args = make_cpp_args(n, k, max_m, n_split, out_dtype, is_arch_support_pdl())
name = "tiny_k_gemm" if use_k_variant else "tiny_n_gemm"
kernel = "TinyKGemmKernel" if use_k_variant else "TinyNGemmKernel"
return load_jit(
"tiny_gemm",
name,
*args,
cuda_files=["gemm/tiny_gemm.cuh"],
cuda_wrappers=[("run", f"TinyNGemmKernel<{args}>::run")],
cuda_wrappers=[("run", f"{kernel}<{args}>::run")],
extra_cuda_cflags=["-O3"],
)
@cache_once
def _jit_tiny_k_gemm_module(
n: int, k: int, max_m: int, n_unroll: int, out_dtype: torch.dtype
) -> Module:
args = make_cpp_args(n, k, max_m, n_unroll, out_dtype, is_arch_support_pdl())
return load_jit(
"tiny_k_gemm",
*args,
cuda_files=["gemm/tiny_gemm.cuh"],
cuda_wrappers=[("run", f"TinyKGemmKernel<{args}>::run")],
extra_cuda_cflags=["-O3"],
def _get_num_sm() -> int:
device = torch.cuda.current_device()
return torch.cuda.get_device_properties(device).multi_processor_count
def _n_gemm_block_threads(k: int) -> int:
return k // _N_GEMM_ELEMS_PER_THREAD
def _k_gemm_reduction_lanes(k: int) -> int:
return k // _K_GEMM_ELEMS_PER_THREAD
def _supports_n_variant(k: int, max_m: int) -> bool:
block = _n_gemm_block_threads(k)
return (
k % _N_GEMM_ELEMS_PER_THREAD == 0
and block % _WARP_THREADS == 0
and block <= _MAX_BLOCK_THREADS
and block >= max_m
)
def _vec_elems() -> int:
"""bf16 elements per vectorized load; mirrors kMaxVecBytes in utils.cuh."""
from sglang.kernels.jit.utils import get_jit_cuda_arch
cuda = tuple(int(v) for v in (torch.version.cuda or "0.0").split(".")[:2])
return 16 if get_jit_cuda_arch().major >= 10 and cuda >= (12, 9) else 8
def _supports_k_variant(n: int, k: int) -> bool:
lanes = _k_gemm_reduction_lanes(k)
return (
k % _K_GEMM_ELEMS_PER_THREAD == 0
and k > _K_GEMM_ELEMS_PER_THREAD
and _WARP_THREADS % lanes == 0
and n % (_WARP_THREADS // lanes) == 0
)
def _default_split_n(n: int, k: int, max_m: int, device: torch.device) -> int:
"""Smallest divisor of n whose n / split_n blocks fit in one wave, subject
to the max_m * split_n <= K / vec_elems block-size constraint; falls back
to the largest split_n satisfying the constraint (multi-wave grid)."""
sm_count = torch.cuda.get_device_properties(device).multi_processor_count
split_cap = (k // _vec_elems()) // max_m
divisors = [d for d in range(1, min(n, split_cap) + 1) if n % d == 0]
if not divisors:
@cache_once
def _default_n_gemm_n_split(n: int, k: int, max_m: int) -> int:
split_cap = _n_gemm_block_threads(k) // max_m
candidates = [d for d in range(1, min(n, split_cap) + 1) if n % d == 0]
if not candidates:
raise RuntimeError(
f"tiny_gemm: no valid split_n for N={n}, K={k}, max_m={max_m};"
" lower max_m"
f"tiny_gemm: no valid n_split for N={n}, K={k}, max_m={max_m};"
" try lower max_m or align the N dimension"
)
for split in divisors:
# try to fit into 1 wave
sm_count = _get_num_sm()
for split in candidates:
if n // split <= sm_count:
return split
return divisors[-1]
return candidates[-1]
def tiny_n_gemm_bf16(
x: torch.Tensor,
w: torch.Tensor,
out: Optional[torch.Tensor] = None,
*,
out_dtype: Optional[torch.dtype] = None,
split_n: Optional[int] = None,
max_m: int = _MAX_M_DEFAULT,
) -> torch.Tensor:
n = w.shape[0]
k = x.shape[1]
if out is None:
out_dtype = out_dtype or torch.bfloat16
out = torch.empty((x.shape[0], n), dtype=out_dtype, device=x.device)
else:
assert out_dtype is None or out_dtype == out.dtype
if split_n is None:
split_n = _default_split_n(n, k, max_m, x.device)
module = _jit_tiny_gemm_module(n, k, max_m, split_n, out.dtype)
module.run(x, w, out)
return out
def _default_k_split_n(n: int, k: int) -> int:
"""Smallest divisor of n whose n / split_n blocks fit one wave, with
split_n * K-lanes whole-warp aligned and within the block-size limit."""
lanes = k // 8 # fixed 16-byte vectors in the K variant
@cache_once
def _default_k_gemm_n_split(n: int, k: int) -> int:
lanes = _k_gemm_reduction_lanes(k)
candidates = [
d
for d in range(1, n + 1)
if n % d == 0 and d * lanes % 32 == 0 and d * lanes <= 1024
for d in range(1, min(n, _MAX_BLOCK_THREADS // lanes) + 1)
if n % d == 0
and d * lanes % _WARP_THREADS == 0
and d * lanes <= _MAX_BLOCK_THREADS
]
# try to fit into 1 wave
if not candidates:
raise RuntimeError(f"tiny_k_gemm: no valid split_n for N={n}, K={k}")
sm_count = torch.cuda.get_device_properties(0).multi_processor_count
raise RuntimeError(f"tiny_gemm: no valid n_split for N={n}, K={k}")
sm_count = _get_num_sm()
for d in candidates:
if n // d <= sm_count:
return d
return candidates[-1]
def tiny_k_gemm_bf16(
@register_custom_op(op_name="tiny_gemm_bf16", mutates_args=["out"])
def _tiny_gemm_custom_op(
x: torch.Tensor,
w: torch.Tensor,
out: torch.Tensor,
max_m: int,
n_split: int,
) -> None:
n, k = w.shape
module = _jit_tiny_gemm_module(n, k, max_m, n_split, out.dtype)
module.run(x, w, out)
@cache_once
def can_use_tiny_gemm(n: int, k: int, max_m: int = 16) -> bool:
"""Whether :func:`tiny_gemm_bf16` can serve ``[m, k] @ [n, k].T`` for
``m <= max_m``. Callers fall back to a general GEMM when this is False."""
if _prefer_k_variant(n, k):
return _supports_k_variant(n, k)
else:
return _supports_n_variant(k, max_m)
@debug_kernel_api
def tiny_gemm_bf16(
x: torch.Tensor,
w: torch.Tensor,
out: Optional[torch.Tensor] = None,
*,
out_dtype: Optional[torch.dtype] = None,
split_n: Optional[int] = None,
max_m: int = _MAX_M_DEFAULT,
n_split: Optional[int] = None,
max_m: int = 16,
) -> torch.Tensor:
"""Small-K / large-N variant: K / 8 lanes of one warp reduce the K
dimension for one output column; each block covers split_n columns and the
exact N / split_n grid fills the SMs (no tail). Requires K / 8 to be a
power of 2 and <= 32 (e.g. K = 128/256). x may be a row-sliced view as
long as rows stay 16-byte aligned.
"""
Equal to `torch.nn.functional.linear(x, w)`.
Call :func:`can_use_tiny_gemm` first: shapes outside the supported set raise
rather than falling back.
split_n trades block count for block size; the default picks the smallest
divisor of N that fits one wave (12 for [1536, 128] on B200: 128 blocks
of 6 warps)."""
n = w.shape[0]
k = x.shape[1]
:param x: Shape [m, k], must be bf16_t, `m <= max_m`
:param w: Shape [n, k], must be bf16_t
"""
n, k = w.shape
if out is None:
out_dtype = out_dtype or torch.bfloat16
out_dtype = torch.bfloat16 if out_dtype is None else out_dtype
out = torch.empty((x.shape[0], n), dtype=out_dtype, device=x.device)
else:
assert out_dtype is None or out_dtype == out.dtype
if split_n is None:
split_n = _default_k_split_n(n, k)
module = _jit_tiny_k_gemm_module(n, k, max_m, split_n, out.dtype)
module.run(x, w, out)
if n_split is None:
n_split = (
_default_k_gemm_n_split(n, k)
if _prefer_k_variant(n, k)
else _default_n_gemm_n_split(n, k, max_m)
)
_tiny_gemm_custom_op(x, w, out, max_m, n_split)
return out
python/sglang/kernels/ops/kimi_k3/__init__.py
+7 -11
View File
@@ -9,11 +9,11 @@ if TYPE_CHECKING:
_is_npu = is_npu()
_K3_N_GEMM_DISPATCH_MAP = {
# (n, k) -> the largest num_tokens where the tiny GEMM still beats cuBLAS.
# Doubles as the compile-time max_m: one kernel is built per m up to it.
_K3_TINY_GEMM_MAX_TOKENS = {
(144, 7168): 16,
(896, 7168): 8,
}
_K3_K_GEMM_DISPATCH_MAP = {
(1536, 128): 12,
}
@@ -65,17 +65,13 @@ def kimi_k3_tiny_gemm(
) -> torch.Tensor:
import torch
from ..gemm.tiny_gemm import tiny_k_gemm_bf16, tiny_n_gemm_bf16
from ..gemm.tiny_gemm import tiny_gemm_bf16
m, k = x.shape
n, _ = w.shape
if not _is_npu:
if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)):
if 0 < m <= max_num_tokens:
return tiny_n_gemm_bf16(x, w)
if max_num_tokens := _K3_K_GEMM_DISPATCH_MAP.get((n, k)):
if 0 < m <= max_num_tokens:
return tiny_k_gemm_bf16(x, w)
max_num_tokens = _K3_TINY_GEMM_MAX_TOKENS.get((n, k))
if not _is_npu and max_num_tokens is not None and 0 < m <= max_num_tokens:
return tiny_gemm_bf16(x, w, max_m=max_num_tokens)
return torch.nn.functional.linear(x, w)
python/sglang/srt/models/deepseek_common/utils.py
+17
View File
@@ -180,3 +180,20 @@ def _get_llama_4_scaling(
1 + torch.floor(positions / original_max_position_embeddings)
)
return scaling[..., None, None]
def tiny_router_gemm_max_tokens(
*, num_experts: int, hidden_size: int, weight_dtype: torch.dtype
) -> int:
"""Rows up to which the tiny GEMM beats cuBLAS for the router, -1 when the
shape or the device rules it out. Doubles as the kernel's compile-time
max_m, so keep it as tight as the measurements allow.
"""
if not _is_cuda or _device_sm < 90 or weight_dtype != torch.bfloat16:
return -1
from sglang.kernels.ops.gemm.tiny_gemm import can_use_tiny_gemm
if not can_use_tiny_gemm(num_experts, hidden_size, max_m=16):
return -1
return 16
python/sglang/srt/models/deepseek_v2.py
+13 -13
View File
@@ -179,7 +179,6 @@ from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
DeepseekV2WeightLoaderMixin,
)
from sglang.srt.models.deepseek_common.utils import (
_device_sm,
_get_llama_4_scaling,
_is_block_scale_fp8,
_is_cpu,
@@ -195,6 +194,7 @@ from sglang.srt.models.deepseek_common.utils import (
_use_aiter_gfx95,
is_wint4afp8_or_wint4a16_config,
quant_blocks_shared_experts_fusion,
tiny_router_gemm_max_tokens,
)
from sglang.srt.runtime_context import (
attention_backends,
@@ -229,9 +229,7 @@ if _use_aiter:
pass
if _is_cuda:
from sglang.kernels.ops.gemm.dsv3_router_gemm import (
dsv3_router_gemm as dsv3_router_gemm,
)
from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16
elif _is_npu:
from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
forward_dsa_core_npu,
@@ -506,6 +504,11 @@ class MoEGate(nn.Module):
self.use_dsa = is_deepseek_dsa(config)
self.dsa_enable_prefill_cp = dsa_enable_prefill_cp
self.mla_enable_prefill_cp = mla_enable_prefill_cp
self.tiny_router_gemm_max_tokens = tiny_router_gemm_max_tokens(
num_experts=config.n_routed_experts,
hidden_size=config.hidden_size,
weight_dtype=self.weight.dtype,
)
def forward(
self,
@@ -538,15 +541,12 @@ class MoEGate(nn.Module):
return linear_bf16_fp32(hidden_states, self.weight)
return F.linear(hidden_states, self.weight, None)
else:
if (
_is_cuda
and hidden_states.shape[0] <= 16
and hidden_states.shape[1] % 1024 == 0
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
and _device_sm >= 90
):
logits = dsv3_router_gemm(
hidden_states, self.weight, out_dtype=torch.float32
if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
logits = tiny_gemm_bf16(
hidden_states,
self.weight,
out_dtype=torch.float32,
max_m=self.tiny_router_gemm_max_tokens,
)
elif _use_aiter:
python/sglang/srt/models/dots3_common/modeling.py
+12 -12
View File
@@ -112,6 +112,7 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
_load_fused_indexer_wk,
)
from sglang.srt.models.deepseek_common.utils import tiny_router_gemm_max_tokens
from sglang.srt.models.dots3_common.fp8 import per_token_group_quant_einsum_fp8
from sglang.srt.runtime_context import (
get_device,
@@ -126,7 +127,6 @@ from sglang.srt.utils import (
ceil_align,
ceil_div,
get_bool_env_var,
get_device_sm,
is_cuda,
is_non_idle_and_non_empty,
log_info_on_rank0,
@@ -135,16 +135,15 @@ from sglang.srt.utils import (
_is_cuda = is_cuda()
_is_fp8_fnuz = is_fp8_fnuz()
_device_sm = get_device_sm()
# Import-time CUDA kernels would block processor imports on CPU CI.
if _is_cuda:
from sgl_kernel import merge_state_v2
from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm
from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16
else:
merge_state_v2 = None
dsv3_router_gemm = None
tiny_gemm_bf16 = None
def _require_cuda() -> None:
@@ -259,17 +258,18 @@ class Dots3MoEGate(nn.Module):
)
else:
self.e_score_correction_bias = None
self.tiny_router_gemm_max_tokens = tiny_router_gemm_max_tokens(
num_experts=config.n_routed_experts,
hidden_size=config.hidden_size,
weight_dtype=self.weight.dtype,
)
def forward(self, hidden_states):
# Use the fused router only for its tuned shapes.
if (
hidden_states.shape[0] <= 16
and hidden_states.shape[1] == 7168
and self.weight.shape[0] == 256
and _device_sm >= 90
):
# router gemm output float32
logits = dsv3_router_gemm(hidden_states, self.weight)
if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
logits = tiny_gemm_bf16(
hidden_states, self.weight, max_m=self.tiny_router_gemm_max_tokens
)
else:
logits = F.linear(hidden_states, self.weight, None)
python/sglang/srt/models/kimi_k3.py
+1 -1
View File
@@ -1057,7 +1057,7 @@ class KimiK3MoE(nn.Module):
topk_output, routed_input = routed_input
else:
# MoEGate produces fp32 router logits on CUDA (via linear_bf16_fp32
# or dsv3_router_gemm); non-CUDA falls back to F.linear (bf16). The
# or tiny_gemm_bf16); non-CUDA falls back to F.linear (bf16). The
# fp32 logits reach the radix router from moe_fused_gate.
router_logits = self.gate(hidden_states)
topk_output = self.topk(hidden_states, router_logits)