[VLM] Support apply qk norm in multi cuda streams (#15720)

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

python/sglang/srt/layers/attention/vision.py

@@ -22,6 +22,10 @@ from sglang.srt.utils import (
     is_npu,
     print_info_once,
 )
+from sglang.srt.utils.multi_stream_utils import (
+    maybe_execute_in_parallel,
+    with_multi_stream,
+)
 
 _is_cuda = is_cuda()
 _is_npu = is_npu()
@@ -532,6 +536,7 @@ class VisionAttention(nn.Module):
             [torch.Tensor, torch.Tensor, Any, Any], Tuple[torch.Tensor, torch.Tensor]
         ] = None,
         use_data_parallel: bool = False,
+        aux_stream: Optional[torch.cuda.Stream] = None,
         **kwargs,
     ):
         super().__init__()
@@ -620,6 +625,8 @@ class VisionAttention(nn.Module):
             tp_size=self.tp_size,
             prefix=add_prefix("proj", prefix),
         )
+        self.aux_stream = aux_stream
+        self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
 
     def _determine_attention_backend(self, passed_backend: Optional[str]) -> str:
         """Decide the multimodal attention backend string.
@@ -655,20 +662,41 @@ class VisionAttention(nn.Module):
 
     def _apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor):
         """apply qk norm for internvl vit attn"""
-        q = q.flatten(1, 2)
-        k = k.flatten(1, 2)
 
-        if self.tp_size > 1:
-            q = tensor_model_parallel_all_gather(q.contiguous())
-            k = tensor_model_parallel_all_gather(k.contiguous())
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-        if self.tp_size > 1:
-            splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)
-            q = splitter(q)[self.tp_rank]
-            k = splitter(k)[self.tp_rank]
-        q = q.unflatten(-1, (-1, self.head_size))
-        k = k.unflatten(-1, (-1, self.head_size))
+        def q_l2norm():
+            q_ = q.flatten(1, 2)
+            if self.tp_size > 1:
+                q_ = tensor_model_parallel_all_gather(q_.contiguous())
+            q_ = self.q_norm(q_)
+            if self.tp_size > 1:
+                splitter = partial(
+                    split_tensor_along_last_dim, num_partitions=self.tp_size
+                )
+                q_ = splitter(q_)[self.tp_rank]
+            q_ = q_.unflatten(-1, (-1, self.head_size))
+            return q_
+
+        def k_l2norm():
+            k_ = k.flatten(1, 2)
+            if self.tp_size > 1:
+                k_ = tensor_model_parallel_all_gather(k_.contiguous())
+            k_ = self.k_norm(k_)
+            if self.tp_size > 1:
+                splitter = partial(
+                    split_tensor_along_last_dim, num_partitions=self.tp_size
+                )
+                k_ = splitter(k_)[self.tp_rank]
+            k_ = k_.unflatten(-1, (-1, self.head_size))
+            return k_
+
+        with with_multi_stream(True):
+            q, k = maybe_execute_in_parallel(
+                q_l2norm,
+                k_l2norm,
+                self.ln_events[0],
+                self.ln_events[1],
+                self.aux_stream,
+            )
         return q, k
 
     def forward(

python/sglang/srt/models/internvl.py

@@ -38,8 +38,11 @@ from sglang.srt.models.qwen3 import Qwen3ForCausalLM
 from sglang.srt.models.qwen3_moe import Qwen3MoeForCausalLM
 from sglang.srt.multimodal.mm_utils import run_dp_sharded_vision_model
 from sglang.srt.server_args import get_global_server_args
+from sglang.srt.utils import is_cuda
 from sglang.utils import logger
 
+_is_cuda = is_cuda()
+
 
 class InternAttention(nn.Module):
     def __init__(
@@ -47,6 +50,7 @@ class InternAttention(nn.Module):
         config,
         quant_config: QuantizationConfig = None,
         use_data_parallel: bool = False,
+        aux_stream: Optional[torch.cuda.Stream] = None,
     ):
         super().__init__()
         self.config = config
@@ -69,6 +73,7 @@ class InternAttention(nn.Module):
             or getattr(config, "use_qk_norm", False),
             flatten_batch=False,
             use_data_parallel=use_data_parallel,
+            aux_stream=aux_stream,
         )
 
         self.proj_drop = nn.Dropout(config.dropout)
@@ -222,6 +227,7 @@ class InternVisionEncoderLayer(nn.Module):
         drop_path_rate: float,
         quant_config: QuantizationConfig = None,
         use_data_parallel: bool = False,
+        aux_stream: Optional[torch.cuda.Stream] = None,
     ):
         super().__init__()
         self.embed_dim = config.hidden_size
@@ -231,6 +237,7 @@ class InternVisionEncoderLayer(nn.Module):
             config=config,
             quant_config=quant_config,
             use_data_parallel=use_data_parallel,
+            aux_stream=aux_stream,
         )
         self.mlp = InternMLP(config, use_data_parallel)
         self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
@@ -296,10 +303,11 @@ class InternVisionEncoder(nn.Module):
             x.item()
             for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)
         ]
+        aux_stream = torch.cuda.Stream() if _is_cuda else None
         self.layers = nn.ModuleList(
             [
                 InternVisionEncoderLayer(
-                    config, dpr[idx], quant_config, use_data_parallel
+                    config, dpr[idx], quant_config, use_data_parallel, aux_stream
                 )
                 for idx in range(config.num_hidden_layers)
             ]

python/sglang/srt/utils/multi_stream_utils.py

@@ -0,0 +1,77 @@
+# Adapted from trtllm.
+
+import threading
+from contextlib import contextmanager
+from typing import Any, Callable, Optional
+
+import torch
+
+
+class do_multi_stream_local(threading.local):
+
+    def __init__(self):
+        self.do_multi_stream = False
+
+
+_local = do_multi_stream_local()
+
+
+def set_do_multi_stream(enable: bool):
+    _local.do_multi_stream = enable
+
+
+def do_multi_stream() -> bool:
+    return _local.do_multi_stream
+
+
+@contextmanager
+def with_multi_stream(enable: bool):
+    prev_do_multi_stream = _local.do_multi_stream
+    set_do_multi_stream(enable)
+    try:
+        yield
+    finally:
+        set_do_multi_stream(prev_do_multi_stream)
+
+
+def maybe_execute_in_parallel(
+    fn0: Callable,
+    fn1: Callable,
+    event0: torch.cuda.Event,
+    event1: torch.cuda.Event,
+    aux_stream: Optional[torch.cuda.Stream] = None,
+) -> tuple[Any, Any]:
+    """Utility function to run two functions in two cuda streams in parallel. Multi-stream is
+    only enabled when cuda graph is turned on because switch stream has extra host overhead.
+
+    This design is mainly for low latency use case. It needs to be improved for max throughput
+    use case.
+    For simplicity, fn0 and fn1 do not support inputs.
+
+    Args:
+        fn0 (Callable): callable for the default stream
+        fn1 (Callable): callable for the second stream, aux_stream
+        event0 (torch.cuda.Event): cuda event for fn0
+        event1 (torch.cuda.Event): cuda event for fn1
+        aux_stream (Optional[torch.cuda.Stream]): the second cuda stream for fn1.
+            Multi-stream is disabled when aux_stream is None.
+
+    Returns:
+        tuple[Any, Any]: the return values of fn0() and fn1()
+    """
+
+    multi_stream = do_multi_stream() and aux_stream is not None
+
+    if multi_stream:
+        event0.record()
+        result0 = fn0()
+
+        with torch.cuda.stream(aux_stream):
+            event0.wait()
+            result1 = fn1()
+            event1.record()
+        event1.wait()
+    else:
+        result0 = fn0()
+        result1 = fn1()
+    return (result0, result1)