From b9af8d2eb9f0486c1578dcdbb1ded9b6350bc641 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Thu, 25 Dec 2025 14:35:07 +0800 Subject: [PATCH] [VLM] Support apply qk norm in multi cuda streams (#15720) Co-authored-by: luoyuan.luo --- python/sglang/srt/layers/attention/vision.py | 54 +++++++++---- python/sglang/srt/models/internvl.py | 10 ++- python/sglang/srt/utils/multi_stream_utils.py | 77 +++++++++++++++++++ 3 files changed, 127 insertions(+), 14 deletions(-) create mode 100644 python/sglang/srt/utils/multi_stream_utils.py diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index e055c66b2..1bec0fe63 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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( diff --git a/python/sglang/srt/models/internvl.py b/python/sglang/srt/models/internvl.py index b23b1b61b..ef17d6008 100644 --- a/python/sglang/srt/models/internvl.py +++ b/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) ] diff --git a/python/sglang/srt/utils/multi_stream_utils.py b/python/sglang/srt/utils/multi_stream_utils.py new file mode 100644 index 000000000..fa5c55837 --- /dev/null +++ b/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)