From fb683be6eb7c6960f3d37ea91ca836c1b8e4120c Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sat, 24 Jan 2026 13:37:44 +0800 Subject: [PATCH] Use attn tp group in embedding for more models (#17570) --- python/sglang/srt/models/bailing_moe.py | 2 +- python/sglang/srt/models/bailing_moe_nextn.py | 2 +- python/sglang/srt/models/falcon_h1.py | 2 +- python/sglang/srt/models/glm4.py | 2 +- python/sglang/srt/models/glm4_moe.py | 2 +- python/sglang/srt/models/glm4_moe_lite.py | 2 +- python/sglang/srt/models/glm4_moe_nextn.py | 2 +- python/sglang/srt/models/gpt_oss.py | 2 +- python/sglang/srt/models/llada2.py | 2 +- python/sglang/srt/models/llama4.py | 2 +- python/sglang/srt/models/longcat_flash.py | 2 +- python/sglang/srt/models/longcat_flash_nextn.py | 2 +- python/sglang/srt/models/mimo_v2_flash.py | 2 +- python/sglang/srt/models/mimo_v2_flash_nextn.py | 2 +- python/sglang/srt/models/qwen2.py | 2 +- python/sglang/srt/models/qwen2_moe.py | 2 +- python/sglang/srt/models/qwen3_next.py | 2 +- python/sglang/srt/models/qwen3_vl.py | 2 +- python/sglang/srt/models/step3_vl.py | 2 +- 19 files changed, 19 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 3a3a20883..fd637492c 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -717,7 +717,7 @@ class BailingMoEModel(nn.Module): self.embed_dim, quant_config=quant_config, prefix=add_prefix("word_embeddings", prefix), - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), ) else: self.word_embeddings = PPMissingLayer() diff --git a/python/sglang/srt/models/bailing_moe_nextn.py b/python/sglang/srt/models/bailing_moe_nextn.py index 76a24f4c9..437ba098c 100644 --- a/python/sglang/srt/models/bailing_moe_nextn.py +++ b/python/sglang/srt/models/bailing_moe_nextn.py @@ -62,7 +62,7 @@ class BailingMoEModelNextN(nn.Module): self.word_embeddings = 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("word_embeddings", prefix), ) diff --git a/python/sglang/srt/models/falcon_h1.py b/python/sglang/srt/models/falcon_h1.py index 551307cb8..628f99c6e 100644 --- a/python/sglang/srt/models/falcon_h1.py +++ b/python/sglang/srt/models/falcon_h1.py @@ -394,7 +394,7 @@ class FalconH1Model(nn.Module): config.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), ) def get_layer(idx: int, prefix: str): diff --git a/python/sglang/srt/models/glm4.py b/python/sglang/srt/models/glm4.py index 75c4ccc69..f5d1c17f8 100644 --- a/python/sglang/srt/models/glm4.py +++ b/python/sglang/srt/models/glm4.py @@ -307,7 +307,7 @@ class Glm4Model(nn.Module): config.vocab_size, config.hidden_size, quant_config=quant_config, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), ) else: diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 784f9c30d..55036ca25 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -895,7 +895,7 @@ class Glm4MoeModel(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() diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index 0490ac47e..e8e0d054b 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -440,7 +440,7 @@ class Glm4MoeLiteModel(DeepseekV2Model): 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() diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py index 5e0c3ac59..1f6e75364 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -58,7 +58,7 @@ class Glm4MoeModelNextN(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/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 451fe2a4f..74ce2e3eb 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/python/sglang/srt/models/gpt_oss.py @@ -499,7 +499,7 @@ class GptOssModel(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), ) else: diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 8cf8902d0..7094be5c5 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -674,7 +674,7 @@ class LLaDA2MoeModel(nn.Module): self.embed_dim, quant_config=quant_config, prefix=add_prefix("word_embeddings", prefix), - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), ) else: self.word_embeddings = PPMissingLayer() diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index 4a4e309bf..46253a0ad 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -486,7 +486,7 @@ class Llama4Model(nn.Module): config.hidden_size, quant_config=quant_config, prefix=add_prefix("embed_tokens", prefix), - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), ) self.layers = make_layers( config.num_hidden_layers, diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 7ef5b4ae1..c16a93797 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -503,7 +503,7 @@ class LongcatFlashModel(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(), ) self.alt_stream = torch.cuda.Stream() diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index f43946600..12c9cb13f 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -213,7 +213,7 @@ class LongcatFlashModelNextN(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/mimo_v2_flash.py b/python/sglang/srt/models/mimo_v2_flash.py index 12a387a96..8a45b3e93 100644 --- a/python/sglang/srt/models/mimo_v2_flash.py +++ b/python/sglang/srt/models/mimo_v2_flash.py @@ -635,7 +635,7 @@ class MiMoV2Model(nn.Module): config.vocab_size, config.hidden_size, quant_config=quant_config, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), ) else: diff --git a/python/sglang/srt/models/mimo_v2_flash_nextn.py b/python/sglang/srt/models/mimo_v2_flash_nextn.py index 6440c3ef1..2408f950d 100644 --- a/python/sglang/srt/models/mimo_v2_flash_nextn.py +++ b/python/sglang/srt/models/mimo_v2_flash_nextn.py @@ -165,7 +165,7 @@ class MiMoV2ModelNextN(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/qwen2.py b/python/sglang/srt/models/qwen2.py index a7dbadec6..30b02ec83 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -277,7 +277,7 @@ class Qwen2Model(nn.Module): config.vocab_size, config.hidden_size, quant_config=quant_config, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), params_dtype=( torch.float32 diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 1e4b53a7d..5e235cfa1 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -570,7 +570,7 @@ class Qwen2MoeModel(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), ) else: diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 6512afa08..017c98a2c 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -793,7 +793,7 @@ class Qwen3NextModel(nn.Module): config.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), ) def get_layer(idx: int, prefix: str): diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 93661fc78..c397b136d 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -307,7 +307,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): self.num_position_embeddings, self.hidden_size, quant_config=quant_config, - enable_tp=not is_dp_attention_enabled(), + use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("pos_embed", prefix), ) else: diff --git a/python/sglang/srt/models/step3_vl.py b/python/sglang/srt/models/step3_vl.py index ad8fb6f9c..5ac9528f9 100644 --- a/python/sglang/srt/models/step3_vl.py +++ b/python/sglang/srt/models/step3_vl.py @@ -443,7 +443,7 @@ class Step3TextModel(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), )