diff --git a/python/sglang/srt/models/apertus.py b/python/sglang/srt/models/apertus.py index 944cf9062441..ee12ae8f2df5 100644 --- a/python/sglang/srt/models/apertus.py +++ b/python/sglang/srt/models/apertus.py @@ -57,7 +57,7 @@ maybe_remap_kv_scale_name, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers logger = logging.getLogger(__name__) @@ -322,13 +322,11 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: ApertusDecoderLayer( config=config, quant_config=quant_config, layer_id=idx, prefix=prefix ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix="model.layers", ) diff --git a/python/sglang/srt/models/arcee.py b/python/sglang/srt/models/arcee.py index 37bd266a4e94..573c6f3ae3c3 100644 --- a/python/sglang/srt/models/arcee.py +++ b/python/sglang/srt/models/arcee.py @@ -51,7 +51,7 @@ maybe_remap_kv_scale_name, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers logger = logging.getLogger(__name__) @@ -296,13 +296,11 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: ArceeDecoderLayer( config=config, quant_config=quant_config, layer_id=idx, prefix=prefix ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix="model.layers", ) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index dbf9a32812f3..8f2694a28a7d 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -83,7 +83,12 @@ enable_fused_set_kv_buffer, ) from sglang.srt.runtime_context import get_exec, get_parallel, get_stream -from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers +from sglang.srt.utils import ( + add_prefix, + is_cuda, + is_non_idle_and_non_empty, + make_pp_layers, +) LoraConfig = None logger = logging.getLogger(__name__) @@ -780,7 +785,7 @@ def __init__( self.embedding_dropout = torch.nn.Dropout(config.embedding_dropout) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: BailingMoEBlock( layer_id=idx, @@ -789,8 +794,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 762608bec54e..89a788548515 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -78,7 +78,7 @@ is_gfx95_supported, is_hip, is_npu, - make_layers, + make_pp_layers, ) from sglang.srt.utils.common import rank0_log @@ -951,11 +951,9 @@ def layer_fn(idx, prefix): alt_stream=self.alt_stream, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( self.num_layers, layer_fn, - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) diff --git a/python/sglang/srt/models/bailing_moe_v3.py b/python/sglang/srt/models/bailing_moe_v3.py index 34d1bcd39aa6..b4cecd562140 100644 --- a/python/sglang/srt/models/bailing_moe_v3.py +++ b/python/sglang/srt/models/bailing_moe_v3.py @@ -99,7 +99,7 @@ is_cuda, is_flashinfer_available, log_info_on_rank0, - make_layers, + make_pp_layers, ) _is_fp8_fnuz = is_fp8_fnuz() @@ -1297,11 +1297,9 @@ def layer_fn(idx, prefix): alt_stream=self.alt_stream, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( self.num_layers, layer_fn, - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 3550446f432d..9d4d4d5640d6 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -209,7 +209,7 @@ add_prefix, is_non_idle_and_non_empty, is_sm90_supported, - make_layers, + make_pp_layers, use_intel_amx_backend, ) from sglang.srt.utils.custom_op import register_custom_op @@ -2923,7 +2923,7 @@ def __init__( else None ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: DeepseekV2DecoderLayer( config=config, @@ -2933,8 +2933,6 @@ def __init__( alt_stream=self.alt_stream, skip_rope=config.qk_rope_head_dim == 0, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), offloader_kwargs=dict( submodule_accessor=lambda layer: ( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index a8c0fb44beaa..5268254c7741 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -196,7 +196,7 @@ is_gfx942_supported, is_gfx1250_supported, log_info_on_rank0, - make_layers, + make_pp_layers, ) from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -4337,7 +4337,7 @@ def __init__( else None ) self.engram_layout = EngramLayout.from_config(config) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: DeepseekV4DecoderLayer( config=config, @@ -4349,8 +4349,6 @@ def __init__( hc_stats_stream=self.hc_stats_stream, moe_routed_quant_stream=self.moe_routed_quant_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/dots3_common/modeling.py b/python/sglang/srt/models/dots3_common/modeling.py index bc7455878a6f..0a3b61eb15f0 100644 --- a/python/sglang/srt/models/dots3_common/modeling.py +++ b/python/sglang/srt/models/dots3_common/modeling.py @@ -131,7 +131,7 @@ is_cuda, is_non_idle_and_non_empty, log_info_on_rank0, - make_layers, + make_pp_layers, ) _is_cuda = is_cuda() @@ -1718,7 +1718,7 @@ def __init__( self.embed_tokens = PPMissingLayer() self.alt_stream = torch.cuda.Stream() if _is_cuda else None - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Dots3DecoderLayer( config=config, @@ -1727,8 +1727,6 @@ def __init__( prefix=prefix, alt_stream=self.alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py index cdadc9314f45..362ed49bd1ba 100644 --- a/python/sglang/srt/models/ernie45_moe_vl.py +++ b/python/sglang/srt/models/ernie45_moe_vl.py @@ -43,7 +43,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4_5_VLMoeMLP from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.srt.utils.hf_transformers.common import get_rope_config logger = logging.getLogger(__name__) @@ -485,7 +485,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Ernie4_5_VLMoeDecoderLayer( layer_id=idx, @@ -493,8 +493,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/exaone4.py b/python/sglang/srt/models/exaone4.py index 736290257f0f..3f0306fff8f3 100644 --- a/python/sglang/srt/models/exaone4.py +++ b/python/sglang/srt/models/exaone4.py @@ -38,7 +38,7 @@ maybe_remap_kv_scale_name, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.utils import get_exception_traceback, logger @@ -338,7 +338,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Exaone4DecoderLayer( config=config, @@ -346,8 +346,6 @@ def __init__( layer_id=idx, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index 41c1ce289b6e..125a96faa2a1 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/python/sglang/srt/models/exaone_moe.py @@ -59,7 +59,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_exec, get_parallel, get_stream -from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers +from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_pp_layers logger = logging.getLogger(__name__) @@ -537,7 +537,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: ExaoneMoEDecoderLayer( layer_id=idx, @@ -546,8 +546,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index 67b6c87a3cc4..65edfc658d9e 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -58,7 +58,7 @@ create_fused_set_kv_buffer_arg, ) from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers logger = logging.getLogger(__name__) @@ -862,7 +862,7 @@ def __init__( self.per_layer_input_scale = None self.per_layer_projection_scale = None - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Gemma4DecoderLayer( layer_id=idx, @@ -870,8 +870,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) diff --git a/python/sglang/srt/models/gemma4_diffusion.py b/python/sglang/srt/models/gemma4_diffusion.py index 08fc93dd0d45..c6a01f48a0f0 100644 --- a/python/sglang/srt/models/gemma4_diffusion.py +++ b/python/sglang/srt/models/gemma4_diffusion.py @@ -57,7 +57,7 @@ ) from sglang.srt.models.gemma4_vision import Gemma4VisionEncoder from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers logger = logging.getLogger(__name__) @@ -362,7 +362,7 @@ def __init__(self, config, quant_config=None, prefix: str = ""): prefix=add_prefix("self_conditioning", prefix), ) self.pp_group = get_parallel().pp_group - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( text_config.num_hidden_layers, lambda idx, prefix: DiffusionGemmaDecoderLayer( layer_id=idx, @@ -370,8 +370,6 @@ def __init__(self, config, quant_config=None, prefix: str = ""): quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) self.norm = RMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) diff --git a/python/sglang/srt/models/gigachat35.py b/python/sglang/srt/models/gigachat35.py index 5e6c66063e34..8da089590e9e 100644 --- a/python/sglang/srt/models/gigachat35.py +++ b/python/sglang/srt/models/gigachat35.py @@ -40,7 +40,7 @@ ) from sglang.srt.models.qwen3_next import Qwen3GatedDeltaNet from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import BumpAllocator, add_prefix, make_layers +from sglang.srt.utils import BumpAllocator, add_prefix, make_pp_layers _GATED_NORM_LOW_RANK = 16 @@ -525,7 +525,7 @@ def __init__( self.alt_stream = torch.cuda.Stream() if torch.cuda.is_available() else None - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: GigaChat35DecoderLayer( config=config, @@ -534,8 +534,6 @@ def __init__( prefix=prefix, alt_stream=self.alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) diff --git a/python/sglang/srt/models/glm4.py b/python/sglang/srt/models/glm4.py index f1feda9e9ec3..13aac1d6f6b7 100644 --- a/python/sglang/srt/models/glm4.py +++ b/python/sglang/srt/models/glm4.py @@ -47,7 +47,7 @@ kv_cache_scales_loader, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config Glm4Config = None @@ -310,7 +310,7 @@ def __init__( # Use the provided decoder layer type or default to Glm4DecoderLayer decoder_layer_type = decoder_layer_type or Glm4DecoderLayer - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: decoder_layer_type( layer_id=idx, @@ -319,8 +319,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 152d95b8fbb5..e47294bb5183 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -97,7 +97,7 @@ is_non_idle_and_non_empty, is_npu, log_info_on_rank0, - make_layers, + make_pp_layers, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -1039,7 +1039,7 @@ def __init__( self.pp_group.rank_in_group, self.pp_group.world_size, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Glm4MoeDecoderLayer( layer_id=idx, @@ -1049,8 +1049,6 @@ def __init__( prefix=prefix, alt_stream=self.alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index 5b756dd154ff..84927139f958 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -82,7 +82,7 @@ add_prefix, is_non_idle_and_non_empty, is_npu, - make_layers, + make_pp_layers, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -740,7 +740,7 @@ def __init__( self.embed_tokens = PPMissingLayer() self.alt_stream = get_stream("alt") if _is_cuda else None - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Glm4MoeLiteDecoderLayer( config=config, @@ -749,8 +749,6 @@ def __init__( prefix=prefix, alt_stream=self.alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/glm5_next.py b/python/sglang/srt/models/glm5_next.py index 159572f3780d..7bd4dc63897b 100644 --- a/python/sglang/srt/models/glm5_next.py +++ b/python/sglang/srt/models/glm5_next.py @@ -122,7 +122,7 @@ LazyValue, add_prefix, log_info_on_rank0, - make_layers, + make_pp_layers, set_weight_attrs, ) @@ -1008,7 +1008,7 @@ def __init__( else None ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Glm5NextDecoderLayer( config=config, @@ -1017,8 +1017,6 @@ def __init__( prefix=prefix, alt_stream=self.alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 9f06a30ec7bb..a7e8f0a54418 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/python/sglang/srt/models/gpt_oss.py @@ -90,7 +90,7 @@ is_flashinfer_available, is_hip, is_npu, - make_layers, + make_pp_layers, ) from sglang.srt.utils.custom_op import register_custom_op @@ -667,7 +667,7 @@ def __init__( # Use the provided decoder layer type or default to GptOssDecoderLayer decoder_layer_type = decoder_layer_type or GptOssDecoderLayer - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: decoder_layer_type( layer_id=idx, @@ -675,8 +675,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/granitemoehybrid.py b/python/sglang/srt/models/granitemoehybrid.py index a0a7804f441a..e67a4c97f4ce 100644 --- a/python/sglang/srt/models/granitemoehybrid.py +++ b/python/sglang/srt/models/granitemoehybrid.py @@ -26,7 +26,7 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.transformers import maybe_prefix from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import make_layers +from sglang.srt.utils import make_pp_layers from .granitemoe import GraniteMoeMoE, GraniteMoeSharedMLP @@ -347,11 +347,9 @@ def get_layer(idx: int, prefix: str): prefix=prefix, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, get_layer, - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index e0339af864f0..b77edd075b61 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -124,7 +124,7 @@ get_parallel, get_platform, ) -from sglang.srt.utils import is_hip, is_npu, make_layers +from sglang.srt.utils import is_hip, is_npu, make_pp_layers from sglang.srt.utils.common import ( BumpAllocator, add_prefix, @@ -2867,7 +2867,7 @@ def __init__( # Disable on HIP code path. self.alt_streams = None if _is_hip else [torch.cuda.Stream() for _ in range(3)] - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: KimiK3DecoderLayer( layer_idx=idx, @@ -2876,8 +2876,6 @@ def __init__( prefix=prefix, alt_streams=self.alt_streams, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index cc7304fcd732..3078c08ae785 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -54,7 +54,7 @@ from sglang.srt.models.llama import LlamaMLP as KimiMLP from sglang.srt.models.transformers import maybe_prefix from sglang.srt.runtime_context import get_parallel, get_stream -from sglang.srt.utils import is_xpu, make_layers +from sglang.srt.utils import is_xpu, make_pp_layers from sglang.srt.utils.common import BumpAllocator, add_prefix, set_weight_attrs @@ -686,7 +686,7 @@ def __init__( self.alt_stream = None if is_xpu() else get_stream("alt") - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: KimiDecoderLayer( layer_idx=idx, @@ -695,8 +695,6 @@ def __init__( prefix=prefix, alt_stream=self.alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) diff --git a/python/sglang/srt/models/laguna.py b/python/sglang/srt/models/laguna.py index fa641d1ee27b..c2d7a7b1940c 100644 --- a/python/sglang/srt/models/laguna.py +++ b/python/sglang/srt/models/laguna.py @@ -55,7 +55,7 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.utils import apply_qk_norm from sglang.srt.runtime_context import get_exec, get_parallel -from sglang.srt.utils import LazyValue, add_prefix, make_layers +from sglang.srt.utils import LazyValue, add_prefix, make_pp_layers logger = logging.getLogger(__name__) @@ -527,7 +527,7 @@ def __init__( self.embed_tokens = PPMissingLayer() decoder_layer_type = decoder_layer_type or LagunaDecoderLayer - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: decoder_layer_type( layer_id=idx, @@ -535,8 +535,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 5c36c9dc6580..e0b2d9dc3c29 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -80,7 +80,7 @@ is_cuda, is_non_idle_and_non_empty, is_npu, - make_layers, + make_pp_layers, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -709,7 +709,7 @@ def __init__( self.embedding_dropout = torch.nn.Dropout(config.embedding_dropout) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: LLaDA2MoeBlock( layer_id=idx, @@ -718,8 +718,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index 0da5cb20a5fe..47670bfbda99 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -53,7 +53,13 @@ ) from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, is_cuda, is_npu, is_xpu, make_layers +from sglang.srt.utils import ( + add_prefix, + is_cuda, + is_npu, + is_xpu, + make_pp_layers, +) from sglang.utils import get_exception_traceback _is_cuda = is_cuda() @@ -394,7 +400,7 @@ def __init__( self.pp_group.rank_in_group, self.pp_group.world_size, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: LlamaDecoderLayer( config=config, @@ -403,8 +409,6 @@ def __init__( start_layer=pp_start_layer, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix="model.layers", ) diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index 886c363a8d4d..da4a6f963dff 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -86,7 +86,7 @@ add_prefix, ceil_align, is_non_idle_and_non_empty, - make_layers, + make_pp_layers, ) MiMoV2Config = None @@ -994,7 +994,7 @@ def __init__( # Use the provided decoder layer type or default to MiMoV2DecoderLayer decoder_layer_type = decoder_layer_type or MiMoV2DecoderLayer - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, layer_fn=lambda idx, prefix: decoder_layer_type( layer_id=idx, @@ -1002,8 +1002,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index df0f14899361..99e188b46b32 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -100,7 +100,7 @@ is_non_idle_and_non_empty, is_npu, is_xpu, - make_layers, + make_pp_layers, ) from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -1110,11 +1110,9 @@ def layer_fn(idx, prefix: str) -> nn.Module: prefix=prefix, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, layer_fn, - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index 1d7ff25a39a2..929b9da8e5b4 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/python/sglang/srt/models/minimax_m3.py @@ -99,7 +99,7 @@ is_hip, is_npu, log_info_on_rank0, - make_layers, + make_pp_layers, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -1443,11 +1443,9 @@ def layer_fn(idx, prefix: str) -> nn.Module: prefix=prefix, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, layer_fn, - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/ministral3.py b/python/sglang/srt/models/ministral3.py index 29137099dcd2..60e29d66449a 100644 --- a/python/sglang/srt/models/ministral3.py +++ b/python/sglang/srt/models/ministral3.py @@ -11,7 +11,7 @@ LlamaForCausalLM, LlamaModel, ) -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers def _get_llama_4_attn_scale( @@ -141,7 +141,7 @@ def __init__( # Override layer creation to use Ministral3Attention super().__init__(config=config, quant_config=quant_config, prefix=prefix) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Ministral3DecoderLayer( config=config, @@ -150,8 +150,6 @@ def __init__( start_layer=self.start_layer, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix="model.layers", ) diff --git a/python/sglang/srt/models/mixtral.py b/python/sglang/srt/models/mixtral.py index 929c13511d84..2b759e98332e 100644 --- a/python/sglang/srt/models/mixtral.py +++ b/python/sglang/srt/models/mixtral.py @@ -54,7 +54,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers logger = logging.getLogger(__name__) @@ -294,13 +294,11 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: MixtralDecoderLayer( config=config, quant_config=quant_config, layer_id=idx, prefix=prefix ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix="layers", return_tuple=True, ) diff --git a/python/sglang/srt/models/nanbeige.py b/python/sglang/srt/models/nanbeige.py index 98648c2d603f..169c5883a234 100644 --- a/python/sglang/srt/models/nanbeige.py +++ b/python/sglang/srt/models/nanbeige.py @@ -33,7 +33,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers logger = logging.getLogger(__name__) @@ -296,7 +296,7 @@ def __init__( # Use the provided decoder layer type or default to NanbeigeDecoderLayer decoder_layer_type = decoder_layer_type or NanbeigeDecoderLayer - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: decoder_layer_type( layer_id=idx, @@ -305,8 +305,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index abb75a27c040..cfaf75c5634d 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -84,7 +84,7 @@ add_prefix, get_current_device_stream_fast, is_cuda, - make_layers, + make_pp_layers, ) from sglang.srt.utils.custom_op import register_custom_op from sglang.utils import logger @@ -701,11 +701,9 @@ def get_layer(idx: int, prefix: str): layer_class = ALL_DECODER_LAYER_TYPES[config.hybrid_override_pattern[idx]] return layer_class(config, idx, quant_config=quant_config, prefix=prefix) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( len(config.hybrid_override_pattern), get_layer, - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/nemotron_nas.py b/python/sglang/srt/models/nemotron_nas.py index aa4e86eb2be7..8e59a70ac81d 100644 --- a/python/sglang/srt/models/nemotron_nas.py +++ b/python/sglang/srt/models/nemotron_nas.py @@ -40,7 +40,7 @@ ) from sglang.srt.models.llama import LlamaAttention, LlamaMLP from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.utils import logger @@ -197,11 +197,9 @@ def get_layer(idx: int, prefix: str): prefix=prefix, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, get_layer, - pp_rank=get_parallel().pp_group.rank_in_group, - pp_size=get_parallel().pp_group.world_size, prefix=add_prefix("layers", prefix), ) if get_parallel().pp_group.is_last_rank: diff --git a/python/sglang/srt/models/opt.py b/python/sglang/srt/models/opt.py index ee09deec8917..01794ae20e43 100644 --- a/python/sglang/srt/models/opt.py +++ b/python/sglang/srt/models/opt.py @@ -43,7 +43,7 @@ kv_cache_scales_loader, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.utils import get_exception_traceback logger = logging.getLogger(__name__) @@ -273,13 +273,11 @@ def __init__( else: self.final_layer_norm = None - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: OPTDecoderLayer( config=config, layer_id=idx, quant_config=quant_config, prefix=prefix ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix="model.layers", ) diff --git a/python/sglang/srt/models/orion.py b/python/sglang/srt/models/orion.py index 95df14375805..349af0b9cbf9 100644 --- a/python/sglang/srt/models/orion.py +++ b/python/sglang/srt/models/orion.py @@ -33,7 +33,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -231,13 +231,11 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: OrionDecoderLayer( config, layer_id=idx, quant_config=quant_config, prefix=prefix ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) diff --git a/python/sglang/srt/models/persimmon.py b/python/sglang/srt/models/persimmon.py index f62a6f28103c..082e42fd2db0 100644 --- a/python/sglang/srt/models/persimmon.py +++ b/python/sglang/srt/models/persimmon.py @@ -23,7 +23,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers class PersimmonMLP(nn.Module): @@ -209,14 +209,12 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: PersimmonDecoderLayer( config, quant_config=quant_config, prefix=prefix, idx=idx ), prefix="model.layers", - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/phi3_small.py b/python/sglang/srt/models/phi3_small.py index ca8b2589421f..b7a3e57f57b1 100644 --- a/python/sglang/srt/models/phi3_small.py +++ b/python/sglang/srt/models/phi3_small.py @@ -25,7 +25,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers @torch.jit.script @@ -308,7 +308,7 @@ def __init__( prefix=add_prefix("embed_tokens", prefix), ) self.mup_embedding_multiplier = config.mup_embedding_multiplier - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Phi3SmallDecoderLayer( config, @@ -316,8 +316,6 @@ def __init__( quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 793b326e4c87..7cc408f9242d 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -55,7 +55,7 @@ ) from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_exec, get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config Qwen2Config = None @@ -353,7 +353,7 @@ def __init__( self.pp_group.rank_in_group, self.pp_group.world_size, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: decoder_layer_type( layer_id=idx, @@ -363,8 +363,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 10284a362079..8ce079a33c63 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -107,7 +107,7 @@ is_cuda, is_hip, is_npu, - make_layers, + make_pp_layers, use_intel_amx_backend, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -1130,7 +1130,7 @@ def __init__( self.pp_group.rank_in_group, self.pp_group.world_size, ) - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: decoder_layer_type( layer_id=idx, @@ -1140,8 +1140,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index c1d17ee1221a..6523472fd2bb 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -134,7 +134,7 @@ is_hip, is_npu, is_xpu, - make_layers, + make_pp_layers, set_weight_attrs, use_intel_amx_backend, ) @@ -1691,11 +1691,9 @@ def get_layer(idx: int, prefix: str): is_nextn=is_nextn, ) - self.layers, self._start_layer, self._end_layer = make_layers( + self.layers, self._start_layer, self._end_layer = make_pp_layers( config.num_hidden_layers, get_layer, - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index 7c6b47b16ccb..bb819d5e29ac 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -72,7 +72,7 @@ bind_or_assign, is_cuda, is_nvidia_cublas_version_ge_12_9, - make_layers, + make_pp_layers, next_power_of_2, ) @@ -1111,7 +1111,7 @@ def __init__( else: self.embed_tokens = nn.Identity() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: SarvamMoEMLADecoderLayer( config=config, @@ -1120,8 +1120,6 @@ def __init__( prefix=prefix, alt_stream=self.alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix="model.layers", ) diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py index 7f2c44d08993..c51ad528baf4 100644 --- a/python/sglang/srt/models/sdar.py +++ b/python/sglang/srt/models/sdar.py @@ -46,7 +46,7 @@ enable_fused_set_kv_buffer, ) from sglang.srt.runtime_context import get_exec, get_parallel, get_stream -from sglang.srt.utils import add_prefix, is_cuda, make_layers +from sglang.srt.utils import add_prefix, is_cuda, make_pp_layers logger = logging.getLogger(__name__) _is_cuda = is_cuda() @@ -357,7 +357,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: SDARBlock( layer_id=idx, @@ -366,8 +366,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py index c3ad34905a7d..3af99a83e397 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -59,7 +59,7 @@ enable_fused_set_kv_buffer, ) from sglang.srt.runtime_context import get_exec, get_parallel, get_stream -from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers +from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_pp_layers logger = logging.getLogger(__name__) _is_cuda = is_cuda() @@ -422,7 +422,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: SDARMoeBlock( config=config, @@ -431,8 +431,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) diff --git a/python/sglang/srt/models/solar.py b/python/sglang/srt/models/solar.py index e9d915e07862..b5d20e08ff29 100644 --- a/python/sglang/srt/models/solar.py +++ b/python/sglang/srt/models/solar.py @@ -54,7 +54,7 @@ # limitations under the License. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/solar.py from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -302,7 +302,7 @@ def __init__( ) else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: SolarDecoderLayer( config=config, @@ -310,8 +310,6 @@ def __init__( layer_id=idx, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) _check_skips_stay_in_stage(config, self.start_layer, self.end_layer) diff --git a/python/sglang/srt/models/spark2_5.py b/python/sglang/srt/models/spark2_5.py index b8012f5f8928..9460da4ce3db 100644 --- a/python/sglang/srt/models/spark2_5.py +++ b/python/sglang/srt/models/spark2_5.py @@ -34,7 +34,7 @@ ) from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers Spark2_5Config = None @@ -322,7 +322,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: Spark2_5DecoderLayer( layer_id=idx, @@ -330,8 +330,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py index f4ab59d4b6ea..9879fafdd37b 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -49,7 +49,12 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_exec, get_parallel, get_stream -from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers +from sglang.srt.utils import ( + add_prefix, + is_cuda, + is_non_idle_and_non_empty, + make_pp_layers, +) Step3p5Config = None @@ -663,7 +668,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, # 1, lambda idx, prefix: Step3p5DecoderLayer( @@ -673,8 +678,6 @@ def __init__( prefix=prefix, alt_stream=alt_stream, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py index 8ff619b703f9..17ba89b54c9e 100644 --- a/python/sglang/srt/models/xllm.py +++ b/python/sglang/srt/models/xllm.py @@ -84,7 +84,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_exec, get_parallel -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, make_pp_layers _XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY = "xllm_source_router_gemm_partitions" _XLLM_SOURCE_ROUTER_PARTITIONS_MISSING = object() @@ -1592,7 +1592,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: XllmDecoderLayer( layer_id=idx, @@ -1600,8 +1600,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/zaya.py b/python/sglang/srt/models/zaya.py index 8b50517d5234..539828ef7f3d 100644 --- a/python/sglang/srt/models/zaya.py +++ b/python/sglang/srt/models/zaya.py @@ -76,7 +76,7 @@ from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs +from sglang.srt.utils import add_prefix, make_pp_layers, set_weight_attrs logger = logging.getLogger(__name__) @@ -1395,7 +1395,7 @@ def __init__( else: self.embed_tokens = PPMissingLayer() - self.layers, self.start_layer, self.end_layer = make_layers( + self.layers, self.start_layer, self.end_layer = make_pp_layers( config.num_hidden_layers, lambda idx, prefix: _build_layer( layer_id=idx, @@ -1403,8 +1403,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index f4e38cf9e896..e2de7af75b41 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1544,6 +1544,29 @@ def make_layers( return modules, start_layer, end_layer +def make_pp_layers( + num_hidden_layers: int, + layer_fn: LayerFn, + prefix: str = "", + return_tuple: bool = False, + offloader_kwargs: Optional[Dict[str, Any]] = None, +) -> Tuple[torch.nn.Module, int, int]: + """Make this pipeline stage's layers, and return them with the stage's range. + + Layers outside ``[start_layer, end_layer)`` are ``PPMissingLayer`` stand-ins. + """ + parallel = get_parallel() + return make_layers( + num_hidden_layers, + layer_fn, + pp_rank=parallel.pp_rank, + pp_size=parallel.pp_size, + prefix=prefix, + return_tuple=return_tuple, + offloader_kwargs=offloader_kwargs, + ) + + def set_random_seed(seed: int) -> None: """Set the random seed for all libraries.""" random.seed(seed) diff --git a/test/registered/unit/utils/test_make_pp_layers.py b/test/registered/unit/utils/test_make_pp_layers.py new file mode 100644 index 000000000000..5526bbdb943b --- /dev/null +++ b/test/registered/unit/utils/test_make_pp_layers.py @@ -0,0 +1,48 @@ +import unittest + +from torch import nn + +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils.common import make_pp_layers +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, published_topology + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +NUM_LAYERS = 6 + + +def _layer(idx, prefix): + return nn.Linear(1, 1) + + +def _stage(): + """Build the layers here; return the stage range and the indices built.""" + layers, start, end = make_pp_layers(NUM_LAYERS, _layer, prefix="model.layers") + built = [ + i for i, layer in enumerate(layers) if not isinstance(layer, PPMissingLayer) + ] + return (start, end), built + + +class TestMakePPLayers(CustomTestCase): + def test_each_stage_builds_its_own_slice(self): + for pp_rank, stage in ((0, (0, 3)), (1, (3, 6))): + with ( + self.subTest(pp_rank=pp_rank), + published_topology(pp_size=2, ranks=dict(world_rank=pp_rank)), + ): + self.assertEqual(_stage(), (stage, list(range(*stage)))) + + def test_a_single_stage_scope_builds_every_layer(self): + # A draft is built with its pipeline narrowed to one stage. + with published_topology(pp_size=2, ranks=dict(world_rank=1)): + with get_parallel().override(pp_size=1, pp_rank=0): + stage, built = _stage() + self.assertEqual(stage, (0, NUM_LAYERS)) + self.assertEqual(built, list(range(NUM_LAYERS))) + + +if __name__ == "__main__": + unittest.main()