[VLM] Adopt jit qk_norm kernel in VLM (#16171)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/jit_kernel/benchmark/bench_qknorm.py

@@ -5,6 +5,10 @@ from typing import Tuple
 import torch
 import triton
 import triton.testing
+from sgl_kernel import rmsnorm
+
+from sglang.jit_kernel.norm import fused_inplace_qknorm
+from sglang.srt.utils import get_current_device_stream_fast
 
 IS_CI = (
     os.getenv("CI", "false").lower() == "true"
@@ -20,13 +24,12 @@ def sglang_aot_qknorm(
     q_weight: torch.Tensor,
     k_weight: torch.Tensor,
 ) -> None:
-    from sgl_kernel import rmsnorm
 
     head_dim = q.shape[-1]
     q = q.view(-1, head_dim)
     k = k.view(-1, head_dim)
 
-    current_stream = torch.cuda.current_stream()
+    current_stream = get_current_device_stream_fast()
     alt_stream.wait_stream(current_stream)
     rmsnorm(q, q_weight, out=q)
     with torch.cuda.stream(alt_stream):
@@ -40,7 +43,6 @@ def sglang_jit_qknorm(
     q_weight: torch.Tensor,
     k_weight: torch.Tensor,
 ) -> None:
-    from sglang.jit_kernel.norm import fused_inplace_qknorm
 
     fused_inplace_qknorm(q, k, q_weight, k_weight)
 
@@ -51,7 +53,6 @@ def flashinfer_qknorm(
     q_weight: torch.Tensor,
     k_weight: torch.Tensor,
 ) -> None:
-    from flashinfer.norm import rmsnorm
 
     rmsnorm(q, q_weight, out=q)
     rmsnorm(k, k_weight, out=k)

python/sglang/jit_kernel/norm.py

@@ -30,7 +30,7 @@ def _jit_norm_module(head_dims: int) -> Module:
 @cache_once
 def can_use_fused_inplace_qknorm(head_dim: int) -> bool:
     logger = logging.getLogger(__name__)
-    if head_dim not in [64, 128, 256]:
+    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: