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

2026-01-08 13:54:33 +08:00
parent 41b434a7e6
commit 48b8dcd42e
9 changed files with 91 additions and 28 deletions
--- 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