A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
Support sequence parallelism combined with pipeline parallelism (#18243)
Signed-off-by: cascade812 <cascade812@outlook.com>
cascade committed
May 17, 2025 at 15:47 UTC
9ab2c02ff8ef05cb6c4b2bdc9a4faaee61892450
3 files changed
+74
-27
tests/distributed/test_sequence_parallel.py
+35
-2
@@ -26,6 +26,7 @@ VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
26
27
class ParallelSetup(NamedTuple):
28
tp_size: int
29
+ pp_size: int
30
sp_enabled: bool
31
eager_mode: bool
32
chunked_prefill: bool
@@ -60,6 +61,7 @@ class SPTestSettings:
61
def detailed(
62
*,
63
tp_base: int = 2,
64
+ pp_base: int = 1,
65
multi_node_only: bool = False,
66
task: TaskOption = "auto",
67
load_format: Optional[str] = None,
@@ -67,18 +69,42 @@ class SPTestSettings:
69
return SPTestSettings(
70
parallel_setups=[
71
ParallelSetup(tp_size=tp_base,
72
+ pp_size=pp_base,
73
sp_enabled=True,
74
eager_mode=False,
75
chunked_prefill=False),
76
ParallelSetup(tp_size=tp_base,
77
+ pp_size=pp_base,
78
sp_enabled=True,
79
eager_mode=False,
80
chunked_prefill=True),
81
ParallelSetup(tp_size=tp_base,
82
+ pp_size=pp_base,
83
sp_enabled=True,
84
eager_mode=True,
85
chunked_prefill=False),
86
ParallelSetup(tp_size=tp_base,
87
+ pp_size=pp_base,
88
+ sp_enabled=True,
89
+ eager_mode=True,
90
+ chunked_prefill=True),
91
+ ParallelSetup(tp_size=tp_base,
92
+ pp_size=2 * pp_base,
93
+ sp_enabled=True,
94
+ eager_mode=False,
95
+ chunked_prefill=False),
96
+ ParallelSetup(tp_size=tp_base,
97
+ pp_size=2 * pp_base,
98
+ sp_enabled=True,
99
+ eager_mode=False,
100
+ chunked_prefill=True),
101
+ ParallelSetup(tp_size=tp_base,
102
+ pp_size=2 * pp_base,
103
+ sp_enabled=True,
104
+ eager_mode=True,
105
+ chunked_prefill=False),
106
+ ParallelSetup(tp_size=tp_base,
107
+ pp_size=2 * pp_base,
108
sp_enabled=True,
109
eager_mode=True,
110
chunked_prefill=True)
@@ -94,6 +120,7 @@ class SPTestSettings:
120
def fast(
121
*,
122
tp_base: int = 2,
123
+ pp_base: int = 1,
124
task: TaskOption = "auto",
125
multi_node_only: bool = False,
126
load_format: Optional[str] = None,
@@ -101,6 +128,12 @@ class SPTestSettings:
128
return SPTestSettings(
129
parallel_setups=[
130
ParallelSetup(tp_size=tp_base,
131
+ pp_size=pp_base,
132
+ sp_enabled=True,
133
+ eager_mode=False,
134
+ chunked_prefill=False),
135
+ ParallelSetup(tp_size=tp_base,
136
+ pp_size=2 * pp_base,
137
sp_enabled=True,
138
eager_mode=False,
139
chunked_prefill=False),
@@ -136,6 +169,7 @@ def _compare_sp(
169
):
170
(
171
tp_size,
172
+ pp_size,
173
sp_enabled,
174
eager_mode,
175
chunked_prefill,
@@ -167,7 +201,6 @@ def _compare_sp(
201
else:
202
model_info.check_available_online(on_fail="skip")
203
170
- pp_size = 1
204
if num_gpus_available < tp_size * pp_size:
205
pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
206
if VLLM_MULTI_NODE and distributed_backend == "mp":
@@ -256,7 +289,7 @@ def _compare_sp(
289
290
SP_TEXT_GENERATION_MODELS = {
291
# [Decoder-only]
259
- "meta-llama/Llama-3.2-1B-Instruct": SPTestSettings.detailed(),
292
+ "meta-llama/Llama-3.2-1B-Instruct": SPTestSettings.fast(),
293
}
294
295
SP_TEST_MODELS = [
vllm/config.py
-12
@@ -4287,18 +4287,6 @@ class VllmConfig:
4287
self.compilation_config.level = CompilationLevel.PIECEWISE
4288
self.compilation_config.set_splitting_ops_for_v1()
4289
4290
- if self.parallel_config is not None and \
4291
- self.parallel_config.tensor_parallel_size > 1 and \
4292
- self.parallel_config.pipeline_parallel_size > 1 and \
4293
- self.compilation_config is not None and \
4294
- self.compilation_config.pass_config is not None and \
4295
- self.compilation_config.pass_config.enable_sequence_parallelism:
4296
- logger.warning_once(
4297
- "Sequence parallelism is not supported with pipeline "
4298
- "parallelism. Disabling sequence parallelism.")
4299
- self.compilation_config.pass_config.\
4300
- enable_sequence_parallelism = False
4301
-
4290
self._set_cudagraph_sizes()
4291
4292
if self.cache_config is not None and \
vllm/v1/worker/gpu_model_runner.py
+39
-13
@@ -1056,6 +1056,40 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1056
indices=out_indices,
1057
)
1058
1059
+ def sync_and_slice_intermediate_tensors(
1060
+ self, num_tokens: int, intermediate_tensors: IntermediateTensors,
1061
+ sync_self: bool) -> IntermediateTensors:
1062
+
1063
+ assert self.intermediate_tensors is not None
1064
+
1065
+ tp = self.vllm_config.parallel_config.tensor_parallel_size
1066
+ enabled_sp = self.vllm_config.compilation_config.pass_config. \
1067
+ enable_sequence_parallelism
1068
+ if enabled_sp:
1069
+ # When sequence parallelism is enabled, we always pad num_tokens
1070
+ # to be a multiple of tensor_parallel_size (tp) earlier
1071
+ assert num_tokens % tp == 0
1072
+ is_residual_scattered = tp > 1 and enabled_sp \
1073
+ and num_tokens % tp == 0
1074
+
1075
+ # When sequence parallelism is enabled, the "residual" tensor is sharded
1076
+ # across tensor parallel ranks, so each rank only needs its own slice.
1077
+ if sync_self:
1078
+ assert intermediate_tensors is not None
1079
+ for k, v in intermediate_tensors.items():
1080
+ is_scattered = "residual" and is_residual_scattered
1081
+ copy_len = num_tokens // tp if is_scattered else \
1082
+ num_tokens
1083
+ self.intermediate_tensors[k][:copy_len].copy_(
1084
+ v[:copy_len], non_blocking=True)
1085
+
1086
+ return IntermediateTensors({
1087
+ k:
1088
+ v[:num_tokens // tp]
1089
+ if k == "residual" and is_residual_scattered else v[:num_tokens]
1090
+ for k, v in self.intermediate_tensors.items()
1091
+ })
1092
+
1093
@torch.inference_mode()
1094
def execute_model(
1095
self,
@@ -1131,15 +1165,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1165
if get_pp_group().is_first_rank:
1166
intermediate_tensors = None
1167
else:
1134
- assert intermediate_tensors is not None
1135
- assert self.intermediate_tensors is not None
1136
- for k, v in intermediate_tensors.items():
1137
- self.intermediate_tensors[k][:num_input_tokens].copy_(
1138
- v[:num_input_tokens], non_blocking=True)
1139
- intermediate_tensors = IntermediateTensors({
1140
- k: v[:num_input_tokens]
1141
- for k, v in self.intermediate_tensors.items()
1142
- })
1168
+ intermediate_tensors = self.sync_and_slice_intermediate_tensors(
1169
+ num_input_tokens, intermediate_tensors, True)
1170
1171
# Run the decoder.
1172
# Use persistent buffers for CUDA graphs.
@@ -1658,10 +1685,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1685
batch_size=self.max_num_tokens,
1686
dtype=self.model_config.dtype,
1687
device=self.device))
1661
- intermediate_tensors = IntermediateTensors({
1662
- k: v[:num_tokens]
1663
- for k, v in self.intermediate_tensors.items()
1664
- })
1688
+
1689
+ intermediate_tensors = self.sync_and_slice_intermediate_tensors(
1690
+ num_tokens, None, False)
1691
1692
with set_forward_context(attn_metadata,
1693
self.vllm_config,