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

python/sglang/jit_kernel/csrc/norm.cuh

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

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

@@ -2,6 +2,8 @@
 
 #include <sgl_kernel/utils.h>
 
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
 #include <dlpack/dlpack.h>
 #include <tvm/ffi/extra/c_env_api.h>
 
@@ -96,6 +98,50 @@ __forceinline__ __device__ void PDLTriggerSecondary() {
 
 namespace host {
 
+// DType
+template <typename DType, int Pack>
+struct PackedDType {
+  static_assert(dependent_false_v<DType>, "Unsupported dtype for Packed");
+};
+
+template <>
+struct PackedDType<float, 2> {
+  using type = float2;
+};
+
+template <>
+struct PackedDType<float, 4> {
+  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<nv_bfloat16, 2> {
+  using type = nv_bfloat162;
+};
+
+struct alignas(8) bf16_4 {
+  nv_bfloat16 x, y, z, w;
+};
+
+template <>
+struct PackedDType<nv_bfloat16, 4> {
+  using type = bf16_4;
+};
+
 inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) {
   if (error != ::cudaSuccess) {
     [[unlikely]];

python/sglang/jit_kernel/include/sgl_kernel/utils.h

@@ -44,6 +44,9 @@
 
 namespace host {
 
+template <typename>
+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) {}
 };

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)

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)

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

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,

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

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