| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | from importlib.util import find_spec |
| 5 | from types import ModuleType |
| 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._aiter_ops import rocm_aiter_ops |
| 16 | from vllm.compilation.passes.fusion.rms_quant_fusion import ( |
| 17 | _rms_input_weight_dtype_match, |
| 18 | ) |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.config.utils import Range |
| 21 | from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce |
| 22 | from vllm.distributed.parallel_state import ( |
| 23 | get_tensor_model_parallel_rank, |
| 24 | get_tensor_model_parallel_world_size, |
| 25 | ) |
| 26 | from vllm.logger import init_logger |
| 27 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 28 | GroupShape, |
| 29 | QuantKey, |
| 30 | ScaleDesc, |
| 31 | kFp8StaticTensorSym, |
| 32 | ) |
| 33 | from vllm.platforms import current_platform |
| 34 | from vllm.utils.torch_utils import ( |
| 35 | direct_register_custom_op, |
| 36 | ) |
| 37 | |
| 38 | from ..inductor_pass import enable_fake_mode |
| 39 | from ..vllm_inductor_pass import ( |
| 40 | VllmFusionPatternMatcherPass, |
| 41 | VllmInductorPass, |
| 42 | VllmPatternMatcherPass, |
| 43 | VllmPatternReplacement, |
| 44 | ) |
| 45 | from .matcher_utils import MatcherQuantFP8 |
| 46 | |
| 47 | FP8_DTYPE = current_platform.fp8_dtype() |
| 48 | |
| 49 | _IR_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default |
| 50 | _IR_FUSED_ADD_RMS_NORM_OP = torch.ops.vllm_ir.fused_add_rms_norm.default |
| 51 | |
| 52 | |
| 53 | def _view_nvfp4_scale_out_for_flashinfer( |
| 54 | scale_out: torch.Tensor, |
| 55 | ) -> torch.Tensor: |
| 56 | """View vLLM's packed NVFP4 scale buffer as FP8 for FlashInfer.""" |
| 57 | return torch.ops.aten.view.dtype(scale_out, FP8_DTYPE) |
| 58 | |
| 59 | |
| 60 | def _view_flashinfer_nvfp4_scale_out_as_int32( |
| 61 | scale_out: torch.Tensor, |
| 62 | ) -> torch.Tensor: |
| 63 | """View FlashInfer's NVFP4 scale buffer back as vLLM's int32 format.""" |
| 64 | return torch.ops.aten.view.dtype(scale_out, torch.int32) |
| 65 | |
| 66 | |
| 67 | def _norm_input_weight_dtype_match(match: pm.Match) -> bool: |
| 68 | """Prevent fusion when the norm input and weight dtypes differ (e.g. a Gemma |
| 69 | fp32 weight.float()+1 gamma), covering rms_norm and fused_add_rms_norm.""" |
| 70 | for node in match.nodes: |
| 71 | if node.target == _IR_RMS_NORM_OP: |
| 72 | x, weight = node.args[0], node.args[1] |
| 73 | elif node.target == _IR_FUSED_ADD_RMS_NORM_OP: |
| 74 | x, weight = node.args[0], node.args[2] |
| 75 | else: |
| 76 | continue |
| 77 | if isinstance(x, fx.Node) and isinstance(weight, fx.Node): |
| 78 | return x.meta["val"].dtype == weight.meta["val"].dtype |
| 79 | return True |
| 80 | |
| 81 | |
| 82 | # The empirical value for small batch |
| 83 | PDL_ADVANCE_LAUNCH_TOKENS = 16 |
| 84 | |
| 85 | logger = init_logger(__name__) |
| 86 | |
| 87 | flashinfer_comm: ModuleType | None = None |
| 88 | if find_spec("flashinfer"): |
| 89 | try: |
| 90 | import flashinfer.comm as _flashinfer_comm |
| 91 | |
| 92 | if hasattr(_flashinfer_comm, "allreduce_fusion") and hasattr( |
| 93 | _flashinfer_comm, "create_allreduce_fusion_workspace" |
| 94 | ): |
| 95 | flashinfer_comm = _flashinfer_comm |
| 96 | except Exception as e: |
| 97 | logger.debug_once("flashinfer.comm import failed: %s", e) |
| 98 | |
| 99 | if hasattr(torch.ops._C, "scaled_fp4_quant"): |
| 100 | STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.out |
| 101 | |
| 102 | # Max size of the input tensor per world size per device capability |
| 103 | # to use flashinfer fused allreduce |
| 104 | FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = { |
| 105 | 90: { |
| 106 | 2: 64, # 64MB |
| 107 | 4: 2, # 2MB |
| 108 | 8: 0.5, # 0.5MB |
| 109 | }, |
| 110 | 100: { |
| 111 | 2: 64, # 64MB |
| 112 | 4: 32, # 32MB |
| 113 | 8: 1, # 1MB |
| 114 | 16: 64, # 64MB (mnnvl multi-node) |
| 115 | }, |
| 116 | 103: { |
| 117 | 2: 64, # 64MB |
| 118 | 4: 64, # 64MB |
| 119 | 8: 4, # 4MB |
| 120 | 16: 64, # 64MB (mnnvl multi-node) |
| 121 | }, |
| 122 | 107: { |
| 123 | 2: 64, # 64MB |
| 124 | 4: 64, # 64MB |
| 125 | 8: 2, # 2MB |
| 126 | }, |
| 127 | } |
| 128 | |
| 129 | # Max size of the input tensor per world size per device capability |
| 130 | # to use flashinfer one shot fused allreduce |
| 131 | # OneShot max size is at most 64MB / world size (FlashInfer restriction) |
| 132 | _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = { |
| 133 | 90: { |
| 134 | 2: 32, # 32MB |
| 135 | 4: 2, # 2MB |
| 136 | 8: 0.5, # 0.5MB |
| 137 | }, |
| 138 | 100: { |
| 139 | 2: 32, # 32MB |
| 140 | 4: 4, # 4MB |
| 141 | 8: 1, # 1MB |
| 142 | }, |
| 143 | 103: { |
| 144 | 2: 32, # 32MB |
| 145 | 4: 4, # 4MB |
| 146 | 8: 2, # 2MB |
| 147 | }, |
| 148 | 107: { |
| 149 | 2: 32, # 32MB |
| 150 | 4: 4, # 4MB |
| 151 | 8: 2, # 2MB |
| 152 | }, |
| 153 | } |
| 154 | |
| 155 | MiB = 1024 * 1024 |
| 156 | |
| 157 | |
| 158 | def _select_flashinfer_allreduce_use_oneshot( |
| 159 | workspace_backend: str, |
| 160 | device_capability: int | None, |
| 161 | world_size: int, |
| 162 | current_tensor_size: int, |
| 163 | ) -> bool | None: |
| 164 | if workspace_backend == "mnnvl": |
| 165 | # FlashInfer sizes MNNVL workspaces around its own AUTO strategy. |
| 166 | # Forcing vLLM's per-rank threshold can request one-shot for tensors |
| 167 | # larger than the MNNVL one-shot workspace. |
| 168 | return None |
| 169 | |
| 170 | if device_capability is None: |
| 171 | max_one_shot_size = None |
| 172 | else: |
| 173 | max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get( |
| 174 | device_capability, {} |
| 175 | ).get(world_size) |
| 176 | return max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB |
| 177 | |
| 178 | |
| 179 | if flashinfer_comm is not None: |
| 180 | from vllm.distributed.device_communicators.flashinfer_all_reduce import ( |
| 181 | destroy_fi_ar_workspace, |
| 182 | get_fi_ar_quant_workspace, |
| 183 | get_fi_ar_workspace, |
| 184 | ) |
| 185 | |
| 186 | ar_fusion_patterns = flashinfer_comm.AllReduceFusionPattern |
| 187 | |
| 188 | def call_trtllm_fused_allreduce_norm( |
| 189 | allreduce_in: torch.Tensor, |
| 190 | residual: torch.Tensor, |
| 191 | rms_gamma: torch.Tensor, |
| 192 | rms_eps: float, |
| 193 | world_size: int, |
| 194 | launch_with_pdl: bool, |
| 195 | fp32_acc: bool, |
| 196 | max_token_num: int, |
| 197 | pattern_code: int, |
| 198 | norm_out: torch.Tensor | None = None, |
| 199 | quant_out: torch.Tensor | None = None, |
| 200 | scale_out: torch.Tensor | None = None, |
| 201 | scale_factor: torch.Tensor | None = None, |
| 202 | weight_bias: float = 0.0, |
| 203 | ) -> None: |
| 204 | # handle transformers backend passing outer batch dim. |
| 205 | if allreduce_in.dim() != 2: |
| 206 | hidden = allreduce_in.shape[-1] |
| 207 | allreduce_in = allreduce_in.view(-1, hidden) |
| 208 | residual = residual.view(-1, hidden) |
| 209 | if norm_out is not None: |
| 210 | norm_out = norm_out.view(-1, hidden) |
| 211 | num_tokens, hidden_size = allreduce_in.shape |
| 212 | element_size = allreduce_in.element_size() |
| 213 | current_tensor_size = num_tokens * hidden_size * element_size |
| 214 | max_tensor_size = max_token_num * hidden_size * element_size |
| 215 | assert current_tensor_size <= max_tensor_size, ( |
| 216 | f"Current tensor size {current_tensor_size} is larger than " |
| 217 | f"max token num {max_token_num} * hidden size {hidden_size} * " |
| 218 | f"element size {element_size}" |
| 219 | ) |
| 220 | curr_device = current_platform.get_device_capability() |
| 221 | device_capability = curr_device.to_int() if curr_device is not None else None |
| 222 | |
| 223 | # Select workspace based on pattern: quant patterns use the quant |
| 224 | # workspace, non-quant patterns use the primary workspace. |
| 225 | is_quant_pattern = pattern_code in ( |
| 226 | ar_fusion_patterns.kARResidualRMSNormFP8Quant, |
| 227 | ar_fusion_patterns.kARResidualRMSNormFP4Quant, |
| 228 | ) |
| 229 | get_workspace_fn = ( |
| 230 | get_fi_ar_quant_workspace if is_quant_pattern else get_fi_ar_workspace |
| 231 | ) |
| 232 | workspace = get_workspace_fn( |
| 233 | world_size=world_size, |
| 234 | rank=get_tensor_model_parallel_rank(), |
| 235 | max_token_num=max_token_num, |
| 236 | hidden_dim=hidden_size, |
| 237 | dtype=allreduce_in.dtype, |
| 238 | group=get_tp_group().cpu_group, |
| 239 | ) |
| 240 | assert workspace is not None, ( |
| 241 | "Flashinfer allreduce workspace must be initialized when using flashinfer" |
| 242 | ) |
| 243 | use_oneshot = _select_flashinfer_allreduce_use_oneshot( |
| 244 | workspace.backend, |
| 245 | device_capability, |
| 246 | world_size, |
| 247 | current_tensor_size, |
| 248 | ) |
| 249 | assert flashinfer_comm is not None |
| 250 | if norm_out is None: |
| 251 | norm_out = allreduce_in |
| 252 | residual_out = residual |
| 253 | else: |
| 254 | # return residual_out as allreduce_out with zeroed residual_in |
| 255 | # as flashinfer does not support rms_norm |
| 256 | # and allreduce_out together |
| 257 | residual_out = allreduce_in |
| 258 | |
| 259 | layout_code = None |
| 260 | # vLLM quant patterns use swizzled scale-factor layout. Non-quant |
| 261 | # patterns ignore layout_code. |
| 262 | if workspace.backend in ("trtllm", "mnnvl"): |
| 263 | layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4 |
| 264 | |
| 265 | flashinfer_comm.allreduce_fusion( |
| 266 | input=allreduce_in, |
| 267 | workspace=workspace, |
| 268 | pattern=pattern_code, |
| 269 | launch_with_pdl=launch_with_pdl, |
| 270 | output=None, |
| 271 | residual_out=residual_out, |
| 272 | norm_out=norm_out, |
| 273 | quant_out=quant_out, |
| 274 | scale_out=scale_out, |
| 275 | residual_in=residual, |
| 276 | rms_gamma=rms_gamma, |
| 277 | rms_eps=rms_eps, |
| 278 | scale_factor=scale_factor, |
| 279 | layout_code=layout_code, |
| 280 | use_oneshot=use_oneshot, |
| 281 | fp32_acc=fp32_acc, |
| 282 | weight_bias=weight_bias, |
| 283 | # The one-shot Lamport all-reduce signals PDL completion before its |
| 284 | # output buffer is committed when trigger_completion_at_end is |
| 285 | # False, so the next PDL-launched kernel can read the uninitialized |
| 286 | # Lamport buffer and produce NaN. This only fires for |
| 287 | # num_tokens <= PDL_ADVANCE_LAUNCH_TOKENS (the batch=1 / spec-decode |
| 288 | # shapes, where the one-shot path is always selected). Complete at |
| 289 | # the end for the one-shot path; the two-shot path is synchronized |
| 290 | # and keeps the early completion. Related one-shot instability in |
| 291 | # the same kernel: flashinfer-ai/flashinfer#1223. |
| 292 | trigger_completion_at_end=(use_oneshot is True) |
| 293 | or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS, |
| 294 | ) |
| 295 | |
| 296 | direct_register_custom_op( |
| 297 | op_name="flashinfer_trtllm_fused_allreduce_norm", |
| 298 | op_func=call_trtllm_fused_allreduce_norm, |
| 299 | mutates_args=[ |
| 300 | "allreduce_in", |
| 301 | "residual", |
| 302 | "norm_out", |
| 303 | "quant_out", |
| 304 | "scale_out", |
| 305 | ], |
| 306 | ) |
| 307 | flashinfer_trtllm_fused_allreduce_norm = ( |
| 308 | torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default |
| 309 | ) |
| 310 | |
| 311 | |
| 312 | class FlashInferFusedAllReduceParams: |
| 313 | """Parameters for FlashInfer fused allreduce operations.""" |
| 314 | |
| 315 | def __init__( |
| 316 | self, |
| 317 | world_size: int, |
| 318 | max_token_num: int = 1024, |
| 319 | ) -> None: |
| 320 | self.world_size = world_size |
| 321 | self.launch_with_pdl = True |
| 322 | self.fp32_acc = True |
| 323 | self.max_token_num = max_token_num |
| 324 | |
| 325 | def get_trtllm_fused_allreduce_kwargs(self) -> dict[str, bool | int]: |
| 326 | return { |
| 327 | "world_size": self.world_size, |
| 328 | "launch_with_pdl": self.launch_with_pdl, |
| 329 | "fp32_acc": self.fp32_acc, |
| 330 | "max_token_num": self.max_token_num, |
| 331 | } |
| 332 | |
| 333 | |
| 334 | # TODO(luka): unify |
| 335 | class BasePattern: |
| 336 | def __init__(self, dtype: torch.dtype, device: str | None) -> None: |
| 337 | self.dtype = dtype |
| 338 | self.device = device |
| 339 | self.tp = get_tp_group() |
| 340 | self.tp_size = get_tensor_model_parallel_world_size() |
| 341 | |
| 342 | def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: |
| 343 | return torch.empty(*args, dtype=self.dtype, device=self.device, **kwargs) |
| 344 | |
| 345 | def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: |
| 346 | return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs) |
| 347 | |
| 348 | |
| 349 | class AllReduceRMSNormPattern(BasePattern): |
| 350 | """This pattern replaces the allreduce + rms norm (without residual) |
| 351 | with fused flashinfer implementation. |
| 352 | Applies to allreduce + rmsnorm before attn in the first Transformer block. |
| 353 | """ |
| 354 | |
| 355 | def __init__( |
| 356 | self, |
| 357 | epsilon: float, |
| 358 | dtype: torch.dtype, |
| 359 | device: str | None, |
| 360 | allreduce_params: FlashInferFusedAllReduceParams, |
| 361 | ) -> None: |
| 362 | super().__init__(dtype, device) |
| 363 | self.epsilon = epsilon |
| 364 | self.allreduce_params = allreduce_params |
| 365 | |
| 366 | def get_inputs(self) -> list[torch.Tensor]: |
| 367 | # input, weight |
| 368 | return [self.empty(5, 16), self.empty(16)] |
| 369 | |
| 370 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 371 | def pattern( |
| 372 | input: torch.Tensor, weight: torch.Tensor |
| 373 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 374 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 375 | rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon) |
| 376 | |
| 377 | return rms, allreduce_output |
| 378 | |
| 379 | def replacement( |
| 380 | input: torch.Tensor, weight: torch.Tensor |
| 381 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 382 | residual = torch.zeros_like(input) |
| 383 | rms_result = torch.empty_like(input) |
| 384 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 385 | allreduce = auto_functionalized( |
| 386 | flashinfer_trtllm_fused_allreduce_norm, |
| 387 | allreduce_in=input, |
| 388 | residual=residual, |
| 389 | norm_out=rms_result, |
| 390 | quant_out=None, |
| 391 | scale_out=None, |
| 392 | rms_gamma=weight, |
| 393 | rms_eps=self.epsilon, |
| 394 | pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, |
| 395 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 396 | ) |
| 397 | # rms_result, allreduce_in |
| 398 | return allreduce[3], allreduce[1] |
| 399 | |
| 400 | pm.register_replacement( |
| 401 | pattern, |
| 402 | replacement, |
| 403 | self.get_inputs(), |
| 404 | pm.fwd_only, |
| 405 | pm_pass, |
| 406 | extra_check=_rms_input_weight_dtype_match, |
| 407 | ) |
| 408 | |
| 409 | |
| 410 | class AllReduceFusedAddRMSNormPattern(BasePattern): |
| 411 | """This pattern replaces the allreduce + rms norm (with residual) |
| 412 | with fused flashinfer implementation. |
| 413 | Applies to o_proj + rmsnorm after attn and mlp + rmsnorm before attn. |
| 414 | """ |
| 415 | |
| 416 | def __init__( |
| 417 | self, |
| 418 | epsilon: float, |
| 419 | dtype: torch.dtype, |
| 420 | device: str | None, |
| 421 | allreduce_params: FlashInferFusedAllReduceParams, |
| 422 | ) -> None: |
| 423 | super().__init__(dtype, device) |
| 424 | self.epsilon = epsilon |
| 425 | self.allreduce_params = allreduce_params |
| 426 | |
| 427 | def get_inputs(self) -> list[torch.Tensor]: |
| 428 | input = self.empty(5, 16) |
| 429 | residual = self.empty(5, 16) |
| 430 | weight = self.empty(16) |
| 431 | |
| 432 | # input goes through allreduce first, always 16-bit |
| 433 | return [residual, input.to(self.dtype), weight] |
| 434 | |
| 435 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 436 | def pattern( |
| 437 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 438 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 439 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 440 | rms, residual = vllm.ir.ops.fused_add_rms_norm( |
| 441 | allreduce_output, residual, weight, self.epsilon |
| 442 | ) |
| 443 | return rms, residual |
| 444 | |
| 445 | def replacement( |
| 446 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 447 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 448 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 449 | allreduce = auto_functionalized( |
| 450 | flashinfer_trtllm_fused_allreduce_norm, |
| 451 | allreduce_in=input, |
| 452 | residual=residual, |
| 453 | norm_out=None, |
| 454 | quant_out=None, |
| 455 | scale_out=None, |
| 456 | rms_gamma=weight, |
| 457 | rms_eps=self.epsilon, |
| 458 | pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, |
| 459 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 460 | ) |
| 461 | # allreduce_in, residual |
| 462 | return allreduce[1], allreduce[2] |
| 463 | |
| 464 | # extra_check routes a Gemma fp32 gamma to AllReduceFusedAddGemmaRMSNormPattern. |
| 465 | pm.register_replacement( |
| 466 | pattern, |
| 467 | replacement, |
| 468 | self.get_inputs(), |
| 469 | pm.fwd_only, |
| 470 | pm_pass, |
| 471 | extra_check=_norm_input_weight_dtype_match, |
| 472 | ) |
| 473 | |
| 474 | # Same pattern, but only return the output and not residual |
| 475 | # (helpful for end of graph where residual is not used again) |
| 476 | first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0] |
| 477 | |
| 478 | pm.register_replacement( |
| 479 | first_return_only(pattern), # type: ignore[no-untyped-call] |
| 480 | first_return_only(replacement), # type: ignore[no-untyped-call] |
| 481 | self.get_inputs(), |
| 482 | pm.fwd_only, |
| 483 | pm_pass, |
| 484 | extra_check=_norm_input_weight_dtype_match, |
| 485 | ) |
| 486 | |
| 487 | |
| 488 | class AllReduceGemmaRMSNormPattern(BasePattern): |
| 489 | """Gemma-style variant of AllReduceRMSNormPattern (no residual).""" |
| 490 | |
| 491 | def __init__( |
| 492 | self, |
| 493 | epsilon: float, |
| 494 | dtype: torch.dtype, |
| 495 | device: str | None, |
| 496 | allreduce_params: FlashInferFusedAllReduceParams, |
| 497 | ) -> None: |
| 498 | super().__init__(dtype, device) |
| 499 | self.epsilon = epsilon |
| 500 | self.allreduce_params = allreduce_params |
| 501 | |
| 502 | def get_inputs(self) -> list[torch.Tensor]: |
| 503 | return [self.empty(5, 16), self.empty(16)] |
| 504 | |
| 505 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 506 | def pattern( |
| 507 | input: torch.Tensor, weight: torch.Tensor |
| 508 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 509 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 510 | rms = vllm.ir.ops.rms_norm( |
| 511 | allreduce_output, weight.float() + 1.0, self.epsilon |
| 512 | ) |
| 513 | return rms, allreduce_output |
| 514 | |
| 515 | def replacement( |
| 516 | input: torch.Tensor, weight: torch.Tensor |
| 517 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 518 | residual = torch.zeros_like(input) |
| 519 | rms_result = torch.empty_like(input) |
| 520 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 521 | allreduce = auto_functionalized( |
| 522 | flashinfer_trtllm_fused_allreduce_norm, |
| 523 | allreduce_in=input, |
| 524 | residual=residual, |
| 525 | norm_out=rms_result, |
| 526 | quant_out=None, |
| 527 | scale_out=None, |
| 528 | rms_gamma=weight, |
| 529 | rms_eps=self.epsilon, |
| 530 | pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, |
| 531 | weight_bias=1.0, |
| 532 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 533 | ) |
| 534 | return allreduce[3], allreduce[1] |
| 535 | |
| 536 | pm.register_replacement( |
| 537 | pattern, |
| 538 | replacement, |
| 539 | self.get_inputs(), |
| 540 | pm.fwd_only, |
| 541 | pm_pass, |
| 542 | ) |
| 543 | |
| 544 | |
| 545 | class AllReduceFusedAddGemmaRMSNormPattern(BasePattern): |
| 546 | """Gemma-style variant of AllReduceFusedAddRMSNormPattern (with residual).""" |
| 547 | |
| 548 | def __init__( |
| 549 | self, |
| 550 | epsilon: float, |
| 551 | dtype: torch.dtype, |
| 552 | device: str | None, |
| 553 | allreduce_params: FlashInferFusedAllReduceParams, |
| 554 | ) -> None: |
| 555 | super().__init__(dtype, device) |
| 556 | self.epsilon = epsilon |
| 557 | self.allreduce_params = allreduce_params |
| 558 | |
| 559 | def get_inputs(self) -> list[torch.Tensor]: |
| 560 | input = self.empty(5, 16) |
| 561 | residual = self.empty(5, 16) |
| 562 | weight = self.empty(16) |
| 563 | return [residual, input.to(self.dtype), weight] |
| 564 | |
| 565 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 566 | def pattern( |
| 567 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 568 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 569 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 570 | rms, residual = vllm.ir.ops.fused_add_rms_norm( |
| 571 | allreduce_output, residual, weight.float() + 1.0, self.epsilon |
| 572 | ) |
| 573 | return rms, residual |
| 574 | |
| 575 | def replacement( |
| 576 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 577 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 578 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 579 | allreduce = auto_functionalized( |
| 580 | flashinfer_trtllm_fused_allreduce_norm, |
| 581 | allreduce_in=input, |
| 582 | residual=residual, |
| 583 | norm_out=None, |
| 584 | quant_out=None, |
| 585 | scale_out=None, |
| 586 | rms_gamma=weight, |
| 587 | rms_eps=self.epsilon, |
| 588 | pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, |
| 589 | weight_bias=1.0, |
| 590 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 591 | ) |
| 592 | return allreduce[1], allreduce[2] |
| 593 | |
| 594 | pm.register_replacement( |
| 595 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 596 | ) |
| 597 | |
| 598 | first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0] |
| 599 | |
| 600 | pm.register_replacement( |
| 601 | first_return_only(pattern), # type: ignore[no-untyped-call] |
| 602 | first_return_only(replacement), # type: ignore[no-untyped-call] |
| 603 | self.get_inputs(), |
| 604 | pm.fwd_only, |
| 605 | pm_pass, |
| 606 | ) |
| 607 | |
| 608 | |
| 609 | class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern): |
| 610 | """This pattern replaces the allreduce + rms norm (without residual) |
| 611 | + static fp8 quant with fused flashinfer implementation. |
| 612 | Applies to allreduce + rmsnorm + quant before attn |
| 613 | in the first Transformer block. |
| 614 | """ |
| 615 | |
| 616 | def __init__( |
| 617 | self, |
| 618 | epsilon: float, |
| 619 | dtype: torch.dtype, |
| 620 | device: str | None, |
| 621 | allreduce_params: FlashInferFusedAllReduceParams, |
| 622 | ) -> None: |
| 623 | super().__init__(dtype, device) |
| 624 | self.epsilon = epsilon |
| 625 | self.allreduce_params = allreduce_params |
| 626 | self.quant_dtype = torch.float8_e4m3fn |
| 627 | self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) |
| 628 | |
| 629 | def get_inputs(self) -> list[torch.Tensor]: |
| 630 | _, scale = self.quant_matcher.inputs() |
| 631 | |
| 632 | # input, weight |
| 633 | return [self.empty(5, 16), self.empty(16), scale] |
| 634 | |
| 635 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 636 | def pattern( |
| 637 | input: torch.Tensor, |
| 638 | weight: torch.Tensor, |
| 639 | scale: torch.Tensor, |
| 640 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 641 | all_reduce = tensor_model_parallel_all_reduce(input) |
| 642 | rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) |
| 643 | quant, _ = self.quant_matcher(rms, scale) |
| 644 | return quant, all_reduce |
| 645 | |
| 646 | def replacement( |
| 647 | input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor |
| 648 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 649 | residual = torch.zeros_like(input) |
| 650 | result_rms = torch.empty_like(input) |
| 651 | result_quant = torch.empty_like(input, dtype=self.quant_dtype) |
| 652 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 653 | allreduce = auto_functionalized( |
| 654 | flashinfer_trtllm_fused_allreduce_norm, |
| 655 | allreduce_in=input, |
| 656 | residual=residual, |
| 657 | norm_out=result_rms, |
| 658 | quant_out=result_quant, |
| 659 | scale_out=None, |
| 660 | rms_gamma=weight, |
| 661 | rms_eps=self.epsilon, |
| 662 | # We don't use norm_out afterwards |
| 663 | pattern_code=( |
| 664 | flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant |
| 665 | ), |
| 666 | scale_factor=scale, |
| 667 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 668 | ) |
| 669 | |
| 670 | # quant_out, allreduce_output |
| 671 | return allreduce[4], allreduce[1] |
| 672 | |
| 673 | pm.register_replacement( |
| 674 | pattern, |
| 675 | replacement, |
| 676 | self.get_inputs(), |
| 677 | pm.fwd_only, |
| 678 | pm_pass, |
| 679 | extra_check=_rms_input_weight_dtype_match, |
| 680 | ) |
| 681 | |
| 682 | |
| 683 | class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern): |
| 684 | """This pattern replaces the allreduce + rms norm (with residual) |
| 685 | + static fp8 quant with fused flashinfer implementation. |
| 686 | Applies to o_proj + rmsnorm after attn + quant and |
| 687 | mlp + rmsnorm + quant before attn. |
| 688 | """ |
| 689 | |
| 690 | def __init__( |
| 691 | self, |
| 692 | epsilon: float, |
| 693 | dtype: torch.dtype, |
| 694 | device: str | None, |
| 695 | allreduce_params: FlashInferFusedAllReduceParams, |
| 696 | ) -> None: |
| 697 | super().__init__(dtype, device) |
| 698 | self.epsilon = epsilon |
| 699 | self.allreduce_params = allreduce_params |
| 700 | self.quant_dtype = torch.float8_e4m3fn |
| 701 | |
| 702 | self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) |
| 703 | |
| 704 | def get_inputs(self) -> list[torch.Tensor]: |
| 705 | input = self.empty(5, 16) |
| 706 | residual = self.empty(5, 16) |
| 707 | weight = self.empty(16) |
| 708 | _, scale = self.quant_matcher.inputs() |
| 709 | |
| 710 | # input goes through allreduce first, always 16-bit |
| 711 | return [residual, input.to(self.dtype), weight, scale] |
| 712 | |
| 713 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 714 | def pattern( |
| 715 | residual: torch.Tensor, |
| 716 | input: torch.Tensor, |
| 717 | weight: torch.Tensor, |
| 718 | scale: torch.Tensor, |
| 719 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 720 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 721 | rms, res = vllm.ir.ops.fused_add_rms_norm( |
| 722 | allreduce_output, residual, weight, self.epsilon |
| 723 | ) |
| 724 | quant, _ = self.quant_matcher(rms, scale) |
| 725 | |
| 726 | return quant, res |
| 727 | |
| 728 | def replacement( |
| 729 | residual: torch.Tensor, |
| 730 | input: torch.Tensor, |
| 731 | weight: torch.Tensor, |
| 732 | scale: torch.Tensor, |
| 733 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 734 | result_quant = torch.empty_like(input, dtype=self.quant_dtype) |
| 735 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 736 | allreduce = auto_functionalized( |
| 737 | flashinfer_trtllm_fused_allreduce_norm, |
| 738 | allreduce_in=input, |
| 739 | residual=residual, |
| 740 | norm_out=None, |
| 741 | quant_out=result_quant, |
| 742 | scale_out=None, |
| 743 | rms_gamma=weight, |
| 744 | rms_eps=self.epsilon, |
| 745 | # We don't use norm_out afterwards |
| 746 | pattern_code=( |
| 747 | flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant |
| 748 | ), |
| 749 | scale_factor=scale, |
| 750 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 751 | ) |
| 752 | # quant_out, rms_norm_residual |
| 753 | return allreduce[4], allreduce[2] |
| 754 | |
| 755 | pm.register_replacement( |
| 756 | pattern, |
| 757 | replacement, |
| 758 | self.get_inputs(), |
| 759 | pm.fwd_only, |
| 760 | pm_pass, |
| 761 | extra_check=_norm_input_weight_dtype_match, |
| 762 | ) |
| 763 | |
| 764 | |
| 765 | class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern): |
| 766 | """This pattern replaces the allreduce + rms norm (without residual) |
| 767 | + static nvfp4 quant with fused flashinfer implementation. |
| 768 | Applies to allreduce + rmsnorm + quant before attn |
| 769 | in the first Transformer block. |
| 770 | """ |
| 771 | |
| 772 | def __init__( |
| 773 | self, |
| 774 | epsilon: float, |
| 775 | dtype: torch.dtype, |
| 776 | device: str | None, |
| 777 | allreduce_params: FlashInferFusedAllReduceParams, |
| 778 | ) -> None: |
| 779 | super().__init__(dtype, device) |
| 780 | self.epsilon = epsilon |
| 781 | self.allreduce_params = allreduce_params |
| 782 | |
| 783 | def get_inputs(self) -> list[torch.Tensor]: |
| 784 | input = torch.empty([1, 16, 16], device=self.device, dtype=self.dtype) |
| 785 | quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8) |
| 786 | input_global_scale = torch.empty( |
| 787 | [1, 1], device=self.device, dtype=torch.float32 |
| 788 | ) |
| 789 | weight = torch.empty([16], device=self.device, dtype=self.dtype) |
| 790 | output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) |
| 791 | |
| 792 | return [input, quant_result, weight, input_global_scale, output_scale] |
| 793 | |
| 794 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 795 | def pattern( |
| 796 | input: torch.Tensor, |
| 797 | quant_result: torch.Tensor, |
| 798 | weight: torch.Tensor, |
| 799 | input_global_scale: torch.Tensor, |
| 800 | output_scale: torch.Tensor, |
| 801 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 802 | all_reduce = tensor_model_parallel_all_reduce(input) |
| 803 | rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) |
| 804 | quant_out_tuple = auto_functionalized( |
| 805 | STATIC_FP4_QUANT_OP, |
| 806 | input=rms, |
| 807 | input_scale=input_global_scale, |
| 808 | is_sf_swizzled_layout=True, |
| 809 | output=quant_result, |
| 810 | output_scale=output_scale, |
| 811 | ) |
| 812 | |
| 813 | # quant_out, allreduce_output, output_scale |
| 814 | return quant_out_tuple[1], all_reduce, quant_out_tuple[2] |
| 815 | |
| 816 | def replacement( |
| 817 | input: torch.Tensor, |
| 818 | quant_result: torch.Tensor, |
| 819 | weight: torch.Tensor, |
| 820 | input_global_scale: torch.Tensor, |
| 821 | output_scale: torch.Tensor, |
| 822 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 823 | residual = torch.zeros_like(input) |
| 824 | result_rms = torch.empty_like(input) |
| 825 | output_scale_fp8 = _view_nvfp4_scale_out_for_flashinfer(output_scale) |
| 826 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 827 | allreduce = auto_functionalized( |
| 828 | flashinfer_trtllm_fused_allreduce_norm, |
| 829 | allreduce_in=input, |
| 830 | residual=residual, |
| 831 | norm_out=result_rms, |
| 832 | quant_out=quant_result, |
| 833 | scale_out=output_scale_fp8, |
| 834 | rms_gamma=weight, |
| 835 | rms_eps=self.epsilon, |
| 836 | # We don't use norm_out afterwards |
| 837 | pattern_code=( |
| 838 | flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant |
| 839 | ), |
| 840 | scale_factor=input_global_scale, |
| 841 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 842 | ) |
| 843 | |
| 844 | # quant_out, allreduce_output, output_scale |
| 845 | return ( |
| 846 | allreduce[4], |
| 847 | allreduce[1], |
| 848 | _view_flashinfer_nvfp4_scale_out_as_int32(allreduce[5]), |
| 849 | ) |
| 850 | |
| 851 | pm.register_replacement( |
| 852 | pattern, |
| 853 | replacement, |
| 854 | self.get_inputs(), |
| 855 | pm.fwd_only, |
| 856 | pm_pass, |
| 857 | extra_check=_rms_input_weight_dtype_match, |
| 858 | ) |
| 859 | |
| 860 | |
| 861 | class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern): |
| 862 | """This pattern replaces the allreduce + rms norm (with residual) |
| 863 | + static nvfp4 quant with fused flashinfer implementation. |
| 864 | Applies to o_proj + rmsnorm after attn + quant and |
| 865 | mlp + rmsnorm + quant before attn. |
| 866 | """ |
| 867 | |
| 868 | def __init__( |
| 869 | self, |
| 870 | epsilon: float, |
| 871 | dtype: torch.dtype, |
| 872 | device: str | None, |
| 873 | allreduce_params: FlashInferFusedAllReduceParams, |
| 874 | ) -> None: |
| 875 | super().__init__(dtype, device) |
| 876 | self.epsilon = epsilon |
| 877 | self.allreduce_params = allreduce_params |
| 878 | |
| 879 | def get_inputs(self) -> list[torch.Tensor]: |
| 880 | input = torch.empty([16, 16], device=self.device, dtype=self.dtype) |
| 881 | |
| 882 | residual = torch.empty([16, 16], device=self.device, dtype=self.dtype) |
| 883 | weight = torch.empty([16, 16], device=self.device, dtype=self.dtype) |
| 884 | quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8) |
| 885 | input_global_scale = torch.empty( |
| 886 | [1, 1], device=self.device, dtype=torch.float32 |
| 887 | ) |
| 888 | output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) |
| 889 | |
| 890 | return [ |
| 891 | quant_result, |
| 892 | residual, |
| 893 | input, |
| 894 | output_scale, |
| 895 | weight, |
| 896 | input_global_scale, |
| 897 | ] |
| 898 | |
| 899 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 900 | def pattern( |
| 901 | quant_result: torch.Tensor, |
| 902 | residual: torch.Tensor, |
| 903 | input: torch.Tensor, |
| 904 | output_scale: torch.Tensor, |
| 905 | weight: torch.Tensor, |
| 906 | input_global_scale: torch.Tensor, |
| 907 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 908 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 909 | rms, residual = vllm.ir.ops.fused_add_rms_norm( |
| 910 | allreduce_output, residual, weight, self.epsilon |
| 911 | ) |
| 912 | quant_out_tuple = auto_functionalized( |
| 913 | STATIC_FP4_QUANT_OP, |
| 914 | input=rms, |
| 915 | input_scale=input_global_scale, |
| 916 | is_sf_swizzled_layout=True, |
| 917 | output=quant_result, |
| 918 | output_scale=output_scale, |
| 919 | ) |
| 920 | |
| 921 | # quant_out, allreduce_output, output_scale |
| 922 | return quant_out_tuple[1], residual, quant_out_tuple[2] |
| 923 | |
| 924 | def replacement( |
| 925 | quant_result: torch.Tensor, |
| 926 | residual: torch.Tensor, |
| 927 | input: torch.Tensor, |
| 928 | output_scale: torch.Tensor, |
| 929 | weight: torch.Tensor, |
| 930 | input_global_scale: torch.Tensor, |
| 931 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 932 | output_scale_fp8 = _view_nvfp4_scale_out_for_flashinfer(output_scale) |
| 933 | assert flashinfer_comm is not None, "FlashInfer must be enabled" |
| 934 | allreduce = auto_functionalized( |
| 935 | flashinfer_trtllm_fused_allreduce_norm, |
| 936 | allreduce_in=input, |
| 937 | residual=residual, |
| 938 | norm_out=None, |
| 939 | quant_out=quant_result, |
| 940 | scale_out=output_scale_fp8, |
| 941 | rms_gamma=weight, |
| 942 | rms_eps=self.epsilon, |
| 943 | # We don't use norm_out afterwards |
| 944 | pattern_code=( |
| 945 | flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant |
| 946 | ), |
| 947 | scale_factor=input_global_scale, |
| 948 | **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(), |
| 949 | ) |
| 950 | # quant_out, rms_norm_residual, output_scale |
| 951 | return ( |
| 952 | allreduce[4], |
| 953 | allreduce[2], |
| 954 | _view_flashinfer_nvfp4_scale_out_as_int32(allreduce[5]), |
| 955 | ) |
| 956 | |
| 957 | pm.register_replacement( |
| 958 | pattern, |
| 959 | replacement, |
| 960 | self.get_inputs(), |
| 961 | pm.fwd_only, |
| 962 | pm_pass, |
| 963 | extra_check=_norm_input_weight_dtype_match, |
| 964 | ) |
| 965 | |
| 966 | |
| 967 | def _fused_ar_workspace_hidden_dim(config: VllmConfig) -> int: |
| 968 | """Widest hidden size across the target and (optional) draft models. |
| 969 | |
| 970 | The FlashInfer allreduce+RMSNorm workspace is a process-global singleton |
| 971 | created eagerly at pass construction and reused by every model. Under |
| 972 | speculative decoding the draft shares it, so it must fit the larger of the |
| 973 | two hidden sizes; a draft wider than the target otherwise overflows the |
| 974 | target-sized buffer (vLLM #52023). |
| 975 | """ |
| 976 | hidden_dim = config.model_config.get_hidden_size() |
| 977 | spec = config.speculative_config |
| 978 | draft_model_config = getattr(spec, "draft_model_config", None) if spec else None |
| 979 | if draft_model_config is not None: |
| 980 | hidden_dim = max(hidden_dim, draft_model_config.get_hidden_size()) |
| 981 | return hidden_dim |
| 982 | |
| 983 | |
| 984 | class AllReduceFusionPass(VllmPatternMatcherPass): |
| 985 | def __init__(self, config: VllmConfig) -> None: |
| 986 | super().__init__(config) |
| 987 | self.disabled = True |
| 988 | self.tp_size = get_tensor_model_parallel_world_size() |
| 989 | if self.tp_size <= 1: |
| 990 | logger.warning_once("AllReduce fusion pass is disabled for tp_size <= 1.") |
| 991 | return |
| 992 | self.patterns: PatternMatcherPass = PatternMatcherPass( |
| 993 | pass_name="all_reduce_fusion_pass" |
| 994 | ) |
| 995 | if config.model_config is None: |
| 996 | logger.warning_once( |
| 997 | "AllReduce fusion pass is disabled for missing model_config." |
| 998 | ) |
| 999 | return |
| 1000 | # Size the shared workspace for the widest model that reuses it (the |
| 1001 | # draft under speculative decoding may be wider than the target). |
| 1002 | self.workspace_hidden_dim = _fused_ar_workspace_hidden_dim(config) |
| 1003 | self.group = get_tp_group().cpu_group |
| 1004 | rank = get_tensor_model_parallel_rank() |
| 1005 | if flashinfer_comm is None: |
| 1006 | logger.warning( |
| 1007 | "Flashinfer is not installed or comm module not found, " |
| 1008 | "skipping allreduce fusion pass" |
| 1009 | ) |
| 1010 | return |
| 1011 | max_size = config.compilation_config.pass_config.flashinfer_max_size( |
| 1012 | self.tp_size |
| 1013 | ) |
| 1014 | if max_size is None: |
| 1015 | # Flashinfer doesn't support current world size |
| 1016 | logger.warning( |
| 1017 | "Flashinfer allreduce fusion is not supported for world size %s" |
| 1018 | " or max size is not provided", |
| 1019 | self.tp_size, |
| 1020 | ) |
| 1021 | return |
| 1022 | element_size = torch.tensor([], dtype=self.model_dtype).element_size() |
| 1023 | self.max_token_num = max_size // (self.workspace_hidden_dim * element_size) |
| 1024 | # take the min to save workspace size and we'll never use more |
| 1025 | # than max_num_batched_tokens anyways |
| 1026 | self.max_token_num = min( |
| 1027 | self.max_token_num, config.scheduler_config.max_num_batched_tokens |
| 1028 | ) |
| 1029 | logger.debug_once( |
| 1030 | f"Flashinfer max size: {max_size // (1024 * 1024)} MB," |
| 1031 | "Maximal number of tokens used by " |
| 1032 | f"Flashinfer Allreduce Fusion: {self.max_token_num}", |
| 1033 | scope="global", |
| 1034 | ) |
| 1035 | |
| 1036 | workspace_kwargs = dict( |
| 1037 | world_size=self.tp_size, |
| 1038 | rank=rank, |
| 1039 | max_token_num=self.max_token_num, |
| 1040 | hidden_dim=self.workspace_hidden_dim, |
| 1041 | dtype=self.model_dtype, |
| 1042 | group=self.group, |
| 1043 | ) |
| 1044 | if get_fi_ar_workspace(**workspace_kwargs) is None: |
| 1045 | logger.warning_once( |
| 1046 | "Failed to initialize Flashinfer allreduce workspace. " |
| 1047 | "Flashinfer allreduce-norm fusion will be disabled." |
| 1048 | ) |
| 1049 | return |
| 1050 | |
| 1051 | self.supports_quant_fusion = ( |
| 1052 | get_fi_ar_quant_workspace(**workspace_kwargs) is not None |
| 1053 | ) |
| 1054 | if not self.supports_quant_fusion: |
| 1055 | logger.warning_once( |
| 1056 | "Failed to initialize Flashinfer allreduce workspace. " |
| 1057 | "Flashinfer allreduce-norm-quant fusion will be disabled." |
| 1058 | ) |
| 1059 | |
| 1060 | self.allreduce_params = FlashInferFusedAllReduceParams( |
| 1061 | world_size=self.tp_size, |
| 1062 | max_token_num=self.max_token_num, |
| 1063 | ) |
| 1064 | |
| 1065 | self.register_patterns() |
| 1066 | self.dump_patterns(config, self.patterns) |
| 1067 | |
| 1068 | @enable_fake_mode |
| 1069 | def register_patterns(self) -> None: |
| 1070 | for epsilon in [1e-5, 1e-6]: |
| 1071 | if self.supports_quant_fusion: |
| 1072 | AllReduceFusedRMSNormStaticQuantFP8Pattern( |
| 1073 | epsilon, |
| 1074 | self.model_dtype, |
| 1075 | self.device, |
| 1076 | self.allreduce_params, |
| 1077 | ).register(self.patterns) |
| 1078 | AllReduceFusedAddRMSNormStaticQuantFP8Pattern( |
| 1079 | epsilon, |
| 1080 | self.model_dtype, |
| 1081 | self.device, |
| 1082 | self.allreduce_params, |
| 1083 | ).register(self.patterns) |
| 1084 | if current_platform.has_device_capability(100): |
| 1085 | AllReduceFusedRMSNormStaticQuantNVFP4Pattern( |
| 1086 | epsilon, |
| 1087 | self.model_dtype, |
| 1088 | self.device, |
| 1089 | self.allreduce_params, |
| 1090 | ).register(self.patterns) |
| 1091 | AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern( |
| 1092 | epsilon, |
| 1093 | self.model_dtype, |
| 1094 | self.device, |
| 1095 | self.allreduce_params, |
| 1096 | ).register(self.patterns) |
| 1097 | AllReduceRMSNormPattern( |
| 1098 | epsilon, |
| 1099 | self.model_dtype, |
| 1100 | self.device, |
| 1101 | self.allreduce_params, |
| 1102 | ).register(self.patterns) |
| 1103 | AllReduceFusedAddRMSNormPattern( |
| 1104 | epsilon, |
| 1105 | self.model_dtype, |
| 1106 | self.device, |
| 1107 | self.allreduce_params, |
| 1108 | ).register(self.patterns) |
| 1109 | AllReduceGemmaRMSNormPattern( |
| 1110 | epsilon, |
| 1111 | self.model_dtype, |
| 1112 | self.device, |
| 1113 | self.allreduce_params, |
| 1114 | ).register(self.patterns) |
| 1115 | AllReduceFusedAddGemmaRMSNormPattern( |
| 1116 | epsilon, |
| 1117 | self.model_dtype, |
| 1118 | self.device, |
| 1119 | self.allreduce_params, |
| 1120 | ).register(self.patterns) |
| 1121 | |
| 1122 | # WARNING: This is a hack to clear the pattern matcher cache |
| 1123 | # and allow multiple values of epsilon. |
| 1124 | torch._inductor.pattern_matcher._seen_patterns.clear() |
| 1125 | |
| 1126 | self.disabled = False |
| 1127 | |
| 1128 | def is_applicable_for_range(self, compile_range: Range) -> bool: |
| 1129 | if self.disabled: |
| 1130 | logger.warning_once("AllReduce fusion pass is disabled.") |
| 1131 | return False |
| 1132 | return bool(compile_range.end <= self.max_token_num) |
| 1133 | |
| 1134 | @VllmInductorPass.time_and_log |
| 1135 | def __call__(self, graph: fx.Graph) -> None: |
| 1136 | if self.disabled: |
| 1137 | logger.debug("AllReduceFusionPass disabled") |
| 1138 | return |
| 1139 | |
| 1140 | self.matched_count = self.patterns.apply(graph) |
| 1141 | logger.debug("Replaced %s patterns", self.matched_count) |
| 1142 | |
| 1143 | def __del__(self) -> None: |
| 1144 | if getattr(self, "disabled", True): |
| 1145 | return |
| 1146 | with contextlib.suppress(Exception): |
| 1147 | destroy_fi_ar_workspace() |
| 1148 | |
| 1149 | |
| 1150 | # TODO: make BasePattern to inherit from VllmPatternReplacement |
| 1151 | class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement): |
| 1152 | def __init__( |
| 1153 | self, |
| 1154 | epsilon: float, |
| 1155 | dtype: torch.dtype, |
| 1156 | device: str | None, |
| 1157 | use_aiter_rmsnorm: bool = True, |
| 1158 | ) -> None: |
| 1159 | super().__init__(dtype, device) |
| 1160 | self.dtype = dtype |
| 1161 | self.epsilon = epsilon |
| 1162 | self.FUSED_AR_RMSNORM_OP = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op() |
| 1163 | |
| 1164 | def get_inputs(self) -> list[torch.Tensor]: |
| 1165 | return [self.empty(5, 16), self.empty(16)] |
| 1166 | |
| 1167 | @property |
| 1168 | def pattern(self): |
| 1169 | def _pattern( |
| 1170 | input: torch.Tensor, weight: torch.Tensor |
| 1171 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1172 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 1173 | rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon) |
| 1174 | |
| 1175 | return rms, allreduce_output |
| 1176 | |
| 1177 | return _pattern |
| 1178 | |
| 1179 | @property |
| 1180 | def replacement(self): |
| 1181 | def _replacement( |
| 1182 | input: torch.Tensor, weight: torch.Tensor |
| 1183 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1184 | residual = torch.zeros_like(input) |
| 1185 | allreduce = self.FUSED_AR_RMSNORM_OP( |
| 1186 | input_=input, |
| 1187 | residual=residual, |
| 1188 | weight=weight.to(input.dtype), |
| 1189 | epsilon=self.epsilon, |
| 1190 | ) |
| 1191 | return allreduce[0], allreduce[1] |
| 1192 | |
| 1193 | return _replacement |
| 1194 | |
| 1195 | |
| 1196 | class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement): |
| 1197 | def __init__( |
| 1198 | self, |
| 1199 | epsilon: float, |
| 1200 | dtype: torch.dtype, |
| 1201 | device: str | None, |
| 1202 | use_aiter_rmsnorm: bool = True, |
| 1203 | ) -> None: |
| 1204 | super().__init__(dtype, device) |
| 1205 | self.epsilon = epsilon |
| 1206 | self.dtype = dtype |
| 1207 | self.FUSED_AR_RMSNORM_OP = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op() |
| 1208 | |
| 1209 | def get_inputs(self) -> list[torch.Tensor]: |
| 1210 | # input, residual, weight |
| 1211 | return [self.empty(5, 16), self.empty(5, 16), self.empty(16)] |
| 1212 | |
| 1213 | @property |
| 1214 | def pattern(self): |
| 1215 | def _pattern( |
| 1216 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 1217 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1218 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 1219 | rms, residual = vllm.ir.ops.fused_add_rms_norm( |
| 1220 | allreduce_output, residual, weight, self.epsilon |
| 1221 | ) |
| 1222 | return rms, residual |
| 1223 | |
| 1224 | return _pattern |
| 1225 | |
| 1226 | @property |
| 1227 | def replacement(self): |
| 1228 | def _replacement( |
| 1229 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 1230 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1231 | allreduce = self.FUSED_AR_RMSNORM_OP( |
| 1232 | input_=input, |
| 1233 | residual=residual, |
| 1234 | weight=weight.to(input.dtype), |
| 1235 | epsilon=self.epsilon, |
| 1236 | ) |
| 1237 | return allreduce[0], allreduce[1] |
| 1238 | |
| 1239 | return _replacement |
| 1240 | |
| 1241 | |
| 1242 | class AiterAllreduceFusedAddRMSNormOutputOnlyPattern( |
| 1243 | AiterAllreduceFusedAddRMSNormPattern |
| 1244 | ): |
| 1245 | """Match the add-RMSNorm form when its residual output is dead.""" |
| 1246 | |
| 1247 | @property |
| 1248 | def pattern(self): |
| 1249 | pattern = super().pattern |
| 1250 | |
| 1251 | def _pattern( |
| 1252 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 1253 | ) -> torch.Tensor: |
| 1254 | return pattern(residual, input, weight)[0] |
| 1255 | |
| 1256 | return _pattern |
| 1257 | |
| 1258 | @property |
| 1259 | def replacement(self): |
| 1260 | replacement = super().replacement |
| 1261 | |
| 1262 | def _replacement( |
| 1263 | residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor |
| 1264 | ) -> torch.Tensor: |
| 1265 | return replacement(residual, input, weight)[0] |
| 1266 | |
| 1267 | return _replacement |
| 1268 | |
| 1269 | |
| 1270 | class AiterAllreduceFusedRMSNormGroupQuantFP8Pattern( |
| 1271 | BasePattern, VllmPatternReplacement |
| 1272 | ): |
| 1273 | """Fuse AllReduce + RMSNorm + per-group FP8 quant into a single AITER |
| 1274 | custom op. |
| 1275 | |
| 1276 | Matches the AR-side analogue of ``AiterRMSFp8GroupQuantPattern`` in |
| 1277 | ``rocm_aiter_fusion.py``: ``all_reduce -> rms_norm -> group_fp8_quant`` |
| 1278 | fans out into ``rocm_aiter_fused_allreduce_rmsnorm_quant_per_group``. |
| 1279 | |
| 1280 | Without this pattern, ``RocmAiterAllReduceFusionPass`` would fuse the |
| 1281 | ``all_reduce + rms_norm`` half (PR #41825 wires that), but the trailing |
| 1282 | ``rocm_aiter_group_fp8_quant`` would still launch as a separate kernel. |
| 1283 | That standalone quant accounts for ~535us / decode step on DSv3.2 MI355X |
| 1284 | TP4 -- this pattern eliminates it by absorbing the quant into the AR |
| 1285 | epilogue. Group size 128 matches the FP8 block-scaled MM kernel used by |
| 1286 | DSv3.2's linear weights. |
| 1287 | """ |
| 1288 | |
| 1289 | def __init__( |
| 1290 | self, |
| 1291 | epsilon: float, |
| 1292 | dtype: torch.dtype, |
| 1293 | device: str | None, |
| 1294 | group_size: int = 128, |
| 1295 | ) -> None: |
| 1296 | super().__init__(dtype, device) |
| 1297 | self.epsilon = epsilon |
| 1298 | self.dtype = dtype |
| 1299 | self.group_size = group_size |
| 1300 | self.FUSED_AR_RMS_QUANT_OP = ( |
| 1301 | rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_op() |
| 1302 | ) |
| 1303 | self.quant_dtype = current_platform.fp8_dtype() |
| 1304 | self.quant_matcher = MatcherQuantFP8( |
| 1305 | QuantKey( |
| 1306 | dtype=self.quant_dtype, |
| 1307 | scale=ScaleDesc(torch.float32, False, GroupShape(1, group_size)), |
| 1308 | symmetric=True, |
| 1309 | ), |
| 1310 | match_rocm_aiter=True, |
| 1311 | ) |
| 1312 | |
| 1313 | def get_inputs(self) -> list[torch.Tensor]: |
| 1314 | # input, weight; hidden dim must be a group_size multiple so the |
| 1315 | # group quant matcher's example trace is well-defined. |
| 1316 | return [self.empty(5, self.group_size), self.empty(self.group_size)] |
| 1317 | |
| 1318 | @property |
| 1319 | def pattern(self): |
| 1320 | def _pattern( |
| 1321 | input: torch.Tensor, weight: torch.Tensor |
| 1322 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1323 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 1324 | rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon) |
| 1325 | quant, scale = self.quant_matcher(rms) |
| 1326 | return quant, scale |
| 1327 | |
| 1328 | return _pattern |
| 1329 | |
| 1330 | @property |
| 1331 | def replacement(self): |
| 1332 | def _replacement( |
| 1333 | input: torch.Tensor, weight: torch.Tensor |
| 1334 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1335 | residual = torch.zeros_like(input) |
| 1336 | result = self.FUSED_AR_RMS_QUANT_OP( |
| 1337 | input_=input, |
| 1338 | residual=residual, |
| 1339 | weight=weight.to(input.dtype), |
| 1340 | epsilon=self.epsilon, |
| 1341 | group_size=self.group_size, |
| 1342 | ) |
| 1343 | # quant_out, scale_out (residual is unused on the no-add path, |
| 1344 | # mirroring how AiterAllreduceFusedRMSNormPattern drops the |
| 1345 | # residual output) |
| 1346 | return result[0], result[2] |
| 1347 | |
| 1348 | return _replacement |
| 1349 | |
| 1350 | |
| 1351 | class AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern( |
| 1352 | BasePattern, VllmPatternReplacement |
| 1353 | ): |
| 1354 | """``fused_add`` variant of ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern``. |
| 1355 | |
| 1356 | Targets the dominant DSv3.2-style post-attention / post-MLP path: |
| 1357 | ``all_reduce -> fused_add_rms_norm -> group_fp8_quant``. Returns the |
| 1358 | FP8 quant output, the residual carry-over, and the per-group scale. |
| 1359 | """ |
| 1360 | |
| 1361 | def __init__( |
| 1362 | self, |
| 1363 | epsilon: float, |
| 1364 | dtype: torch.dtype, |
| 1365 | device: str | None, |
| 1366 | group_size: int = 128, |
| 1367 | ) -> None: |
| 1368 | super().__init__(dtype, device) |
| 1369 | self.epsilon = epsilon |
| 1370 | self.dtype = dtype |
| 1371 | self.group_size = group_size |
| 1372 | self.FUSED_AR_RMS_QUANT_OP = ( |
| 1373 | rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_op() |
| 1374 | ) |
| 1375 | self.quant_dtype = current_platform.fp8_dtype() |
| 1376 | self.quant_matcher = MatcherQuantFP8( |
| 1377 | QuantKey( |
| 1378 | dtype=self.quant_dtype, |
| 1379 | scale=ScaleDesc(torch.float32, False, GroupShape(1, group_size)), |
| 1380 | symmetric=True, |
| 1381 | ), |
| 1382 | match_rocm_aiter=True, |
| 1383 | ) |
| 1384 | |
| 1385 | def get_inputs(self) -> list[torch.Tensor]: |
| 1386 | # residual, input, weight |
| 1387 | return [ |
| 1388 | self.empty(5, self.group_size), |
| 1389 | self.empty(5, self.group_size), |
| 1390 | self.empty(self.group_size), |
| 1391 | ] |
| 1392 | |
| 1393 | @property |
| 1394 | def pattern(self): |
| 1395 | def _pattern( |
| 1396 | residual: torch.Tensor, |
| 1397 | input: torch.Tensor, |
| 1398 | weight: torch.Tensor, |
| 1399 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 1400 | allreduce_output = tensor_model_parallel_all_reduce(input) |
| 1401 | rms, residual_out = vllm.ir.ops.fused_add_rms_norm( |
| 1402 | allreduce_output, residual, weight, self.epsilon |
| 1403 | ) |
| 1404 | quant, scale = self.quant_matcher(rms) |
| 1405 | return quant, scale, residual_out |
| 1406 | |
| 1407 | return _pattern |
| 1408 | |
| 1409 | @property |
| 1410 | def replacement(self): |
| 1411 | def _replacement( |
| 1412 | residual: torch.Tensor, |
| 1413 | input: torch.Tensor, |
| 1414 | weight: torch.Tensor, |
| 1415 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 1416 | result = self.FUSED_AR_RMS_QUANT_OP( |
| 1417 | input_=input, |
| 1418 | residual=residual, |
| 1419 | weight=weight.to(input.dtype), |
| 1420 | epsilon=self.epsilon, |
| 1421 | group_size=self.group_size, |
| 1422 | ) |
| 1423 | # quant_out, scale_out, residual_out |
| 1424 | return result[0], result[2], result[1] |
| 1425 | |
| 1426 | return _replacement |
| 1427 | |
| 1428 | |
| 1429 | class AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( |
| 1430 | BasePattern, VllmPatternReplacement |
| 1431 | ): |
| 1432 | """Indexer-fan-out variant of ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern``. |
| 1433 | |
| 1434 | Targets the DSv3.2 post-attention / post-MLP path where the post-AR normed |
| 1435 | activation has two consumers: a per-group FP8 quant for ``fused_qkv_a_proj`` |
| 1436 | *and* a bf16 ``rocm_unquantized_gemm`` for the indexer ``wk_weights_proj``. |
| 1437 | The single-consumer pattern above cannot fire when this fan-out is present, |
| 1438 | so without this pattern the standalone FP8 quant kernel survives unfused |
| 1439 | (~535us / decode step on DSv3.2 MI355X TP4). |
| 1440 | |
| 1441 | Lowers to ``rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm`` |
| 1442 | (the ``emit_bf16=True`` variant of the AR+RMS+QUANT launcher, which returns |
| 1443 | FP8 quant + scales + bf16 normed activations in one kernel) and rewires the |
| 1444 | indexer GEMM onto the emitted bf16 norm output. The RMS output is also a |
| 1445 | graph output in DSv3.2's residual carry; it is returned as a pattern output |
| 1446 | so the matcher can substitute the bf16 norm in its place. |
| 1447 | |
| 1448 | The trailing FP8 group-quant is matched via ``MatcherQuantFP8`` (consistent |
| 1449 | with the sibling patterns above), which traces both ``QuantFP8.forward_hip`` |
| 1450 | and ``forward_native`` paths and so matches whichever op the call site |
| 1451 | lowers to (``vllm.rocm_aiter_group_fp8_quant``). |
| 1452 | """ |
| 1453 | |
| 1454 | def __init__( |
| 1455 | self, |
| 1456 | epsilon: float, |
| 1457 | dtype: torch.dtype, |
| 1458 | device: str | None, |
| 1459 | group_size: int = 128, |
| 1460 | ) -> None: |
| 1461 | super().__init__(dtype, device) |
| 1462 | self.epsilon = epsilon |
| 1463 | self.dtype = dtype |
| 1464 | self.group_size = group_size |
| 1465 | self.FUSED_AR_RMS_QUANT_BF16_OP = ( |
| 1466 | rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() # noqa: E501 |
| 1467 | ) |
| 1468 | self.quant_dtype = current_platform.fp8_dtype() |
| 1469 | self.quant_matcher = MatcherQuantFP8( |
| 1470 | QuantKey( |
| 1471 | dtype=self.quant_dtype, |
| 1472 | scale=ScaleDesc(torch.float32, False, GroupShape(1, group_size)), |
| 1473 | symmetric=True, |
| 1474 | ), |
| 1475 | match_rocm_aiter=True, |
| 1476 | ) |
| 1477 | |
| 1478 | def get_inputs(self) -> list[torch.Tensor]: |
| 1479 | h = self.group_size |
| 1480 | indexer_out = 8 |
| 1481 | return [ |
| 1482 | self.empty(5, h), |
| 1483 | self.empty(5, h), |
| 1484 | self.empty(h), |
| 1485 | self.empty(indexer_out, h), |
| 1486 | ] |
| 1487 | |
| 1488 | @property |
| 1489 | def pattern(self): |
| 1490 | eps = self.epsilon |
| 1491 | |
| 1492 | def _pattern( |
| 1493 | residual: torch.Tensor, |
| 1494 | input_: torch.Tensor, |
| 1495 | norm_weight: torch.Tensor, |
| 1496 | indexer_weight: torch.Tensor, |
| 1497 | ) -> tuple[ |
| 1498 | torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor |
| 1499 | ]: |
| 1500 | ar_out = tensor_model_parallel_all_reduce(input_) |
| 1501 | rms, res_out = vllm.ir.ops.fused_add_rms_norm( |
| 1502 | ar_out, residual, norm_weight, eps |
| 1503 | ) |
| 1504 | q, s = self.quant_matcher(rms) |
| 1505 | idx = torch.ops.vllm.rocm_unquantized_gemm(rms, indexer_weight) |
| 1506 | return q, s, res_out, idx, rms |
| 1507 | |
| 1508 | return _pattern |
| 1509 | |
| 1510 | @property |
| 1511 | def replacement(self): |
| 1512 | gs = self.group_size |
| 1513 | eps = self.epsilon |
| 1514 | |
| 1515 | def _replacement( |
| 1516 | residual: torch.Tensor, |
| 1517 | input_: torch.Tensor, |
| 1518 | norm_weight: torch.Tensor, |
| 1519 | indexer_weight: torch.Tensor, |
| 1520 | ) -> tuple[ |
| 1521 | torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor |
| 1522 | ]: |
| 1523 | fused = self.FUSED_AR_RMS_QUANT_BF16_OP( |
| 1524 | input_=input_, |
| 1525 | residual=residual, |
| 1526 | weight=norm_weight.to(input_.dtype), |
| 1527 | epsilon=eps, |
| 1528 | group_size=gs, |
| 1529 | ) |
| 1530 | quant_out, residual_out, scale_out, bf16_norm = ( |
| 1531 | fused[0], |
| 1532 | fused[1], |
| 1533 | fused[2], |
| 1534 | fused[3], |
| 1535 | ) |
| 1536 | idx = torch.ops.vllm.rocm_unquantized_gemm(bf16_norm, indexer_weight) |
| 1537 | return quant_out, scale_out, residual_out, idx, bf16_norm |
| 1538 | |
| 1539 | return _replacement |
| 1540 | |
| 1541 | |
| 1542 | class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): |
| 1543 | def __init__(self, config: VllmConfig) -> None: |
| 1544 | super().__init__(config, "rocm_aiter_allreduce_fusion_pass") |
| 1545 | self.disabled = True |
| 1546 | self.tp_size = get_tensor_model_parallel_world_size() |
| 1547 | if self.tp_size <= 1: |
| 1548 | logger.warning_once("AllReduce fusion pass is disabled for tp_size <= 1.") |
| 1549 | return |
| 1550 | |
| 1551 | if config.model_config is None: |
| 1552 | logger.warning_once( |
| 1553 | "AllReduce fusion pass is disabled for missing model_config." |
| 1554 | ) |
| 1555 | return |
| 1556 | |
| 1557 | ca_comm = rocm_aiter_ops.get_aiter_allreduce() |
| 1558 | if ca_comm is None: |
| 1559 | logger.warning_once( |
| 1560 | "AITER allreduce fusions are disabled " |
| 1561 | "because AITER Custom All Reduce is not enabled. " |
| 1562 | "Set VLLM_ROCM_USE_AITER_CUSTOM_AR=1 " |
| 1563 | "to enable it." |
| 1564 | ) |
| 1565 | return |
| 1566 | self.ca_comm = ca_comm |
| 1567 | |
| 1568 | hidden_dim = config.model_config.get_hidden_size() |
| 1569 | element_size = torch.tensor([], dtype=self.model_dtype).element_size() |
| 1570 | max_size = ca_comm.effective_max_size() |
| 1571 | _AITER_OLD_FUSED_AR_RMS_HIDDEN = (512, 1024, 2048, 4096) |
| 1572 | if ( |
| 1573 | not ca_comm.supports_dynamic_hidden_dim |
| 1574 | and hidden_dim not in _AITER_OLD_FUSED_AR_RMS_HIDDEN |
| 1575 | ): |
| 1576 | logger.warning_once( |
| 1577 | "AITER allreduce-rmsnorm fusion disabled: aiter<0.1.12 " |
| 1578 | "only supports hidden_dim in %s; got %d. Upgrade aiter to " |
| 1579 | ">=0.1.12 to enable fusion for this model.", |
| 1580 | _AITER_OLD_FUSED_AR_RMS_HIDDEN, |
| 1581 | hidden_dim, |
| 1582 | ) |
| 1583 | return |
| 1584 | |
| 1585 | max_token_num = max_size // (hidden_dim * element_size) |
| 1586 | self.max_token_num = min( |
| 1587 | max_token_num, |
| 1588 | config.scheduler_config.max_num_batched_tokens, |
| 1589 | ) |
| 1590 | |
| 1591 | # Only register the AR+RMS+per-group-FP8-quant patterns when the |
| 1592 | # running aiter exposes the kernel. Older aiter builds (pre PR #2823) |
| 1593 | # fall back to the AR+RMS-only fusion paired with PR #41825's |
| 1594 | # standalone RMS+quant fusion -- still correct, just leaves the |
| 1595 | # post-AR quant as a standalone kernel. |
| 1596 | supports_per_group_quant = ca_comm.supports_per_group_quant |
| 1597 | if not supports_per_group_quant: |
| 1598 | logger.warning_once( |
| 1599 | "AITER AR+RMS+per-group-FP8-quant fusion disabled: aiter " |
| 1600 | "build is missing 'fused_ar_rms_per_group_quant'. Upgrade " |
| 1601 | "aiter past PR #2823 to enable the trailing per-group " |
| 1602 | "FP8 quant fusion." |
| 1603 | ) |
| 1604 | |
| 1605 | for epsilon in [1e-5, 1e-6]: |
| 1606 | # Quant-fused variants must register first so the pattern matcher |
| 1607 | # tries them before the AR+RMS-only variants. Otherwise the |
| 1608 | # AR+RMS-only fusion runs first and consumes the all_reduce node, |
| 1609 | # leaving the trailing quant op stranded as an unfused kernel. |
| 1610 | # Register larger subgraphs first (DeepSeek indexer fan-out, then |
| 1611 | # quant-only AR+RMS+quant, then AR+RMS-only). |
| 1612 | if supports_per_group_quant: |
| 1613 | self.register( |
| 1614 | AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( |
| 1615 | epsilon, |
| 1616 | self.model_dtype, |
| 1617 | self.device, |
| 1618 | ) |
| 1619 | ) |
| 1620 | self.register( |
| 1621 | AiterAllreduceFusedRMSNormGroupQuantFP8Pattern( |
| 1622 | epsilon, |
| 1623 | self.model_dtype, |
| 1624 | self.device, |
| 1625 | ) |
| 1626 | ) |
| 1627 | self.register( |
| 1628 | AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern( |
| 1629 | epsilon, |
| 1630 | self.model_dtype, |
| 1631 | self.device, |
| 1632 | ) |
| 1633 | ) |
| 1634 | |
| 1635 | self.register( |
| 1636 | AiterAllreduceFusedRMSNormPattern( |
| 1637 | epsilon, |
| 1638 | self.model_dtype, |
| 1639 | self.device, |
| 1640 | ) |
| 1641 | ) |
| 1642 | self.register( |
| 1643 | AiterAllreduceFusedAddRMSNormPattern( |
| 1644 | epsilon, |
| 1645 | self.model_dtype, |
| 1646 | self.device, |
| 1647 | ) |
| 1648 | ) |
| 1649 | self.register( |
| 1650 | AiterAllreduceFusedAddRMSNormOutputOnlyPattern( |
| 1651 | epsilon, |
| 1652 | self.model_dtype, |
| 1653 | self.device, |
| 1654 | ) |
| 1655 | ) |
| 1656 | |
| 1657 | # WARNING: This is a hack to clear the pattern matcher cache |
| 1658 | # and allow multiple values of epsilon. |
| 1659 | torch._inductor.pattern_matcher._seen_patterns.clear() |
| 1660 | |
| 1661 | self.disabled = False |
| 1662 | |
| 1663 | self.dump_patterns(config, self.pm_pass) |
| 1664 | |
| 1665 | def is_applicable_for_range(self, compile_range: Range) -> bool: |
| 1666 | if self.disabled: |
| 1667 | logger.warning_once("AllReduce fusion pass is disabled.") |
| 1668 | return False |
| 1669 | return bool(compile_range.end <= self.max_token_num) |
| 1670 | |
| 1671 | @VllmInductorPass.time_and_log |
| 1672 | def __call__(self, graph: fx.Graph) -> None: |
| 1673 | self.matched_count = self.pm_pass.apply(graph) |
| 1674 | VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count |
| 1675 | logger.debug( |
| 1676 | "%s Replaced %s patterns", self.__class__.__name__, self.matched_count |
| 1677 | ) |