diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 04d2d0652..6a7aa2aee 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -736,6 +736,13 @@ class Envs: # Distributed SGLANG_DSV4_FIX_TP_ATTN_A2A_SCATTER = EnvBool(True) SGLANG_SHARED_EXPERT_TP1 = EnvBool(False) + # Replicate the input embedding across TP ranks instead of sharding it + # along the vocab dimension (saves an all-reduce/all-gather in the embed + # lookup at the cost of replicated embedding weights). Drives both the + # target and every draft that shares its embedding (see + # get_embedding_tp_kwargs); they must stay in lock-step. Currently only + # applies to the Deepseek-V2 family (Deepseek V3.1, Kimi K2.5) + drafts. + SGLANG_ENABLE_EMBED_REPLICATION = EnvBool(False) # Symmetric Memory SGLANG_SYMM_MEM_PREALLOC_GB_SIZE = EnvInt(-1) SGLANG_DEBUG_SYMM_MEM = EnvBool(False) diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index cb35f091a..3deb1ad0e 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -19,6 +19,7 @@ from sglang.srt.distributed import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) +from sglang.srt.environ import envs 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 ( @@ -26,6 +27,7 @@ from sglang.srt.layers.dp_attention import ( get_attention_tp_rank, get_attention_tp_size, is_allocation_symmetric, + is_dp_attention_enabled, ) from sglang.srt.layers.parameter import BasevLLMParameter from sglang.srt.layers.quantization.base_config import ( @@ -160,6 +162,28 @@ def get_masked_input_and_mask( return input_, ~vocab_mask +def get_embedding_tp_kwargs() -> dict: + """Vocab-parallel layout kwargs for the *input embedding* of models that + support embedding replication (the DeepSeek-V2 target family: DeepSeek + V3.1 / Kimi K2.5, plus their EAGLE3 / NextN drafts). + + EAGLE / NextN share the target's ``embed_tokens.weight`` tensor with the + draft (``set_embed`` / ``set_embed_and_head``), so the target and every + draft that shares it MUST use the same vocab-parallel layout -- otherwise + the draft's masking/index math runs against a tensor with a different + layout and accept_len silently drops. Route all of them through this one + helper so they can never drift. + """ + if envs.SGLANG_ENABLE_EMBED_REPLICATION.get(): + # Replicate the full table on every rank: skips the embed all-reduce + # at the cost of duplicated embedding weights. + return {"enable_tp": False} + # Shard along the vocab dim. Under DP attention each rank owns only its + # local tokens, so reduce within the attention-TP group, not the full TP + # group. + return {"enable_tp": True, "use_attn_tp_group": is_dp_attention_enabled()} + + class VocabParallelEmbedding(torch.nn.Module): """Embedding parallelized in the vocabulary dimension. diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 8b668cf56..517789a8d 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -36,7 +36,6 @@ from sglang.srt.layers.attention.dsa.utils import ( from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, get_attention_cp_size, - is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ReplicatedLinear @@ -55,6 +54,7 @@ from sglang.srt.layers.utils.cp_utils import ( from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, + get_embedding_tp_kwargs, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.deepseek_common.utils import enable_nextn_moe_bf16_cast_to_fp8 @@ -99,8 +99,8 @@ class DeepseekModelNextN(nn.Module): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, - use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), + **get_embedding_tp_kwargs(), ) self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 75da94ee3..b6bb780eb 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -78,7 +78,6 @@ from sglang.srt.layers.dp_attention import ( get_attention_tp_group, get_attention_tp_rank, get_attention_tp_size, - is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( @@ -134,6 +133,7 @@ from sglang.srt.layers.utils.cp_utils import ( from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, + get_embedding_tp_kwargs, ) from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors @@ -2223,7 +2223,7 @@ class DeepseekV2Model(nn.Module): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, - use_attn_tp_group=is_dp_attention_enabled(), + **get_embedding_tp_kwargs(), ) else: self.embed_tokens = PPMissingLayer() diff --git a/python/sglang/srt/models/kimi_k25_eagle3.py b/python/sglang/srt/models/kimi_k25_eagle3.py index 24c8586b0..621f554b4 100644 --- a/python/sglang/srt/models/kimi_k25_eagle3.py +++ b/python/sglang/srt/models/kimi_k25_eagle3.py @@ -33,6 +33,7 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, + get_embedding_tp_kwargs, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader @@ -199,6 +200,7 @@ class Eagle3MLAModel(nn.Module): config.vocab_size, config.hidden_size, prefix=add_prefix("embed_tokens", prefix), + **get_embedding_tp_kwargs(), ) target_hidden_size = (