From 48b8dcd42e55a0826fbba4acc36bdc0a84f35bb6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E4=B8=80=E6=B6=B5?= Date: Thu, 8 Jan 2026 13:54:33 +0800 Subject: [PATCH] [jit kernel] support dtype as a cpp template parameter (#16452) --- python/sglang/jit_kernel/csrc/norm.cuh | 42 +++++++++-------- .../jit_kernel/include/sgl_kernel/utils.cuh | 46 +++++++++++++++++++ .../jit_kernel/include/sgl_kernel/utils.h | 3 ++ python/sglang/jit_kernel/norm.py | 10 ++-- python/sglang/jit_kernel/utils.py | 10 ++++ .../runtime/layers/layernorm.py | 2 +- .../runtime/models/dits/zimage.py | 2 +- python/sglang/srt/layers/attention/vision.py | 2 +- python/sglang/srt/models/utils.py | 2 +- 9 files changed, 91 insertions(+), 28 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/norm.cuh b/python/sglang/jit_kernel/csrc/norm.cuh index c5e369df4..a44ecba4c 100644 --- a/python/sglang/jit_kernel/csrc/norm.cuh +++ b/python/sglang/jit_kernel/csrc/norm.cuh @@ -119,10 +119,15 @@ __global__ void fused_qknorm(const QKNormParams __grid_constant__ params) { PDLTriggerSecondary(); // launch secondary kernel } -template +template struct QKNormKernel { - template - static constexpr auto qknorm_kernel = fused_qknorm; + static_assert( + std::is_same_v || std::is_same_v, + "Unsupported DType: QKNormKernel only supports half and nv_bfloat16."); + using DType2 = host::PackedDType::type; + + // only initialize once (static variable) to avoid overhead + static constexpr auto kernel = fused_qknorm; 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(...) (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 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(dtype) - .with_device(device) + .with_dtype(dtype) + .template with_device(device) .verify(q); TensorMatcher({N, K, D}) // k input .with_strides({Sk, D, 1}) - .with_dtype(dtype) - .with_device(device) + .with_dtype(dtype) + .template with_device(device) .verify(k); TensorMatcher({D}) // weight - .with_dtype(dtype) - .with_device(device) + .with_dtype(dtype) + .template with_device(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; - static constexpr auto fp16_kernel = qknorm_kernel; - 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(); - 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); diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index 8d7da5ee3..0e503a3c0 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -2,6 +2,8 @@ #include +#include +#include #include #include @@ -96,6 +98,50 @@ __forceinline__ __device__ void PDLTriggerSecondary() { namespace host { +// DType +template +struct PackedDType { + static_assert(dependent_false_v, "Unsupported dtype for Packed"); +}; + +template <> +struct PackedDType { + using type = float2; +}; + +template <> +struct PackedDType { + using type = float4; +}; + +template <> +struct PackedDType<__half, 2> { + using type = __half2; +}; + +struct alignas(8) half4 { + __half x, y, z, w; +}; + +template <> +struct PackedDType<__half, 4> { + using type = half4; +}; + +template <> +struct PackedDType { + using type = nv_bfloat162; +}; + +struct alignas(8) bf16_4 { + nv_bfloat16 x, y, z, w; +}; + +template <> +struct PackedDType { + using type = bf16_4; +}; + inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) { if (error != ::cudaSuccess) { [[unlikely]]; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/python/sglang/jit_kernel/include/sgl_kernel/utils.h index c3cc7c1e5..106ae14ee 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.h @@ -44,6 +44,9 @@ namespace host { +template +inline constexpr bool dependent_false_v = false; + struct DebugInfo : public source_location_t { DebugInfo(source_location_t loc = source_location_t::current()) : source_location_t(loc) {} }; diff --git a/python/sglang/jit_kernel/norm.py b/python/sglang/jit_kernel/norm.py index 50f2315f8..963e5ed06 100644 --- a/python/sglang/jit_kernel/norm.py +++ b/python/sglang/jit_kernel/norm.py @@ -17,8 +17,8 @@ if TYPE_CHECKING: @cache_once -def _jit_norm_module(head_dims: int) -> Module: - args = make_cpp_args(head_dims, is_arch_support_pdl()) +def _jit_norm_module(head_dims: int, dtype: torch.dtype) -> Module: + args = make_cpp_args(head_dims, is_arch_support_pdl(), dtype) return load_jit( "norm", *args, @@ -28,13 +28,13 @@ def _jit_norm_module(head_dims: int) -> Module: @cache_once -def can_use_fused_inplace_qknorm(head_dim: int) -> bool: +def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool: logger = logging.getLogger(__name__) if head_dim not in [64, 128, 256, 512, 1024]: logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel") return False try: - _jit_norm_module(head_dim) + _jit_norm_module(head_dim, dtype) return True except Exception as e: logger.warning(f"Failed to load JIT QK-Norm kernel: {e}") @@ -51,5 +51,5 @@ def fused_inplace_qknorm( head_dim: int = 0, ) -> None: head_dim = head_dim or q.size(-1) - module = _jit_norm_module(head_dim) + module = _jit_norm_module(head_dim, q.dtype) module.qknorm(q, k, q_weight, k_weight, eps) diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index 4de953497..fa326db00 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -5,6 +5,8 @@ import pathlib from functools import lru_cache from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar, Union +import torch + if TYPE_CHECKING: from tvm_ffi import Module @@ -54,6 +56,14 @@ def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList: return "true" if arg else "false" if isinstance(arg, (int, float)): return str(arg) + if isinstance(arg, torch.dtype): + if arg == torch.float: + return "float" + if arg == torch.float16: + return "__half" + if arg == torch.bfloat16: + return "nv_bfloat16" + raise TypeError(f"Not implement this arg wrapper yet: {arg}") raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}") return CPPArgList(_convert(arg) for arg in args) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 78ed8099d..3cf61d54f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -445,7 +445,7 @@ def apply_qk_norm( q.is_cuda and allow_inplace and (q_eps == k_eps) - and can_use_fused_inplace_qknorm(head_dim) + and can_use_fused_inplace_qknorm(head_dim, q.dtype) ): fused_inplace_qknorm( q=q.view(batch_size, -1, head_dim), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index b14b63496..e2ddfeaf8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -159,7 +159,7 @@ class ZImageAttention(nn.Module): if ( q.is_cuda and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon) - and can_use_fused_inplace_qknorm(self.head_dim) + and can_use_fused_inplace_qknorm(self.head_dim, q.dtype) ): q, k = apply_qk_norm( q=q, diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index ff603e8f8..9dd35e5e3 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -825,7 +825,7 @@ class VisionAttention(nn.Module): # internvl if self.qk_normalization: # jit kernel - if can_use_jit_qk_norm(self.head_size): + if can_use_jit_qk_norm(self.head_size, q.dtype): # q: [tokens, head, head_size] -> [tokens, embed_dim] head_dim_for_norm = head * self.head_size diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 3ebd82448..e5211e833 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -233,7 +233,7 @@ def apply_qk_norm( and allow_inplace # TODO(dark): this can be relaxed if needed and (q_eps == k_eps) # TODO(dark): this can also be relaxed and not envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() - and can_use_fused_inplace_qknorm(head_dim) + and can_use_fused_inplace_qknorm(head_dim, q.dtype) ): fused_inplace_qknorm( q=q.view(batch_size, -1, head_dim),