Use attn_tp_group for all reduce in token embedding (#17403)

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)
 

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,6 +489,9 @@ 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:
+                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

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),
         )
 

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()