From 0c5a81acb85b963f533cc85024736724a1182884 Mon Sep 17 00:00:00 2001 From: Zheng Li Date: Fri, 30 Jan 2026 20:44:25 +0800 Subject: [PATCH] [BUGFIX] Fix dp size > 1 for qwen3 vl model (#17624) Co-authored-by: yizhang2077 <1109276519@qq.com> --- python/sglang/srt/layers/attention/vision.py | 2 ++ python/sglang/srt/layers/linear.py | 12 +++++++-- .../srt/model_executor/forward_batch_info.py | 10 ++++++- python/sglang/srt/models/qwen3_vl.py | 27 ++++++++++--------- python/sglang/srt/multimodal/mm_utils.py | 16 ++++++++--- 5 files changed, 48 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 648ff091a..11092f4fa 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -596,6 +596,7 @@ class VisionAttention(nn.Module): [torch.Tensor, torch.Tensor, Any, Any], Tuple[torch.Tensor, torch.Tensor] ] = None, use_data_parallel: bool = False, + use_dp_attention_reduce: bool = False, aux_stream: Optional[torch.cuda.Stream] = None, **kwargs, ): @@ -688,6 +689,7 @@ class VisionAttention(nn.Module): tp_size=self.tp_size, reduce_results=False, prefix=add_prefix("proj", prefix), + use_dp_attention_reduce=use_dp_attention_reduce, ) self.aux_stream = aux_stream self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] if aux_stream else [] diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 98d97f20e..ebe7bb581 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -21,7 +21,10 @@ from sglang.srt.distributed import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) -from sglang.srt.layers.dp_attention import is_allocation_symmetric +from sglang.srt.layers.dp_attention import ( + get_attention_tp_group, + is_allocation_symmetric, +) from sglang.srt.layers.parameter import ( BasevLLMParameter, BlockQuantScaleParameter, @@ -1263,6 +1266,7 @@ class RowParallelLinear(LinearBase): tp_rank: Optional[int] = None, tp_size: Optional[int] = None, use_presharded_weights: bool = False, + use_dp_attention_reduce: bool = False, ): quant_config = None if _disable_hip_linear_quant else quant_config super().__init__( @@ -1271,6 +1275,7 @@ class RowParallelLinear(LinearBase): self.input_is_parallel = input_is_parallel self.reduce_results = reduce_results + self.use_dp_attention_reduce = use_dp_attention_reduce # Divide the weight matrix along the last dimension. if tp_rank is None: @@ -1419,7 +1424,10 @@ class RowParallelLinear(LinearBase): output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_) if self.reduce_results and self.tp_size > 1 and not skip_all_reduce: - output = tensor_model_parallel_all_reduce(output_parallel) + if self.use_dp_attention_reduce: + output = get_attention_tp_group().all_reduce(output_parallel) + else: + output = tensor_model_parallel_all_reduce(output_parallel) else: output = output_parallel diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index efd9f07d3..bcd1fda10 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -860,7 +860,15 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ) if self.mrope_positions is not None: - self.mrope_positions = self._pad_tensor_to_size(self.mrope_positions, bs) + self.mrope_positions = torch.cat( + [ + self.mrope_positions, + self.mrope_positions.new_zeros( + 3, num_tokens - self.mrope_positions.shape[1] + ), + ], + dim=1, + ) # TODO: check if we need to pad other tensors if self.extend_seq_lens is not None: diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index c397b136d..54d224b24 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -25,14 +25,15 @@ from einops import rearrange from transformers.activations import ACT2FN from sglang.srt.configs.qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig -from sglang.srt.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) +from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.distributed.parallel_state import get_pp_group from sglang.srt.environ import envs from sglang.srt.layers.attention.vision import VisionAttention -from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.dp_attention import ( + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, +) from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.pooler import Pooler, PoolingType @@ -85,10 +86,8 @@ class Qwen3_VisionMLP(nn.Module): use_data_parallel: bool = False, ): super().__init__() - self.tp_size = ( - 1 if use_data_parallel else get_tensor_model_parallel_world_size() - ) - self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + self.tp_size = 1 if use_data_parallel else get_attention_tp_size() + self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank() self.linear_fc1 = ColumnParallelLinear( in_features, hidden_features, @@ -106,6 +105,7 @@ class Qwen3_VisionMLP(nn.Module): prefix=add_prefix("linear_fc2", prefix), tp_size=self.tp_size, tp_rank=self.tp_rank, + use_dp_attention_reduce=is_dp_attention_enabled(), ) self.act = ACT2FN[hidden_act] @@ -176,6 +176,7 @@ class Qwen3_VisionBlock(nn.Module): quant_config=quant_config, prefix=add_prefix("attn", prefix), use_data_parallel=use_data_parallel, + use_dp_attention_reduce=is_dp_attention_enabled(), ) self.mlp = Qwen3_VisionMLP( dim, @@ -235,10 +236,8 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module): self.norm = norm_layer( self.hidden_size if use_postshuffle_norm else context_dim ) - self.tp_size = ( - 1 if use_data_parallel else get_tensor_model_parallel_world_size() - ) - self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + self.tp_size = 1 if use_data_parallel else get_attention_tp_size() + self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank() self.linear_fc1 = ColumnParallelLinear( self.hidden_size, self.hidden_size, @@ -257,6 +256,7 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module): prefix=add_prefix("linear_fc2", prefix), tp_size=self.tp_size, tp_rank=self.tp_rank, + use_dp_attention_reduce=is_dp_attention_enabled(), ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -713,6 +713,7 @@ class Qwen3VLForConditionalGeneration(nn.Module): self.config.vocab_size, self.config.hidden_size, quant_config=quant_config, + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, prefix=add_prefix("lm_head", prefix), ) else: diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index afce1079b..bb7802f7f 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/python/sglang/srt/multimodal/mm_utils.py @@ -495,11 +495,19 @@ def run_dp_sharded_mrope_vision_model( ``` """ - tp_size = get_tensor_model_parallel_world_size() + from sglang.srt.layers.dp_attention import ( + get_attention_tp_group, + get_attention_tp_rank, + get_attention_tp_size, + ) + + tp_size = get_attention_tp_size() + if tp_size == 1: + return vision_model(pixel_values, grid_thw=torch.tensor(grid_thw_list)) # GPU_0 tp_rank_local = 0 # GPU_1 tp_rank_local = 1 - tp_rank_local = get_tensor_model_parallel_rank() + tp_rank_local = get_attention_tp_rank() # patches_per_image = [1000, 100, 200, 50] patches_per_image = [math.prod(grid_thw) for grid_thw in grid_thw_list] @@ -611,7 +619,9 @@ def run_dp_sharded_mrope_vision_model( image_embeds_local_padded = image_embeds_local # Do all_gather to collect embeddings from all ranks - gathered_embeds = tensor_model_parallel_all_gather(image_embeds_local_padded, dim=0) + gathered_embeds = get_attention_tp_group().all_gather( + image_embeds_local_padded, dim=0 + ) # Remove padding and reconstruct per-rank embeddings rank_embeddings = list[torch.Tensor]()