chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[jit kernel] support dtype as a cpp template parameter (#16452)

Browse Source
This commit is contained in:
陈一涵
2026-01-08 13:54:33 +08:00
committed by GitHub
parent 41b434a7e6
commit 48b8dcd42e
9 changed files with 91 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
python/sglang/jit_kernel/csrc/norm.cuh
View File

@@ -119,10 +119,15 @@ __global__ void fused_qknorm(const QKNormParams __grid_constant__ params) {
PDLTriggerSecondary<kUsePDL>(); // launch secondary kernel
}
template <int64_t kHeadDim, bool kUsePDL>
template <int64_t kHeadDim, bool kUsePDL, typename DType>
struct QKNormKernel {
template <typename PackedFloat, typename Float>
static constexpr auto qknorm_kernel = fused_qknorm<kHeadDim, kUsePDL, PackedFloat, Float>;
static_assert(
std::is_same_v<DType, half> || std::is_same_v<DType, nv_bfloat16>,
"Unsupported DType: QKNormKernel only supports half and nv_bfloat16.");
using DType2 = host::PackedDType<DType, 2>::type;
// only initialize once (static variable) to avoid overhead
static constexpr auto kernel = fused_qknorm<kHeadDim, kUsePDL, DType2, DType>;
static void
run(const tvm::ffi::TensorView q,
@@ -141,19 +146,27 @@ struct QKNormKernel {
auto dtype = SymbolicDType{};
auto device = SymbolicDevice{};
/*
* We need the .template disambiguator here because this call happens in a dependent context.
* After switching to with_dtype<DType>(...) (where DType is a template parameter), the chained expression becomes
* dependent. In C++, when calling a member function template via ./-> on a dependent expression, the compiler may
* otherwise parse <kDLCUDA> as the < operator instead of template arguments. Adding .template forces correct
* parsing and fixes compilation errors (often seen with NVCC/clang). Ref:
* https://en.cppreference.com/w/cpp/language/dependent_name
*/
TensorMatcher({N, Q, D}) // q input
.with_strides({Sq, D, 1})
.with_dtype<nv_bfloat16, half>(dtype)
.with_device<kDLCUDA>(device)
.with_dtype<DType>(dtype)
.template with_device<kDLCUDA>(device)
.verify(q);
TensorMatcher({N, K, D}) // k input
.with_strides({Sk, D, 1})
.with_dtype<nv_bfloat16, half>(dtype)
.with_device<kDLCUDA>(device)
.with_dtype<DType>(dtype)
.template with_device<kDLCUDA>(device)
.verify(k);
TensorMatcher({D}) // weight
.with_dtype<nv_bfloat16, half>(dtype)
.with_device<kDLCUDA>(device)
.with_dtype<DType>(dtype)
.template with_device<kDLCUDA>(device)
.verify(q_weight)
.verify(k_weight);
@@ -177,19 +190,10 @@ struct QKNormKernel {
.num_tokens = num_tokens,
};
// only initialize once (static variable) to avoid overhead
static constexpr auto bf16_kernel = qknorm_kernel<nv_bfloat162, nv_bfloat16>;
static constexpr auto fp16_kernel = qknorm_kernel<half2, half>;
static const uint32_t kMaxOccupancyTable[2] = {
runtime::get_blocks_per_sm(fp16_kernel, kThreadsPerBlock),
runtime::get_blocks_per_sm(bf16_kernel, kThreadsPerBlock),
};
static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kThreadsPerBlock);
static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
// choose kernel based on dtype
const bool use_bf16 = dtype.is_type<nv_bfloat16>();
const auto kernel = use_bf16 ? bf16_kernel : fp16_kernel;
const auto max_occupancy = kMaxOccupancyTable[use_bf16 ? 1 : 0];
const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock);