| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import functools |
| 5 | from collections.abc import Callable, Sequence |
| 6 | from typing import Any |
| 7 | |
| 8 | import torch |
| 9 | import torch._inductor.pattern_matcher as pm |
| 10 | import torch.fx as fx |
| 11 | from torch._higher_order_ops.auto_functionalize import auto_functionalized |
| 12 | from torch._inductor.pattern_matcher import PatternMatcherPass |
| 13 | |
| 14 | import vllm.ir.ops |
| 15 | from vllm.config import VllmConfig |
| 16 | from vllm.config.utils import Range |
| 17 | from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce |
| 18 | from vllm.distributed.parallel_state import ( |
| 19 | get_tensor_model_parallel_rank, |
| 20 | get_tensor_model_parallel_world_size, |
| 21 | ) |
| 22 | from vllm.logger import init_logger |
| 23 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 24 | kFp8StaticTensorSym, |
| 25 | ) |
| 26 | |
| 27 | from ..inductor_pass import enable_fake_mode |
| 28 | from ..utility.noop_elimination import NoOpEliminationPass |
| 29 | from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass |
| 30 | from .matcher_utils import MatcherQuantFP8 |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | if hasattr(torch.ops._C, "scaled_fp4_quant"): |
| 35 | SCALED_FP4_QUANT_OUT_OVERLOAD = torch.ops._C.scaled_fp4_quant.out |
| 36 | SCALED_FP4_QUANT_DEFAULT_OVERLOAD = torch.ops._C.scaled_fp4_quant.default |
| 37 | |
| 38 | # Min hidden size per device capability for sequence parallelism |
| 39 | # Only apply sequence parallelism for models with hidden_size >= threshold |
| 40 | SP_MIN_HIDDEN_SIZE: dict[int, int] = { |
| 41 | 90: 8192, # H100: only for models with hidden_size >= 8192 |
| 42 | 100: 8192, # Blackwell family: only for models with hidden_size >= 8192 |
| 43 | } |
| 44 | |
| 45 | # Min size per GPU per device capability for sequence parallelism |
| 46 | # Total min size = min_per_gpu_size * tp_size |
| 47 | # This ensures the threshold scales appropriately with tensor parallelism |
| 48 | SP_MIN_PER_GPU_SIZE_MB: dict[int, float] = { |
| 49 | 90: 8, # 8MB per GPU for H100 |
| 50 | # Use a more conservative threshold on Blackwell so TP8 starts later. |
| 51 | 100: 32, |
| 52 | } |
| 53 | |
| 54 | |
| 55 | def get_sequence_parallelism_threshold( |
| 56 | hidden_size: int, |
| 57 | tp_size: int, |
| 58 | element_size: int, |
| 59 | ) -> int | None: |
| 60 | """Calculate the minimum token threshold for applying sequence parallelism. |
| 61 | |
| 62 | Returns None if sequence parallelism should not be applied based on model size. |
| 63 | |
| 64 | Branching logic based on device capability: |
| 65 | - Check if hidden_size >= SP_MIN_HIDDEN_SIZE[device_capability] |
| 66 | - If not, returns None (SP disabled for small models on this device) |
| 67 | - If yes, calculates threshold based on per-GPU size |
| 68 | |
| 69 | Formula: min_token_num = (min_per_gpu_size_mb * tp_size * MiB) // |
| 70 | (hidden_size * element_size) |
| 71 | """ |
| 72 | from vllm.platforms import current_platform |
| 73 | |
| 74 | if current_platform.is_xpu(): |
| 75 | min_hidden_size = 4096 |
| 76 | min_per_gpu_size_mb = 8.0 |
| 77 | elif current_platform.is_cuda(): |
| 78 | capability = current_platform.get_device_capability() |
| 79 | if capability is None: |
| 80 | return None |
| 81 | |
| 82 | # Collapse Blackwell variants (sm100/sm103/...) into one policy bucket. |
| 83 | if current_platform.is_device_capability_family(100): |
| 84 | device_capability = 100 |
| 85 | else: |
| 86 | device_capability = capability.to_int() |
| 87 | |
| 88 | # Check if device has configured thresholds |
| 89 | _hidden = SP_MIN_HIDDEN_SIZE.get(device_capability) |
| 90 | _gpu_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability) |
| 91 | if _hidden is None or _gpu_mb is None: |
| 92 | return None |
| 93 | min_hidden_size, min_per_gpu_size_mb = _hidden, _gpu_mb |
| 94 | else: |
| 95 | return None |
| 96 | |
| 97 | # Only apply sequence parallelism for models meeting the size threshold |
| 98 | if hidden_size < min_hidden_size: |
| 99 | return None |
| 100 | |
| 101 | MiB = 1024 * 1024 |
| 102 | min_size = min_per_gpu_size_mb * MiB * tp_size |
| 103 | return int(min_size // (hidden_size * element_size)) |
| 104 | |
| 105 | |
| 106 | def get_first_out_wrapper( |
| 107 | fn: Callable[..., Sequence[torch.Tensor]], |
| 108 | ) -> Callable[..., torch.Tensor]: |
| 109 | @functools.wraps(fn) |
| 110 | def wrapper(*args: Any) -> torch.Tensor: |
| 111 | return fn(*args)[0] |
| 112 | |
| 113 | return wrapper |
| 114 | |
| 115 | |
| 116 | class _SequenceParallelPatternHelper: |
| 117 | """Helper for sequence parallelism patterns.""" |
| 118 | |
| 119 | def __init__( |
| 120 | self, |
| 121 | epsilon: float, |
| 122 | dtype: torch.dtype, |
| 123 | device: str | None, |
| 124 | ) -> None: |
| 125 | self.epsilon = epsilon |
| 126 | self.dtype = dtype |
| 127 | self.device = device |
| 128 | self.tp_group = get_tp_group() |
| 129 | self.tp_size = get_tensor_model_parallel_world_size() |
| 130 | self.tp_rank = get_tensor_model_parallel_rank() |
| 131 | |
| 132 | def _all_reduce(self, x: torch.Tensor) -> torch.Tensor: |
| 133 | return tensor_model_parallel_all_reduce(x) |
| 134 | |
| 135 | def _reduce_scatter(self, x: torch.Tensor) -> torch.Tensor: |
| 136 | return torch.ops.vllm.reduce_scatter.default( |
| 137 | x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name |
| 138 | ) |
| 139 | |
| 140 | def _all_gather(self, x: torch.Tensor) -> torch.Tensor: |
| 141 | return torch.ops.vllm.all_gather.default( |
| 142 | x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name |
| 143 | ) |
| 144 | |
| 145 | def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: |
| 146 | return torch.empty(*args, dtype=self.dtype, device=self.device, **kwargs) |
| 147 | |
| 148 | def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: |
| 149 | return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs) |
| 150 | |
| 151 | |
| 152 | class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper): |
| 153 | def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: |
| 154 | super().__init__(epsilon, dtype, device) |
| 155 | |
| 156 | def get_inputs(self) -> list[torch.Tensor]: |
| 157 | # input, weight |
| 158 | return [self.empty([1, 8, 4]), self.empty([4])] |
| 159 | |
| 160 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 161 | def pattern( |
| 162 | input: torch.Tensor, |
| 163 | weight: torch.Tensor, |
| 164 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 165 | all_reduce = self._all_reduce(input) |
| 166 | rmsnorm = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) |
| 167 | |
| 168 | return rmsnorm, all_reduce |
| 169 | |
| 170 | def replacement( |
| 171 | input: torch.Tensor, |
| 172 | weight: torch.Tensor, |
| 173 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 174 | reduce_scatter = self._reduce_scatter(input) |
| 175 | |
| 176 | rmsnorm = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) |
| 177 | all_gather = self._all_gather(rmsnorm) |
| 178 | return all_gather, reduce_scatter |
| 179 | |
| 180 | pm.register_replacement( |
| 181 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 182 | ) |
| 183 | |
| 184 | |
| 185 | class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper): |
| 186 | def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: |
| 187 | super().__init__(epsilon, dtype, device) |
| 188 | |
| 189 | def get_inputs(self) -> list[torch.Tensor]: |
| 190 | mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) |
| 191 | |
| 192 | residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) |
| 193 | rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) |
| 194 | |
| 195 | return [ |
| 196 | residual, |
| 197 | mm_1, |
| 198 | rms_norm_weights, |
| 199 | ] |
| 200 | |
| 201 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 202 | def pattern( |
| 203 | residual: torch.Tensor, |
| 204 | mm_1: torch.Tensor, |
| 205 | rms_norm_weights: torch.Tensor, |
| 206 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 207 | all_reduce = self._all_reduce(mm_1) |
| 208 | rmsnorm = vllm.ir.ops.fused_add_rms_norm( |
| 209 | all_reduce, residual, rms_norm_weights, self.epsilon |
| 210 | ) |
| 211 | return rmsnorm[0], rmsnorm[1] |
| 212 | |
| 213 | def replacement( |
| 214 | residual: torch.Tensor, |
| 215 | mm_1: torch.Tensor, |
| 216 | rms_norm_weights: torch.Tensor, |
| 217 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 218 | # The pattern matcher replaces from the end of the graph |
| 219 | # (last layer first). At the time each match is replaced, |
| 220 | # the preceding layer has NOT been replaced yet, so |
| 221 | # `residual` is still full-size and the slice below is |
| 222 | # correct. Once the preceding layer IS replaced, its |
| 223 | # residual output shrinks to [local_len, H], and this |
| 224 | # slice becomes semantically incorrect (e.g. for rank > 0, |
| 225 | # the indices would be out of bounds). However, since the |
| 226 | # symbolic output shape equals the input shape, |
| 227 | # NoOpEliminationPass (called at the end of |
| 228 | # SequenceParallelismPass.__call__) removes these slices |
| 229 | # before the graph is ever executed or compiled. |
| 230 | reduce_scatter = self._reduce_scatter(mm_1) |
| 231 | local_len = reduce_scatter.size(0) |
| 232 | # when the preceding VocabParallelEmbedding is excluded |
| 233 | # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), |
| 234 | # the FirstAllReduceRMSNorm pattern is never matched. we must |
| 235 | # perform a proper TP-aware slice here. simply using `[0:local_len]` |
| 236 | # would incorrectly cause all ranks to process rank 0's chunk. |
| 237 | residual = residual[ |
| 238 | self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... |
| 239 | ] |
| 240 | rmsnorm = vllm.ir.ops.fused_add_rms_norm( |
| 241 | reduce_scatter, residual, rms_norm_weights, self.epsilon |
| 242 | ) |
| 243 | all_gather = self._all_gather(rmsnorm[0]) |
| 244 | # residual output is now [local_len, H]; the next layer's |
| 245 | # slice on it is semantically incorrect until |
| 246 | # NoOpEliminationPass removes it. |
| 247 | return all_gather, rmsnorm[1] |
| 248 | |
| 249 | pm.register_replacement( |
| 250 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 251 | ) |
| 252 | pm.register_replacement( |
| 253 | get_first_out_wrapper(pattern), |
| 254 | get_first_out_wrapper(replacement), |
| 255 | self.get_inputs(), |
| 256 | pm.fwd_only, |
| 257 | pm_pass, |
| 258 | ) |
| 259 | |
| 260 | |
| 261 | class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): |
| 262 | def __init__( |
| 263 | self, |
| 264 | epsilon: float, |
| 265 | dtype: torch.dtype, |
| 266 | device: str | None, |
| 267 | ) -> None: |
| 268 | super().__init__(epsilon, dtype, device) |
| 269 | self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) |
| 270 | |
| 271 | def get_inputs(self) -> list[torch.Tensor]: |
| 272 | # input, weight, scale |
| 273 | return [self.empty([1, 8, 4]), self.empty([4]), self.empty_f32([1, 1])] |
| 274 | |
| 275 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 276 | def pattern( |
| 277 | input: torch.Tensor, |
| 278 | weight: torch.Tensor, |
| 279 | scale: torch.Tensor, |
| 280 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 281 | all_reduce = self._all_reduce(input) |
| 282 | rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) |
| 283 | quant, _ = self.quant_matcher(rms, scale) |
| 284 | return quant, all_reduce |
| 285 | |
| 286 | def replacement( |
| 287 | input: torch.Tensor, |
| 288 | weight: torch.Tensor, |
| 289 | scale: torch.Tensor, |
| 290 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 291 | reduce_scatter = self._reduce_scatter(input) |
| 292 | rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) |
| 293 | quant, _ = self.quant_matcher(rms, scale) |
| 294 | all_gather = self._all_gather(quant) |
| 295 | |
| 296 | return all_gather, reduce_scatter |
| 297 | |
| 298 | pm.register_replacement( |
| 299 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 300 | ) |
| 301 | |
| 302 | |
| 303 | class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): |
| 304 | def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: |
| 305 | super().__init__(epsilon, dtype, device) |
| 306 | self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) |
| 307 | |
| 308 | def get_inputs(self) -> list[torch.Tensor]: |
| 309 | mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) |
| 310 | residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) |
| 311 | rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) |
| 312 | scale = torch.empty([1, 1], device=self.device, dtype=torch.float32) |
| 313 | |
| 314 | return [residual, mm_1, rms_norm_weights, scale] |
| 315 | |
| 316 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 317 | def pattern( |
| 318 | residual: torch.Tensor, |
| 319 | mm_1: torch.Tensor, |
| 320 | rms_norm_weights: torch.Tensor, |
| 321 | scale: torch.Tensor, |
| 322 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 323 | all_reduce = self._all_reduce(mm_1) |
| 324 | rms, residual_out = vllm.ir.ops.fused_add_rms_norm( |
| 325 | all_reduce, residual, rms_norm_weights, self.epsilon |
| 326 | ) |
| 327 | quant, _ = self.quant_matcher(rms, scale) |
| 328 | return quant, residual_out |
| 329 | |
| 330 | def replacement( |
| 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 | # See MiddleAllReduceRMSNormPattern.replacement for a |
| 337 | # detailed explanation of the temporary slice below: |
| 338 | # it is correct when first inserted, becomes semantically |
| 339 | # incorrect after the preceding layer is replaced, and is |
| 340 | # removed by NoOpEliminationPass before the graph is compiled. |
| 341 | reduce_scatter = self._reduce_scatter(mm_1) |
| 342 | local_len = reduce_scatter.size(0) |
| 343 | # when the preceding VocabParallelEmbedding is excluded |
| 344 | # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), |
| 345 | # the FirstAllReduceRMSNorm pattern is never matched. we must |
| 346 | # perform a proper TP-aware slice here. simply using `[0:local_len]` |
| 347 | # would incorrectly cause all ranks to process rank 0's chunk. |
| 348 | residual = residual[ |
| 349 | self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... |
| 350 | ] |
| 351 | rms, residual_out = vllm.ir.ops.fused_add_rms_norm( |
| 352 | reduce_scatter, residual, rms_norm_weights, self.epsilon |
| 353 | ) |
| 354 | quant, _ = self.quant_matcher(rms, scale) |
| 355 | all_gather = self._all_gather(quant) |
| 356 | # residual output is now [local_len, H]; the next layer's |
| 357 | # slice on it is semantically incorrect until |
| 358 | # NoOpEliminationPass removes it. |
| 359 | return all_gather, residual_out |
| 360 | |
| 361 | pm.register_replacement( |
| 362 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 363 | ) |
| 364 | |
| 365 | pm.register_replacement( |
| 366 | get_first_out_wrapper(pattern), |
| 367 | get_first_out_wrapper(replacement), |
| 368 | self.get_inputs(), |
| 369 | pm.fwd_only, |
| 370 | pm_pass, |
| 371 | ) |
| 372 | |
| 373 | |
| 374 | class FirstAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): |
| 375 | def get_inputs(self) -> list[torch.Tensor]: |
| 376 | input = self.empty([8, 16]) |
| 377 | weight = self.empty([16]) |
| 378 | input_global_scale = self.empty_f32([1, 1]) |
| 379 | quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) |
| 380 | output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) |
| 381 | return [input, weight, input_global_scale, quant_output, output_scale] |
| 382 | |
| 383 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 384 | def pattern( |
| 385 | input: torch.Tensor, |
| 386 | weight: torch.Tensor, |
| 387 | input_global_scale: torch.Tensor, |
| 388 | quant_output: torch.Tensor, |
| 389 | output_scale: torch.Tensor, |
| 390 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 391 | all_reduce = self._all_reduce(input) |
| 392 | rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) |
| 393 | quant = auto_functionalized( |
| 394 | SCALED_FP4_QUANT_OUT_OVERLOAD, |
| 395 | input=rms, |
| 396 | input_scale=input_global_scale, |
| 397 | is_sf_swizzled_layout=True, |
| 398 | output=quant_output, |
| 399 | output_scale=output_scale, |
| 400 | ) |
| 401 | return quant[1], all_reduce, quant[2] |
| 402 | |
| 403 | def replacement( |
| 404 | input: torch.Tensor, |
| 405 | weight: torch.Tensor, |
| 406 | input_global_scale: torch.Tensor, |
| 407 | quant_output: torch.Tensor, |
| 408 | output_scale: torch.Tensor, |
| 409 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 410 | reduce_scatter = self._reduce_scatter(input) |
| 411 | rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) |
| 412 | rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) |
| 413 | quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( |
| 414 | rms, |
| 415 | input_global_scale, |
| 416 | True, |
| 417 | ) |
| 418 | return ( |
| 419 | self._all_gather(quant[0]), |
| 420 | reduce_scatter, |
| 421 | self._all_gather(quant[1]), |
| 422 | ) |
| 423 | |
| 424 | pm.register_replacement( |
| 425 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 426 | ) |
| 427 | |
| 428 | |
| 429 | class MiddleAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): |
| 430 | def get_inputs(self) -> list[torch.Tensor]: |
| 431 | mm_1 = self.empty([8, 16]) |
| 432 | residual = self.empty([8, 16]) |
| 433 | rms_norm_weights = self.empty([16]) |
| 434 | input_global_scale = self.empty_f32([1, 1]) |
| 435 | quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) |
| 436 | output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) |
| 437 | return [ |
| 438 | residual, |
| 439 | mm_1, |
| 440 | rms_norm_weights, |
| 441 | input_global_scale, |
| 442 | quant_output, |
| 443 | output_scale, |
| 444 | ] |
| 445 | |
| 446 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 447 | def pattern( |
| 448 | residual: torch.Tensor, |
| 449 | mm_1: torch.Tensor, |
| 450 | rms_norm_weights: torch.Tensor, |
| 451 | input_global_scale: torch.Tensor, |
| 452 | quant_output: torch.Tensor, |
| 453 | output_scale: torch.Tensor, |
| 454 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 455 | all_reduce = self._all_reduce(mm_1) |
| 456 | rms, residual_out = vllm.ir.ops.fused_add_rms_norm( |
| 457 | all_reduce, residual, rms_norm_weights, self.epsilon |
| 458 | ) |
| 459 | quant = auto_functionalized( |
| 460 | SCALED_FP4_QUANT_OUT_OVERLOAD, |
| 461 | input=rms, |
| 462 | input_scale=input_global_scale, |
| 463 | is_sf_swizzled_layout=True, |
| 464 | output=quant_output, |
| 465 | output_scale=output_scale, |
| 466 | ) |
| 467 | return quant[1], residual_out, quant[2] |
| 468 | |
| 469 | def replacement( |
| 470 | residual: torch.Tensor, |
| 471 | mm_1: torch.Tensor, |
| 472 | rms_norm_weights: torch.Tensor, |
| 473 | input_global_scale: torch.Tensor, |
| 474 | quant_output: torch.Tensor, |
| 475 | output_scale: torch.Tensor, |
| 476 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 477 | # Keep this slice in sync with the non-quantized SP replacement: |
| 478 | # once the previous SP pattern fires, it becomes a no-op. |
| 479 | reduce_scatter = self._reduce_scatter(mm_1) |
| 480 | residual = residual[0 : reduce_scatter.size(0), ...] |
| 481 | rms, residual_out = vllm.ir.ops.fused_add_rms_norm( |
| 482 | reduce_scatter, residual, rms_norm_weights, self.epsilon |
| 483 | ) |
| 484 | rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) |
| 485 | quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( |
| 486 | rms, |
| 487 | input_global_scale, |
| 488 | True, |
| 489 | ) |
| 490 | return self._all_gather(quant[0]), residual_out, self._all_gather(quant[1]) |
| 491 | |
| 492 | pm.register_replacement( |
| 493 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 494 | ) |
| 495 | |
| 496 | |
| 497 | class SequenceParallelismPass(VllmPatternMatcherPass): |
| 498 | """This pass enables sequence parallelism for models. |
| 499 | It identifies patterns where an AllReduce operation is followed by |
| 500 | an RMSNorm (or RMSNorm and then Quantization) operation. |
| 501 | These patterns are replaced with a ReduceScatter operation, followed by |
| 502 | a local RMSNorm/Quantization, and then an AllGather operation. |
| 503 | |
| 504 | The general transformation is: |
| 505 | Input -> AllReduce -> RMSNorm -> Output |
| 506 | becomes |
| 507 | Input -> ReduceScatter -> RMSNorm -> AllGather -> Output |
| 508 | |
| 509 | While this pass itself does not directly yield performance improvements, |
| 510 | it lays the groundwork for subsequent fusion passes, such as |
| 511 | GEMM + ReduceScatter and AllGather + GEMM fusions. These fusions can |
| 512 | significantly reduce communication overhead and improve overall model |
| 513 | performance. |
| 514 | |
| 515 | This pass is only supported when compiling the whole graph (fullgraph |
| 516 | mode, i.e. using Inductor graph partition or empty splitting_ops). |
| 517 | Piecewise compilation is not supported because the residual tensor |
| 518 | gets split across TP ranks, causing size mismatches at subgraph |
| 519 | boundaries. |
| 520 | |
| 521 | This pass splits up the residual tensor across TP ranks and hence divides |
| 522 | its size. The pattern matcher starts at the end of the graph (last layer |
| 523 | first), so when each replacement inserts a residual slice, the preceding |
| 524 | layer has not been replaced yet and the slice is correct. Once the |
| 525 | preceding layer IS replaced, its residual output shrinks and the slice |
| 526 | becomes semantically incorrect (out-of-bounds indices for rank > 0). |
| 527 | The graph is never executed in this intermediate state — |
| 528 | NoOpEliminationPass removes these slices based on symbolic shape equality |
| 529 | (input shape == output shape) before the graph is compiled. |
| 530 | """ |
| 531 | |
| 532 | @enable_fake_mode |
| 533 | def __init__(self, config: VllmConfig) -> None: |
| 534 | super().__init__(config) |
| 535 | |
| 536 | # Get min_token_num threshold |
| 537 | # Read min_token_num from config (calculated during config init) |
| 538 | self.min_token_num = None |
| 539 | if config.model_config is not None: |
| 540 | pass_config = config.compilation_config.pass_config |
| 541 | self.min_token_num = pass_config.sp_min_token_num |
| 542 | |
| 543 | if self.min_token_num is not None: |
| 544 | # Take the min to avoid exceeding max_num_batched_tokens |
| 545 | max_batched = config.scheduler_config.max_num_batched_tokens |
| 546 | if max_batched is not None: |
| 547 | self.min_token_num = min(self.min_token_num, max_batched) |
| 548 | logger.debug_once( |
| 549 | f"Sequence parallelism min token threshold: {self.min_token_num}", |
| 550 | scope="global", |
| 551 | ) |
| 552 | |
| 553 | # Used to clean up redundant views created temporarily |
| 554 | # to circumvent residual shape change issues |
| 555 | self.noop_cleanup = NoOpEliminationPass(config) |
| 556 | self.noop_cleanup.pass_name = f"{self.pass_name}.{self.noop_cleanup.pass_name}" |
| 557 | |
| 558 | self.patterns: PatternMatcherPass = PatternMatcherPass( |
| 559 | pass_name="sequence_parallelism_pass" |
| 560 | ) |
| 561 | |
| 562 | for epsilon in [1e-5, 1e-6]: |
| 563 | # RMSNorm + Static FP8 quantization patterns |
| 564 | FirstAllReduceRMSNormStaticFP8Pattern( |
| 565 | epsilon, self.model_dtype, self.device |
| 566 | ).register(self.patterns) |
| 567 | MiddleAllReduceRMSNormStaticFP8Pattern( |
| 568 | epsilon, self.model_dtype, self.device |
| 569 | ).register(self.patterns) |
| 570 | |
| 571 | if "SCALED_FP4_QUANT_OUT_OVERLOAD" in globals(): |
| 572 | FirstAllReduceRMSNormStaticNVFP4Pattern( |
| 573 | epsilon, self.model_dtype, self.device |
| 574 | ).register(self.patterns) |
| 575 | MiddleAllReduceRMSNormStaticNVFP4Pattern( |
| 576 | epsilon, self.model_dtype, self.device |
| 577 | ).register(self.patterns) |
| 578 | |
| 579 | # Normal RMSNorm patterns |
| 580 | FirstAllReduceRMSNormPattern( |
| 581 | epsilon, self.model_dtype, self.device |
| 582 | ).register(self.patterns) |
| 583 | |
| 584 | MiddleAllReduceRMSNormPattern( |
| 585 | epsilon, self.model_dtype, self.device |
| 586 | ).register(self.patterns) |
| 587 | |
| 588 | self.dump_patterns(config, self.patterns) |
| 589 | |
| 590 | def is_applicable_for_range(self, compile_range: Range) -> bool: |
| 591 | """Determines if sequence parallelism should be applied for the given |
| 592 | compile range. |
| 593 | |
| 594 | SP is only beneficial for larger batch sizes where the communication |
| 595 | overhead is amortized. For small batches, the overhead of splitting |
| 596 | and gathering tensors across TP ranks outweighs the benefits. |
| 597 | |
| 598 | Returns False (SP disabled) when: |
| 599 | - min_token_num is None (SP disabled for this device/config) |
| 600 | - The compile range starts below the minimum token threshold |
| 601 | """ |
| 602 | assert ( |
| 603 | self.compilation_config.use_inductor_graph_partition |
| 604 | or not self.compilation_config.splitting_ops |
| 605 | ), "SequenceParallelismPass requires full-graph compilation" |
| 606 | |
| 607 | # min_token_num is None when SP is disabled for this device/config |
| 608 | # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size) |
| 609 | if self.min_token_num is None: |
| 610 | return False |
| 611 | |
| 612 | # Only apply SP when batch size meets the minimum threshold |
| 613 | return compile_range.start >= self.min_token_num |
| 614 | |
| 615 | @VllmInductorPass.time_and_log |
| 616 | def __call__(self, graph: fx.Graph) -> None: |
| 617 | self.matched_count = self.patterns.apply(graph) |
| 618 | logger.debug("Replaced %s patterns", self.matched_count) |
| 619 | # Clean up reshape nodes |
| 620 | self.noop_cleanup(graph) |