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

[Refactor] Pass PP start_layer via model constructor instead of forward_batch.token_to_kv_pool (#25825)

Browse Source 
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-19 22:16:07 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent a4b51d35ef
commit 052abcc0dd
9 changed files with 59 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/glm4_moe.py
+12 -1
View File
@@ -28,6 +28,7 @@ from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import (
get_moe_expert_parallel_world_size,
get_pp_group,
get_pp_indices,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
parallel_state,
@@ -187,6 +188,7 @@ class Glm4MoeAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
layer_id: int = 0,
start_layer: int = 0,
rope_theta: float = 1000000,
partial_rotary_factor: float = 0.5,
rope_scaling: Optional[Dict[str, Any]] = None,
@@ -201,6 +203,7 @@ class Glm4MoeAttention(nn.Module):
) -> None:
super().__init__()
self.hidden_size = hidden_size
self.start_layer = start_layer
attn_tp_rank = get_attention_tp_rank()
attn_tp_size = get_attention_tp_size()
@@ -312,7 +315,7 @@ class Glm4MoeAttention(nn.Module):
)
q, k = self.rotary_emb(positions, q, k)
else:
if self.attn.layer_id == forward_batch.token_to_kv_pool.start_layer:
if self.attn.layer_id == self.start_layer:
self.rotary_emb.get_cos_sin_with_position(positions)
if self.use_qk_norm:
eps = self.q_norm.variance_epsilon
@@ -788,6 +791,7 @@ class Glm4MoeDecoderLayer(nn.Module):
self,
config: PretrainedConfig,
layer_id: int,
start_layer: int = 0,
quant_config: Optional[QuantizationConfig] = None,
is_nextn: bool = False,
prefix: str = "",
@@ -815,6 +819,7 @@ class Glm4MoeDecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
layer_id=layer_id,
start_layer=start_layer,
rope_theta=rope_theta,
rope_scaling=rope_scaling,
partial_rotary_factor=partial_rotary_factor,
@@ -1073,10 +1078,16 @@ class Glm4MoeModel(nn.Module):
self.embed_tokens = PPMissingLayer()
self.alt_stream = torch.cuda.Stream() if _is_cuda else None
pp_start_layer, _ = get_pp_indices(
config.num_hidden_layers,
self.pp_group.rank_in_group,
self.pp_group.world_size,
)
self.layers, self.start_layer, self.end_layer = make_layers(
config.num_hidden_layers,
lambda idx, prefix: Glm4MoeDecoderLayer(
layer_id=idx,
start_layer=pp_start_layer,
config=config,
quant_config=quant_config,
prefix=prefix,
python/sglang/srt/models/llama.py
+16 -2
View File
@@ -27,6 +27,7 @@ from transformers import LlamaConfig
from sglang.srt.distributed import (
get_pp_group,
get_pp_indices,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
@@ -131,6 +132,7 @@ class LlamaAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
layer_id: int = 0,
start_layer: int = 0,
rope_theta: float = 10000,
rope_scaling: Optional[Dict[str, Any]] = None,
rope_is_neox_style: bool = True,
@@ -141,6 +143,7 @@ class LlamaAttention(nn.Module):
) -> None:
super().__init__()
self.hidden_size = hidden_size
self.start_layer = start_layer
tp_size = get_tensor_model_parallel_world_size()
self.total_num_heads = num_heads
assert self.total_num_heads % tp_size == 0
@@ -210,7 +213,7 @@ class LlamaAttention(nn.Module):
def forward_prepare_npu(self, positions, hidden_states, forward_batch):
qkv, _ = self.qkv_proj(hidden_states)
if self.attn.layer_id == forward_batch.token_to_kv_pool.start_layer:
if self.attn.layer_id == self.start_layer:
self.rotary_emb.get_cos_sin_with_position(positions)
q, k, v = split_qkv_rmsnorm_rope(
qkv,
@@ -254,6 +257,7 @@ class LlamaDecoderLayer(nn.Module):
self,
config: LlamaConfig,
layer_id: int = 0,
start_layer: int = 0,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
@@ -285,6 +289,7 @@ class LlamaDecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
layer_id=layer_id,
start_layer=start_layer,
rope_theta=rope_theta,
rope_scaling=rope_scaling,
rope_is_neox_style=rope_is_neox_style,
@@ -352,10 +357,19 @@ class LlamaModel(nn.Module):
else:
self.embed_tokens = PPMissingLayer()
pp_start_layer, _ = get_pp_indices(
config.num_hidden_layers,
self.pp_group.rank_in_group,
self.pp_group.world_size,
)
self.layers, self.start_layer, self.end_layer = make_layers(
config.num_hidden_layers,
lambda idx, prefix: LlamaDecoderLayer(
config=config, quant_config=quant_config, layer_id=idx, prefix=prefix
config=config,
quant_config=quant_config,
layer_id=idx,
start_layer=pp_start_layer,
prefix=prefix,
),
pp_rank=self.pp_group.rank_in_group,
pp_size=self.pp_group.world_size,
python/sglang/srt/models/llama_eagle.py
+1 -1
View File
@@ -44,7 +44,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__(config, layer_id, quant_config, prefix)
super().__init__(config, layer_id, quant_config=quant_config, prefix=prefix)
# Skip the input_layernorm
# https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
python/sglang/srt/models/llama_eagle3.py
+1 -1
View File
@@ -49,7 +49,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__(config, layer_id, quant_config, prefix)
super().__init__(config, layer_id, quant_config=quant_config, prefix=prefix)
# Input layer concats embeds + target_hidden before qkv (input dim 2x).
self.is_input_layer = layer_id == 0
python/sglang/srt/models/qwen2.py
+9
View File
@@ -24,6 +24,7 @@ from torch import nn
from sglang.srt.distributed import (
get_pp_group,
get_pp_indices,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
@@ -200,12 +201,14 @@ class Qwen2DecoderLayer(nn.Module):
self,
config: Qwen2Config,
layer_id: int = 0,
start_layer: int = 0,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
self.start_layer = start_layer
rope_theta, rope_scaling = get_rope_config(config)
max_position_embeddings = getattr(config, "max_position_embeddings", 32768)
head_dim = getattr(config, "head_dim", None)
@@ -295,10 +298,16 @@ class Qwen2Model(nn.Module):
# Use the provided decoder layer type or default to Qwen2DecoderLayer
decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
pp_start_layer, _ = get_pp_indices(
config.num_hidden_layers,
self.pp_group.rank_in_group,
self.pp_group.world_size,
)
self.layers, self.start_layer, self.end_layer = make_layers(
config.num_hidden_layers,
lambda idx, prefix: decoder_layer_type(
layer_id=idx,
start_layer=pp_start_layer,
config=config,
quant_config=quant_config,
prefix=prefix,
python/sglang/srt/models/qwen2_eagle.py
+1 -1
View File
@@ -45,7 +45,7 @@ class Qwen2DecoderLayer(Qwen2DecoderLayer):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__(config, layer_id, quant_config, prefix=prefix)
super().__init__(config, layer_id, quant_config=quant_config, prefix=prefix)
# Skip the input_layernorm
# https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
python/sglang/srt/models/qwen2_moe.py
+9
View File
@@ -32,6 +32,7 @@ from sglang.srt.distributed import (
get_moe_data_parallel_world_size,
get_moe_expert_parallel_world_size,
get_pp_group,
get_pp_indices,
get_tensor_model_parallel_world_size,
tensor_model_parallel_all_reduce,
)
@@ -600,6 +601,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
self,
config: PretrainedConfig,
layer_id: int,
start_layer: int = 0,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
@@ -607,6 +609,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.start_layer = start_layer
rope_theta, rope_scaling = get_rope_config(config)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
qkv_bias = getattr(config, "qkv_bias", True)
@@ -747,10 +750,16 @@ class Qwen2MoeModel(nn.Module):
# Use the provided decoder layer type or default to Qwen2MoeDecoderLayer
decoder_layer_type = decoder_layer_type or Qwen2MoeDecoderLayer
pp_start_layer, _ = get_pp_indices(
config.num_hidden_layers,
self.pp_group.rank_in_group,
self.pp_group.world_size,
)
self.layers, self.start_layer, self.end_layer = make_layers(
config.num_hidden_layers,
lambda idx, prefix: decoder_layer_type(
layer_id=idx,
start_layer=pp_start_layer,
config=config,
quant_config=quant_config,
prefix=prefix,
python/sglang/srt/models/qwen3.py
+5 -1
View File
@@ -64,6 +64,7 @@ class Qwen3Attention(nn.Module):
num_heads: int,
num_kv_heads: int,
layer_id: int = 0,
start_layer: int = 0,
rope_theta: float = 1000000,
rope_scaling: Optional[Dict[str, Any]] = None,
head_dim: Optional[int] = None,
@@ -76,6 +77,7 @@ class Qwen3Attention(nn.Module):
) -> None:
super().__init__()
self.hidden_size = hidden_size
self.start_layer = start_layer
self.tp_size = get_tensor_model_parallel_world_size()
self.total_num_heads = num_heads
attn_tp_rank = get_attention_tp_rank()
@@ -181,7 +183,7 @@ class Qwen3Attention(nn.Module):
def forward_prepare_npu(self, positions, hidden_states, forward_batch):
qkv, _ = self.qkv_proj(hidden_states)
if self.attn.layer_id == forward_batch.token_to_kv_pool.start_layer:
if self.attn.layer_id == self.start_layer:
self.rotary_emb.get_cos_sin_with_position(positions)
q, k, v = split_qkv_rmsnorm_rope(
qkv,
@@ -310,6 +312,7 @@ class Qwen3DecoderLayer(nn.Module):
self,
config: Qwen3Config,
layer_id: int = 0,
start_layer: int = 0,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
@@ -333,6 +336,7 @@ class Qwen3DecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
layer_id=layer_id,
start_layer=start_layer,
rope_theta=rope_theta,
rope_scaling=rope_scaling,
head_dim=head_dim,
python/sglang/srt/models/qwen3_moe.py
+5 -1
View File
@@ -445,6 +445,7 @@ class Qwen3MoeAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
layer_id: int = 0,
start_layer: int = 0,
rope_theta: float = 10000,
rope_scaling: Optional[Dict[str, Any]] = None,
max_position_embeddings: int = 8192,
@@ -459,6 +460,7 @@ class Qwen3MoeAttention(nn.Module):
) -> None:
super().__init__()
self.hidden_size = hidden_size
self.start_layer = start_layer
attn_tp_rank = get_attention_tp_rank()
attn_tp_size = get_attention_tp_size()
@@ -568,7 +570,7 @@ class Qwen3MoeAttention(nn.Module):
forward_batch: ForwardBatch,
):
qkv, _ = self.qkv_proj(hidden_states)
if self.attn.layer_id == forward_batch.token_to_kv_pool.start_layer:
if self.attn.layer_id == self.start_layer:
self.rotary_emb.get_cos_sin_with_position(positions)
q, k, v = split_qkv_rmsnorm_rope(
qkv,
@@ -721,6 +723,7 @@ class Qwen3MoeDecoderLayer(nn.Module):
self,
config: Qwen3MoeConfig,
layer_id: int,
start_layer: int = 0,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
@@ -744,6 +747,7 @@ class Qwen3MoeDecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
layer_id=layer_id,
start_layer=start_layer,
rope_theta=rope_theta,
rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,