@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e6327c9b3e
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, 04:45 UTC

[Feature] Support sequence parallelism for static fp8 quantization (#19181)

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

cascade committed Jun 23, 2025 at 13:09 UTC e6327c9b3eb2111305d6b19bb3e5d1fd852eea55
7 files changed +534 -198
tests/compile/test_sequence_parallelism.py
+144 -17
@@ -6,7 +6,9 @@ import torch
6
7 import vllm.envs as envs
8 from vllm.compilation.fix_functionalization import FixFunctionalizationPass
9 +from vllm.compilation.fusion import FusionPass
10 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
11 +from vllm.compilation.noop_elimination import NoOpEliminationPass
12 from vllm.compilation.sequence_parallelism import SequenceParallelismPass
13 from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,
14 PassConfig, VllmConfig)
@@ -14,12 +16,15 @@ from vllm.distributed import tensor_model_parallel_all_reduce
16 from vllm.distributed.parallel_state import (init_distributed_environment,
17 initialize_model_parallel)
18 from vllm.model_executor.layers.layernorm import RMSNorm
19 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20 + Fp8LinearOp)
21 from vllm.platforms import current_platform
22 from vllm.utils import update_environment_variables
23
24 from ..utils import multi_gpu_test
25 from .backend import TestBackend
26
27 +FP8_DTYPE = current_platform.fp8_dtype()
28 prompts = [
29 "Hello, my name is",
30 "The president of the United States is",
@@ -30,13 +35,16 @@ prompts = [
35
36 class TestModel(torch.nn.Module):
37
33 - def __init__(self, hidden_size=16, intermediate_size=32):
38 + def __init__(self,
39 + hidden_size=16,
40 + intermediate_size=32,
41 + vllm_config: VllmConfig = None):
42 super().__init__()
43 self.hidden_size = hidden_size
44 self.intermediate_size = intermediate_size
45 self.gate_proj = torch.nn.Parameter(
46 torch.empty((intermediate_size, hidden_size)))
39 - self.norm = RMSNorm(hidden_size, 1e-05)
47 + self.norm = RMSNorm(intermediate_size, 1e-05)
48 # Initialize weights
49 torch.nn.init.normal_(self.gate_proj, std=0.02)
50
@@ -79,32 +87,138 @@ class TestModel(torch.nn.Module):
87 return [torch.ops._C.fused_add_rms_norm.default]
88
89
90 +class TestQuantModel(torch.nn.Module):
91 +
92 + def __init__(self,
93 + hidden_size=16,
94 + intermediate_size=32,
95 + vllm_config: VllmConfig = None):
96 + super().__init__()
97 + self.hidden_size = hidden_size
98 + self.intermediate_size = intermediate_size
99 + self.vllm_config = vllm_config
100 + self.gate_proj = torch.nn.Parameter(torch.empty(
101 + (intermediate_size, hidden_size)),
102 + requires_grad=False)
103 + self.norm = RMSNorm(intermediate_size, 1e-05)
104 + # Initialize weights
105 + torch.nn.init.normal_(self.gate_proj, std=0.02)
106 +
107 + self.fp8_linear = Fp8LinearOp(cutlass_fp8_supported=True,
108 + use_per_token_if_dynamic=False)
109 +
110 + self.scale = torch.rand(1, dtype=torch.float32)
111 + # Create a weight that is compatible with torch._scaled_mm,
112 + # which expects a column-major layout.
113 + self.w = torch.rand(hidden_size,
114 + intermediate_size).to(dtype=FP8_DTYPE).t()
115 + self.wscale = torch.rand(1, dtype=torch.float32)
116 +
117 + def forward(self, hidden_states, residual):
118 + """
119 + Forward pass implementing the operations in the FX graph
120 +
121 + Args:
122 + hidden_states: Input tensor
123 + residual: Residual tensor from previous layer
124 +
125 + Returns:
126 + Tuple containing the output tensor
127 + """
128 + # Reshape input
129 + view = hidden_states.reshape(-1, self.hidden_size)
130 +
131 + #matrix multiplication
132 + permute = self.gate_proj.permute(1, 0)
133 + mm = torch.mm(view, permute)
134 +
135 + # Tensor parallel all-reduce
136 + all_reduce = tensor_model_parallel_all_reduce(mm)
137 +
138 + # layer normalization
139 + norm_output, residual_output = self.norm(all_reduce, residual)
140 +
141 + # for static input quantization
142 + # self.fp8_linear is initialized with use_per_token_if_dynamic=False
143 + fp8_linear_result = self.fp8_linear.apply(norm_output,
144 + self.w,
145 + self.wscale,
146 + input_scale=self.scale.to(
147 + norm_output.device))
148 +
149 + return fp8_linear_result, residual_output
150 +
151 + def ops_in_model_before(self):
152 + ops_to_remove = [torch.ops.vllm.all_reduce.default
153 + ] # Always removed by SP
154 + # The following are only removed if fusion happens
155 + if self.vllm_config and self.vllm_config.compilation_config \
156 + .pass_config.enable_fusion:
157 + ops_to_remove.extend([
158 + torch.ops._C.fused_add_rms_norm.default,
159 + torch.ops._C.static_scaled_fp8_quant.default,
160 + ])
161 + return ops_to_remove
162 +
163 + def ops_in_model_after(self):
164 + ops_to_add = [
165 + torch.ops.vllm.reduce_scatter.default,
166 + torch.ops.vllm.all_gather.default
167 + ]
168 + # The following is only added if fusion happens
169 + if self.vllm_config and self.vllm_config.compilation_config \
170 + .pass_config.enable_fusion:
171 + ops_to_add.append(
172 + torch.ops._C.fused_add_rms_norm_static_fp8_quant.default)
173 + return ops_to_add
174 +
175 + def ops_in_model(self):
176 + if self.vllm_config and self.vllm_config.compilation_config \
177 + .pass_config.enable_fusion:
178 + # If fusion happens, the fused op is the one
179 + # we check for (de)functionalization
180 + return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default
181 + ] # noqa: E501
182 + else:
183 + # If no fusion, the original ops are checked
184 + return [
185 + torch.ops._C.fused_add_rms_norm.default,
186 + # TODO functionalization pass does not handle this yet
187 + # torch.ops._C.static_scaled_fp8_quant.default,
188 + ]
189 +
190 +
191 @multi_gpu_test(num_gpus=2)
192 +@pytest.mark.parametrize("test_model_cls", [TestModel, TestQuantModel])
193 @pytest.mark.parametrize("batch_size", [8])
194 @pytest.mark.parametrize("seq_len", [16])
195 @pytest.mark.parametrize("hidden_size", [16])
196 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
197 +@pytest.mark.parametrize("enable_fusion", [True, False])
198 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
199 reason="Only test on CUDA")
89 -def test_sequence_parallelism_pass(batch_size: int, seq_len: int,
90 - hidden_size: int, dtype: torch.dtype):
200 +def test_sequence_parallelism_pass(test_model_cls: type[torch.nn.Module],
201 + batch_size: int, seq_len: int,
202 + hidden_size: int, dtype: torch.dtype,
203 + enable_fusion: bool):
204 num_processes = 2
205
206 def run_torch_spawn(fn, nprocs):
207 # need to use torch.mp.spawn otherwise will have problems with
208 # torch.distributed and cuda
209 torch.multiprocessing.spawn(fn,
97 - args=(num_processes, batch_size, seq_len,
98 - hidden_size, dtype),
210 + args=(num_processes, test_model_cls,
211 + batch_size, seq_len, hidden_size,
212 + dtype, enable_fusion),
213 nprocs=nprocs)
214
215 run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
216
217
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):
218 +def sequence_parallelism_pass_on_test_model(
219 + local_rank: int, world_size: int,
220 + test_model_cls: type[torch.nn.Module], batch_size: int, seq_len: int,
221 + hidden_size: int, dtype: torch.dtype, enable_fusion: bool):
222 current_platform.seed_everything(0)
223
224 device = torch.device(f"cuda:{local_rank}")
@@ -127,26 +241,39 @@ def sequence_parallelism_pass_on_test_model(local_rank: int, world_size: int,
241 # configure vllm config for SequenceParallelismPass
242 vllm_config = VllmConfig()
243 vllm_config.compilation_config = CompilationConfig(pass_config=PassConfig(
130 - enable_sequence_parallelism=True))
244 + enable_sequence_parallelism=True,
245 + enable_fusion=enable_fusion,
246 + enable_noop=True)) # NoOp needed for fusion
247 vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
248
249 # this is a fake model name to construct the model config
250 # in the vllm_config, it's not really used.
135 - model = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
136 - vllm_config.model_config = ModelConfig(model=model,
251 + model_name = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
252 + vllm_config.model_config = ModelConfig(model=model_name,
253 task="auto",
138 - tokenizer=model,
254 + tokenizer=model_name,
255 tokenizer_mode="auto",
256 trust_remote_code=True,
257 dtype=dtype,
258 seed=42)
259
260 sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
145 - backend_no_func = TestBackend(sequence_parallelism_pass)
261 + noop_pass = NoOpEliminationPass(vllm_config)
262 func_pass = FixFunctionalizationPass(vllm_config)
147 - backend_func = TestBackend(sequence_parallelism_pass, func_pass)
263
149 - model = TestModel(hidden_size, hidden_size * 2)
264 + passes_for_backend = [noop_pass, sequence_parallelism_pass]
265 +
266 + if enable_fusion:
267 + fusion_pass = FusionPass.instance(vllm_config)
268 + passes_for_backend.append(fusion_pass)
269 +
270 + backend_no_func = TestBackend(*passes_for_backend)
271 + backend_func = TestBackend(*passes_for_backend, func_pass)
272 +
273 + model = test_model_cls(hidden_size,
274 + hidden_size * 2,
275 + vllm_config=vllm_config)
276 +
277 hidden_states = torch.randn((batch_size * seq_len, hidden_size),
278 dtype=dtype)
279 residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
tests/distributed/test_sequence_parallel.py
+52 -56
@@ -28,7 +28,7 @@ VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
28 class ParallelSetup(NamedTuple):
29 tp_size: int
30 pp_size: int
31 - sp_enabled: bool
31 + enable_fusion: bool
32 eager_mode: bool
33 chunked_prefill: bool
34
@@ -67,49 +67,18 @@ class SPTestSettings:
67 task: TaskOption = "auto",
68 load_format: Optional[str] = None,
69 ):
70 + parallel_setups = []
71 + for eager_mode_val in [False, True]:
72 + for pp_multiplier in [1, 2]:
73 + for chunked_prefill_val in [False, True]:
74 + parallel_setups.append(
75 + ParallelSetup(tp_size=tp_base,
76 + pp_size=pp_multiplier * pp_base,
77 + enable_fusion=False,
78 + eager_mode=eager_mode_val,
79 + chunked_prefill=chunked_prefill_val))
80 return SPTestSettings(
71 - parallel_setups=[
72 - ParallelSetup(tp_size=tp_base,
73 - pp_size=pp_base,
74 - sp_enabled=True,
75 - eager_mode=False,
76 - chunked_prefill=False),
77 - ParallelSetup(tp_size=tp_base,
78 - pp_size=pp_base,
79 - sp_enabled=True,
80 - eager_mode=False,
81 - chunked_prefill=True),
82 - ParallelSetup(tp_size=tp_base,
83 - pp_size=pp_base,
84 - sp_enabled=True,
85 - eager_mode=True,
86 - chunked_prefill=False),
87 - ParallelSetup(tp_size=tp_base,
88 - pp_size=pp_base,
89 - sp_enabled=True,
90 - eager_mode=True,
91 - chunked_prefill=True),
92 - ParallelSetup(tp_size=tp_base,
93 - pp_size=2 * pp_base,
94 - sp_enabled=True,
95 - eager_mode=False,
96 - chunked_prefill=False),
97 - ParallelSetup(tp_size=tp_base,
98 - pp_size=2 * pp_base,
99 - sp_enabled=True,
100 - eager_mode=False,
101 - chunked_prefill=True),
102 - ParallelSetup(tp_size=tp_base,
103 - pp_size=2 * pp_base,
104 - sp_enabled=True,
105 - eager_mode=True,
106 - chunked_prefill=False),
107 - ParallelSetup(tp_size=tp_base,
108 - pp_size=2 * pp_base,
109 - sp_enabled=True,
110 - eager_mode=True,
111 - chunked_prefill=True)
112 - ],
81 + parallel_setups=parallel_setups,
82 distributed_backends=["mp", "ray"],
83 vllm_major_versions=["1", "1"],
84 task=task,
@@ -126,19 +95,44 @@ class SPTestSettings:
95 multi_node_only: bool = False,
96 load_format: Optional[str] = None,
97 ):
98 + parallel_setups = []
99 + for eager_mode_val in [False, True]:
100 + for pp_multiplier in [1, 2]:
101 + for chunked_prefill_val in [False, True]:
102 + parallel_setups.append(
103 + ParallelSetup(tp_size=tp_base,
104 + pp_size=pp_multiplier * pp_base,
105 + enable_fusion=False,
106 + eager_mode=eager_mode_val,
107 + chunked_prefill=chunked_prefill_val))
108 return SPTestSettings(
130 - parallel_setups=[
109 + parallel_setups=parallel_setups,
110 + distributed_backends=["mp", "ray"],
111 + vllm_major_versions=["1", "1"],
112 + task=task,
113 + test_options=SPTestOptions(multi_node_only=multi_node_only,
114 + load_format=load_format),
115 + )
116 +
117 + @staticmethod
118 + def fp8_quant(
119 + *,
120 + tp_base: int = 2,
121 + pp_base: int = 1,
122 + task: TaskOption = "auto",
123 + multi_node_only: bool = False,
124 + load_format: Optional[str] = None,
125 + ):
126 + parallel_setups = []
127 + for fusion_val in [False, True]:
128 + parallel_setups.append(
129 ParallelSetup(tp_size=tp_base,
130 pp_size=pp_base,
133 - sp_enabled=True,
134 - eager_mode=False,
135 - chunked_prefill=False),
136 - ParallelSetup(tp_size=tp_base,
137 - pp_size=2 * pp_base,
138 - sp_enabled=True,
139 - eager_mode=False,
140 - chunked_prefill=False),
141 - ],
131 + enable_fusion=fusion_val,
132 + eager_mode=True,
133 + chunked_prefill=False))
134 + return SPTestSettings(
135 + parallel_setups=parallel_setups,
136 distributed_backends=["mp", "ray"],
137 vllm_major_versions=["1", "1"],
138 task=task,
@@ -171,7 +165,7 @@ def _compare_sp(
165 (
166 tp_size,
167 pp_size,
174 - sp_enabled,
168 + enable_fusion,
169 eager_mode,
170 chunked_prefill,
171 ) = parallel_setup
@@ -240,9 +234,9 @@ def _compare_sp(
234 'compile_sizes': [4, 8],
235 'splitting_ops': [],
236 'pass_config': {
243 - 'enable_sequence_parallelism': sp_enabled,
237 + 'enable_sequence_parallelism': True,
238 + 'enable_fusion': enable_fusion,
239 'enable_noop': True,
245 - 'enable_fusion': True,
240 },
241 }
242
@@ -291,12 +285,14 @@ def _compare_sp(
285 SP_TEXT_GENERATION_MODELS = {
286 # [Decoder-only]
287 "meta-llama/Llama-3.2-1B-Instruct": SPTestSettings.fast(),
288 + "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8": SPTestSettings.fp8_quant(),
289 }
290
291 SP_TEST_MODELS = [
292 # TODO support other models
293 # [LANGUAGE GENERATION]
294 "meta-llama/Llama-3.2-1B-Instruct",
295 + "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"
296 ]
297
298
tests/models/registry.py
+2 -1
@@ -193,7 +193,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
193 extras={"tiny": "ai21labs/Jamba-tiny-dev"}), # noqa: E501
194 "LlamaForCausalLM": _HfExamplesInfo("meta-llama/Llama-3.2-1B-Instruct",
195 extras={"guard": "meta-llama/Llama-Guard-3-1B", # noqa: E501
196 - "hermes": "NousResearch/Hermes-3-Llama-3.1-8B"}), # noqa: E501
196 + "hermes": "NousResearch/Hermes-3-Llama-3.1-8B", # noqa: E501
197 + "fp8": "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"}), # noqa: E501
198 "LLaMAForCausalLM": _HfExamplesInfo("decapoda-research/llama-7b-hf",
199 is_available_online=False),
200 "MambaForCausalLM": _HfExamplesInfo("state-spaces/mamba-130m-hf"),
vllm/compilation/fusion.py
+2 -2
@@ -345,8 +345,8 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
345 # 0 is always None
346 fused_return_mapping = {1: (quant_node, 1), 2: (rms_node, 2)}
347 self.insert_fused_node(fused_return_mapping,
348 - epsilon=rms_node.kwargs["epsilon"],
349 - **kwargs)
348 + **kwargs,
349 + epsilon=rms_node.kwargs["epsilon"])
350
351
352 class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
vllm/compilation/pass_manager.py
+4 -4
@@ -51,15 +51,15 @@ class PostGradPassManager(CustomGraphPass):
51 if self.pass_config.enable_noop:
52 self.passes += [NoOpEliminationPass(config)]
53
54 - if self.pass_config.enable_fusion:
55 - self.passes += [FusionPass.instance(config)]
56 - self.passes += [ActivationQuantFusionPass(config)]
57 -
54 if self.pass_config.enable_sequence_parallelism:
55 self.passes += [SequenceParallelismPass(config)]
56 if self.pass_config.enable_async_tp:
57 self.passes += [AsyncTPPass(config)]
58
59 + if self.pass_config.enable_fusion:
60 + self.passes += [FusionPass.instance(config)]
61 + self.passes += [ActivationQuantFusionPass(config)]
62 +
63 if self.pass_config.enable_attn_fusion:
64 self.passes += [AttnFusionPass(config)]
65
vllm/compilation/sequence_parallelism.py
+328 -114
@@ -12,91 +12,142 @@ from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
12 from vllm.distributed.parallel_state import (
13 get_tensor_model_parallel_world_size)
14 from vllm.logger import init_logger
15 +from vllm.platforms import current_platform
16
17 from .vllm_inductor_pass import VllmInductorPass
18
19 logger = init_logger(__name__)
20
21
21 -class AllReduceRMSNormPattern:
22 +class _RMSNormAndQuantOpHelper:
23 + """Base helper for RMSNorm and RMSNorm + Quantization functionalization."""
24
23 - def __init__(self, epsilon: float, dtype: torch.dtype, device: str):
25 + def __init__(self,
26 + epsilon: float,
27 + dtype: torch.dtype,
28 + device: str,
29 + quant_op: Optional[torch._ops.OpOverload] = None,
30 + **kwargs):
31 self.epsilon = epsilon
32 self.dtype = dtype
33 self.device = device
27 -
28 -
29 -class EmbeddingAllReduceRMSNormPattern(AllReduceRMSNormPattern):
34 + self.quant_op = quant_op
35 +
36 + def _functional_rmsnorm(self, result_buffer, input_tensor, weight_tensor):
37 + return torch.ops.higher_order.auto_functionalized(
38 + torch.ops._C.rms_norm.default,
39 + result=result_buffer,
40 + input=input_tensor,
41 + weight=weight_tensor,
42 + epsilon=self.epsilon)
43 +
44 + def _functional_fused_add_rmsnorm(self, input_tensor, residual_tensor,
45 + weight_tensor):
46 + return torch.ops.higher_order.auto_functionalized(
47 + torch.ops._C.fused_add_rms_norm.default,
48 + input=input_tensor,
49 + residual=residual_tensor,
50 + weight=weight_tensor,
51 + epsilon=self.epsilon)
52 +
53 + def _functional_rmsnorm_then_quant(self, rmsnorm_result_buffer,
54 + quant_result_buffer, input_tensor,
55 + weight_tensor, scale_tensor):
56 + if self.quant_op is None:
57 + raise RuntimeError(
58 + "_RMSNormAndQuantOpHelper was not initialized with a quant_op."
59 + )
60 + rmsnorm_out_tuple = self._functional_rmsnorm(rmsnorm_result_buffer,
61 + input_tensor,
62 + weight_tensor)
63 + quant_out_tuple = torch.ops.higher_order.auto_functionalized(
64 + self.quant_op,
65 + result=quant_result_buffer,
66 + input=rmsnorm_out_tuple[1],
67 + scale=scale_tensor)
68 + return quant_out_tuple
69 +
70 + def _functional_fused_add_rmsnorm_then_quant(self, quant_result_buffer,
71 + input_tensor, residual_tensor,
72 + weight_tensor, scale_tensor):
73 + if self.quant_op is None:
74 + raise RuntimeError(
75 + "_RMSNormAndQuantOpHelper was not initialized with a quant_op."
76 + )
77 + fused_add_rmsnorm_out_tuple = self._functional_fused_add_rmsnorm(
78 + input_tensor, residual_tensor, weight_tensor)
79 + quant_out_tuple = torch.ops.higher_order.auto_functionalized(
80 + self.quant_op,
81 + result=quant_result_buffer,
82 + input=fused_add_rmsnorm_out_tuple[1],
83 + scale=scale_tensor)
84 + return quant_out_tuple, fused_add_rmsnorm_out_tuple[2]
85 +
86 +
87 +class _SequenceParallelPatternHelper(_RMSNormAndQuantOpHelper):
88 + """Helper for sequence parallelism patterns."""
89 +
90 + def __init__(self,
91 + epsilon: float,
92 + dtype: torch.dtype,
93 + device: str,
94 + quant_op: Optional[torch._ops.OpOverload] = None,
95 + **kwargs):
96 + super().__init__(epsilon, dtype, device, quant_op=quant_op, **kwargs)
97 + self.tp_group = get_tp_group()
98 + self.tp_size = get_tensor_model_parallel_world_size()
99 +
100 + def _all_reduce(self, x: torch.Tensor) -> torch.Tensor:
101 + return tensor_model_parallel_all_reduce(x)
102 +
103 + def _reduce_scatter(self, x: torch.Tensor) -> torch.Tensor:
104 + return torch.ops.vllm.reduce_scatter.default(
105 + x,
106 + dim=0,
107 + world_size=self.tp_size,
108 + group_name=self.tp_group.unique_name)
109 +
110 + def _all_gather(self, x: torch.Tensor) -> torch.Tensor:
111 + return torch.ops.vllm.all_gather.default(
112 + x,
113 + dim=0,
114 + world_size=self.tp_size,
115 + group_name=self.tp_group.unique_name)
116 +
117 +
118 +class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
119
120 def get_inputs(self):
32 - arg2_1 = torch.empty([16, 4], device=self.device, dtype=self.dtype)
33 - mul_6 = torch.tensor([[3, 7, 1, 4, 9, 2, 5, 0]],
34 - device=self.device,
35 - dtype=torch.long)
36 - unsqueeze = torch.rand([1, 8, 1], device=self.device, \
37 - dtype=self.dtype) > 0.5
38 - full_default = torch.zeros([1, 8, 4], device=self.device, \
39 - dtype=self.dtype)
121 + input = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
122 permute = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
123 arg3_1 = torch.empty([4], device=self.device, dtype=self.dtype)
124
43 - return [arg2_1, mul_6, unsqueeze, full_default, permute, arg3_1]
125 + return [input, permute, arg3_1]
126
127 def register(self, pm_pass: PatternMatcherPass):
128
129 def pattern(
48 - arg2_1: torch.Tensor,
49 - mul_6: torch.Tensor,
50 - unsqueeze: torch.Tensor,
51 - full_default: torch.Tensor,
130 + input: torch.Tensor,
131 permute: torch.Tensor,
132 arg3_1: torch.Tensor,
133 ):
55 - embedding = torch.ops.aten.embedding.default(arg2_1, mul_6)
56 - where = torch.ops.aten.where.self(unsqueeze, full_default,
57 - embedding)
58 - all_reduce = tensor_model_parallel_all_reduce(where)
59 - rmsnorm = torch.ops.higher_order.auto_functionalized(
60 - torch.ops._C.rms_norm.default,
61 - result=permute,
62 - input=all_reduce,
63 - weight=arg3_1,
64 - epsilon=self.epsilon,
65 - )
134 + all_reduce = self._all_reduce(input)
135 + rmsnorm = self._functional_rmsnorm(permute, all_reduce, arg3_1)
136
137 return rmsnorm[1], all_reduce
138
139 def replacement(
70 - arg2_1: torch.Tensor,
71 - mul_6: torch.Tensor,
72 - unsqueeze: torch.Tensor,
73 - full_default: torch.Tensor,
140 + input: torch.Tensor,
141 permute: torch.Tensor,
142 arg3_1: torch.Tensor,
143 ):
77 - embedding = torch.ops.aten.embedding.default(arg2_1, mul_6)
78 - where = torch.ops.aten.where.self(unsqueeze, full_default,
79 - embedding)
80 -
81 - tp = get_tp_group()
82 - tp_size = get_tensor_model_parallel_world_size()
83 - reduce_scatter = torch.ops.vllm.reduce_scatter.default(
84 - where, dim=0, world_size=tp_size, group_name=tp.unique_name)
144 + reduce_scatter = self._reduce_scatter(input)
145
146 rmsnorm_result = torch.empty_like(reduce_scatter)
87 - rmsnorm = torch.ops.higher_order.auto_functionalized(
88 - torch.ops._C.rms_norm.default,
89 - result=rmsnorm_result,
90 - input=reduce_scatter,
91 - weight=arg3_1,
92 - epsilon=self.epsilon,
93 - )
147 + rmsnorm = self._functional_rmsnorm(rmsnorm_result, reduce_scatter,
148 + arg3_1)
149
95 - all_gather = torch.ops.vllm.all_gather.default(
96 - rmsnorm[1],
97 - dim=0,
98 - world_size=tp_size,
99 - group_name=tp.unique_name)
150 + all_gather = self._all_gather(rmsnorm[1])
151
152 return all_gather, reduce_scatter
153
@@ -104,7 +155,7 @@ class EmbeddingAllReduceRMSNormPattern(AllReduceRMSNormPattern):
155 pm.fwd_only, pm_pass)
156
157
107 -class MiddleAllReduceRMSNormPattern(AllReduceRMSNormPattern):
158 +class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
159
160 def get_inputs(self):
161 mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -127,16 +178,9 @@ class MiddleAllReduceRMSNormPattern(AllReduceRMSNormPattern):
178 mm_1: torch.Tensor,
179 rms_norm_weights: torch.Tensor,
180 ) -> tuple[torch.Tensor, torch.Tensor]:
130 - all_reduce = tensor_model_parallel_all_reduce(mm_1)
131 -
132 - rmsnorm = torch.ops.higher_order.auto_functionalized(
133 - torch.ops._C.fused_add_rms_norm.default,
134 - input=all_reduce,
135 - residual=residual,
136 - weight=rms_norm_weights,
137 - epsilon=self.epsilon,
138 - )
139 -
181 + all_reduce = self._all_reduce(mm_1)
182 + rmsnorm = self._functional_fused_add_rmsnorm(
183 + all_reduce, residual, rms_norm_weights)
184 return rmsnorm[1], rmsnorm[2]
185
186 def replacement(
@@ -144,32 +188,17 @@ class MiddleAllReduceRMSNormPattern(AllReduceRMSNormPattern):
188 mm_1: torch.Tensor,
189 rms_norm_weights: torch.Tensor,
190 ) -> tuple[torch.Tensor, torch.Tensor]:
147 - tp = get_tp_group()
148 - tp_size = get_tensor_model_parallel_world_size()
149 - reduce_scatter = torch.ops.vllm.reduce_scatter.default(
150 - mm_1, dim=0, world_size=tp_size, group_name=tp.unique_name)
151 -
152 - # TODO is it possible to extract epsilon from somewhere
153 - rmsnorm = torch.ops.higher_order.auto_functionalized(
154 - torch.ops._C.fused_add_rms_norm.default,
155 - input=reduce_scatter,
156 - residual=residual,
157 - weight=rms_norm_weights,
158 - epsilon=self.epsilon,
159 - )
160 -
161 - all_gather = torch.ops.vllm.all_gather.default(
162 - rmsnorm[1],
163 - dim=0,
164 - world_size=tp_size,
165 - group_name=tp.unique_name)
191 + reduce_scatter = self._reduce_scatter(mm_1)
192 + rmsnorm = self._functional_fused_add_rmsnorm(
193 + reduce_scatter, residual, rms_norm_weights)
194 + all_gather = self._all_gather(rmsnorm[1])
195 return all_gather, rmsnorm[2]
196
197 pm.register_replacement(pattern, replacement, self.get_inputs(),
198 pm.fwd_only, pm_pass)
199
200
172 -class LastAllReduceRMSNormPattern(AllReduceRMSNormPattern):
201 +class LastAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
202
203 def get_inputs(self):
204 mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -192,16 +221,9 @@ class LastAllReduceRMSNormPattern(AllReduceRMSNormPattern):
221 mm_1: torch.Tensor,
222 rms_norm_weights: torch.Tensor,
223 ) -> tuple[torch.Tensor, torch.Tensor]:
195 - all_reduce = tensor_model_parallel_all_reduce(mm_1)
196 -
197 - rmsnorm = torch.ops.higher_order.auto_functionalized(
198 - torch.ops._C.fused_add_rms_norm.default,
199 - input=all_reduce,
200 - residual=residual,
201 - weight=rms_norm_weights,
202 - epsilon=self.epsilon,
203 - )
204 -
224 + all_reduce = self._all_reduce(mm_1)
225 + rmsnorm = self._functional_fused_add_rmsnorm(
226 + all_reduce, residual, rms_norm_weights)
227 return rmsnorm[1]
228
229 def replacement(
@@ -209,26 +231,185 @@ class LastAllReduceRMSNormPattern(AllReduceRMSNormPattern):
231 mm_1: torch.Tensor,
232 rms_norm_weights: torch.Tensor,
233 ) -> tuple[torch.Tensor, torch.Tensor]:
212 - tp = get_tp_group()
213 - tp_size = get_tensor_model_parallel_world_size()
214 - reduce_scatter = torch.ops.vllm.reduce_scatter.default(
215 - mm_1, dim=0, world_size=tp_size, group_name=tp.unique_name)
216 -
217 - # TODO is it possible to extract epsilon from somewhere
218 - rmsnorm = torch.ops.higher_order.auto_functionalized(
219 - torch.ops._C.fused_add_rms_norm.default,
220 - input=reduce_scatter,
221 - residual=residual,
222 - weight=rms_norm_weights,
223 - epsilon=self.epsilon,
224 - )
234 + reduce_scatter = self._reduce_scatter(mm_1)
235 + rmsnorm = self._functional_fused_add_rmsnorm(
236 + reduce_scatter, residual, rms_norm_weights)
237 + normalized = self._all_gather(rmsnorm[1])
238 + return normalized
239 +
240 + pm.register_replacement(pattern, replacement, self.get_inputs(),
241 + pm.fwd_only, pm_pass)
242 +
243 +
244 +FP8_DTYPE = current_platform.fp8_dtype()
245 +
246 +
247 +class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
248 +
249 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str,
250 + op: torch._ops.OpOverload):
251 + super().__init__(epsilon, dtype, device, quant_op=op)
252 +
253 + def get_inputs(self):
254 + input = torch.zeros([1, 8, 4], device=self.device, dtype=self.dtype)
255 + rmsnorm_result = torch.empty([1, 8, 4],
256 + device=self.device,
257 + dtype=self.dtype)
258 + quant_result = torch.empty([1, 8, 4],
259 + device=self.device,
260 + dtype=FP8_DTYPE)
261 + weight = torch.empty([4], device=self.device, dtype=self.dtype)
262 + scale = torch.tensor(1.0, device=self.device, dtype=torch.float32)
263 + return [input, rmsnorm_result, quant_result, weight, scale]
264 +
265 + def register(self, pm_pass: PatternMatcherPass):
266 +
267 + def pattern(
268 + input: torch.Tensor,
269 + rmsnorm_result: torch.Tensor,
270 + quant_result: torch.Tensor,
271 + weight: torch.Tensor,
272 + scale: torch.Tensor,
273 + ):
274 + all_reduce = self._all_reduce(input)
275 + static_fp8 = self._functional_rmsnorm_then_quant(
276 + rmsnorm_result, quant_result, all_reduce, weight, scale)
277 + return static_fp8[1], all_reduce
278 +
279 + def replacement(
280 + input: torch.Tensor,
281 + rmsnorm_result: torch.Tensor,
282 + quant_result: torch.Tensor,
283 + weight: torch.Tensor,
284 + scale: torch.Tensor,
285 + ):
286 + reduce_scatter = self._reduce_scatter(input)
287 +
288 + rmsnorm_result = torch.empty_like(reduce_scatter,
289 + dtype=rmsnorm_result.dtype)
290 + quant_result = torch.empty_like(
291 + rmsnorm_result, # Output of RMSNorm
292 + dtype=quant_result.dtype)
293 + static_fp8 = self._functional_rmsnorm_then_quant(
294 + rmsnorm_result, quant_result, reduce_scatter, weight, scale)
295 + all_gather = self._all_gather(static_fp8[1])
296 +
297 + return all_gather, reduce_scatter
298 +
299 + pm.register_replacement(pattern, replacement, self.get_inputs(),
300 + pm.fwd_only, pm_pass)
301 +
302 +
303 +class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
304 +
305 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str,
306 + op: torch._ops.OpOverload):
307 + super().__init__(epsilon, dtype, device, quant_op=op)
308 +
309 + def get_inputs(self):
310 + mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
311 +
312 + residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
313 + rms_norm_weights = torch.empty([4, 4],
314 + device=self.device,
315 + dtype=self.dtype)
316 + result = torch.empty([4, 4], device=self.device, dtype=FP8_DTYPE)
317 + scale = torch.empty([1, 1], device=self.device, dtype=torch.float32)
318 +
319 + return [
320 + result,
321 + residual,
322 + mm_1,
323 + rms_norm_weights,
324 + scale,
325 + ]
326 +
327 + def register(self, pm_pass: PatternMatcherPass):
328 +
329 + def pattern(
330 + result: torch.Tensor,
331 + residual: torch.Tensor,
332 + mm_1: torch.Tensor,
333 + rms_norm_weights: torch.Tensor,
334 + scale: torch.Tensor,
335 + ) -> tuple[torch.Tensor, torch.Tensor]:
336 + all_reduce = self._all_reduce(mm_1)
337 + static_fp8, rmsnorm_residual_out = self._functional_fused_add_rmsnorm_then_quant( # noqa: E501
338 + result, all_reduce, residual, rms_norm_weights, scale)
339 + return static_fp8[1], rmsnorm_residual_out
340 +
341 + def replacement(
342 + result: torch.Tensor,
343 + residual: torch.Tensor,
344 + mm_1: torch.Tensor,
345 + rms_norm_weights: torch.Tensor,
346 + scale: torch.Tensor,
347 + ) -> tuple[torch.Tensor, torch.Tensor]:
348 + reduce_scatter = self._reduce_scatter(mm_1)
349 + quant_result_buf = torch.empty_like(reduce_scatter,
350 + dtype=result.dtype)
351 + static_fp8, rmsnorm_residual_out = self._functional_fused_add_rmsnorm_then_quant( # noqa: E501
352 + quant_result_buf, reduce_scatter, residual, rms_norm_weights,
353 + scale)
354 + all_gather = self._all_gather(static_fp8[1])
355 + return all_gather, rmsnorm_residual_out
356 +
357 + pm.register_replacement(pattern, replacement, self.get_inputs(),
358 + pm.fwd_only, pm_pass)
359 +
360 +
361 +class LastAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
362 +
363 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str,
364 + op: torch._ops.OpOverload):
365 + super().__init__(epsilon, dtype, device, quant_op=op)
366 +
367 + def get_inputs(self):
368 + mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
369 +
370 + residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
371 + rms_norm_weights = torch.empty([4, 4],
372 + device=self.device,
373 + dtype=self.dtype)
374 + result = torch.empty([4, 4], device=self.device, dtype=FP8_DTYPE)
375 + scale = torch.empty([1, 1], device=self.device, dtype=torch.float32)
376 +
377 + return [
378 + result,
379 + residual,
380 + mm_1,
381 + rms_norm_weights,
382 + scale,
383 + ]
384 +
385 + def register(self, pm_pass: PatternMatcherPass):
386
226 - normalized = torch.ops.vllm.all_gather.default(
227 - rmsnorm[1],
228 - dim=0,
229 - world_size=tp_size,
230 - group_name=tp.unique_name)
387 + def pattern(
388 + result: torch.Tensor,
389 + residual: torch.Tensor,
390 + mm_1: torch.Tensor,
391 + rms_norm_weights: torch.Tensor,
392 + scale: torch.Tensor,
393 + ) -> tuple[torch.Tensor, torch.Tensor]:
394 + all_reduce = self._all_reduce(mm_1)
395 + static_fp8, _ = self._functional_fused_add_rmsnorm_then_quant(
396 + result, all_reduce, residual, rms_norm_weights, scale)
397 + return static_fp8[1]
398
399 + def replacement(
400 + result: torch.Tensor,
401 + residual: torch.Tensor,
402 + mm_1: torch.Tensor,
403 + rms_norm_weights: torch.Tensor,
404 + scale: torch.Tensor,
405 + ) -> tuple[torch.Tensor, torch.Tensor]:
406 + reduce_scatter = self._reduce_scatter(mm_1)
407 + quant_result_buf = torch.empty_like(reduce_scatter,
408 + dtype=result.dtype)
409 + static_fp8, _ = self._functional_fused_add_rmsnorm_then_quant(
410 + quant_result_buf, reduce_scatter, residual, rms_norm_weights,
411 + scale)
412 + normalized = self._all_gather(static_fp8[1])
413 return normalized
414
415 pm.register_replacement(pattern, replacement, self.get_inputs(),
@@ -236,21 +417,54 @@ class LastAllReduceRMSNormPattern(AllReduceRMSNormPattern):
417
418
419 class SequenceParallelismPass(VllmInductorPass):
420 + """
421 + This pass enables sequence parallelism for models.
422 + It identifies patterns where an AllReduce operation is followed by
423 + an RMSNorm (or RMSNorm and then Quantization) operation.
424 + These patterns are replaced with a ReduceScatter operation, followed by
425 + a local RMSNorm/Quantization, and then an AllGather operation.
426 +
427 + The general transformation is:
428 + Input -> AllReduce -> RMSNorm -> Output
429 + becomes
430 + Input -> ReduceScatter -> RMSNorm -> AllGather -> Output
431 +
432 + While this pass itself does not directly yield performance improvements,
433 + it lays the groundwork for subsequent fusion passes, such as
434 + GEMM + ReduceScatter and AllGather + GEMM fusions. These fusions can
435 + significantly reduce communication overhead and improve overall model
436 + performance.
437 + """
438
439 def __init__(self, config: VllmConfig):
440 super().__init__(config)
441
442 self.patterns: PatternMatcherPass = PatternMatcherPass(
443 pass_name="sequence_parallelism_pass")
444 +
445 for epsilon in [1e-5, 1e-6]:
246 - EmbeddingAllReduceRMSNormPattern(
247 - epsilon, self.model_dtype, self.device).register(self.patterns)
446 + # RMSNorm + Static FP8 quantization patterns
447 + fp8_quant_op = torch.ops._C.static_scaled_fp8_quant.default
448 + FirstAllReduceRMSNormStaticFP8Pattern(
449 + epsilon, self.model_dtype, self.device,
450 + fp8_quant_op).register(self.patterns)
451 + MiddleAllReduceRMSNormStaticFP8Pattern(
452 + epsilon, self.model_dtype, self.device,
453 + fp8_quant_op).register(self.patterns)
454 + LastAllReduceRMSNormStaticFP8Pattern(
455 + epsilon, self.model_dtype, self.device,
456 + fp8_quant_op).register(self.patterns)
457 +
458 + # Normal RMSNorm patterns
459 + FirstAllReduceRMSNormPattern(epsilon, self.model_dtype,
460 + self.device).register(self.patterns)
461
462 MiddleAllReduceRMSNormPattern(epsilon, self.model_dtype,
463 self.device).register(self.patterns)
464
465 LastAllReduceRMSNormPattern(epsilon, self.model_dtype,
466 self.device).register(self.patterns)
467 +
468 # WARNING: This is a hack to clear the pattern matcher cache
469 # and allow multiple values of epsilon.
470 torch._inductor.pattern_matcher._seen_patterns.clear()
vllm/config.py
+2 -4
@@ -3802,11 +3802,11 @@ class PassConfig:
3802 its own stages (before, after, maybe in-between)."""
3803 dump_graph_dir: Path = Path(".")
3804 """Directory to dump the graphs."""
3805 - enable_fusion: bool = True
3805 + enable_fusion: bool = field(default_factory=lambda: not envs.VLLM_USE_V1)
3806 """Whether to enable the custom fusion (RMSNorm/SiluMul+quant) pass."""
3807 enable_attn_fusion: bool = False
3808 """Whether to enable the custom attention+quant fusion pass."""
3809 - enable_noop: bool = True
3809 + enable_noop: bool = field(default_factory=lambda: not envs.VLLM_USE_V1)
3810 """Whether to enable the custom no-op elimination pass."""
3811 enable_sequence_parallelism: bool = False
3812 """Whether to enable sequence parallelism."""
@@ -4451,8 +4451,6 @@ class VllmConfig:
4451 # By default, V1 uses piecewise CUDA graphs. If full_cuda_graph
4452 # is set to True, full CUDA graphs will be used.
4453 self.compilation_config.cudagraph_num_of_warmups = 1
4454 - self.compilation_config.pass_config.enable_fusion = False
4455 - self.compilation_config.pass_config.enable_noop = False
4454 self.compilation_config.level = CompilationLevel.PIECEWISE
4455 self.compilation_config.set_splitting_ops_for_v1()
4456