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

Optimize uneven PP layer distribution logic to improve PP performance (#13977)

Browse Source
This commit is contained in:
Shangming Cai
2025-11-26 22:05:03 +08:00
committed by GitHub
parent 15729dbc8e
commit b704b0a922
1 changed files with 6 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/utils.py
+6 -7
View File
@@ -65,7 +65,7 @@ def get_pp_indices(
) -> Tuple[int, int]: ) -> Tuple[int, int]:
"""Try to evenly distribute layers across partitions. """Try to evenly distribute layers across partitions.
If the number of layers is not divisible by the number of partitions, If the number of layers is not divisible by the number of partitions,
the first N partitions will have one extra layer, where N = remainder. the last N partitions will have one extra layer, where N = remainder.
""" """
# partition_list_str can be set to None in sglang # partition_list_str can be set to None in sglang
partition_list_str = os.getenv("SGLANG_PP_LAYER_PARTITION", None) partition_list_str = os.getenv("SGLANG_PP_LAYER_PARTITION", None)
@@ -85,16 +85,15 @@ def get_pp_indices(
else: else:
base_layers = num_hidden_layers // pp_size base_layers = num_hidden_layers // pp_size
remainder = num_hidden_layers % pp_size remainder = num_hidden_layers % pp_size
# Distribute the extra layers to the first 'remainder' partitions # Distribute the extra layers to the last 'remainder' partitions
if pp_rank < remainder: if pp_rank >= pp_size - remainder:
partitions_without_extra_layer = pp_size - remainder
# This partition gets one extra layer # This partition gets one extra layer
start_layer = pp_rank * (base_layers + 1) start_layer = pp_rank * (base_layers + 1) - partitions_without_extra_layer
end_layer = start_layer + (base_layers + 1) end_layer = start_layer + (base_layers + 1)
else: else:
# This partition gets only base layers # This partition gets only base layers
start_layer = ( start_layer = pp_rank * base_layers
remainder * (base_layers + 1) + (pp_rank - remainder) * base_layers
)
end_layer = start_layer + base_layers end_layer = start_layer + base_layers
return (start_layer, end_layer) return (start_layer, end_layer)