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

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