@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 71ea614d4a
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, 08:30 UTC

[Feature]Add async tensor parallelism using compilation pass (#17882)

Signed-off-by: cascade812 <cascade812@outlook.com>

cascade committed May 23, 2025 at 01:03 UTC 71ea614d4ab2e36a953f0123d63de95bac07f7d8
11 files changed +472 -56
.buildkite/test-pipeline.yaml
+1
@@ -316,6 +316,7 @@ steps:
316 - pytest -v -s compile/test_fusion.py
317 - pytest -v -s compile/test_silu_mul_quant_fusion.py
318 - pytest -v -s compile/test_sequence_parallelism.py
319 + - pytest -v -s compile/test_async_tp.py
320
321 - label: PyTorch Fullgraph Smoke Test # 9min
322 mirror_hardwares: [amdexperimental, amdproduction]
tests/compile/backend.py
+18
@@ -5,6 +5,8 @@ from typing import Callable, Union
5
6 from torch import fx
7
8 +from vllm.compilation.fx_utils import (find_specified_fn,
9 + find_specified_fn_maybe)
10 from vllm.compilation.inductor_pass import InductorPass
11 from vllm.config import get_current_vllm_config
12
@@ -44,3 +46,19 @@ class TestBackend:
46 self.graph_post_pass = deepcopy(graph)
47 # assign by reference, will reflect the final state of the graph
48 self.final_graph = graph
49 +
50 + def check_before_ops(self, ops,
51 + find_fn=find_specified_fn, \
52 + find_fn_maybe=find_specified_fn_maybe, \
53 + ops_fully_replaced=True):
54 + for op in ops:
55 + find_fn(self.graph_pre_pass.nodes, op)
56 + if ops_fully_replaced:
57 + assert find_fn_maybe(self.graph_post_pass.nodes, op) is None
58 +
59 + def check_after_ops(self, ops,
60 + find_fn=find_specified_fn, \
61 + find_fn_maybe=find_specified_fn_maybe):
62 + for op in ops:
63 + find_fn(self.graph_post_pass.nodes, op)
64 + assert find_fn_maybe(self.graph_pre_pass.nodes, op) is None
tests/compile/test_async_tp.py new
+248
@@ -0,0 +1,248 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import json
4 +
5 +import pytest
6 +import torch
7 +
8 +import vllm.envs as envs
9 +from vllm.compilation.collective_fusion import AsyncTPPass
10 +from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,
11 + PassConfig, VllmConfig)
12 +from vllm.distributed import (tensor_model_parallel_all_gather,
13 + tensor_model_parallel_reduce_scatter)
14 +from vllm.distributed.parallel_state import (init_distributed_environment,
15 + initialize_model_parallel)
16 +from vllm.platforms import current_platform
17 +from vllm.utils import update_environment_variables
18 +
19 +from ..models.registry import HF_EXAMPLE_MODELS
20 +from ..utils import (compare_two_settings, create_new_process_for_each_test,
21 + multi_gpu_test)
22 +from .backend import TestBackend
23 +
24 +prompts = [
25 + "Hello, my name is",
26 + "The president of the United States is",
27 + "The capital of France is",
28 + "The future of AI is",
29 +]
30 +
31 +
32 +class TestMMRSModel(torch.nn.Module):
33 +
34 + def __init__(self, hidden_size=16):
35 + super().__init__()
36 + self.hidden_size = hidden_size
37 + self.gate_proj = torch.nn.Parameter(torch.empty(
38 + (self.hidden_size * 2, hidden_size)),
39 + requires_grad=False)
40 + # Initialize weights
41 + torch.nn.init.normal_(self.gate_proj, std=0.02)
42 +
43 + def forward(self, hidden_states):
44 + """
45 + Forward pass implementing the mm + reduce scatter in the FX graph
46 +
47 + """
48 + # Reshape input
49 + view = hidden_states.reshape(-1, self.hidden_size)
50 +
51 + # matrix multiplication
52 + permute = self.gate_proj.permute(1, 0)
53 + mm = torch.mm(view, permute)
54 + reduce_scatter = tensor_model_parallel_reduce_scatter(mm, dim=0)
55 + return reduce_scatter
56 +
57 + def ops_in_model_before(self):
58 + return [torch.ops.vllm.reduce_scatter.default]
59 +
60 + def ops_in_model_after(self):
61 + return [torch.ops.symm_mem.fused_matmul_reduce_scatter.default]
62 +
63 +
64 +class TestAGMMModel(torch.nn.Module):
65 +
66 + def __init__(self, hidden_size=16):
67 + super().__init__()
68 + self.hidden_size = hidden_size
69 + self.weight = torch.nn.Parameter(torch.empty(
70 + (hidden_size, hidden_size)),
71 + requires_grad=False)
72 + # Initialize weights
73 + torch.nn.init.normal_(self.weight, std=0.02)
74 +
75 + def forward(self, hidden_states):
76 + """
77 + Forward pass implementing the mm + all gather in the FX graph
78 + """
79 + # Reshape input
80 + view = hidden_states.reshape(-1, self.hidden_size)
81 + all_gather = tensor_model_parallel_all_gather(view, dim=0)
82 + permute = self.weight.permute(1, 0)
83 + mm = torch.mm(all_gather, permute)
84 + return mm
85 +
86 + def ops_in_model_before(self):
87 + return [torch.ops.vllm.all_gather.default]
88 +
89 + def ops_in_model_after(self):
90 + return [torch.ops.symm_mem.fused_all_gather_matmul.default]
91 +
92 +
93 +@multi_gpu_test(num_gpus=2)
94 +@pytest.mark.parametrize("test_model", [TestMMRSModel, TestAGMMModel])
95 +@pytest.mark.parametrize("batch_size", [8])
96 +@pytest.mark.parametrize("seq_len", [16])
97 +@pytest.mark.parametrize("hidden_size", [16])
98 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
99 +@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
100 + reason="Only test on CUDA")
101 +def test_async_tp_pass_replace(test_model: str, batch_size: int, seq_len: int,
102 + hidden_size: int, dtype: torch.dtype):
103 + num_processes = 2
104 +
105 + def run_torch_spawn(fn, nprocs):
106 + # need to use torch.mp.spawn otherwise will have problems with
107 + # torch.distributed and cuda
108 + torch.multiprocessing.spawn(fn,
109 + args=(num_processes, test_model,
110 + batch_size, seq_len, hidden_size,
111 + dtype),
112 + nprocs=nprocs)
113 +
114 + run_torch_spawn(async_tp_pass_on_test_model, num_processes)
115 +
116 +
117 +def async_tp_pass_on_test_model(local_rank: int, world_size: int,
118 + test_model_cls: torch.nn.Module,
119 + batch_size: int, seq_len: int,
120 + hidden_size: int, dtype: torch.dtype):
121 + current_platform.seed_everything(0)
122 +
123 + device = torch.device(f"cuda:{local_rank}")
124 + torch.cuda.set_device(device)
125 + torch.set_default_device(device)
126 + torch.set_default_dtype(dtype)
127 +
128 + update_environment_variables({
129 + 'RANK': str(local_rank),
130 + 'LOCAL_RANK': str(local_rank),
131 + 'WORLD_SIZE': str(world_size),
132 + 'MASTER_ADDR': 'localhost',
133 + 'MASTER_PORT': '12345',
134 + })
135 +
136 + # initialize distributed
137 + init_distributed_environment()
138 + initialize_model_parallel(tensor_model_parallel_size=world_size)
139 +
140 + # configure vllm config for SequenceParallelismPass
141 + vllm_config = VllmConfig()
142 + vllm_config.compilation_config = CompilationConfig(pass_config=PassConfig(
143 + enable_async_tp=True, ), )
144 + vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
145 +
146 + # this is a fake model name to construct the model config
147 + # in the vllm_config, it's not really used.
148 + model_name = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
149 + vllm_config.model_config = ModelConfig(model=model_name,
150 + task="auto",
151 + tokenizer=model_name,
152 + tokenizer_mode="auto",
153 + trust_remote_code=True,
154 + dtype=dtype,
155 + seed=42)
156 +
157 + async_tp_pass = AsyncTPPass(vllm_config)
158 + backend = TestBackend(async_tp_pass)
159 +
160 + model = test_model_cls(hidden_size)
161 +
162 + hidden_states = torch.randn((batch_size * seq_len, hidden_size),
163 + dtype=dtype,
164 + requires_grad=False)
165 +
166 + compiled_model = torch.compile(model, backend=backend)
167 + compiled_model(hidden_states)
168 +
169 + # In pre-nodes, all gather or reduce scatter should exist,
170 + # fused_matmul_reduce_scatter or fused_all_gather_matmul should not
171 + backend.check_before_ops(model.ops_in_model_before(),
172 + ops_fully_replaced=False)
173 +
174 + # In post-nodes, fused_matmul_reduce_scatter or \
175 + # fused_all_gather_matmul should exist
176 + backend.check_after_ops(model.ops_in_model_after())
177 +
178 +
179 +@create_new_process_for_each_test()
180 +@pytest.mark.parametrize("model_id", ["meta-llama/Llama-3.2-1B-Instruct"])
181 +@pytest.mark.parametrize("tp_size", [2])
182 +@pytest.mark.parametrize("async_tp_enabled", [True])
183 +@pytest.mark.parametrize("distributed_backend", ["mp"])
184 +@pytest.mark.parametrize("eager_mode", [False, True])
185 +def test_async_tp_pass_correctness(
186 + model_id: str,
187 + tp_size: int,
188 + async_tp_enabled: bool,
189 + distributed_backend: str,
190 + eager_mode: bool,
191 + num_gpus_available: int,
192 +):
193 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
194 + model_info.check_transformers_version(on_fail="skip")
195 + model_info.check_available_online(on_fail="skip")
196 +
197 + pp_size = 1
198 + if num_gpus_available < tp_size:
199 + pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
200 +
201 + common_args = [
202 + "--dtype",
203 + "bfloat16",
204 + "--max-model-len",
205 + "2048",
206 + "--max-num-seqs",
207 + "8",
208 + ]
209 + if eager_mode:
210 + common_args.append("--enforce-eager")
211 +
212 + compilation_config = {
213 + 'level': 3,
214 + 'compile_sizes': [2, 4, 8],
215 + 'splitting_ops': [],
216 + 'pass_config': {
217 + 'enable_async_tp': async_tp_enabled
218 + },
219 + }
220 +
221 + async_tp_env = tp_env = {
222 + "VLLM_USE_V1": "1",
223 + }
224 +
225 + aysnc_tp_args = [
226 + *common_args,
227 + "--tensor-parallel-size",
228 + str(tp_size),
229 + "--distributed-executor-backend",
230 + distributed_backend,
231 + "--compilation_config",
232 + json.dumps(compilation_config),
233 + ]
234 +
235 + tp_args = [
236 + *common_args,
237 + "--tensor-parallel-size",
238 + str(tp_size),
239 + "--distributed-executor-backend",
240 + "mp",
241 + ]
242 +
243 + compare_two_settings(model_id,
244 + aysnc_tp_args,
245 + tp_args,
246 + async_tp_env,
247 + tp_env,
248 + method="generate")
tests/compile/test_fusion.py
+17 -19
@@ -29,6 +29,10 @@ class TestModel(torch.nn.Module):
29 self.cutlass_fp8_enabled = cutlass_fp8_enabled
30 self.norm = [RMSNorm(hidden_size, eps) for _ in range(3)]
31 self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]
32 + self.key = QuantKey(dtype=FP8_DTYPE,
33 + static=static,
34 + per_tensor=static,
35 + symmetric=True)
36 if static:
37 self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]
38 else:
@@ -59,6 +63,15 @@ class TestModel(torch.nn.Module):
63 y3, resid = self.norm[2](x3, resid) # use resid here
64 return y3
65
66 + def ops_in_model_before(self):
67 + return [QUANT_OPS[self.key]]
68 +
69 + def ops_in_model_after(self):
70 + return [
71 + FUSED_OPS[FusedRMSQuantKey(self.key, False)],
72 + FUSED_OPS[FusedRMSQuantKey(self.key, True)]
73 + ]
74 +
75
76 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
77 @pytest.mark.parametrize("hidden_size", [64, 3392, 4096])
@@ -107,25 +120,10 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
120
121 torch.testing.assert_close(result, result2, atol=ATOL, rtol=RTOL)
122
110 - # Check substitution worked
111 - pre_nodes = backend.graph_pre_pass.nodes
112 - post_nodes = backend.graph_post_pass.nodes
113 -
114 - # static is per-tensor, dynamic is per-token
115 - key = QuantKey(dtype=FP8_DTYPE,
116 - static=static,
117 - per_tensor=static,
118 - symmetric=True)
119 - rms_quant = FUSED_OPS[FusedRMSQuantKey(key, False)]
120 - add_rms_quant = FUSED_OPS[FusedRMSQuantKey(key, True)]
121 - fp8_quant = QUANT_OPS[key]
122 -
123 # In pre-nodes, fp8 quant should be there and fused kernels should not
124 - assert find_auto_fn_maybe(pre_nodes, rms_quant) is None
125 - assert find_auto_fn_maybe(pre_nodes, add_rms_quant) is None
126 - find_auto_fn(pre_nodes, fp8_quant)
124 + backend.check_before_ops(model.ops_in_model_before(), find_auto_fn,
125 + find_auto_fn_maybe)
126
127 # In post-nodes, fused kernels should be there and fp8 quant should not
129 - find_auto_fn(post_nodes, rms_quant)
130 - find_auto_fn(post_nodes, add_rms_quant)
131 - assert find_auto_fn_maybe(post_nodes, fp8_quant) is None
128 + backend.check_after_ops(model.ops_in_model_after(), find_auto_fn,
129 + find_auto_fn_maybe)
tests/compile/test_sequence_parallelism.py
+18 -29
@@ -5,9 +5,7 @@ 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)
8 +from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
9 from vllm.compilation.sequence_parallelism import SequenceParallelismPass
10 from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,
11 PassConfig, VllmConfig)
@@ -21,17 +19,6 @@ from vllm.utils import update_environment_variables
19 from ..utils import multi_gpu_test
20 from .backend import TestBackend
21
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 -
22 prompts = [
23 "Hello, my name is",
24 "The president of the United States is",
@@ -78,6 +65,18 @@ class TestModel(torch.nn.Module):
65
66 return norm_output, residual_output
67
68 + def ops_in_model_before(self):
69 + return [torch.ops.vllm.all_reduce.default]
70 +
71 + def ops_in_model_after(self):
72 + return [
73 + torch.ops.vllm.reduce_scatter.default,
74 + torch.ops.vllm.all_gather.default
75 + ]
76 +
77 + def ops_in_model(self):
78 + return [torch.ops._C.fused_add_rms_norm.default]
79 +
80
81 @multi_gpu_test(num_gpus=2)
82 @pytest.mark.parametrize("batch_size", [8])
@@ -156,26 +155,16 @@ def sequence_parallelism_pass_on_test_model(local_rank: int, world_size: int,
155 compiled_model_func = torch.compile(model, backend=backend_func)
156 compiled_model_func(hidden_states, residual)
157
159 - # Check substitution worked
160 - pre_nodes = backend_no_func.graph_pre_pass.nodes
161 - post_nodes = backend_no_func.graph_post_pass.nodes
162 -
158 # In pre-nodes, all reduce should be there,
159 # reduce scatter and all gather should not
165 - for op in OPS_IN_MODEL_BEFORE:
166 - find_specified_fn(pre_nodes, op)
167 - for op in OPS_IN_MODEL_AFTER:
168 - assert find_specified_fn_maybe(pre_nodes, op) is None
160 + backend_no_func.check_before_ops(model.ops_in_model_before())
161
162 # In post-nodes, reduce scatter and all gather should be there,
163 # all reduce should not
172 - for op in OPS_IN_MODEL_AFTER:
173 - find_specified_fn(post_nodes, op)
174 - for op in OPS_IN_MODEL_BEFORE:
175 - assert find_specified_fn_maybe(post_nodes, op) is None
164 + backend_no_func.check_after_ops(model.ops_in_model_after())
165
166 # check if the functionalization pass is applied
178 - for op in OPS_IN_MODEL:
167 + for op in model.ops_in_model():
168 find_auto_fn(backend_no_func.graph_post_pass.nodes, op)
169 assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes,
170 op) is None # noqa: E501
@@ -183,7 +172,7 @@ def sequence_parallelism_pass_on_test_model(local_rank: int, world_size: int,
172 # make sure the ops were all de-functionalized
173 found = dict()
174 for node in backend_func.graph_post_pass.nodes:
186 - for op in OPS_IN_MODEL:
175 + for op in model.ops_in_model():
176 if is_func(node, op):
177 found[op] = True
189 - assert all(found[op] for op in OPS_IN_MODEL)
178 + assert all(found[op] for op in model.ops_in_model())
vllm/compilation/collective_fusion.py new
+126
@@ -0,0 +1,126 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional
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 +from torch.distributed._symmetric_memory import enable_symm_mem_for_group
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.distributed import get_tp_group
12 +from vllm.distributed.parallel_state import (
13 + get_tensor_model_parallel_world_size)
14 +from vllm.logger import init_logger
15 +
16 +from .vllm_inductor_pass import VllmInductorPass
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +class BasePattern:
22 +
23 + def __init__(self, dtype: torch.dtype, device: str):
24 + self.dtype = dtype
25 + self.device = device
26 + self.tp = get_tp_group()
27 + self.tp_size = get_tensor_model_parallel_world_size()
28 +
29 +
30 +class GEMMReduceScatterPattern(BasePattern):
31 +
32 + def get_inputs(self):
33 + mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
34 + mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
35 + return [mul, mm_weight]
36 +
37 + def register(self, pm_pass: PatternMatcherPass):
38 +
39 + def pattern(mul: torch.Tensor, mm_weight: torch.Tensor):
40 + mm = torch.ops.aten.mm.default(mul, mm_weight)
41 + reduce_scatter = torch.ops.vllm.reduce_scatter.default(
42 + mm,
43 + dim=0,
44 + world_size=self.tp_size,
45 + group_name=self.tp.unique_name)
46 + return reduce_scatter
47 +
48 + def replacement(mul: torch.Tensor, mm_weight: torch.Tensor):
49 + gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
50 + mul,
51 + mm_weight,
52 + "avg",
53 + scatter_dim=0,
54 + group_name=self.tp.device_group.group_name,
55 + )
56 +
57 + return gemm_rs
58 +
59 + pm.register_replacement(pattern, replacement, self.get_inputs(),
60 + pm.fwd_only, pm_pass)
61 +
62 +
63 +class AllGatherGEMMPattern(BasePattern):
64 +
65 + def get_inputs(self):
66 + x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
67 + weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
68 +
69 + return [x, weight]
70 +
71 + def register(self, pm_pass: PatternMatcherPass):
72 +
73 + def pattern(
74 + x: torch.Tensor,
75 + weight: torch.Tensor,
76 + ) -> tuple[torch.Tensor, torch.Tensor]:
77 + all_gather = torch.ops.vllm.all_gather.default(
78 + x,
79 + dim=0,
80 + world_size=self.tp_size,
81 + group_name=self.tp.unique_name)
82 +
83 + return torch.ops.aten.mm.default(all_gather, weight)
84 +
85 + def replacement(
86 + x: torch.Tensor,
87 + weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
88 + ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
89 + x,
90 + [weight],
91 + gather_dim=0,
92 + group_name=self.tp.device_group.group_name,
93 + )
94 + return mm_outputs
95 +
96 + pm.register_replacement(pattern, replacement, self.get_inputs(),
97 + pm.fwd_only, pm_pass)
98 +
99 +
100 +class AsyncTPPass(VllmInductorPass):
101 +
102 + def __init__(self, config: VllmConfig):
103 + super().__init__(config)
104 +
105 + # Enable symmetric memory for the TP process group
106 + enable_symm_mem_for_group(get_tp_group().device_group.group_name)
107 + self.patterns: PatternMatcherPass = PatternMatcherPass(
108 + pass_name="async_tp_pass")
109 + GEMMReduceScatterPattern(self.model_dtype,
110 + self.device).register(self.patterns)
111 +
112 + AllGatherGEMMPattern(self.model_dtype,
113 + self.device).register(self.patterns)
114 +
115 + def is_applicable_for_shape(self, shape: Optional[int]) -> bool:
116 + # only do replace for specific shapes
117 + tp_size = get_tensor_model_parallel_world_size()
118 + return shape is not None and shape % tp_size == 0
119 +
120 + def __call__(self, graph: fx.Graph):
121 + self.begin()
122 + self.dump_graph(graph, "before_async_tp_pass")
123 + count = self.patterns.apply(graph)
124 + logger.debug("Replaced %s patterns", count)
125 + self.dump_graph(graph, "after_async_tp_pass")
126 + self.end_and_log()
vllm/compilation/pass_manager.py
+3
@@ -6,6 +6,7 @@ from vllm.config import VllmConfig
6 from vllm.logger import init_logger
7
8 from .activation_quant_fusion import ActivationQuantFusionPass
9 +from .collective_fusion import AsyncTPPass
10 from .fix_functionalization import FixFunctionalizationPass
11 from .fusion import FusionPass
12 from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
@@ -54,6 +55,8 @@ class PostGradPassManager(CustomGraphPass):
55
56 if self.pass_config.enable_sequence_parallelism:
57 self.passes += [SequenceParallelismPass(config)]
58 + if self.pass_config.enable_async_tp:
59 + self.passes += [AsyncTPPass(config)]
60
61 self.fix_functionalization = FixFunctionalizationPass(config)
62
vllm/compilation/sequence_parallelism.py
+5 -4
@@ -243,24 +243,25 @@ class SequenceParallelismPass(VllmInductorPass):
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)
246 + epsilon, self.model_dtype, self.device).register(self.patterns)
247
248 - MiddleAllReduceRMSNormPattern(epsilon, self.dtype,
248 + MiddleAllReduceRMSNormPattern(epsilon, self.model_dtype,
249 self.device).register(self.patterns)
250
251 - LastAllReduceRMSNormPattern(epsilon, self.dtype,
251 + LastAllReduceRMSNormPattern(epsilon, self.model_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
258 tp_size = get_tensor_model_parallel_world_size()
259 return shape is not None and shape % tp_size == 0
260
261 def __call__(self, graph: fx.Graph):
262 + self.begin()
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")
267 + self.end_and_log()
vllm/compilation/vllm_inductor_pass.py
+2 -1
@@ -26,7 +26,8 @@ class VllmInductorPass(InductorPass):
26
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
29 + self.model_dtype = config.model_config.dtype if config.model_config \
30 + else None
31 self.device = config.device_config.device if config.device_config \
32 else None
33 self.pass_name = self.__class__.__name__
vllm/config.py
+10 -1
@@ -3652,6 +3652,8 @@ class PassConfig:
3652 """Whether to enable the custom no-op elimination pass."""
3653 enable_sequence_parallelism: bool = False
3654 """Whether to enable sequence parallelism."""
3655 + enable_async_tp: bool = False
3656 + """Whether to enable async TP."""
3657
3658 def uuid(self):
3659 """
@@ -3661,7 +3663,8 @@ class PassConfig:
3663 compilation.
3664 """
3665 include = {
3664 - "enable_fusion", "enable_noop", "enable_sequence_parallelism"
3666 + "enable_fusion", "enable_noop", "enable_sequence_parallelism",
3667 + "enable_async_tp"
3668 }
3669 dict_ = {k: v for k, v in asdict(self).items() if k in include}
3670 return InductorPass.hash_dict(dict_)
@@ -4274,6 +4277,12 @@ class VllmConfig:
4277
4278 if self.compilation_config is None:
4279 self.compilation_config = CompilationConfig()
4280 +
4281 + # async tp is built on top of sequence parallelism
4282 + # and requires it to be enabled.
4283 + if self.compilation_config.pass_config.enable_async_tp:
4284 + self.compilation_config.pass_config.enable_sequence_parallelism = \
4285 + True
4286 if self.compilation_config.pass_config.enable_sequence_parallelism:
4287 self.compilation_config.custom_ops.append("+rms_norm")
4288 if envs.VLLM_USE_V1 and self.model_config is not None and \
vllm/distributed/parallel_state.py
+24 -2
@@ -120,7 +120,7 @@ def reduce_scatter(tensor: torch.Tensor, dim: int, world_size: int,
120 group = _groups[group_name]()
121 if group is None:
122 raise ValueError(f"Group {group_name} is destroyed.")
123 - return group.reduce_scatter(tensor, dim)
123 + return group._reduce_scatter_out_place(tensor, dim)
124
125
126 def reduce_scatter_fake(tensor: torch.Tensor, dim: int, world_size: int,
@@ -136,7 +136,7 @@ def all_gather(tensor: torch.Tensor, dim: int, world_size: int,
136 group = _groups[group_name]()
137 if group is None:
138 raise ValueError(f"Group {group_name} is destroyed.")
139 - return group.all_gather(tensor, dim)
139 + return group._all_gather_out_place(tensor, dim)
140
141
142 def all_gather_fake(tensor: torch.Tensor, dim: int, world_size: int,
@@ -161,6 +161,7 @@ if supports_custom_op():
161 op_func=reduce_scatter,
162 mutates_args=[],
163 fake_impl=reduce_scatter_fake,
164 + dispatch_key=current_platform.dispatch_key,
165 )
166
167 direct_register_custom_op(
@@ -168,6 +169,7 @@ if supports_custom_op():
169 op_func=all_gather,
170 mutates_args=[],
171 fake_impl=all_gather_fake,
172 + dispatch_key=current_platform.dispatch_key,
173 )
174
175
@@ -367,6 +369,16 @@ class GroupCoordinator:
369 assert -input_.dim() <= dim < input_.dim(), (
370 f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
371
372 + if self.use_custom_op_call:
373 + return torch.ops.vllm.all_gather(input_,
374 + dim,
375 + world_size,
376 + group_name=self.unique_name)
377 + else:
378 + return self._all_gather_out_place(input_, dim)
379 +
380 + def _all_gather_out_place(self, input_: torch.Tensor,
381 + dim: int) -> torch.Tensor:
382 return self.device_communicator.all_gather(input_, dim)
383
384 def reduce_scatter(self,
@@ -379,6 +391,16 @@ class GroupCoordinator:
391 assert -input_.dim() <= dim < input_.dim(), (
392 f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
393
394 + if self.use_custom_op_call:
395 + return torch.ops.vllm.reduce_scatter(input_,
396 + dim,
397 + world_size,
398 + group_name=self.unique_name)
399 + else:
400 + return self._reduce_scatter_out_place(input_, dim)
401 +
402 + def _reduce_scatter_out_place(self, input_: torch.Tensor,
403 + dim: int) -> torch.Tensor:
404 return self.device_communicator.reduce_scatter(input_, dim)
405
406 def gather(self,