@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 690fe019f0
Fork 0

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, 13:30 UTC

[Feature] support sequence parallelism using compilation pass (#16155)

Signed-off-by: cascade812 <cascade812@outlook.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

cascade committed Apr 27, 2025 at 06:29 UTC 690fe019f046c03978ebe728770e5200d4e9d5f2
21 files changed +1072 -44
.buildkite/test-pipeline.yaml
+3
@@ -299,6 +299,7 @@ steps:
299 commands:
300 - pytest -v -s compile/test_pass_manager.py
301 - pytest -v -s compile/test_fusion.py
302 + - pytest -v -s compile/test_sequence_parallelism.py
303
304 - label: PyTorch Fullgraph Smoke Test # 9min
305 source_file_dependencies:
@@ -583,6 +584,8 @@ steps:
584 - pytest models/encoder_decoder/language/test_bart.py -v -s -m 'distributed(num_gpus=2)'
585 - pytest models/encoder_decoder/vision_language/test_broadcast.py -v -s -m 'distributed(num_gpus=2)'
586 - pytest models/decoder_only/vision_language/test_models.py -v -s -m 'distributed(num_gpus=2)'
587 + # test sequence parallel
588 + - pytest -v -s distributed/test_sequence_parallel.py
589 # this test fails consistently.
590 # TODO: investigate and fix
591 # - pytest -v -s spec_decode/e2e/test_integration_dist_tp2.py
tests/compile/test_functionalization.py
+8 -6
@@ -10,7 +10,7 @@ from vllm.compilation.fusion import (FUSED_OPS, FusionPass, QuantKey,
10 kFp8DynamicTokenSym, kFp8StaticTensorSym)
11 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
12 from vllm.compilation.noop_elimination import NoOpEliminationPass
13 -from vllm.config import CompilationConfig
13 +from vllm.config import CompilationConfig, VllmConfig
14
15 from .backend import TestBackend
16
@@ -49,13 +49,15 @@ def test_fix_functionalization(model: str, quant_key: QuantKey,
49 do_fusion: bool):
50 torch.set_default_device("cuda")
51
52 - config = CompilationConfig.PassConfig(enable_fusion=do_fusion,
53 - enable_noop=True)
54 - noop_pass = NoOpEliminationPass(config)
55 - fusion_pass = FusionPass.instance(config)
52 + vllm_config = VllmConfig()
53 + vllm_config.compilation_config = CompilationConfig(pass_config= \
54 + CompilationConfig.PassConfig(enable_fusion=do_fusion,
55 + enable_noop=True))
56 + noop_pass = NoOpEliminationPass(vllm_config)
57 + fusion_pass = FusionPass.instance(vllm_config)
58
59 passes = [noop_pass, fusion_pass] if do_fusion else [noop_pass]
58 - func_pass = FixFunctionalizationPass(config)
60 + func_pass = FixFunctionalizationPass(vllm_config)
61 backend_func = TestBackend(*passes, func_pass)
62 backend_no_func = TestBackend(*passes)
63
tests/compile/test_fusion.py
+5 -4
@@ -77,12 +77,13 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
77
78 vllm_config = VllmConfig(compilation_config=CompilationConfig(
79 level=CompilationLevel.PIECEWISE, custom_ops=["+rms_norm"]))
80 + vllm_config.compilation_config.pass_config = \
81 + CompilationConfig.PassConfig(enable_fusion=True,
82 + enable_noop=True)
83 with vllm.config.set_current_vllm_config(vllm_config):
84 # Reshape pass is needed for the fusion pass to work
82 - config = CompilationConfig.PassConfig(enable_fusion=True,
83 - enable_noop=True)
84 - noop_pass = NoOpEliminationPass(config)
85 - fusion_pass = FusionPass.instance(config)
85 + noop_pass = NoOpEliminationPass(vllm_config)
86 + fusion_pass = FusionPass.instance(vllm_config)
87
88 backend = TestBackend(noop_pass, fusion_pass)
89 model = TestModel(hidden_size, eps, static, cutlass_fp8_enabled)
tests/compile/test_pass_manager.py
+5 -4
@@ -6,7 +6,7 @@ import torch
6
7 from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
8 from vllm.compilation.pass_manager import PostGradPassManager
9 -from vllm.config import CompilationConfig
9 +from vllm.config import VllmConfig
10
11
12 # dummy custom pass that doesn't inherit
@@ -16,7 +16,7 @@ def simple_callable(graph: torch.fx.Graph):
16
17 # Should fail to add directly to the pass manager
18 def test_bad_callable():
19 - config = CompilationConfig().pass_config
19 + config = VllmConfig()
20
21 pass_manager = PostGradPassManager()
22 pass_manager.configure(config)
@@ -43,7 +43,7 @@ class ProperPass(InductorPass):
43 ],
44 )
45 def test_pass_manager_uuid(callable):
46 - config = CompilationConfig().pass_config
46 + config = VllmConfig()
47
48 pass_manager = PostGradPassManager()
49 pass_manager.configure(config)
@@ -64,7 +64,8 @@ def test_pass_manager_uuid(callable):
64
65 # UUID should be different due to config change
66 config2 = copy.deepcopy(config)
67 - config2.enable_fusion = not config2.enable_fusion
67 + config2.compilation_config.pass_config.enable_fusion = not \
68 + config2.compilation_config.pass_config.enable_fusion
69 pass_manager3 = PostGradPassManager()
70 pass_manager3.configure(config2)
71 pass_manager3.add(callable)
tests/compile/test_sequence_parallelism.py new
+190
@@ -0,0 +1,190 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import pytest
4 +import torch
5 +
6 +import vllm.envs as envs
7 +from vllm.compilation.fix_functionalization import FixFunctionalizationPass
8 +from vllm.compilation.fx_utils import (find_auto_fn, find_auto_fn_maybe,
9 + find_specified_fn,
10 + find_specified_fn_maybe, is_func)
11 +from vllm.compilation.sequence_parallelism import SequenceParallelismPass
12 +from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,
13 + VllmConfig)
14 +from vllm.distributed import tensor_model_parallel_all_reduce
15 +from vllm.distributed.parallel_state import (init_distributed_environment,
16 + initialize_model_parallel)
17 +from vllm.model_executor.layers.layernorm import RMSNorm
18 +from vllm.platforms import current_platform
19 +from vllm.utils import update_environment_variables
20 +
21 +from ..utils import multi_gpu_test
22 +from .backend import TestBackend
23 +
24 +OPS_IN_MODEL_BEFORE = [
25 + torch.ops.vllm.all_reduce.default,
26 +]
27 +
28 +OPS_IN_MODEL_AFTER = [
29 + torch.ops.vllm.reduce_scatter.default,
30 + torch.ops.vllm.all_gather.default,
31 +]
32 +
33 +OPS_IN_MODEL = [torch.ops._C.fused_add_rms_norm.default]
34 +
35 +prompts = [
36 + "Hello, my name is",
37 + "The president of the United States is",
38 + "The capital of France is",
39 + "The future of AI is",
40 +]
41 +
42 +
43 +class TestModel(torch.nn.Module):
44 +
45 + def __init__(self, hidden_size=16, intermediate_size=32):
46 + super().__init__()
47 + self.hidden_size = hidden_size
48 + self.intermediate_size = intermediate_size
49 + self.gate_proj = torch.nn.Parameter(
50 + torch.empty((intermediate_size, hidden_size)))
51 + self.norm = RMSNorm(hidden_size, 1e-05)
52 + # Initialize weights
53 + torch.nn.init.normal_(self.gate_proj, std=0.02)
54 +
55 + def forward(self, hidden_states, residual):
56 + """
57 + Forward pass implementing the operations in the FX graph
58 +
59 + Args:
60 + hidden_states: Input tensor
61 + residual: Residual tensor from previous layer
62 +
63 + Returns:
64 + Tuple containing the output tensor
65 + """
66 + # Reshape input
67 + view = hidden_states.reshape(-1, self.hidden_size)
68 +
69 + #matrix multiplication
70 + permute = self.gate_proj.permute(1, 0)
71 + mm = torch.mm(view, permute)
72 +
73 + # Tensor parallel all-reduce
74 + all_reduce = tensor_model_parallel_all_reduce(mm)
75 +
76 + # layer normalization
77 + norm_output, residual_output = self.norm(all_reduce, residual)
78 +
79 + return norm_output, residual_output
80 +
81 +
82 +@multi_gpu_test(num_gpus=2)
83 +@pytest.mark.parametrize("batch_size", [8])
84 +@pytest.mark.parametrize("seq_len", [16])
85 +@pytest.mark.parametrize("hidden_size", [16])
86 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
87 +@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
88 + reason="Only test on CUDA")
89 +def test_sequence_parallelism_pass(batch_size: int, seq_len: int,
90 + hidden_size: int, dtype: torch.dtype):
91 + num_processes = 2
92 +
93 + def run_torch_spawn(fn, nprocs):
94 + # need to use torch.mp.spawn otherwise will have problems with
95 + # torch.distributed and cuda
96 + torch.multiprocessing.spawn(fn,
97 + args=(num_processes, batch_size, seq_len,
98 + hidden_size, dtype),
99 + nprocs=nprocs)
100 +
101 + run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
102 +
103 +
104 +def sequence_parallelism_pass_on_test_model(local_rank: int, world_size: int,
105 + batch_size: int, seq_len: int,
106 + hidden_size: int,
107 + dtype: torch.dtype):
108 + current_platform.seed_everything(0)
109 +
110 + device = torch.device(f"cuda:{local_rank}")
111 + torch.cuda.set_device(device)
112 + torch.set_default_device(device)
113 + torch.set_default_dtype(dtype)
114 +
115 + update_environment_variables({
116 + 'RANK': str(local_rank),
117 + 'LOCAL_RANK': str(local_rank),
118 + 'WORLD_SIZE': str(world_size),
119 + 'MASTER_ADDR': 'localhost',
120 + 'MASTER_PORT': '12345',
121 + })
122 +
123 + # initialize distributed
124 + init_distributed_environment()
125 + initialize_model_parallel(tensor_model_parallel_size=world_size)
126 +
127 + # configure vllm config for SequenceParallelismPass
128 + vllm_config = VllmConfig()
129 + vllm_config.compilation_config = CompilationConfig(
130 + pass_config=CompilationConfig.PassConfig(
131 + enable_sequence_parallelism=True, ), )
132 + vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
133 +
134 + # this is a fake model name to construct the model config
135 + # in the vllm_config, it's not really used.
136 + model = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
137 + vllm_config.model_config = ModelConfig(model=model,
138 + task="auto",
139 + tokenizer=model,
140 + tokenizer_mode="auto",
141 + trust_remote_code=True,
142 + dtype=dtype,
143 + seed=42)
144 +
145 + sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
146 + backend_no_func = TestBackend(sequence_parallelism_pass)
147 + func_pass = FixFunctionalizationPass(vllm_config)
148 + backend_func = TestBackend(sequence_parallelism_pass, func_pass)
149 +
150 + model = TestModel(hidden_size, hidden_size * 2)
151 + hidden_states = torch.randn((batch_size * seq_len, hidden_size),
152 + dtype=dtype)
153 + residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
154 +
155 + compiled_model_no_func = torch.compile(model, backend=backend_no_func)
156 + compiled_model_no_func(hidden_states, residual)
157 + compiled_model_func = torch.compile(model, backend=backend_func)
158 + compiled_model_func(hidden_states, residual)
159 +
160 + # Check substitution worked
161 + pre_nodes = backend_no_func.graph_pre_pass.nodes
162 + post_nodes = backend_no_func.graph_post_pass.nodes
163 +
164 + # In pre-nodes, all reduce should be there,
165 + # reduce scatter and all gather should not
166 + for op in OPS_IN_MODEL_BEFORE:
167 + find_specified_fn(pre_nodes, op)
168 + for op in OPS_IN_MODEL_AFTER:
169 + assert find_specified_fn_maybe(pre_nodes, op) is None
170 +
171 + # In post-nodes, reduce scatter and all gather should be there,
172 + # all reduce should not
173 + for op in OPS_IN_MODEL_AFTER:
174 + find_specified_fn(post_nodes, op)
175 + for op in OPS_IN_MODEL_BEFORE:
176 + assert find_specified_fn_maybe(post_nodes, op) is None
177 +
178 + # check if the functionalization pass is applied
179 + for op in OPS_IN_MODEL:
180 + find_auto_fn(backend_no_func.graph_post_pass.nodes, op)
181 + assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes,
182 + op) is None # noqa: E501
183 +
184 + # make sure the ops were all de-functionalized
185 + found = dict()
186 + for node in backend_func.graph_post_pass.nodes:
187 + for op in OPS_IN_MODEL:
188 + if is_func(node, op):
189 + found[op] = True
190 + assert all(found[op] for op in OPS_IN_MODEL)
tests/distributed/test_comm_ops.py
+30 -1
@@ -14,7 +14,8 @@ import torch
14
15 from vllm.distributed import (broadcast_tensor_dict, get_pp_group,
16 tensor_model_parallel_all_gather,
17 - tensor_model_parallel_all_reduce)
17 + tensor_model_parallel_all_reduce,
18 + tensor_model_parallel_reduce_scatter)
19
20 from ..utils import init_test_distributed_environment, multi_process_parallel
21
@@ -47,6 +48,34 @@ def all_reduce_test_worker(
48 torch.testing.assert_close(t, expected)
49
50
51 +@ray.remote(num_gpus=1, max_calls=1)
52 +def reduce_scatter_test_worker(monkeypatch: pytest.MonkeyPatch, tp_size: int,
53 + pp_size: int, rank: int,
54 + distributed_init_port: str):
55 + # it is important to delete the CUDA_VISIBLE_DEVICES environment variable
56 + # so that each worker can see all the GPUs
57 + # they will be able to set the device to the correct GPU
58 + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
59 + device = torch.device(f"cuda:{rank}")
60 + torch.cuda.set_device(device)
61 + init_test_distributed_environment(tp_size, pp_size, rank,
62 + distributed_init_port)
63 +
64 + num_elements = 8
65 + all_tensors = [
66 + torch.arange(num_elements, dtype=torch.float32, device="cuda") *
67 + (r + 1) for r in range(tp_size)
68 + ]
69 +
70 + index = rank % tp_size
71 + partition_size = num_elements // tp_size
72 + all_reduce = torch.sum(torch.stack(all_tensors, dim=0), dim=0)
73 + expected = all_reduce[index * partition_size:(index + 1) * partition_size]
74 + t = all_tensors[index]
75 + t = tensor_model_parallel_reduce_scatter(t, 0)
76 + torch.testing.assert_close(t, expected)
77 +
78 +
79 @ray.remote(num_gpus=1, max_calls=1)
80 def all_gather_test_worker(
81 monkeypatch: pytest.MonkeyPatch,
tests/distributed/test_sequence_parallel.py new
+296
@@ -0,0 +1,296 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +WARNING: This test runs in both single-node (4 GPUs) and multi-node
4 + (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
5 + important to set the distributed backend to "mp" to avoid Ray scheduling
6 + all workers in a node other than the head node, which can cause the test
7 + to fail.
8 +"""
9 +import json
10 +import os
11 +from dataclasses import dataclass
12 +from typing import Literal, NamedTuple, Optional
13 +
14 +import pytest
15 +
16 +from vllm.config import TaskOption
17 +from vllm.logger import init_logger
18 +
19 +from ..models.registry import HF_EXAMPLE_MODELS
20 +from ..utils import compare_two_settings, create_new_process_for_each_test
21 +
22 +logger = init_logger("test_sequence_parallel")
23 +
24 +VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
25 +
26 +
27 +class ParallelSetup(NamedTuple):
28 + tp_size: int
29 + sp_enabled: bool
30 + eager_mode: bool
31 + chunked_prefill: bool
32 +
33 +
34 +class SPTestOptions(NamedTuple):
35 + multi_node_only: bool
36 + load_format: Optional[str] = None
37 +
38 +
39 +@dataclass
40 +class SPTestSettings:
41 + parallel_setups: list[ParallelSetup]
42 + # NOTE: the length of distributed_backends and
43 + # vllm_major_versions should be the same, and they
44 + # are first zipped together to iterate over all
45 + # test settings.
46 + distributed_backends: list[str]
47 + # vllm major version: "0" for V0, "1" for V1
48 + vllm_major_versions: list[str]
49 + task: TaskOption
50 + test_options: SPTestOptions
51 +
52 + def __post_init__(self):
53 + if len(self.distributed_backends) != len(self.vllm_major_versions):
54 + raise ValueError(
55 + f"Length mismatch: distributed_backends "
56 + f"({len(self.distributed_backends)}) != "
57 + f"vllm_major_versions ({len(self.vllm_major_versions)})")
58 +
59 + @staticmethod
60 + def detailed(
61 + *,
62 + tp_base: int = 2,
63 + multi_node_only: bool = False,
64 + task: TaskOption = "auto",
65 + load_format: Optional[str] = None,
66 + ):
67 + return SPTestSettings(
68 + parallel_setups=[
69 + ParallelSetup(tp_size=tp_base,
70 + sp_enabled=True,
71 + eager_mode=False,
72 + chunked_prefill=False),
73 + ParallelSetup(tp_size=tp_base,
74 + sp_enabled=True,
75 + eager_mode=False,
76 + chunked_prefill=True),
77 + ParallelSetup(tp_size=tp_base,
78 + sp_enabled=True,
79 + eager_mode=True,
80 + chunked_prefill=False),
81 + ParallelSetup(tp_size=tp_base,
82 + sp_enabled=True,
83 + eager_mode=True,
84 + chunked_prefill=True)
85 + ],
86 + distributed_backends=["mp", "ray"],
87 + vllm_major_versions=["1", "1"],
88 + task=task,
89 + test_options=SPTestOptions(multi_node_only=multi_node_only,
90 + load_format=load_format),
91 + )
92 +
93 + @staticmethod
94 + def fast(
95 + *,
96 + tp_base: int = 2,
97 + task: TaskOption = "auto",
98 + multi_node_only: bool = False,
99 + load_format: Optional[str] = None,
100 + ):
101 + return SPTestSettings(
102 + parallel_setups=[
103 + ParallelSetup(tp_size=tp_base,
104 + sp_enabled=True,
105 + eager_mode=False,
106 + chunked_prefill=False),
107 + ],
108 + distributed_backends=["mp", "ray"],
109 + vllm_major_versions=["1", "1"],
110 + task=task,
111 + test_options=SPTestOptions(multi_node_only=multi_node_only,
112 + load_format=load_format),
113 + )
114 +
115 + def iter_params(self, model_id: str):
116 + opts = self.test_options
117 +
118 + for parallel_setup in self.parallel_setups:
119 + for backend, vllm_major_version in zip(self.distributed_backends,
120 + self.vllm_major_versions):
121 + yield (model_id, parallel_setup, backend, vllm_major_version,
122 + self.task, opts)
123 +
124 +
125 +def _compare_sp(
126 + model_id: str,
127 + parallel_setup: ParallelSetup,
128 + distributed_backend: str,
129 + vllm_major_version: str,
130 + task: TaskOption,
131 + test_options: SPTestOptions,
132 + num_gpus_available: int,
133 + *,
134 + method: Literal["generate", "encode"],
135 + is_multimodal: bool,
136 +):
137 + (
138 + tp_size,
139 + sp_enabled,
140 + eager_mode,
141 + chunked_prefill,
142 + ) = parallel_setup
143 +
144 + multi_node_only, load_format = test_options
145 +
146 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
147 + model_info.check_transformers_version(on_fail="skip")
148 +
149 + trust_remote_code = model_info.trust_remote_code
150 + tokenizer_mode = model_info.tokenizer_mode
151 + hf_overrides = model_info.hf_overrides
152 +
153 + if load_format == "dummy":
154 + # Avoid OOM
155 + text_overrides = {
156 + "num_hidden_layers": 4,
157 + "hidden_size": 512,
158 + "intermediate_size": 800,
159 + "num_attention_heads": 4,
160 + "num_key_value_heads": 1,
161 + }
162 +
163 + if is_multimodal:
164 + hf_overrides.update({"text_config": text_overrides})
165 + else:
166 + hf_overrides.update(text_overrides)
167 + else:
168 + model_info.check_available_online(on_fail="skip")
169 +
170 + pp_size = 1
171 + if num_gpus_available < tp_size * pp_size:
172 + pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
173 + if VLLM_MULTI_NODE and distributed_backend == "mp":
174 + pytest.skip("Skipping multi-node pipeline parallel test for "
175 + "multiprocessing distributed backend")
176 + if multi_node_only and not VLLM_MULTI_NODE:
177 + pytest.skip("Not in multi-node setting")
178 +
179 + common_args = [
180 + # use half precision for speed and memory savings in CI environment
181 + "--dtype",
182 + "float16",
183 + "--max-model-len",
184 + "2048",
185 + "--max-num-seqs",
186 + "8",
187 + ]
188 + if chunked_prefill:
189 + common_args.append("--enable-chunked-prefill")
190 + if eager_mode:
191 + common_args.append("--enforce-eager")
192 + if task != "auto":
193 + common_args.extend(["--task", task])
194 + if trust_remote_code:
195 + common_args.append("--trust-remote-code")
196 + if tokenizer_mode:
197 + common_args.extend(["--tokenizer-mode", tokenizer_mode])
198 + if load_format:
199 + common_args.extend(["--load-format", load_format])
200 + if hf_overrides:
201 + common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
202 +
203 + compilation_config = {
204 + 'level': 3,
205 + 'custom_ops': ["+rms_norm"],
206 + 'compile_sizes': [4, 8],
207 + 'splitting_ops': [],
208 + 'pass_config': {
209 + 'enable_sequence_parallism': sp_enabled,
210 + 'enable_noop': True,
211 + 'enable_fusion': True,
212 + },
213 + }
214 +
215 + tp_sp_env = tp_env = {
216 + "VLLM_USE_V1": vllm_major_version,
217 + }
218 +
219 + tp_sp_args = [
220 + *common_args,
221 + "--tensor-parallel-size",
222 + str(tp_size),
223 + "--distributed-executor-backend",
224 + distributed_backend,
225 + "--compilation_config",
226 + str(compilation_config),
227 + ]
228 +
229 + tp_env = {
230 + "VLLM_USE_V1": vllm_major_version,
231 + }
232 + tp_args = [
233 + *common_args,
234 + "--tensor-parallel-size",
235 + str(tp_size),
236 + "--distributed-executor-backend",
237 + "mp",
238 + ]
239 +
240 + try:
241 + compare_two_settings(model_id,
242 + tp_sp_args,
243 + tp_args,
244 + tp_sp_env,
245 + tp_env,
246 + method=method)
247 + except Exception:
248 + testing_ray_compiled_graph = tp_sp_env is not None
249 + if testing_ray_compiled_graph and vllm_major_version == "0":
250 + # Ray Compiled Graph tests are flaky for V0,
251 + # so we don't want to fail the test
252 + logger.exception("Ray Compiled Graph tests failed")
253 + else:
254 + raise
255 +
256 +
257 +SP_TEXT_GENERATION_MODELS = {
258 + # [Decoder-only]
259 + "meta-llama/Llama-3.2-1B-Instruct": SPTestSettings.detailed(),
260 +}
261 +
262 +SP_TEST_MODELS = [
263 + # TODO support other models
264 + # [LANGUAGE GENERATION]
265 + "meta-llama/Llama-3.2-1B-Instruct",
266 +]
267 +
268 +
269 +@pytest.mark.parametrize(
270 + ("model_id", "parallel_setup", "distributed_backend", "vllm_major_version",
271 + "task", "test_options"),
272 + [
273 + params for model_id, settings in SP_TEXT_GENERATION_MODELS.items()
274 + for params in settings.iter_params(model_id)
275 + if model_id in SP_TEST_MODELS
276 + ],
277 +)
278 +@create_new_process_for_each_test()
279 +def test_tp_sp_generation(
280 + model_id: str,
281 + parallel_setup: ParallelSetup,
282 + distributed_backend: str,
283 + vllm_major_version: str,
284 + task: TaskOption,
285 + test_options: SPTestOptions,
286 + num_gpus_available,
287 +):
288 + _compare_sp(model_id,
289 + parallel_setup,
290 + distributed_backend,
291 + vllm_major_version,
292 + task,
293 + test_options,
294 + num_gpus_available,
295 + method="generate",
296 + is_multimodal=False)
vllm/compilation/backends.py
+1 -1
@@ -339,7 +339,7 @@ class VllmBackend:
339
340 def configure_post_pass(self):
341 config = self.compilation_config
342 - self.post_grad_pass_manager.configure(config.pass_config)
342 + self.post_grad_pass_manager.configure(self.vllm_config)
343
344 # Post-grad custom passes are run using the post_grad_custom_post_pass
345 # hook. If a pass for that hook exists, add it to the pass manager.
vllm/compilation/compiler_interface.py
+8 -5
@@ -15,6 +15,8 @@ import vllm.envs as envs
15 from vllm.config import VllmConfig
16 from vllm.utils import is_torch_equal_or_newer
17
18 +from .inductor_pass import pass_context
19 +
20
21 class CompilerInterface:
22 """
@@ -312,11 +314,12 @@ class InductorAdaptor(CompilerInterface):
314 torch._functorch.config.patch(
315 enable_remote_autograd_cache=False))
316
315 - compiled_graph = compile_fx(
316 - graph,
317 - example_inputs,
318 - inner_compile=hijacked_compile_fx_inner,
319 - config_patches=current_config)
317 + with pass_context(runtime_shape):
318 + compiled_graph = compile_fx(
319 + graph,
320 + example_inputs,
321 + inner_compile=hijacked_compile_fx_inner,
322 + config_patches=current_config)
323
324 # We treat VLLM_DISABLE_COMPILE_CACHE as the overall switch for torch
325 # compilation cache. So turn off the checks if we disable the
vllm/compilation/fusion.py
+4 -4
@@ -9,7 +9,7 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
9 from torch._inductor.pattern_matcher import PatternMatcherPass
10 from torch._ops import OpOverload
11
12 -from vllm.config import CompilationConfig
12 +from vllm.config import VllmConfig
13 from vllm.logger import init_logger
14 from vllm.platforms import current_platform
15
@@ -531,7 +531,7 @@ class FusionPass(VllmInductorPass):
531 _instance: 'Optional[FusionPass]' = None
532
533 @classmethod
534 - def instance(cls, config: CompilationConfig.PassConfig):
534 + def instance(cls, config: VllmConfig):
535 """
536 Get the singleton instance of the FusionPass.
537 If the instance exists, the config is updated but
@@ -540,10 +540,10 @@ class FusionPass(VllmInductorPass):
540 if cls._instance is None:
541 cls._instance = FusionPass(config)
542 else:
543 - cls._instance.config = config
543 + cls._instance.pass_config = config.compilation_config.pass_config
544 return cls._instance
545
546 - def __init__(self, config: CompilationConfig.PassConfig):
546 + def __init__(self, config: VllmConfig):
547 assert self.__class__._instance is None, \
548 "FusionPass singleton instance already exists"
549 super().__init__(config)
vllm/compilation/fx_utils.py
+16
@@ -12,6 +12,22 @@ def is_func(node: fx.Node, target) -> bool:
12 return node.op == "call_function" and node.target == target
13
14
15 +# Returns the first specified node with the given op (if it exists)
16 +def find_specified_fn_maybe(nodes: Iterable[fx.Node],
17 + op: OpOverload) -> Optional[fx.Node]:
18 + for node in nodes:
19 + if node.target == op:
20 + return node
21 + return None
22 +
23 +
24 +# Returns the first specified node with the given op
25 +def find_specified_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:
26 + node = find_specified_fn_maybe(nodes, op)
27 + assert node is not None, f"Could not find {op} in nodes {nodes}"
28 + return node
29 +
30 +
31 # Returns the first auto_functionalized node with the given op (if it exists)
32 def find_auto_fn_maybe(nodes: Iterable[fx.Node],
33 op: OpOverload) -> Optional[fx.Node]:
vllm/compilation/inductor_pass.py
+32
@@ -4,6 +4,7 @@ import hashlib
4 import inspect
5 import json
6 import types
7 +from contextlib import contextmanager
8 from typing import Any, Callable, Dict, Optional, Union
9
10 import torch
@@ -18,6 +19,34 @@ else:
19 from .torch25_custom_graph_pass import ( # noqa: yapf
20 Torch25CustomGraphPass as CustomGraphPass)
21
22 +_pass_context = None
23 +
24 +
25 +class PassContext:
26 +
27 + def __init__(self, runtime_shape: Optional[int]):
28 + self.runtime_shape = runtime_shape
29 +
30 +
31 +def get_pass_context() -> PassContext:
32 + """Get the current pass context."""
33 + assert _pass_context is not None
34 + return _pass_context
35 +
36 +
37 +@contextmanager
38 +def pass_context(runtime_shape: Optional[int]):
39 + """A context manager that stores the current pass context,
40 + usually it is a list of sizes to specialize.
41 + """
42 + global _pass_context
43 + prev_context = _pass_context
44 + _pass_context = PassContext(runtime_shape)
45 + try:
46 + yield
47 + finally:
48 + _pass_context = prev_context
49 +
50
51 class InductorPass(CustomGraphPass):
52 """
@@ -62,6 +91,9 @@ class InductorPass(CustomGraphPass):
91 encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
92 return hashlib.sha256(encoded).hexdigest()
93
94 + def is_applicable_for_shape(self, shape: Optional[int]):
95 + return True
96 +
97
98 class CallableInductorPass(InductorPass):
99 """
vllm/compilation/pass_manager.py
+18 -11
@@ -4,13 +4,15 @@ from typing import List
4
5 from torch import fx as fx
6
7 -from vllm.config import CompilationConfig
7 +from vllm.config import VllmConfig
8 from vllm.logger import init_logger
9
10 from .fix_functionalization import FixFunctionalizationPass
11 from .fusion import FusionPass
12 -from .inductor_pass import CustomGraphPass, InductorPass
12 +from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
13 from .noop_elimination import NoOpEliminationPass
14 +from .sequence_parallelism import SequenceParallelismPass
15 +from .vllm_inductor_pass import VllmInductorPass
16
17 logger = init_logger(__name__)
18
@@ -31,24 +33,29 @@ class PostGradPassManager(CustomGraphPass):
33 """
34
35 def __init__(self):
34 - self.passes: List[InductorPass] = []
36 + self.passes: List[VllmInductorPass] = []
37
38 def __call__(self, graph: fx.Graph):
39 + shape = get_pass_context().runtime_shape
40 for pass_ in self.passes:
38 - pass_(graph)
41 + if pass_.is_applicable_for_shape(shape):
42 + pass_(graph)
43
44 # always run fix_functionalization last
45 self.fix_functionalization(graph)
46
43 - def configure(self, pass_config: CompilationConfig.PassConfig):
44 - self.pass_config = pass_config
45 - if pass_config.enable_noop:
46 - self.passes += [NoOpEliminationPass(pass_config)]
47 + def configure(self, config: VllmConfig):
48 + self.pass_config = config.compilation_config.pass_config
49 + if self.pass_config.enable_noop:
50 + self.passes += [NoOpEliminationPass(config)]
51
48 - if pass_config.enable_fusion:
49 - self.passes += [FusionPass.instance(pass_config)]
52 + if self.pass_config.enable_fusion:
53 + self.passes += [FusionPass.instance(config)]
54
51 - self.fix_functionalization = FixFunctionalizationPass(pass_config)
55 + if self.pass_config.enable_sequence_parallelism:
56 + self.passes += [SequenceParallelismPass(config)]
57 +
58 + self.fix_functionalization = FixFunctionalizationPass(config)
59
60 def add(self, pass_: InductorPass):
61 assert isinstance(pass_, InductorPass)
vllm/compilation/sequence_parallelism.py new
+266
@@ -0,0 +1,266 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional, Tuple
3 +
4 +import torch
5 +import torch._inductor.pattern_matcher as pm
6 +import torch.fx as fx
7 +from torch._inductor.pattern_matcher import PatternMatcherPass
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
11 +from vllm.distributed.parallel_state import (
12 + get_tensor_model_parallel_world_size)
13 +from vllm.logger import init_logger
14 +
15 +from .vllm_inductor_pass import VllmInductorPass
16 +
17 +logger = init_logger(__name__)
18 +
19 +
20 +class AllReduceRMSNormPattern:
21 +
22 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str):
23 + self.epsilon = epsilon
24 + self.dtype = dtype
25 + self.device = device
26 +
27 +
28 +class EmbeddingAllReduceRMSNormPattern(AllReduceRMSNormPattern):
29 +
30 + def get_inputs(self):
31 + arg2_1 = torch.empty([16, 4], device=self.device, dtype=self.dtype)
32 + mul_6 = torch.tensor([[3, 7, 1, 4, 9, 2, 5, 0]],
33 + device=self.device,
34 + dtype=torch.long)
35 + unsqueeze = torch.rand([1, 8, 1], device=self.device, \
36 + dtype=self.dtype) > 0.5
37 + full_default = torch.zeros([1, 8, 4], device=self.device, \
38 + dtype=self.dtype)
39 + permute = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
40 + arg3_1 = torch.empty([4], device=self.device, dtype=self.dtype)
41 +
42 + return [arg2_1, mul_6, unsqueeze, full_default, permute, arg3_1]
43 +
44 + def register(self, pm_pass: PatternMatcherPass):
45 +
46 + def pattern(
47 + arg2_1: torch.Tensor,
48 + mul_6: torch.Tensor,
49 + unsqueeze: torch.Tensor,
50 + full_default: torch.Tensor,
51 + permute: torch.Tensor,
52 + arg3_1: torch.Tensor,
53 + ):
54 + embedding = torch.ops.aten.embedding.default(arg2_1, mul_6)
55 + where = torch.ops.aten.where.self(unsqueeze, full_default,
56 + embedding)
57 + all_reduce = tensor_model_parallel_all_reduce(where)
58 + rmsnorm = torch.ops.higher_order.auto_functionalized(
59 + torch.ops._C.rms_norm.default,
60 + result=permute,
61 + input=all_reduce,
62 + weight=arg3_1,
63 + epsilon=self.epsilon,
64 + )
65 +
66 + return rmsnorm[1], all_reduce
67 +
68 + def replacement(
69 + arg2_1: torch.Tensor,
70 + mul_6: torch.Tensor,
71 + unsqueeze: torch.Tensor,
72 + full_default: torch.Tensor,
73 + permute: torch.Tensor,
74 + arg3_1: torch.Tensor,
75 + ):
76 + embedding = torch.ops.aten.embedding.default(arg2_1, mul_6)
77 + where = torch.ops.aten.where.self(unsqueeze, full_default,
78 + embedding)
79 +
80 + tp = get_tp_group()
81 + tp_size = get_tensor_model_parallel_world_size()
82 + reduce_scatter = torch.ops.vllm.reduce_scatter.default(
83 + where, dim=0, world_size=tp_size, group_name=tp.unique_name)
84 +
85 + rmsnorm_result = torch.empty_like(reduce_scatter)
86 + rmsnorm = torch.ops.higher_order.auto_functionalized(
87 + torch.ops._C.rms_norm.default,
88 + result=rmsnorm_result,
89 + input=reduce_scatter,
90 + weight=arg3_1,
91 + epsilon=self.epsilon,
92 + )
93 +
94 + all_gather = torch.ops.vllm.all_gather.default(
95 + rmsnorm[1],
96 + dim=0,
97 + world_size=tp_size,
98 + group_name=tp.unique_name)
99 +
100 + return all_gather, reduce_scatter
101 +
102 + pm.register_replacement(pattern, replacement, self.get_inputs(),
103 + pm.fwd_only, pm_pass)
104 +
105 +
106 +class MiddleAllReduceRMSNormPattern(AllReduceRMSNormPattern):
107 +
108 + def get_inputs(self):
109 + mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
110 +
111 + residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
112 + rms_norm_weights = torch.empty([4, 4],
113 + device=self.device,
114 + dtype=self.dtype)
115 +
116 + return [
117 + residual,
118 + mm_1,
119 + rms_norm_weights,
120 + ]
121 +
122 + def register(self, pm_pass: PatternMatcherPass):
123 +
124 + def pattern(
125 + residual: torch.Tensor,
126 + mm_1: torch.Tensor,
127 + rms_norm_weights: torch.Tensor,
128 + ) -> Tuple[torch.Tensor, torch.Tensor]:
129 + all_reduce = tensor_model_parallel_all_reduce(mm_1)
130 +
131 + rmsnorm = torch.ops.higher_order.auto_functionalized(
132 + torch.ops._C.fused_add_rms_norm.default,
133 + input=all_reduce,
134 + residual=residual,
135 + weight=rms_norm_weights,
136 + epsilon=self.epsilon,
137 + )
138 +
139 + return rmsnorm[1], rmsnorm[2]
140 +
141 + def replacement(
142 + residual: torch.Tensor,
143 + mm_1: torch.Tensor,
144 + rms_norm_weights: torch.Tensor,
145 + ) -> Tuple[torch.Tensor, torch.Tensor]:
146 + tp = get_tp_group()
147 + tp_size = get_tensor_model_parallel_world_size()
148 + reduce_scatter = torch.ops.vllm.reduce_scatter.default(
149 + mm_1, dim=0, world_size=tp_size, group_name=tp.unique_name)
150 +
151 + # TODO is it possible to extract epsilon from somewhere
152 + rmsnorm = torch.ops.higher_order.auto_functionalized(
153 + torch.ops._C.fused_add_rms_norm.default,
154 + input=reduce_scatter,
155 + residual=residual,
156 + weight=rms_norm_weights,
157 + epsilon=self.epsilon,
158 + )
159 +
160 + all_gather = torch.ops.vllm.all_gather.default(
161 + rmsnorm[1],
162 + dim=0,
163 + world_size=tp_size,
164 + group_name=tp.unique_name)
165 + return all_gather, rmsnorm[2]
166 +
167 + pm.register_replacement(pattern, replacement, self.get_inputs(),
168 + pm.fwd_only, pm_pass)
169 +
170 +
171 +class LastAllReduceRMSNormPattern(AllReduceRMSNormPattern):
172 +
173 + def get_inputs(self):
174 + mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
175 +
176 + residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
177 + rms_norm_weights = torch.empty([4, 4],
178 + device=self.device,
179 + dtype=self.dtype)
180 +
181 + return [
182 + residual,
183 + mm_1,
184 + rms_norm_weights,
185 + ]
186 +
187 + def register(self, pm_pass: PatternMatcherPass):
188 +
189 + def pattern(
190 + residual: torch.Tensor,
191 + mm_1: torch.Tensor,
192 + rms_norm_weights: torch.Tensor,
193 + ) -> Tuple[torch.Tensor, torch.Tensor]:
194 + all_reduce = tensor_model_parallel_all_reduce(mm_1)
195 +
196 + rmsnorm = torch.ops.higher_order.auto_functionalized(
197 + torch.ops._C.fused_add_rms_norm.default,
198 + input=all_reduce,
199 + residual=residual,
200 + weight=rms_norm_weights,
201 + epsilon=self.epsilon,
202 + )
203 +
204 + return rmsnorm[1]
205 +
206 + def replacement(
207 + residual: torch.Tensor,
208 + mm_1: torch.Tensor,
209 + rms_norm_weights: torch.Tensor,
210 + ) -> Tuple[torch.Tensor, torch.Tensor]:
211 + tp = get_tp_group()
212 + tp_size = get_tensor_model_parallel_world_size()
213 + reduce_scatter = torch.ops.vllm.reduce_scatter.default(
214 + mm_1, dim=0, world_size=tp_size, group_name=tp.unique_name)
215 +
216 + # TODO is it possible to extract epsilon from somewhere
217 + rmsnorm = torch.ops.higher_order.auto_functionalized(
218 + torch.ops._C.fused_add_rms_norm.default,
219 + input=reduce_scatter,
220 + residual=residual,
221 + weight=rms_norm_weights,
222 + epsilon=self.epsilon,
223 + )
224 +
225 + normalized = torch.ops.vllm.all_gather.default(
226 + rmsnorm[1],
227 + dim=0,
228 + world_size=tp_size,
229 + group_name=tp.unique_name)
230 +
231 + return normalized
232 +
233 + pm.register_replacement(pattern, replacement, self.get_inputs(),
234 + pm.fwd_only, pm_pass)
235 +
236 +
237 +class SequenceParallelismPass(VllmInductorPass):
238 +
239 + def __init__(self, config: VllmConfig):
240 + super().__init__(config)
241 +
242 + self.patterns: PatternMatcherPass = PatternMatcherPass(
243 + pass_name="sequence_parallelism_pass")
244 + for epsilon in [1e-5, 1e-6]:
245 + EmbeddingAllReduceRMSNormPattern(
246 + epsilon, self.dtype, self.device).register(self.patterns)
247 +
248 + MiddleAllReduceRMSNormPattern(epsilon, self.dtype,
249 + self.device).register(self.patterns)
250 +
251 + LastAllReduceRMSNormPattern(epsilon, self.dtype,
252 + self.device).register(self.patterns)
253 + # WARNING: This is a hack to clear the pattern matcher cache
254 + # and allow multiple values of epsilon.
255 + torch._inductor.pattern_matcher._seen_patterns.clear()
256 +
257 + def is_applicable_for_shape(self, shape: Optional[int]) -> bool:
258 + # only do replace for specific shapes
259 + tp_size = get_tensor_model_parallel_world_size()
260 + return shape is not None and shape % tp_size == 0
261 +
262 + def __call__(self, graph: fx.Graph):
263 + self.dump_graph(graph, "before_sequence_parallelism_pass")
264 + count = self.patterns.apply(graph)
265 + logger.debug("Replaced %s patterns", count)
266 + self.dump_graph(graph, "after_sequence_parallelism_pass")
vllm/compilation/vllm_inductor_pass.py
+8 -5
@@ -4,7 +4,7 @@ import time
4
5 import torch
6
7 -from vllm.config import CompilationConfig
7 +from vllm.config import CompilationConfig, VllmConfig
8 # yapf: disable
9 from vllm.distributed import get_tensor_model_parallel_rank as get_tp_rank
10 from vllm.distributed import (
@@ -24,16 +24,19 @@ class VllmInductorPass(InductorPass):
24 It provides timing, logging, and dumping utilities.
25 """
26
27 - def __init__(self, config: CompilationConfig.PassConfig):
28 - self.config = config
27 + def __init__(self, config: VllmConfig):
28 + self.pass_config = config.compilation_config.pass_config
29 + self.dtype = config.model_config.dtype if config.model_config else None
30 + self.device = config.device_config.device if config.device_config \
31 + else None
32 self.pass_name = self.__class__.__name__
33
34 def dump_graph(self, graph: torch.fx.Graph, stage: str, always=False):
32 - if stage in self.config.dump_graph_stages or always:
35 + if stage in self.pass_config.dump_graph_stages or always:
36 # Make sure filename includes rank in the distributed setting
37 parallel = p_is_init() and get_tp_world_size() > 1
38 rank = f"-{get_tp_rank()}" if parallel else ""
36 - filepath = self.config.dump_graph_dir / f"{stage}{rank}.py"
39 + filepath = self.pass_config.dump_graph_dir / f"{stage}{rank}.py"
40
41 logger.info("%s printing graph to %s", self.pass_name, filepath)
42 with open(filepath, "w") as f:
vllm/config.py
+50 -2
@@ -3405,11 +3405,13 @@ class CompilationConfig(BaseModel):
3405 - enable_fusion: whether to enable the custom fusion pass.
3406 - enable_noop: whether to enable the custom no-op elimination pass.
3407 TODO(luka) better pass enabling system.
3408 + - enable_sequence_parallelism: whether to enable sequence parallelism.
3409 """
3410 dump_graph_stages: list[str] = Field(default_factory=list)
3411 dump_graph_dir: Path = Field(default=Path("."))
3412 enable_fusion: bool = True
3413 enable_noop: bool = True
3414 + enable_sequence_parallelism: bool = False
3415
3416 def uuid(self):
3417 """
@@ -3418,7 +3420,8 @@ class CompilationConfig(BaseModel):
3420 Do not include dump_graph_* in the hash - they don't affect
3421 compilation.
3422 """
3421 - dict_ = self.model_dump(include={"enable_fusion", "enable_noop"})
3423 + dict_ = self.model_dump(include={"enable_fusion", "enable_noop", \
3424 + "enable_sequence_parallelism"})
3425 return InductorPass.hash_dict(dict_)
3426
3427 def model_post_init(self, __context: Any) -> None:
@@ -3840,6 +3843,8 @@ class VllmConfig:
3843
3844 if self.compilation_config is None:
3845 self.compilation_config = CompilationConfig()
3846 + if self.compilation_config.pass_config.enable_sequence_parallelism:
3847 + self.compilation_config.custom_ops.append("+rms_norm")
3848 if envs.VLLM_USE_V1 and self.model_config is not None and \
3849 not self.model_config.enforce_eager:
3850 # NOTE(woosuk): Currently, we use inductor because the piecewise
@@ -3847,7 +3852,8 @@ class VllmConfig:
3852 # FIXME(woosuk): Disable inductor to reduce the compilation time
3853 # and avoid any potential issues with the inductor.
3854 # FIXME(rob): Add function to set all of these.
3850 - self.compilation_config.custom_ops = ["none"]
3855 + if not self.compilation_config.custom_ops:
3856 + self.compilation_config.custom_ops = ["none"]
3857 self.compilation_config.use_cudagraph = True
3858 self.compilation_config.use_inductor = True
3859 self.compilation_config.cudagraph_num_of_warmups = 1
@@ -3856,6 +3862,18 @@ class VllmConfig:
3862 self.compilation_config.level = CompilationLevel.PIECEWISE
3863 self.compilation_config.set_splitting_ops_for_v1()
3864
3865 + if self.parallel_config is not None and \
3866 + self.parallel_config.tensor_parallel_size > 1 and \
3867 + self.parallel_config.pipeline_parallel_size > 1 and \
3868 + self.compilation_config is not None and \
3869 + self.compilation_config.pass_config is not None and \
3870 + self.compilation_config.pass_config.enable_sequence_parallelism:
3871 + logger.warning_once(
3872 + "Sequence parallelism is not supported with pipeline "
3873 + "parallelism. Disabling sequence parallelism.")
3874 + self.compilation_config.pass_config.\
3875 + enable_sequence_parallelism = False
3876 +
3877 self._set_cudagraph_sizes()
3878
3879 if self.cache_config is not None and \
@@ -3895,6 +3913,26 @@ class VllmConfig:
3913 if not self.instance_id:
3914 self.instance_id = random_uuid()[:5]
3915
3916 + def update_sizes_for_sequence_parallelism(self,
3917 + possible_sizes: list) -> list:
3918 + # remove the sizes that not multiple of tp_size when
3919 + # enable sequence parallelism
3920 + removed_sizes = [
3921 + size for size in possible_sizes
3922 + if size % self.parallel_config.tensor_parallel_size != 0
3923 + ]
3924 + if removed_sizes:
3925 + logger.warning(
3926 + "Batch sizes %s are removed because they are not "
3927 + "multiple of tp_size %d when "
3928 + "sequence parallelism is enabled", removed_sizes,
3929 + self.parallel_config.tensor_parallel_size)
3930 +
3931 + return [
3932 + size for size in possible_sizes
3933 + if size % self.parallel_config.tensor_parallel_size == 0
3934 + ]
3935 +
3936 def _set_cudagraph_sizes(self):
3937 """
3938 cudagraph batchsize padding logic:
@@ -3932,6 +3970,11 @@ class VllmConfig:
3970 not self.model_config.enforce_eager:
3971
3972 possible_sizes = [1, 2, 4] + [8 * i for i in range(1, 1025)]
3973 + if self.parallel_config.tensor_parallel_size > 1 and \
3974 + self.compilation_config.pass_config.enable_sequence_parallelism:
3975 + possible_sizes = self.update_sizes_for_sequence_parallelism(
3976 + possible_sizes)
3977 +
3978 # find the minimum size that is larger than max_num_seqs,
3979 # which then becomes the max_batchsize_to_capture
3980 larger_sizes = [
@@ -3955,6 +3998,11 @@ class VllmConfig:
3998 not self.model_config.enforce_eager:
3999 batch_size_capture_list = [1, 2, 4
4000 ] + [i for i in range(8, 513, 8)]
4001 + if self.parallel_config.tensor_parallel_size > 1 and \
4002 + self.compilation_config.pass_config.enable_sequence_parallelism:
4003 + batch_size_capture_list = \
4004 + self.update_sizes_for_sequence_parallelism(batch_size_capture_list)
4005 +
4006 max_num_tokens = self.scheduler_config.max_num_batched_tokens
4007 batch_size_capture_list = [
4008 size for size in batch_size_capture_list
vllm/distributed/communication_op.py
+6
@@ -19,6 +19,12 @@ def tensor_model_parallel_all_gather(input_: torch.Tensor,
19 return get_tp_group().all_gather(input_, dim)
20
21
22 +def tensor_model_parallel_reduce_scatter(input_: torch.Tensor,
23 + dim: int = -1) -> torch.Tensor:
24 + """Reduce-Scatter the input tensor across model parallel group."""
25 + return get_tp_group().reduce_scatter(input_, dim)
26 +
27 +
28 def tensor_model_parallel_gather(input_: torch.Tensor,
29 dst: int = 0,
30 dim: int = -1) -> Optional[torch.Tensor]:
vllm/distributed/device_communicators/base_device_communicator.py
+34
@@ -61,6 +61,40 @@ class DeviceCommunicatorBase:
61 input_size[dim + 1:])
62 return output_tensor
63
64 + def reduce_scatter(self,
65 + input_: torch.Tensor,
66 + dim: int = -1) -> torch.Tensor:
67 + world_size = self.world_size
68 + # Bypass the function if we are using only 1 GPU.
69 + if world_size == 1:
70 + return input_
71 + assert -input_.dim() <= dim < input_.dim(), (
72 + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
73 +
74 + if dim < 0:
75 + # Convert negative dim to positive.
76 + dim += input_.dim()
77 +
78 + # Note: This will produce an incorrect answer if we don't make
79 + # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
80 + input_tensor = input_.movedim(0, dim).contiguous()
81 +
82 + assert input_tensor.shape[0] % world_size == 0
83 + chunk_size = input_tensor.shape[0] // world_size
84 + output_shape = (chunk_size, ) + input_tensor.shape[1:]
85 +
86 + output_tensor = torch.empty(output_shape,
87 + dtype=input_tensor.dtype,
88 + device=input_tensor.device)
89 +
90 + # Perform reduce-scatter operation
91 + torch.distributed.reduce_scatter_tensor(output_tensor,
92 + input_tensor,
93 + group=self.device_group)
94 +
95 + # Reshape before returning
96 + return output_tensor.movedim(0, dim).contiguous()
97 +
98 def gather(self,
99 input_: torch.Tensor,
100 dst: int = 0,
vllm/distributed/device_communicators/cuda_communicator.py
+25
@@ -70,6 +70,31 @@ class CudaCommunicator(DeviceCommunicatorBase):
70 torch.distributed.all_reduce(out, group=self.device_group)
71 return out
72
73 + def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
74 + world_size = self.world_size
75 + pynccl_comm = self.pynccl_comm
76 + assert pynccl_comm is not None
77 + if dim < 0:
78 + # Convert negative dim to positive.
79 + dim += input_.dim()
80 +
81 + # Note: This will produce an incorrect answer if we don't make
82 + # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
83 + input_tensor = input_.movedim(0, dim).contiguous()
84 +
85 + assert input_tensor.shape[0] % world_size == 0
86 + chunk_size = input_tensor.shape[0] // world_size
87 + output_shape = (chunk_size, ) + input_tensor.shape[1:]
88 +
89 + output = torch.empty(output_shape,
90 + dtype=input_tensor.dtype,
91 + device=input_tensor.device)
92 +
93 + pynccl_comm.reduce_scatter(output, input_)
94 +
95 + # Reshape before returning
96 + return output.movedim(0, dim).contiguous()
97 +
98 def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
99 """Sends a tensor to the destination rank in a non-blocking way"""
100 """NOTE: `dst` is the local rank of the destination rank."""
vllm/distributed/parallel_state.py
+58
@@ -113,6 +113,38 @@ def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
113 return torch.empty_like(tensor)
114
115
116 +def reduce_scatter(tensor: torch.Tensor, dim: int, world_size: int,
117 + group_name: str) -> torch.Tensor:
118 + assert group_name in _groups, f"Group {group_name} is not found."
119 + group = _groups[group_name]()
120 + if group is None:
121 + raise ValueError(f"Group {group_name} is destroyed.")
122 + return group.reduce_scatter(tensor, dim)
123 +
124 +
125 +def reduce_scatter_fake(tensor: torch.Tensor, dim: int, world_size: int,
126 + group_name: str) -> torch.Tensor:
127 + new_shape = list(tensor.shape)
128 + new_shape[dim] = tensor.shape[dim] // world_size
129 + return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)
130 +
131 +
132 +def all_gather(tensor: torch.Tensor, dim: int, world_size: int,
133 + group_name: str) -> torch.Tensor:
134 + assert group_name in _groups, f"Group {group_name} is not found."
135 + group = _groups[group_name]()
136 + if group is None:
137 + raise ValueError(f"Group {group_name} is destroyed.")
138 + return group.all_gather(tensor, dim)
139 +
140 +
141 +def all_gather_fake(tensor: torch.Tensor, dim: int, world_size: int,
142 + group_name: str) -> torch.Tensor:
143 + new_shape = list(tensor.shape)
144 + new_shape[dim] = tensor.shape[dim] * world_size
145 + return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)
146 +
147 +
148 if supports_custom_op():
149 from vllm.platforms import current_platform
150 direct_register_custom_op(
@@ -123,6 +155,20 @@ if supports_custom_op():
155 dispatch_key=current_platform.dispatch_key,
156 )
157
158 + direct_register_custom_op(
159 + op_name="reduce_scatter",
160 + op_func=reduce_scatter,
161 + mutates_args=[],
162 + fake_impl=reduce_scatter_fake,
163 + )
164 +
165 + direct_register_custom_op(
166 + op_name="all_gather",
167 + op_func=all_gather,
168 + mutates_args=[],
169 + fake_impl=all_gather_fake,
170 + )
171 +
172
173 class GroupCoordinator:
174 """
@@ -322,6 +368,18 @@ class GroupCoordinator:
368
369 return self.device_communicator.all_gather(input_, dim)
370
371 + def reduce_scatter(self,
372 + input_: torch.Tensor,
373 + dim: int = -1) -> torch.Tensor:
374 + world_size = self.world_size
375 + # Bypass the function if we are using only 1 GPU.
376 + if world_size == 1:
377 + return input_
378 + assert -input_.dim() <= dim < input_.dim(), (
379 + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
380 +
381 + return self.device_communicator.reduce_scatter(input_, dim)
382 +
383 def gather(self,
384 input_: torch.Tensor,
385 dst: int = 0,
vllm/v1/worker/gpu_model_runner.py
+9 -1
@@ -1027,7 +1027,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1027 num_scheduled_tokens)
1028 else:
1029 # Eager mode.
1030 - num_input_tokens = num_scheduled_tokens
1030 + # Pad tokens to multiple of tensor_parallel_size when
1031 + # enabled collective fusion for SP
1032 + tp_size = self.vllm_config.parallel_config.tensor_parallel_size
1033 + if self.vllm_config.compilation_config.pass_config. \
1034 + enable_sequence_parallelism and tp_size > 1:
1035 + from vllm.utils import round_up
1036 + num_input_tokens = round_up(num_scheduled_tokens, tp_size)
1037 + else:
1038 + num_input_tokens = num_scheduled_tokens
1039 attn_metadata.num_input_tokens = num_input_tokens
1040
1041 # _prepare_inputs may reorder the batch, so we must gather multi