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

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh

@@ -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

@@ -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>
+  static_assert(N % N_SPLIT == 0, "N must be divisible by n_split");
-__global__ __launch_bounds__(K / kTinyNGemmVecSize, 1)  // 1 block per SM
+  static_assert((K * sizeof(bf16_t)) % kBytes == 0, "K must be divisible by kBytes");
-    void tiny_n_gemm_kernel(OutT* __restrict__ out, const bf16_t* __restrict__ x, const bf16_t* __restrict__ w) {
+  static_assert(kBytes % device::kMaxVecBytes == 0);
-  constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize;
+  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<bf16x2_t, kVecSize / 2>;
-  using vec_t = AlignedVector<bf16_t, kTinyNGemmVecSize>;
 
   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.
+  // prefetch weight before PDL
-  vec_t wv[N_SPLIT];
+  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) {
+      for (uint32_t u = 0; u < kUnroll; ++u) {
-        acc = device::math::fma_f32_bf16(xv[m][i], wv[n][i], acc);
+        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) {
+template <typename Trait, uint32_t M, typename Out, bool kUsePDL>
-  const uint32_t offset = static_cast<uint32_t>(vec_offset * 16);
+TINY_GEMM_KERNEL void tiny_k_gemm_kernel(const TinyGEMMParams params) {
-#if defined(USE_ROCM)
+  using namespace device;
-  *reinterpret_cast<uint4*>(static_cast<char*>(smem_dst) + offset) =
+  constexpr uint32_t N = Trait::N;
-      *reinterpret_cast<const uint4*>(static_cast<const char*>(gmem_src) + offset);
+  constexpr uint32_t K = Trait::K;
-#else
+  constexpr uint32_t N_SPLIT = Trait::N_SPLIT;
-  const uint32_t smem_addr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst)) + offset;
+  constexpr uint32_t kVecSize = Trait::kVecSize;
-  const uint64_t gmem_addr = static_cast<uint64_t>(__cvta_generic_to_global(gmem_src)) + offset;
+  constexpr uint32_t kNumKLanes = Trait::kNumKLanes;
-  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" : : "r"(smem_addr), "l"(gmem_addr) : "memory");
+  using vec_t = AlignedVector<bf16x2_t, kVecSize / 2>;
-#endif
-}
 
-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
+    dot_product(xv[m], wv, acc);
-    for (uint32_t i = 0; i < kTinyKGemmVecSize; ++i) {
-      acc = device::math::fma_f32_bf16(xv[m][i], wv[i], 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;
+  using Trait = GEMMTraitN<N, K, N_SPLIT, 32>;
-  static constexpr uint32_t kNumBlocks = N / N_SPLIT;
+  using KernelFn = void (*)(TinyGEMMParams);
-  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*);
 
   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);
+    if (num_tokens == 0) return;
-    LaunchKernel(kNumBlocks, kBlockSize, device.unwrap())
+    CHECK_HOST(num_tokens >= 1 && num_tokens <= kMaxM);
-        .enable_pdl(kUsePDL)(
+    // x may be a row-sliced view of a wider buffer, but the rows are loaded as
-            kTable[num_tokens],
+    // whole vectors, so every row start must stay vector-aligned.
-            static_cast<OutT*>(out.data_ptr()),
+    CHECK_HOST(x.stride(0) % Trait::kVecSize == 0)
-            static_cast<const bf16_t*>(x.data_ptr()),
+        << "x rows must stay aligned to the vector width, got stride " << x.stride(0);
-            static_cast<const bf16_t*>(w.data_ptr()));
+    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;
+  using Trait = GEMMTraitK<N, K, N_SPLIT, 16>;
-  static constexpr uint32_t kBlockSize = N_SPLIT * kNumKLanes;
+  using KernelFn = void (*)(TinyGEMMParams);
-  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);
 
   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_strides({-1, 1}).with_dtype<bf16_t>().with_device(device).verify(x);
-    TensorMatcher({M, K}).with_dtype<bf16_t>().with_strides({-1, 1}).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));
+    if (num_tokens == 0) return;
-    RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM);
+    CHECK_HOST(num_tokens >= 1 && num_tokens <= kMaxM);
-    RuntimeCheck(
+    // x may be a row-sliced view of a wider buffer, but the rows are loaded as
-        x_stride * sizeof(bf16_t) % (kTinyKGemmVecSize * sizeof(bf16_t)) == 0,
+    // whole vectors, so every row start must stay vector-aligned.
-        "x rows must stay aligned to the vector width, got stride ",
+    CHECK_HOST(x.stride(0) % Trait::kVecSize == 0)
-        x_stride);
+        << "x rows must stay aligned to the vector width, got stride " << x.stride(0);
-    LaunchKernel(kNumBlocks, kBlockSize, device.unwrap())
+    const auto params = TinyGEMMParams{
-        .enable_pdl(kUsePDL)(
+        .out = out.data_ptr(),
-            kTable[num_tokens],
+        .x = static_cast<const bf16_t*>(x.data_ptr()),
-            static_cast<OutT*>(out.data_ptr()),
+        .w = static_cast<const bf16_t*>(w.data_ptr()),
-            static_cast<const bf16_t*>(x.data_ptr()),
+        .stride_x = static_cast<int64_t>(x.stride(0)),
-            static_cast<const bf16_t*>(w.data_ptr()),
+    };
-            x_stride);
+    LaunchKernel(Trait::kNumBlocks, Trait::kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kTable[num_tokens], params);
   }
 };
 

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(
+# NOTE: named after the kernel entry, not the op -- a bare `tiny_gemm` here would
-    hidden_states: torch.Tensor,
+# be shadowed by the submodule of the same name as soon as it is imported.
-    router_weights: torch.Tensor,
+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."""
+    """Tiny bf16 GEMM ``x[m, k] @ w[n, k].T``, support bf16/fp32 out"""
-    impl = get_kernel("gemm.dsv3_router_gemm", KernelBackend.JIT)
+    impl = get_kernel("gemm.tiny_gemm", KernelBackend.JIT)
-    if out_dtype is None:
+    return impl(x, w, out, out_dtype=out_dtype, max_m=max_m)
-        return impl(hidden_states, router_weights, output=output)
-    return impl(hidden_states, router_weights, out_dtype, output)
 
 
 __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

@@ -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

@@ -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(
+def _get_num_sm() -> int:
-    n: int, k: int, max_m: int, n_unroll: int, out_dtype: torch.dtype
+    device = torch.cuda.current_device()
-) -> Module:
+    return torch.cuda.get_device_properties(device).multi_processor_count
-    args = make_cpp_args(n, k, max_m, n_unroll, out_dtype, is_arch_support_pdl())
+
-    return load_jit(
+
-        "tiny_k_gemm",
+def _n_gemm_block_threads(k: int) -> int:
-        *args,
+    return k // _N_GEMM_ELEMS_PER_THREAD
-        cuda_files=["gemm/tiny_gemm.cuh"],
+
-        cuda_wrappers=[("run", f"TinyKGemmKernel<{args}>::run")],
+
-        extra_cuda_cflags=["-O3"],
+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:
+def _supports_k_variant(n: int, k: int) -> bool:
-    """bf16 elements per vectorized load; mirrors kMaxVecBytes in utils.cuh."""
+    lanes = _k_gemm_reduction_lanes(k)
-    from sglang.kernels.jit.utils import get_jit_cuda_arch
+    return (
-
+        k % _K_GEMM_ELEMS_PER_THREAD == 0
-    cuda = tuple(int(v) for v in (torch.version.cuda or "0.0").split(".")[:2])
+        and k > _K_GEMM_ELEMS_PER_THREAD
-    return 16 if get_jit_cuda_arch().major >= 10 and cuda >= (12, 9) else 8
+        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:
+@cache_once
-    """Smallest divisor of n whose n / split_n blocks fit in one wave, subject
+def _default_n_gemm_n_split(n: int, k: int, max_m: int) -> int:
-    to the max_m * split_n <= K / vec_elems block-size constraint; falls back
+    split_cap = _n_gemm_block_threads(k) // max_m
-    to the largest split_n satisfying the constraint (multi-wave grid)."""
+    candidates = [d for d in range(1, min(n, split_cap) + 1) if n % d == 0]
-    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    if not candidates:
-    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:
         raise RuntimeError(
-            f"tiny_gemm: no valid split_n for N={n}, K={k}, max_m={max_m};"
+            f"tiny_gemm: no valid n_split for N={n}, K={k}, max_m={max_m};"
-            " lower 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(
+@cache_once
-    x: torch.Tensor,
+def _default_k_gemm_n_split(n: int, k: int) -> int:
-    w: torch.Tensor,
+    lanes = _k_gemm_reduction_lanes(k)
-    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
     candidates = [
         d
-        for d in range(1, n + 1)
+        for d in range(1, min(n, _MAX_BLOCK_THREADS // lanes) + 1)
-        if n % d == 0 and d * lanes % 32 == 0 and d * lanes <= 1024
+        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}")
+        raise RuntimeError(f"tiny_gemm: no valid n_split for N={n}, K={k}")
-    sm_count = torch.cuda.get_device_properties(0).multi_processor_count
+    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,
+    n_split: Optional[int] = None,
-    max_m: int = _MAX_M_DEFAULT,
+    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
+    Equal to `torch.nn.functional.linear(x, w)`.
-    exact N / split_n grid fills the SMs (no tail). Requires K / 8 to be a
+    Call :func:`can_use_tiny_gemm` first: shapes outside the supported set raise
-    power of 2 and <= 32 (e.g. K = 128/256). x may be a row-sliced view as
+    rather than falling back.
-    long as rows stay 16-byte aligned.
 
-    split_n trades block count for block size; the default picks the smallest
+    :param x: Shape [m, k], must be bf16_t, `m <= max_m`
-    divisor of N that fits one wave (12 for [1536, 128] on B200: 128 blocks
+    :param w: Shape [n, k], must be bf16_t
-    of 6 warps)."""
+    """
-    n = w.shape[0]
+    n, k = w.shape
-    k = x.shape[1]
     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:
+    if n_split is None:
-        split_n = _default_k_split_n(n, k)
+        n_split = (
-    module = _jit_tiny_k_gemm_module(n, k, max_m, split_n, out.dtype)
+            _default_k_gemm_n_split(n, k)
-    module.run(x, w, out)
+            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

@@ -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:
+    max_num_tokens = _K3_TINY_GEMM_MAX_TOKENS.get((n, k))
-        if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)):
+    if not _is_npu and max_num_tokens is not None and 0 < m <= max_num_tokens:
-            if 0 < m <= max_num_tokens:
+        return tiny_gemm_bf16(x, w, max_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)
     return torch.nn.functional.linear(x, w)
 
 

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

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 (
+    from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16
-        dsv3_router_gemm as dsv3_router_gemm,
-    )
 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 (
+            if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
-                _is_cuda
+                logits = tiny_gemm_bf16(
-                and hidden_states.shape[0] <= 16
+                    hidden_states,
-                and hidden_states.shape[1] % 1024 == 0
+                    self.weight,
-                and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
+                    out_dtype=torch.float32,
-                and _device_sm >= 90
+                    max_m=self.tiny_router_gemm_max_tokens,
-            ):
-                logits = dsv3_router_gemm(
-                    hidden_states, self.weight, out_dtype=torch.float32
                 )
 
             elif _use_aiter:

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 (
+        if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
-            hidden_states.shape[0] <= 16
+            logits = tiny_gemm_bf16(
-            and hidden_states.shape[1] == 7168
+                hidden_states, self.weight, max_m=self.tiny_router_gemm_max_tokens
-            and self.weight.shape[0] == 256
+            )
-            and _device_sm >= 90
-        ):
-            # router gemm output float32
-            logits = dsv3_router_gemm(hidden_states, self.weight)
         else:
             logits = F.linear(hidden_states, self.weight, None)
 

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)

test/registered/kernels/benchmark/gemm/bench_dsv3_router_gemm.py

@@ -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()

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()

test/registered/kernels/ops/gemm/test_dsv3_router_gemm.py

@@ -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"]))

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"]))

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 (
+from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16
-    tiny_k_gemm_bf16,
-    tiny_n_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
         )