diff --git a/python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh deleted file mode 100644 index a7b1ca3d3..000000000 --- a/python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh +++ /dev/null @@ -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 -#include - -#include -#include -#include -#include -#include - -#include - -#include - -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(); - - int k_base = tid * VPT; -#pragma unroll - for (int ki = 0; ki < kIters; ++ki, k_base += kElemsPerKIter) { - AlignedVector b_vec; - b_vec.load(b_col + k_base); -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; ++m_idx) { - AlignedVector 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(a_vec[k]) * cast(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(final_sum); - } - - PDLTriggerSecondary(); -} - -template -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; - 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 -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(output, mat_a, mat_b, device); - } else { - RouterGemmDispatcher::run( - num_tokens, output, mat_a, mat_b, device); - } - } -}; - -// Base case: kBegin == kEnd -template -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 == kEnd) { - invokeRouterGemm(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 -struct DSV3RouterGemmKernel { - static_assert( - kNumExperts == kDefaultNumExperts || kNumExperts == kKimiK2NumExperts, - "required num_experts == 256 or num_experts == 384"); - - using OutT = std::conditional_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(); - - TensorMatcher({M, K}).with_dtype().with_device(device).verify(mat_a); - TensorMatcher({N, K}).with_dtype().with_device(device).verify(mat_b); - TensorMatcher({M, N}).with_dtype().with_device(device).verify(output); - - const int num_tokens = static_cast(M.unwrap()); - - RouterGemmDispatcher<1, 16, OutT, kNumExperts, kHiddenDim, kUsePDL>::run( - num_tokens, - static_cast(output.data_ptr()), - static_cast(mat_a.data_ptr()), - static_cast(mat_b.data_ptr()), - device.unwrap()); - } -}; - -} // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh index e019a36a8..1cabaddf3 100644 --- a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh @@ -9,41 +9,115 @@ #include #include -#include +#include #include #include namespace sglang { -using namespace device; +constexpr uint32_t kMaxBlockThreads = 1024; -constexpr uint32_t kTinyNGemmVecSize = kMaxVecBytes / sizeof(bf16_t); +template +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 -__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 +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 +SGL_DEVICE void dot_product(device::AlignedVector a, device::AlignedVector 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(a[i]); + const auto [b0, b1] = cast(b[i]); + acc += a0 * b0; + acc += a1 * b1; +#endif + } +} + +template +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; + using vec_t = AlignedVector; 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(); - 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(bf16x2_t{xv[m][2 * i], xv[m][2 * i + 1]}); - const auto [w0, w1] = cast(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(); __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(acc[0]); + static_cast(params.out)[m * N + bx * N_SPLIT + n] = cast(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(vec_offset * 16); -#if defined(USE_ROCM) - *reinterpret_cast(static_cast(smem_dst) + offset) = - *reinterpret_cast(static_cast(gmem_src) + offset); -#else - const uint32_t smem_addr = static_cast(__cvta_generic_to_shared(smem_dst)) + offset; - const uint64_t gmem_addr = static_cast(__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 +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; -constexpr uint32_t kTinyKGemmVecSize = 16 / sizeof(bf16_t); // NOTE: no need to be large - -template -__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; - 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(warp::reduce_sum(acc)); + const auto sum = warp::reduce_sum(acc); + static_cast(params.out)[m * N + n_idx] = cast(sum); } PDLTriggerSecondary(); } template 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; + using KernelFn = void (*)(TinyGEMMParams); template static constexpr auto make_table(std::index_sequence) { - return std::array{nullptr, tiny_n_gemm_kernel...}; + return std::array{nullptr, tiny_n_gemm_kernel...}; } + static constexpr auto kTable = make_table(std::make_index_sequence{}); 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(); - TensorMatcher({M, K}).with_dtype().with_device(device).verify(x); + TensorMatcher({M, K}).with_strides({-1, 1}).with_dtype().with_device(device).verify(x); TensorMatcher({N, K}).with_dtype().with_device(device).verify(w); TensorMatcher({M, N}).with_dtype().with_device(device).verify(out); const auto num_tokens = static_cast(M.unwrap()); - RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM); - LaunchKernel(kNumBlocks, kBlockSize, device.unwrap()) - .enable_pdl(kUsePDL)( - kTable[num_tokens], - static_cast(out.data_ptr()), - static_cast(x.data_ptr()), - static_cast(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(x.data_ptr()), + .w = static_cast(w.data_ptr()), + .stride_x = static_cast(x.stride(0)), + }; + LaunchKernel(Trait::kNumBlocks, Trait::kBlockSize, device.unwrap()) // + .enable_pdl(kUsePDL)(kTable[num_tokens], params); } }; template 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; + using KernelFn = void (*)(TinyGEMMParams); template static constexpr auto make_table(std::index_sequence) { - return std::array{nullptr, tiny_k_gemm_kernel...}; + return std::array{nullptr, tiny_k_gemm_kernel...}; } + static constexpr auto kTable = make_table(std::make_index_sequence{}); 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(); - // x may be a row-sliced view of a wider fused buffer: allow stride != K. - TensorMatcher({M, K}).with_dtype().with_strides({-1, 1}).with_device(device).verify(x); + TensorMatcher({M, K}).with_strides({-1, 1}).with_dtype().with_device(device).verify(x); TensorMatcher({N, K}).with_dtype().with_device(device).verify(w); TensorMatcher({M, N}).with_dtype().with_device(device).verify(out); const auto num_tokens = static_cast(M.unwrap()); - const auto x_stride = static_cast(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(out.data_ptr()), - static_cast(x.data_ptr()), - static_cast(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(x.data_ptr()), + .w = static_cast(w.data_ptr()), + .stride_x = static_cast(x.stride(0)), + }; + LaunchKernel(Trait::kNumBlocks, Trait::kBlockSize, device.unwrap()) // + .enable_pdl(kUsePDL)(kTable[num_tokens], params); } }; diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index f588f6956..33623a5d5 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -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", ] diff --git a/python/sglang/kernels/ops/gemm/dsv3_router_gemm.py b/python/sglang/kernels/ops/gemm/dsv3_router_gemm.py deleted file mode 100644 index 6a49bbdbb..000000000 --- a/python/sglang/kernels/ops/gemm/dsv3_router_gemm.py +++ /dev/null @@ -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 diff --git a/python/sglang/kernels/ops/gemm/tiny_gemm.py b/python/sglang/kernels/ops/gemm/tiny_gemm.py index 379ff3c5f..1b429732f 100644 --- a/python/sglang/kernels/ops/gemm/tiny_gemm.py +++ b/python/sglang/kernels/ops/gemm/tiny_gemm.py @@ -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 diff --git a/python/sglang/kernels/ops/kimi_k3/__init__.py b/python/sglang/kernels/ops/kimi_k3/__init__.py index c95fe21de..f6aba29b2 100644 --- a/python/sglang/kernels/ops/kimi_k3/__init__.py +++ b/python/sglang/kernels/ops/kimi_k3/__init__.py @@ -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) diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index b4079630d..45064be29 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index f3ed43efe..65c880d48 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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: diff --git a/python/sglang/srt/models/dots3_common/modeling.py b/python/sglang/srt/models/dots3_common/modeling.py index 9a64c59a6..360bd114f 100644 --- a/python/sglang/srt/models/dots3_common/modeling.py +++ b/python/sglang/srt/models/dots3_common/modeling.py @@ -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) diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index afe3d9c30..f72789117 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -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) diff --git a/test/registered/kernels/benchmark/gemm/bench_dsv3_router_gemm.py b/test/registered/kernels/benchmark/gemm/bench_dsv3_router_gemm.py deleted file mode 100644 index 1a4442bab..000000000 --- a/test/registered/kernels/benchmark/gemm/bench_dsv3_router_gemm.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Benchmark for DeepSeek V3 router GEMM (JIT kernel vs torch). - -Run on a Hopper (SM90+) GPU: - python -m sglang.kernels.jit.benchmark.bench_dsv3_router_gemm -""" - -import torch -import torch.nn.functional as F - -from sglang.kernels.jit.benchmark import marker -from sglang.kernels.jit.benchmark.utils import create_random -from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime -from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci - -register_cuda_ci( - est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" -) -register_amd_ci(est_time=5, stage="jit-kernel-benchmark", runner_config="amd") - - -def _torch(mat_a, mat_b, out_dtype): - return F.linear(mat_a, mat_b).to(out_dtype) - - -FN_MAP = { - "jit": dsv3_router_gemm, - "torch": _torch, -} - - -@marker.parametrize("num_experts", [256, 384], [256]) -@marker.parametrize("hidden_dim", [6144, 7168], [7168]) -@marker.parametrize("num_tokens", list(range(1, 17)), [1, 8, 16]) -@marker.parametrize("out_dtype", [torch.bfloat16, torch.float32]) -@marker.benchmark("provider", ["jit", "torch"]) -def benchmark(num_experts, hidden_dim, num_tokens, out_dtype, provider): - mat_a = create_random(num_tokens, hidden_dim) - mat_b = create_random(num_experts, hidden_dim) - return marker.do_bench( - FN_MAP[provider], - input_args=(mat_a, mat_b), - input_kwargs={"out_dtype": out_dtype}, - ) - - -if __name__ == "__main__": - if is_hip_runtime() or get_jit_cuda_arch().major < 9: - print( - "dsv3_router_gemm JIT kernel requires SM90+ (Hopper). Skipping benchmark." - ) - else: - benchmark.run() diff --git a/test/registered/kernels/benchmark/gemm/bench_tiny_gemm.py b/test/registered/kernels/benchmark/gemm/bench_tiny_gemm.py new file mode 100644 index 000000000..9375344a7 --- /dev/null +++ b/test/registered/kernels/benchmark/gemm/bench_tiny_gemm.py @@ -0,0 +1,49 @@ +"""Benchmark for the tiny GEMM (JIT kernel vs torch). + +Run on a Hopper (SM90+) GPU: + python -m sglang.kernels.jit.benchmark.bench_tiny_gemm +""" + +import torch +import torch.nn.functional as F + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random +from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + + +def _torch(x, w, out_dtype): + return F.linear(x, w).to(out_dtype) + + +def _jit(x, w, out_dtype): + return tiny_gemm_bf16(x, w, out_dtype=out_dtype, max_m=16) + + +FN_MAP = {"jit": _jit, "torch": _torch} + +SHAPES = [(256, 7168), (384, 7168), (256, 4096), (896, 7168), (144, 7168), (1536, 128)] + + +@marker.parametrize("out_dtype", [torch.bfloat16, torch.float32]) +@marker.parametrize("shape", SHAPES, [(384, 7168)]) +@marker.parametrize("num_tokens", list(range(1, 17)), [1, 8, 16]) +@marker.benchmark("provider", ["jit", "torch"]) +def benchmark(shape, num_tokens, out_dtype, provider): + n, k = shape + x = create_random(num_tokens, k) + w = create_random(n, k) + return marker.do_bench( + FN_MAP[provider], + input_args=(x, w), + input_kwargs={"out_dtype": out_dtype}, + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/gemm/test_dsv3_router_gemm.py b/test/registered/kernels/ops/gemm/test_dsv3_router_gemm.py deleted file mode 100644 index 981fbd959..000000000 --- a/test/registered/kernels/ops/gemm/test_dsv3_router_gemm.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Tests for JIT dsv3_router_gemm kernel.""" - -import itertools -import sys - -import pytest -import torch - -from sglang.kernels.jit.utils import ( - get_ci_test_range, - get_jit_cuda_arch, - is_hip_runtime, -) -from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=37, stage="base-b-kernel-unit", runner_config="1-gpu-large") -# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps. -register_cuda_ci(est_time=110, stage="nightly", runner_config="1-gpu-large") - -HIDDEN_DIMS = [1024, 4096, 5120, 6144, 7168] -ROUTER_GEMM_CASES = get_ci_test_range( - list( - itertools.product( - [256, 384], - HIDDEN_DIMS, - list(range(1, 17)), - [torch.bfloat16, torch.float32], - ) - ), - [ - (256, 1024, 1, torch.bfloat16), - (256, 7168, 6, torch.bfloat16), - (256, 6144, 4, torch.float32), - (384, 7168, 8, torch.bfloat16), - (256, 7168, 16, torch.float32), - (384, 5120, 16, torch.float32), - ], -) -ATOL = 1e-2 -RTOL = 1e-2 - - -def _ref(hidden_states, router_weights, out_dtype): - return (hidden_states.float() @ router_weights.float().T).to(out_dtype) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "num_experts,hidden_dim,num_tokens,out_dtype", ROUTER_GEMM_CASES -) -def test_dsv3_router_gemm(num_experts, hidden_dim, num_tokens, out_dtype): - if is_hip_runtime() or get_jit_cuda_arch().major < 9: - pytest.skip("SM90+ required") - - mat_a = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda") - mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.bfloat16, device="cuda") - - ref = _ref(mat_a, mat_b, out_dtype) - out = dsv3_router_gemm(mat_a, mat_b, out_dtype=out_dtype) - - assert out.shape == (num_tokens, num_experts) - assert out.dtype == out_dtype - torch.testing.assert_close(out.float(), ref.float(), atol=ATOL, rtol=RTOL) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/gemm/test_tiny_gemm.py b/test/registered/kernels/ops/gemm/test_tiny_gemm.py new file mode 100644 index 000000000..dcfd43dc1 --- /dev/null +++ b/test/registered/kernels/ops/gemm/test_tiny_gemm.py @@ -0,0 +1,56 @@ +"""Tests for the JIT tiny_gemm kernels.""" + +import sys + +import pytest +import torch + +from sglang.kernels.jit.utils import ( + get_ci_test_range, + get_jit_cuda_arch, + is_hip_runtime, +) +from sglang.kernels.ops.gemm.tiny_gemm import can_use_tiny_gemm, tiny_gemm_bf16 +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=300, stage="nightly", runner_config="1-gpu-large") + +# One kernel is built per m in [1, MAX_M], so hold max_m fixed across the sweep: +# every num_tokens of a shape then shares one JIT module. +MAX_M = 16 +SHAPES = [(256, 7168), (384, 7168), (256, 4096), (896, 7168), (144, 7168), (1536, 128)] + +TINY_GEMM_CASES = get_ci_test_range( + [ + (n, k, num_tokens, dtype) + for n, k in SHAPES + for num_tokens in range(1, MAX_M + 1) + for dtype in (torch.bfloat16, torch.float32) + ], + [ + (384, 7168, 1, torch.float32), + (384, 7168, 4, torch.float32), + (896, 7168, 8, torch.float32), + (1536, 128, 16, torch.bfloat16), + ], +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("n,k,num_tokens,out_dtype", TINY_GEMM_CASES) +def test_tiny_gemm(n, k, num_tokens, out_dtype): + if is_hip_runtime() or get_jit_cuda_arch().major < 9: + pytest.skip("SM90+ required") + + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + assert can_use_tiny_gemm(n, k, MAX_M) + out = tiny_gemm_bf16(x, w, out_dtype=out_dtype, max_m=MAX_M) + ref = torch.nn.functional.linear(x, w) + torch.testing.assert_close(out.float(), ref.float(), atol=1e-2, rtol=1e-2) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py b/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py index cf1cac046..e4b6088f3 100644 --- a/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py +++ b/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py @@ -24,10 +24,7 @@ from sglang.kernels.ops.attention.vision_rope import ( prepare_fused_qk_complex_rope_inplace, ) from sglang.kernels.ops.elementwise import add3 -from sglang.kernels.ops.gemm.tiny_gemm import ( - tiny_k_gemm_bf16, - tiny_n_gemm_bf16, -) +from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16 from sglang.kernels.ops.kvcache.set_mla_kv_buffer import set_mla_kv_buffer from sglang.kernels.ops.mm.process.image import ( _normalize_and_patchify_torch, @@ -379,17 +376,19 @@ class TestKimiK3PrerequisiteOps(CustomTestCase): ) def test_tiny_gemm_variants(self): + """Both K3 gate-projection shapes, one per kernel variant: N=144 is the + tiny dimension for the first, K=128 for the second.""" torch.manual_seed(2) x = torch.randn(2, 7168, device="cuda", dtype=torch.bfloat16) / 8 weight = torch.randn(144, 7168, device="cuda", dtype=torch.bfloat16) / 8 - actual = tiny_n_gemm_bf16(x, weight, out_dtype=torch.float32) + actual = tiny_gemm_bf16(x, weight, out_dtype=torch.float32) torch.testing.assert_close( actual.double(), x.double() @ weight.double().t(), rtol=1e-3, atol=1e-3 ) x = torch.randn(7, 128, device="cuda", dtype=torch.bfloat16) / 4 weight = torch.randn(1536, 128, device="cuda", dtype=torch.bfloat16) / 4 - actual = tiny_k_gemm_bf16(x, weight) + actual = tiny_gemm_bf16(x, weight) torch.testing.assert_close( actual.double(), x.double() @ weight.double().t(), rtol=2e-2, atol=2e-2 )