chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Use attn tp group in embedding for more models (#17570)

Browse Source
This commit is contained in:
Ke Bao
2026-01-24 13:37:44 +08:00
committed by GitHub
parent 176da1bbdd
commit fb683be6eb
19 changed files with 19 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/models/bailing_moe.py
View File

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

2
python/sglang/srt/models/bailing_moe_nextn.py
View File

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

2
python/sglang/srt/models/falcon_h1.py
View File

@@ -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):

2
python/sglang/srt/models/glm4.py
View File

@@ -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:

2
python/sglang/srt/models/glm4_moe.py
View File

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

2
python/sglang/srt/models/glm4_moe_lite.py
View File

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

2
python/sglang/srt/models/glm4_moe_nextn.py
View File

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

2
python/sglang/srt/models/gpt_oss.py
View File

@@ -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:

2
python/sglang/srt/models/llada2.py
View File

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

2
python/sglang/srt/models/llama4.py
View File

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

2
python/sglang/srt/models/longcat_flash.py
View File

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

2
python/sglang/srt/models/longcat_flash_nextn.py
View File

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

2
python/sglang/srt/models/mimo_v2_flash.py
View File

@@ -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:

2
python/sglang/srt/models/mimo_v2_flash_nextn.py
View File

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

2
python/sglang/srt/models/qwen2.py
View File

@@ -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

2
python/sglang/srt/models/qwen2_moe.py
View File

@@ -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:

2
python/sglang/srt/models/qwen3_next.py
View File

@@ -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):

2
python/sglang/srt/models/qwen3_vl.py
View File

@@ -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:

2
python/sglang/srt/models/step3_vl.py
View File

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