diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 6c2db5f26..0b5e2765d 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -564,6 +564,10 @@ def attn_tp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): return get_attention_tp_group().reduce_scatter_tensor(output, input) +def attn_tp_all_reduce(input: torch.Tensor): + return get_attention_tp_group().all_reduce(input) + + def attn_tp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor): return get_attention_tp_group().all_gather_into_tensor(output, input) diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index 192970904..1df12f66e 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -19,7 +19,11 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( ) from sglang.srt.layers.amx_utils import PackWeightMethod from sglang.srt.layers.communicator import get_attn_tp_context -from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size +from sglang.srt.layers.dp_attention import ( + attn_tp_all_reduce, + get_attention_tp_rank, + get_attention_tp_size, +) from sglang.srt.layers.parameter import BasevLLMParameter from sglang.srt.layers.quantization.base_config import ( QuantizationConfig, @@ -210,6 +214,7 @@ class VocabParallelEmbedding(torch.nn.Module): self.quant_config = quant_config self.enable_tp = enable_tp + self.use_attn_tp_group = use_attn_tp_group if self.enable_tp: if use_attn_tp_group: tp_rank = get_attention_tp_rank() @@ -484,8 +489,11 @@ class VocabParallelEmbedding(torch.nn.Module): # Mask the output embedding. output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) if not get_attn_tp_context().input_scattered: - # Reduce across all the model parallel GPUs. - output_parallel = tensor_model_parallel_all_reduce(output_parallel) + if self.use_attn_tp_group: + output_parallel = attn_tp_all_reduce(output_parallel) + else: + # Reduce across all the model parallel GPUs. + output_parallel = tensor_model_parallel_all_reduce(output_parallel) return output_parallel def extra_repr(self) -> str: diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 2fa3d2943..a64239031 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -86,7 +86,7 @@ class DeepseekModelNextN(nn.Module): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), ) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index d91bc5024..b5d274a3c 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -2474,7 +2474,7 @@ class DeepseekV2Model(nn.Module): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), ) else: self.embed_tokens = PPMissingLayer()