| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | from typing import Any |
| 5 | |
| 6 | import torch |
| 7 | import torch._inductor.pattern_matcher as pm |
| 8 | from torch import fx |
| 9 | from torch._inductor.fx_passes.post_grad import view_to_reshape |
| 10 | from torch._inductor.pattern_matcher import PatternMatcherPass |
| 11 | |
| 12 | import vllm.ir.ops |
| 13 | import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401 |
| 14 | from vllm._aiter_ops import rocm_aiter_ops |
| 15 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 18 | GroupShape, |
| 19 | QuantKey, |
| 20 | ScaleDesc, |
| 21 | kFp8Dynamic128Sym, |
| 22 | ) |
| 23 | from vllm.platforms import current_platform |
| 24 | |
| 25 | from ..inductor_pass import enable_fake_mode |
| 26 | from ..vllm_inductor_pass import ( |
| 27 | VllmFusionPatternMatcherPass, |
| 28 | VllmInductorPass, |
| 29 | VllmPatternMatcherPass, |
| 30 | VllmPatternReplacement, |
| 31 | _fx_view_to_reshape, |
| 32 | fold_consecutive_reshapes, |
| 33 | remove_noop_reshapes, |
| 34 | ) |
| 35 | from .matcher_utils import ( |
| 36 | MatcherQuantFP8, |
| 37 | MatcherRMSNormGated, |
| 38 | MatcherSiluAndMul, |
| 39 | ) |
| 40 | from .rms_quant_fusion import ( |
| 41 | FusedRMSQuantKey, |
| 42 | ) |
| 43 | |
| 44 | logger = init_logger(__name__) |
| 45 | FP8_DTYPE = current_platform.fp8_dtype() |
| 46 | |
| 47 | |
| 48 | class AiterRMSNormQuantPattern: |
| 49 | def __init__( |
| 50 | self, epsilon: float, key: FusedRMSQuantKey, match_aiter_quant: bool = True |
| 51 | ): |
| 52 | self.epsilon = epsilon |
| 53 | self.quant_dtype = key.quant.dtype |
| 54 | self.device = torch.device("cuda") |
| 55 | |
| 56 | self.quant_matcher = MatcherQuantFP8( |
| 57 | key.quant, |
| 58 | match_rocm_aiter=match_aiter_quant, |
| 59 | ) |
| 60 | |
| 61 | def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: |
| 62 | return torch.empty(*args, dtype=torch.bfloat16, device=self.device, **kwargs) |
| 63 | |
| 64 | def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: |
| 65 | return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs) |
| 66 | |
| 67 | |
| 68 | class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern): |
| 69 | """AITER RMSNorm + Dynamic Quantization pattern.""" |
| 70 | |
| 71 | FUSED_OP = rocm_aiter_ops.get_rmsnorm_fused_dynamic_quant_op() |
| 72 | |
| 73 | def __init__( |
| 74 | self, |
| 75 | epsilon: float, |
| 76 | quant_dtype: torch.dtype, |
| 77 | match_aiter_quant: bool = True, |
| 78 | group_shape: GroupShape = GroupShape.PER_TOKEN, |
| 79 | symmetric: bool = True, |
| 80 | ) -> None: |
| 81 | scale = ScaleDesc(torch.float32, False, group_shape) |
| 82 | key = FusedRMSQuantKey( |
| 83 | fused_add=False, |
| 84 | quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric), |
| 85 | ) |
| 86 | |
| 87 | super().__init__(epsilon, key, match_aiter_quant) |
| 88 | |
| 89 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 90 | def pattern( |
| 91 | input: torch.Tensor, |
| 92 | weight: torch.Tensor, |
| 93 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 94 | result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon) |
| 95 | result, scale = self.quant_matcher(result_rms) |
| 96 | return result, scale |
| 97 | |
| 98 | def replacement( |
| 99 | input: torch.Tensor, |
| 100 | weight: torch.Tensor, |
| 101 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 102 | result = self.FUSED_OP( |
| 103 | x=input, |
| 104 | weight=weight, |
| 105 | epsilon=self.epsilon, |
| 106 | quant_dtype=self.quant_dtype, |
| 107 | ) |
| 108 | |
| 109 | return result[0], result[1] |
| 110 | |
| 111 | pm.register_replacement( |
| 112 | pattern, |
| 113 | replacement, |
| 114 | # input, weight |
| 115 | [self.empty(5, 16), self.empty(16)], |
| 116 | pm.fwd_only, |
| 117 | pm_pass, |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern): |
| 122 | """AITER RMSNorm Fused Add + Dynamic Quantization pattern.""" |
| 123 | |
| 124 | FUSED_OP = rocm_aiter_ops.get_rmsnorm_fused_add_dynamic_quant_op() |
| 125 | |
| 126 | def __init__( |
| 127 | self, |
| 128 | epsilon: float, |
| 129 | quant_dtype: torch.dtype, |
| 130 | match_aiter_quant: bool = True, |
| 131 | group_shape: GroupShape = GroupShape.PER_TOKEN, |
| 132 | symmetric: bool = True, |
| 133 | ) -> None: |
| 134 | scale = ScaleDesc(torch.float32, False, group_shape) |
| 135 | key = FusedRMSQuantKey( |
| 136 | fused_add=True, |
| 137 | quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric), |
| 138 | ) |
| 139 | |
| 140 | super().__init__(epsilon, key, match_aiter_quant) |
| 141 | |
| 142 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 143 | def pattern( |
| 144 | input: torch.Tensor, |
| 145 | weight: torch.Tensor, |
| 146 | residual: torch.Tensor, |
| 147 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 148 | result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( |
| 149 | input, residual, weight, self.epsilon |
| 150 | ) |
| 151 | result, scale = self.quant_matcher(result_rms) |
| 152 | |
| 153 | return result, residual_out, scale |
| 154 | |
| 155 | def replacement( |
| 156 | input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor |
| 157 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 158 | result = self.FUSED_OP( |
| 159 | x=input, |
| 160 | residual=residual, |
| 161 | weight=weight, |
| 162 | epsilon=self.epsilon, |
| 163 | quant_dtype=self.quant_dtype, |
| 164 | ) |
| 165 | |
| 166 | return result[0], result[1], result[2] |
| 167 | |
| 168 | inputs = [ |
| 169 | self.empty(5, 16), # input |
| 170 | self.empty(16), # weight |
| 171 | self.empty(5, 16), # residual |
| 172 | ] |
| 173 | |
| 174 | pm.register_replacement( |
| 175 | pattern, |
| 176 | replacement, |
| 177 | inputs, |
| 178 | pm.fwd_only, |
| 179 | pm_pass, |
| 180 | ) |
| 181 | |
| 182 | |
| 183 | class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern): |
| 184 | """This pattern fuses aiter rms_norm & group fp8 quant custom |
| 185 | ops into an aiter rms_norm_group_fp8_quant op. |
| 186 | """ |
| 187 | |
| 188 | FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op() |
| 189 | |
| 190 | def __init__( |
| 191 | self, |
| 192 | epsilon: float, |
| 193 | quant_dtype: torch.dtype, |
| 194 | group_shape: GroupShape, |
| 195 | match_aiter_quant: bool = True, |
| 196 | symmetric: bool = True, |
| 197 | ) -> None: |
| 198 | scale = ScaleDesc(torch.float32, False, group_shape) |
| 199 | key = FusedRMSQuantKey( |
| 200 | fused_add=False, |
| 201 | quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric), |
| 202 | ) |
| 203 | |
| 204 | super().__init__(epsilon, key, match_aiter_quant) |
| 205 | |
| 206 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 207 | def pattern( |
| 208 | input: torch.Tensor, |
| 209 | weight: torch.Tensor, |
| 210 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 211 | result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon) |
| 212 | result, scale = self.quant_matcher(result_rms) |
| 213 | return result, scale |
| 214 | |
| 215 | def replacement( |
| 216 | input: torch.Tensor, |
| 217 | weight: torch.Tensor, |
| 218 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 219 | at = self.FUSED_OP( |
| 220 | x=input, |
| 221 | weight=weight, |
| 222 | variance_epsilon=self.epsilon, |
| 223 | group_size=128, |
| 224 | ) |
| 225 | |
| 226 | return at[0], at[1] |
| 227 | |
| 228 | pm.register_replacement( |
| 229 | pattern, |
| 230 | replacement, |
| 231 | # input, weight |
| 232 | [self.empty(5, 16), self.empty(16)], |
| 233 | pm.fwd_only, |
| 234 | pm_pass, |
| 235 | ) |
| 236 | |
| 237 | |
| 238 | class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern): |
| 239 | """This pattern fuses aiter rms_norm_with_add & group fp8 quant custom ops |
| 240 | into a aiter rms_norm_with_add_group_fp8_quant op. |
| 241 | """ |
| 242 | |
| 243 | FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_add_fused_quant_op() |
| 244 | |
| 245 | def __init__( |
| 246 | self, |
| 247 | epsilon: float, |
| 248 | quant_dtype: torch.dtype, |
| 249 | group_shape: GroupShape, |
| 250 | match_aiter_quant: bool = True, |
| 251 | symmetric: bool = True, |
| 252 | ) -> None: |
| 253 | scale = ScaleDesc(torch.float32, False, group_shape) |
| 254 | key = FusedRMSQuantKey( |
| 255 | fused_add=True, |
| 256 | quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric), |
| 257 | ) |
| 258 | |
| 259 | super().__init__(epsilon, key, match_aiter_quant) |
| 260 | |
| 261 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 262 | def pattern( |
| 263 | input: torch.Tensor, |
| 264 | weight: torch.Tensor, |
| 265 | residual: torch.Tensor, |
| 266 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 267 | result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( |
| 268 | input, residual, weight, self.epsilon |
| 269 | ) |
| 270 | result, scale = self.quant_matcher(result_rms) |
| 271 | |
| 272 | return result, residual_out, scale |
| 273 | |
| 274 | def replacement( |
| 275 | input: torch.Tensor, |
| 276 | weight: torch.Tensor, |
| 277 | residual: torch.Tensor, |
| 278 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 279 | at = self.FUSED_OP( |
| 280 | x=input, |
| 281 | residual=residual, |
| 282 | weight=weight, |
| 283 | variance_epsilon=self.epsilon, |
| 284 | group_size=128, |
| 285 | ) |
| 286 | |
| 287 | # result, scale, residual |
| 288 | return at[0], at[1], at[2] |
| 289 | |
| 290 | inputs = [ |
| 291 | self.empty(5, 16), # input |
| 292 | self.empty(16), # weight |
| 293 | self.empty(5, 16), # residual |
| 294 | ] |
| 295 | |
| 296 | pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass) |
| 297 | |
| 298 | |
| 299 | class DoubleAiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern): |
| 300 | """Pattern matching ``rms_norm`` whose output feeds *two* distinct |
| 301 | ``rocm_aiter_group_fp8_quant`` consumers, replacing it with two |
| 302 | independent fused ``rms_norm_group_fp8_quant`` ops. |
| 303 | |
| 304 | Repeating the rms_norm in the replacement is preferable to leaving |
| 305 | the fused 16-bit rms output materialized for two unfused quant |
| 306 | consumers, and matches what the previous manual graph surgery |
| 307 | achieved by cloning the rms_norm node. |
| 308 | """ |
| 309 | |
| 310 | FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op() |
| 311 | |
| 312 | def __init__( |
| 313 | self, |
| 314 | epsilon: float, |
| 315 | quant_dtype: torch.dtype, |
| 316 | group_shape: GroupShape, |
| 317 | match_aiter_quant: bool = True, |
| 318 | symmetric: bool = True, |
| 319 | ) -> None: |
| 320 | scale = ScaleDesc(torch.float32, False, group_shape) |
| 321 | key = FusedRMSQuantKey( |
| 322 | fused_add=False, |
| 323 | quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric), |
| 324 | ) |
| 325 | |
| 326 | super().__init__(epsilon, key, match_aiter_quant) |
| 327 | |
| 328 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 329 | def pattern( |
| 330 | input: torch.Tensor, |
| 331 | weight: torch.Tensor, |
| 332 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 333 | result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon) |
| 334 | result1, scale1 = self.quant_matcher(result_rms) |
| 335 | result2, scale2 = self.quant_matcher(result_rms) |
| 336 | return result1, scale1, result2, scale2 |
| 337 | |
| 338 | def replacement( |
| 339 | input: torch.Tensor, |
| 340 | weight: torch.Tensor, |
| 341 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 342 | at1 = self.FUSED_OP( |
| 343 | x=input, |
| 344 | weight=weight, |
| 345 | variance_epsilon=self.epsilon, |
| 346 | group_size=128, |
| 347 | ) |
| 348 | at2 = self.FUSED_OP( |
| 349 | x=input, |
| 350 | weight=weight, |
| 351 | variance_epsilon=self.epsilon, |
| 352 | group_size=128, |
| 353 | ) |
| 354 | |
| 355 | return at1[0], at1[1], at2[0], at2[1] |
| 356 | |
| 357 | pm.register_replacement( |
| 358 | pattern, |
| 359 | replacement, |
| 360 | # input, weight |
| 361 | [self.empty(5, 16), self.empty(16)], |
| 362 | pm.fwd_only, |
| 363 | pm_pass, |
| 364 | ) |
| 365 | |
| 366 | |
| 367 | class DoubleAiterRMSFp8GroupQuantViewPattern(AiterRMSNormQuantPattern): |
| 368 | """View-tolerant variant of ``DoubleAiterRMSFp8GroupQuantPattern``. |
| 369 | |
| 370 | Matches the same 1-to-2 fan-out, but with a ``view``/``reshape`` between |
| 371 | the ``rms_norm`` output and the two ``rocm_aiter_group_fp8_quant`` |
| 372 | consumers:: |
| 373 | |
| 374 | rms_norm -> view -> rocm_aiter_group_fp8_quant |
| 375 | \\-> view -> rocm_aiter_group_fp8_quant |
| 376 | |
| 377 | This shape arises in DeepSeek-V3.2's MLA indexer q_c norm, where the |
| 378 | FP8 linear path's 2D-flatten boilerplate |
| 379 | (``Fp8BlockScaledMMLinearKernel.apply_weights``) inserts a view between |
| 380 | the rms_norm output and each FP8 group quant op. The non-view sibling |
| 381 | pattern silently no-ops on this graph because the pattern matcher |
| 382 | requires the in-graph and in-pattern node shapes to align. |
| 383 | |
| 384 | The trace_fn runs Inductor's ``view_to_reshape`` post-grad pass to |
| 385 | normalize ``view`` to ``reshape`` in both the pattern and the input |
| 386 | graph, widening the match without touching the no-view sibling. |
| 387 | """ |
| 388 | |
| 389 | FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op() |
| 390 | |
| 391 | def __init__( |
| 392 | self, |
| 393 | epsilon: float, |
| 394 | quant_dtype: torch.dtype, |
| 395 | group_shape: GroupShape, |
| 396 | match_aiter_quant: bool = True, |
| 397 | symmetric: bool = True, |
| 398 | ) -> None: |
| 399 | scale = ScaleDesc(torch.float32, False, group_shape) |
| 400 | key = FusedRMSQuantKey( |
| 401 | fused_add=False, |
| 402 | quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric), |
| 403 | ) |
| 404 | |
| 405 | super().__init__(epsilon, key, match_aiter_quant) |
| 406 | |
| 407 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 408 | def pattern( |
| 409 | input: torch.Tensor, |
| 410 | weight: torch.Tensor, |
| 411 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 412 | result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon) |
| 413 | view_rms = result_rms.view(-1, result_rms.shape[-1]) |
| 414 | result1, scale1 = self.quant_matcher(view_rms) |
| 415 | result2, scale2 = self.quant_matcher(view_rms) |
| 416 | return result1, scale1, result2, scale2 |
| 417 | |
| 418 | def replacement( |
| 419 | input: torch.Tensor, |
| 420 | weight: torch.Tensor, |
| 421 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 422 | at1 = self.FUSED_OP( |
| 423 | x=input, |
| 424 | weight=weight, |
| 425 | variance_epsilon=self.epsilon, |
| 426 | group_size=128, |
| 427 | ) |
| 428 | at2 = self.FUSED_OP( |
| 429 | x=input, |
| 430 | weight=weight, |
| 431 | variance_epsilon=self.epsilon, |
| 432 | group_size=128, |
| 433 | ) |
| 434 | |
| 435 | return at1[0], at1[1], at2[0], at2[1] |
| 436 | |
| 437 | def trace_with_view_to_reshape(*args: Any, **kwargs: Any) -> fx.GraphModule: |
| 438 | gm = pm.fwd_only(*args, **kwargs) |
| 439 | view_to_reshape(gm) |
| 440 | return gm |
| 441 | |
| 442 | pm.register_replacement( |
| 443 | pattern, |
| 444 | replacement, |
| 445 | # input, weight |
| 446 | [self.empty(5, 16), self.empty(16)], |
| 447 | trace_with_view_to_reshape, |
| 448 | pm_pass, |
| 449 | ) |
| 450 | |
| 451 | |
| 452 | class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern): |
| 453 | """Matches decomposed RMSNormGated + reshape + group FP8 quant and replaces |
| 454 | with rocm_aiter_fused_rms_gated_fp8_group_quant. |
| 455 | |
| 456 | The norm operates per-head, on either (N*H, D) or (N, H, D). The compiler |
| 457 | folds the reshape chain so after norm the result goes through |
| 458 | reshape->merge->quant. The pattern reshapes to (N, H*D) before calling |
| 459 | MatcherQuantFP8 so that _quantize_group_native sees the full hidden dim |
| 460 | and computes the correct num_groups. |
| 461 | """ |
| 462 | |
| 463 | FUSED_OP = rocm_aiter_ops.get_fused_rms_gated_fp8_group_quant_op() |
| 464 | |
| 465 | def __init__( |
| 466 | self, |
| 467 | epsilon: float, |
| 468 | quant_dtype: torch.dtype, |
| 469 | group_shape: GroupShape, |
| 470 | num_heads: int, |
| 471 | head_dim: int, |
| 472 | match_aiter_quant: bool = True, |
| 473 | symmetric: bool = True, |
| 474 | ) -> None: |
| 475 | scale = ScaleDesc(torch.float32, False, group_shape) |
| 476 | key = FusedRMSQuantKey( |
| 477 | fused_add=False, |
| 478 | quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric), |
| 479 | ) |
| 480 | super().__init__(epsilon, key, match_aiter_quant) |
| 481 | self.rmsnorm_gated_matcher = MatcherRMSNormGated(epsilon) |
| 482 | self.num_heads = num_heads |
| 483 | self.head_dim = head_dim |
| 484 | |
| 485 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 486 | self._register_layout(pm_pass, flatten_heads=True) |
| 487 | # Both layouts reach the norm. With an opaque quant op they trace to |
| 488 | # the same pattern; with a native quant they differ, because its group |
| 489 | # reshape back to (N, H, D) is a no-op the compiler already dropped. |
| 490 | if not self.quant_matcher.enabled: |
| 491 | self._register_layout(pm_pass, flatten_heads=False) |
| 492 | |
| 493 | def _register_layout( |
| 494 | self, pm_pass: PatternMatcherPass, *, flatten_heads: bool |
| 495 | ) -> None: |
| 496 | num_heads = self.num_heads |
| 497 | head_dim = self.head_dim |
| 498 | hidden_dim = num_heads * head_dim |
| 499 | quant_matcher = self.quant_matcher |
| 500 | |
| 501 | def pattern( |
| 502 | x: torch.Tensor, |
| 503 | z: torch.Tensor, |
| 504 | weight: torch.Tensor, |
| 505 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 506 | normed = self.rmsnorm_gated_matcher(x, z, weight) |
| 507 | merged = normed.reshape(-1, hidden_dim) |
| 508 | quant_out, scales_out = quant_matcher(merged) |
| 509 | return quant_out, scales_out |
| 510 | |
| 511 | def replacement( |
| 512 | x: torch.Tensor, |
| 513 | z: torch.Tensor, |
| 514 | weight: torch.Tensor, |
| 515 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 516 | # The pattern matches on op structure rather than rank, so either |
| 517 | # per-head layout can land here. The fused op takes (M, D), and |
| 518 | # both operands are contiguous, so collapsing the heads is a view. |
| 519 | x = x.reshape(-1, head_dim) |
| 520 | z = z.reshape(-1, head_dim) |
| 521 | fused = self.FUSED_OP( |
| 522 | x=x, |
| 523 | weight=weight, |
| 524 | bias=None, |
| 525 | z=z, |
| 526 | eps=self.epsilon, |
| 527 | norm_before_gate=True, |
| 528 | activation="silu", |
| 529 | group_size=head_dim, |
| 530 | ) |
| 531 | fp8_out = fused[0] |
| 532 | scales_out = fused[1] |
| 533 | fp8_reshaped = fp8_out.reshape(-1, hidden_dim) |
| 534 | scales_reshaped = scales_out.reshape(-1, num_heads) |
| 535 | return fp8_reshaped, scales_reshaped |
| 536 | |
| 537 | n_tokens = 2 |
| 538 | shape = ( |
| 539 | (n_tokens * num_heads, head_dim) |
| 540 | if flatten_heads |
| 541 | else (n_tokens, num_heads, head_dim) |
| 542 | ) |
| 543 | x = self.empty(*shape) |
| 544 | z = self.empty(*shape) |
| 545 | w = self.empty(head_dim) |
| 546 | |
| 547 | def trace_fn(*args, **kwargs): |
| 548 | gm = pm.fwd_only(*args, **kwargs) |
| 549 | _fx_view_to_reshape(gm) |
| 550 | fold_consecutive_reshapes(gm) |
| 551 | if not flatten_heads: |
| 552 | remove_noop_reshapes(gm) |
| 553 | return gm |
| 554 | |
| 555 | pm.register_replacement( |
| 556 | pattern, |
| 557 | replacement, |
| 558 | [x, z, w], |
| 559 | trace_fn, |
| 560 | pm_pass, |
| 561 | ) |
| 562 | |
| 563 | |
| 564 | class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): |
| 565 | """This pass fuses aiter rms_norm & vllm/aiter quant custom ops |
| 566 | into a fused rms_norm_quant op. |
| 567 | It also supports fused_add_rms_norm. |
| 568 | """ |
| 569 | |
| 570 | @enable_fake_mode |
| 571 | def __init__(self, config: VllmConfig) -> None: |
| 572 | super().__init__(config) |
| 573 | |
| 574 | self.patterns: PatternMatcherPass = PatternMatcherPass( |
| 575 | pass_name="rocm_aiter_rms_norm_quant_fusion_pass" |
| 576 | ) |
| 577 | |
| 578 | # Discover (num_heads, head_dim) pairs for gated RMSNorm patterns |
| 579 | # from GatedDeltaNetAttention layers in static_forward_context. |
| 580 | from vllm.model_executor.layers.mamba.gdn.base import ( |
| 581 | GatedDeltaNetAttention, |
| 582 | ) |
| 583 | |
| 584 | gdn_layers = get_layers_from_vllm_config( |
| 585 | config, |
| 586 | GatedDeltaNetAttention, # type: ignore[type-abstract] |
| 587 | ) |
| 588 | gated_norm_shapes: set[tuple[int, int]] = set() |
| 589 | for layer in gdn_layers.values(): |
| 590 | num_v_heads = getattr(layer, "num_v_heads", None) or getattr( |
| 591 | layer, "num_heads", None |
| 592 | ) |
| 593 | head_v_dim = getattr(layer, "head_v_dim", None) or getattr( |
| 594 | layer, "head_dim", None |
| 595 | ) |
| 596 | |
| 597 | assert num_v_heads is not None and head_v_dim is not None |
| 598 | |
| 599 | gated_norm_shapes.add((num_v_heads // layer.tp_size, head_v_dim)) |
| 600 | |
| 601 | # RDNA4 uses native quant ops and supports only Triton replacements. |
| 602 | match_aiter_quant_op = not rocm_aiter_ops.is_rdna_aiter_enabled() |
| 603 | |
| 604 | # Make sure fused add patterns are before simple rms norm, |
| 605 | # as the latter is a subset of the former in torch ops. |
| 606 | # The DoubleQuant patterns handle 1 rms_norm -> 2 group_fp8_quant |
| 607 | # fan-out (e.g. DSv3.2) and must be registered before the single |
| 608 | # group-quant pattern so they match first. The view-tolerant variant |
| 609 | # additionally covers the rms_norm -> view -> 2x quant shape that |
| 610 | # appears when the FP8 linear path inserts a 2D-flatten boilerplate |
| 611 | # (DSv3.2 MLA indexer q_c norm). |
| 612 | for epsilon in [1e-5, 1e-6]: |
| 613 | # Fuse aiter rms_norm + 2x aiter group fp8 quant |
| 614 | DoubleAiterRMSFp8GroupQuantPattern( |
| 615 | epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op |
| 616 | ).register(self.patterns) |
| 617 | |
| 618 | # View-tolerant sibling for DSv3.2 q_c norm fan-out |
| 619 | DoubleAiterRMSFp8GroupQuantViewPattern( |
| 620 | epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op |
| 621 | ).register(self.patterns) |
| 622 | |
| 623 | # Fuse aiter rms_norm + aiter dynamic group fp8 quant |
| 624 | AiterRMSFp8GroupQuantPattern( |
| 625 | epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op |
| 626 | ).register(self.patterns) |
| 627 | |
| 628 | # Fuse aiter fused_add_rms_norm + aiter dynamic group fp8 quant |
| 629 | AiterFusedAddRMSFp8GroupQuantPattern( |
| 630 | epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op |
| 631 | ).register(self.patterns) |
| 632 | |
| 633 | # When quant_fp8 custom ops are disabled, both AITER and native |
| 634 | # quant matchers trace through QuantFP8's native implementation. |
| 635 | # Registering both variants would create duplicate Inductor |
| 636 | # patterns. |
| 637 | is_quant_fp8_enabled = config.compilation_config.is_custom_op_enabled( |
| 638 | "quant_fp8" |
| 639 | ) |
| 640 | match_aiter_quant_options = ( |
| 641 | [True, False] if is_quant_fp8_enabled else [False] |
| 642 | ) |
| 643 | if not match_aiter_quant_op: |
| 644 | match_aiter_quant_options = [] |
| 645 | |
| 646 | for match_aiter_quant in match_aiter_quant_options: |
| 647 | # Fuse aiter rms_norm + (aiter / vllm built-in) |
| 648 | # dynamic per-token fp8 quant |
| 649 | AiterRMSNormDynamicQuantPattern( |
| 650 | epsilon, FP8_DTYPE, match_aiter_quant=match_aiter_quant |
| 651 | ).register(self.patterns) |
| 652 | |
| 653 | # Fuse aiter fused_add_rms_norm + (aiter / vllm built-in) |
| 654 | # dynamic per-token fp8 quant |
| 655 | AiterFusedAddRMSNormDynamicQuantPattern( |
| 656 | epsilon, FP8_DTYPE, match_aiter_quant=match_aiter_quant |
| 657 | ).register(self.patterns) |
| 658 | |
| 659 | # Fuse decomposed RMSNormGated + group fp8 quant. |
| 660 | # The replacement op (fused_rms_gated_fp8_group_quant) requires |
| 661 | # an aiter version that includes the GDN triton kernel renames. |
| 662 | if gated_norm_shapes and ( |
| 663 | rocm_aiter_ops.are_gdn_triton_kernels_available() |
| 664 | or rocm_aiter_ops.is_rdna_gdn_triton_kernels_available() |
| 665 | ): |
| 666 | for num_heads, head_dim in gated_norm_shapes: |
| 667 | if head_dim != 128: |
| 668 | continue |
| 669 | AiterRMSNormGatedFp8GroupQuantPattern( |
| 670 | epsilon, |
| 671 | FP8_DTYPE, |
| 672 | GroupShape(1, 128), |
| 673 | num_heads=num_heads, |
| 674 | head_dim=head_dim, |
| 675 | match_aiter_quant=match_aiter_quant_op, |
| 676 | ).register(self.patterns) |
| 677 | |
| 678 | self.dump_patterns(config, self.patterns) |
| 679 | |
| 680 | @VllmInductorPass.time_and_log |
| 681 | def __call__(self, graph: fx.Graph) -> None: |
| 682 | self.matched_count = self.patterns.apply(graph) |
| 683 | logger.debug( |
| 684 | "%s Replaced %s patterns", self.__class__.__name__, self.matched_count |
| 685 | ) |
| 686 | |
| 687 | def uuid(self) -> str: |
| 688 | fusion_patterns = [ |
| 689 | AiterRMSNormDynamicQuantPattern, |
| 690 | AiterFusedAddRMSNormDynamicQuantPattern, |
| 691 | AiterRMSFp8GroupQuantPattern, |
| 692 | AiterFusedAddRMSFp8GroupQuantPattern, |
| 693 | DoubleAiterRMSFp8GroupQuantPattern, |
| 694 | DoubleAiterRMSFp8GroupQuantViewPattern, |
| 695 | AiterRMSNormGatedFp8GroupQuantPattern, |
| 696 | ] |
| 697 | return self.hash_source(self, *fusion_patterns) |
| 698 | |
| 699 | |
| 700 | class AiterSiluMulFp8GroupQuantPattern(VllmPatternReplacement): |
| 701 | """This pattern fuses aiter silu_and_mul & group fp8 quant custom |
| 702 | ops into an aiter silu_and_mul_group_fp8_quant op. |
| 703 | """ |
| 704 | |
| 705 | FUSED_SILU_MUL_QUANT_OP = rocm_aiter_ops.get_act_mul_fused_fp8_group_quant_op() |
| 706 | |
| 707 | def __init__(self, match_aiter_quant_op: bool = True) -> None: |
| 708 | self.silu_and_mul_matcher = MatcherSiluAndMul() |
| 709 | self.quant_matcher = MatcherQuantFP8( |
| 710 | quant_key=kFp8Dynamic128Sym, match_rocm_aiter=match_aiter_quant_op |
| 711 | ) |
| 712 | |
| 713 | def get_inputs(self) -> list[torch.Tensor]: |
| 714 | return [ |
| 715 | self.silu_and_mul_matcher.inputs()[0], |
| 716 | ] |
| 717 | |
| 718 | @property |
| 719 | def pattern(self): |
| 720 | def _pattern( |
| 721 | input: torch.Tensor, |
| 722 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 723 | at1 = self.silu_and_mul_matcher(input) |
| 724 | at2 = self.quant_matcher(at1) |
| 725 | return at2[0], at2[1] |
| 726 | |
| 727 | return _pattern |
| 728 | |
| 729 | @property |
| 730 | def replacement(self): |
| 731 | def _replacement( |
| 732 | input: torch.Tensor, |
| 733 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 734 | at = self.FUSED_SILU_MUL_QUANT_OP(x=input, group_size=128) |
| 735 | return at[0], at[1] |
| 736 | |
| 737 | return _replacement |
| 738 | |
| 739 | |
| 740 | class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmFusionPatternMatcherPass): |
| 741 | """This pass fuses a pre-defined set of custom ops into fused ops. |
| 742 | It uses the torch pattern matcher to find the patterns and replace them. |
| 743 | |
| 744 | Because patterns can only be registered once, the pass is a singleton. |
| 745 | This will be addressed in a future version of PyTorch: |
| 746 | https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980 |
| 747 | """ |
| 748 | |
| 749 | def __init__(self, config: VllmConfig) -> None: |
| 750 | super().__init__(config, "rocm_aiter_silu_mul_fp8_group_quant_fusion_pass") |
| 751 | |
| 752 | match_aiter_quant_op = not rocm_aiter_ops.is_rdna_aiter_enabled() |
| 753 | |
| 754 | self.register( |
| 755 | AiterSiluMulFp8GroupQuantPattern(match_aiter_quant_op=match_aiter_quant_op) |
| 756 | ) |
| 757 | |
| 758 | self.dump_patterns(config, self.pm_pass) |
| 759 | |
| 760 | |
| 761 | class AddAiterRMSNormPadPattern: |
| 762 | """This pattern replaces an aiter_rmsnorm_with_add & a pad op |
| 763 | with a custom triton_add_rmsnorm_pad op from AITER. |
| 764 | """ |
| 765 | |
| 766 | AITER_TRITON_ADD_RMSNORM_PAD_OP = rocm_aiter_ops.get_triton_add_rmsnorm_pad_op() |
| 767 | |
| 768 | def __init__( |
| 769 | self, |
| 770 | epsilon: float, |
| 771 | hidden_size: int, |
| 772 | x_pad_to_multiple: int, |
| 773 | ): |
| 774 | self.epsilon = epsilon |
| 775 | self.hidden_size = hidden_size |
| 776 | self.x_pad_to_multiple = x_pad_to_multiple |
| 777 | |
| 778 | def get_inputs(self) -> list[torch.Tensor]: |
| 779 | device = torch.device("cuda") |
| 780 | dtype = torch.bfloat16 |
| 781 | input = torch.empty(5, 16, dtype=dtype, device=device) |
| 782 | weight = torch.empty(16, dtype=dtype, device=device) |
| 783 | residual = torch.empty(5, 16, dtype=dtype, device=device) |
| 784 | router_weight = torch.empty([8, 16], dtype=dtype, device=device) |
| 785 | router_bias = torch.empty([8], dtype=dtype, device=device) |
| 786 | return [input, weight, residual, router_weight, router_bias] |
| 787 | |
| 788 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 789 | def pattern( |
| 790 | input: torch.Tensor, |
| 791 | weight: torch.Tensor, |
| 792 | residual: torch.Tensor, |
| 793 | router_weight: torch.Tensor, |
| 794 | router_bias: torch.Tensor, |
| 795 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 796 | pad_size = self.x_pad_to_multiple - ( |
| 797 | self.hidden_size % self.x_pad_to_multiple |
| 798 | ) |
| 799 | result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( |
| 800 | input, residual, weight, self.epsilon |
| 801 | ) |
| 802 | router_logits = torch.ops.vllm.rocm_unquantized_gemm( |
| 803 | result_rms, router_weight, router_bias |
| 804 | ) |
| 805 | result = torch.nn.functional.pad( |
| 806 | result_rms, (0, pad_size), mode="constant", value=0.0 |
| 807 | ) |
| 808 | return result, residual_out, router_logits |
| 809 | |
| 810 | def replacement( |
| 811 | input: torch.Tensor, |
| 812 | weight: torch.Tensor, |
| 813 | residual: torch.Tensor, |
| 814 | router_weight: torch.Tensor, |
| 815 | router_bias: torch.Tensor, |
| 816 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 817 | at = self.AITER_TRITON_ADD_RMSNORM_PAD_OP( |
| 818 | x=input, |
| 819 | weight=weight, |
| 820 | variance_epsilon=self.epsilon, |
| 821 | residual=residual, |
| 822 | x_pad_to_multiple=self.x_pad_to_multiple, |
| 823 | ) |
| 824 | result_padded = at[0] |
| 825 | router_logits = torch.ops.vllm.rocm_unquantized_gemm( |
| 826 | result_padded[:, : self.hidden_size], router_weight, router_bias |
| 827 | ) |
| 828 | residual_out = at[1] |
| 829 | return result_padded, residual_out, router_logits |
| 830 | |
| 831 | pm.register_replacement( |
| 832 | pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass |
| 833 | ) |
| 834 | |
| 835 | |
| 836 | class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass): |
| 837 | """This pass replaces an AITER CK RMSNorm + residual add and a pad op |
| 838 | with an triton_add_rmsnorm_pad op from AITER. |
| 839 | """ |
| 840 | |
| 841 | def __init__(self, config: VllmConfig): |
| 842 | super().__init__(config) |
| 843 | self.patterns: PatternMatcherPass = PatternMatcherPass( |
| 844 | pass_name="rocm_aiter_triton_add_rmsnorm_pad_fusion_pass" |
| 845 | ) |
| 846 | |
| 847 | # gpt-oss has hidden size 2880 |
| 848 | # padded to a multiple of 128 on gfx942 and 256 on gfx950 respectively |
| 849 | hidden_size = 2880 |
| 850 | for epsilon in [1e-5, 1e-6]: |
| 851 | for x_pad_to_multiple in [128, 256]: |
| 852 | AddAiterRMSNormPadPattern( |
| 853 | epsilon, hidden_size, x_pad_to_multiple |
| 854 | ).register(self.patterns) |
| 855 | |
| 856 | self.dump_patterns(config, self.patterns) |
| 857 | |
| 858 | @VllmInductorPass.time_and_log |
| 859 | def __call__(self, graph: torch.fx.Graph) -> None: |
| 860 | self.matched_count = self.patterns.apply(graph) |
| 861 | logger.debug("Replaced %s patterns", self.matched_count) |
| 862 | |
| 863 | def uuid(self) -> str: |
| 864 | return VllmInductorPass.hash_source(self, AddAiterRMSNormPadPattern) |
| 865 | |
| 866 | |
| 867 | class MLADualRMSNormPattern( |
| 868 | VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]] |
| 869 | ): |
| 870 | """Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into |
| 871 | AITER's ``fused_qk_rmsnorm`` HIP kernel. |
| 872 | |
| 873 | Target FX-graph pattern (unfused, ``vllm_ir`` stage):: |
| 874 | |
| 875 | gemm -> split_with_sizes([q_dim, kv_dim]) |
| 876 | +-- q_c -> vllm_ir.rms_norm(q_c, q_w, eps) |
| 877 | +-- kv_lora -> split_with_sizes([kv_c_dim, k_pe_dim]) |
| 878 | +-- kv_c -> vllm_ir.rms_norm(kv_c, kv_w, eps) |
| 879 | +-- k_pe |
| 880 | |
| 881 | The pattern covers the connected subgraph rooted at the first |
| 882 | ``split_with_sizes`` (which produces ``q_c`` and ``kv_lora``), |
| 883 | through the two ``rms_norm`` calls, and the ``k_pe`` passthrough. |
| 884 | """ |
| 885 | |
| 886 | def __init__(self, epsilon: float) -> None: |
| 887 | self._epsilon = epsilon |
| 888 | |
| 889 | def get_inputs(self) -> list[torch.Tensor]: |
| 890 | q_dim, kv_c_dim, k_pe_dim = 8, 4, 2 |
| 891 | return [ |
| 892 | self.empty_bf16(5, q_dim + kv_c_dim + k_pe_dim), |
| 893 | self.empty_bf16(q_dim), |
| 894 | self.empty_bf16(kv_c_dim), |
| 895 | ] |
| 896 | |
| 897 | @property |
| 898 | def pattern( |
| 899 | self, |
| 900 | ) -> Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: |
| 901 | eps = self._epsilon |
| 902 | |
| 903 | def _pattern( |
| 904 | projected: torch.Tensor, |
| 905 | q_weight: torch.Tensor, |
| 906 | kv_weight: torch.Tensor, |
| 907 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 908 | q_dim = q_weight.shape[0] |
| 909 | kv_dim = projected.shape[-1] - q_dim |
| 910 | kv_c_dim = kv_weight.shape[0] |
| 911 | k_pe_dim = kv_dim - kv_c_dim |
| 912 | q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) |
| 913 | kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) |
| 914 | q_normed = vllm.ir.ops.rms_norm(q_c, q_weight, eps) |
| 915 | kv_normed = vllm.ir.ops.rms_norm(kv_c, kv_weight, eps) |
| 916 | return q_normed, kv_normed, k_pe |
| 917 | |
| 918 | return _pattern |
| 919 | |
| 920 | @property |
| 921 | def replacement( |
| 922 | self, |
| 923 | ) -> Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: |
| 924 | eps = self._epsilon |
| 925 | |
| 926 | def _replacement( |
| 927 | projected: torch.Tensor, |
| 928 | q_weight: torch.Tensor, |
| 929 | kv_weight: torch.Tensor, |
| 930 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 931 | q_dim = q_weight.shape[0] |
| 932 | kv_dim = projected.shape[-1] - q_dim |
| 933 | kv_c_dim = kv_weight.shape[0] |
| 934 | k_pe_dim = kv_dim - kv_c_dim |
| 935 | q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) |
| 936 | kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) |
| 937 | q_normed, kv_normed = torch.ops.vllm.fused_mla_dual_rms_norm( |
| 938 | q_c, |
| 939 | q_weight, |
| 940 | kv_c, |
| 941 | kv_weight, |
| 942 | eps, |
| 943 | eps, |
| 944 | ) |
| 945 | return q_normed, kv_normed, k_pe |
| 946 | |
| 947 | return _replacement |
| 948 | |
| 949 | |
| 950 | class MLADualRMSPerTokenQuantPattern( |
| 951 | VllmPatternReplacement[ |
| 952 | ..., |
| 953 | tuple[ |
| 954 | torch.Tensor, |
| 955 | torch.Tensor, |
| 956 | torch.Tensor, |
| 957 | torch.Tensor, |
| 958 | ], |
| 959 | ] |
| 960 | ): |
| 961 | """Fuse the MLA FP8 attention path -- q-latent RMSNorm + FP8 *per-token* quant |
| 962 | plus kv-latent RMSNorm -- into AITER's ``fused_qk_rmsnorm_per_token_quant``. |
| 963 | |
| 964 | With a per-token FP8 ``q_b_proj`` (Quark / ModelOpt), the earlier |
| 965 | ``RocmAiterRMSNormQuantFusionPass`` folds the q side into |
| 966 | ``rocm_aiter_rmsnorm_fused_dynamic_quant`` and leaves the kv side a plain |
| 967 | ``vllm_ir.rms_norm``. This pattern matches that asymmetric pair:: |
| 968 | |
| 969 | gemm -> split_with_sizes([q_dim, kv_dim]) |
| 970 | +-- q_c -> rocm_aiter_rmsnorm_fused_dynamic_quant -> (q_fp8, q_scale) |
| 971 | +-- kv_lora -> split_with_sizes([kv_c_dim, k_pe_dim]) |
| 972 | +-- kv_c -> vllm_ir.rms_norm -> kv_normed (bf16) |
| 973 | +-- k_pe |
| 974 | """ |
| 975 | |
| 976 | DYNAMIC_QUANT_OP = rocm_aiter_ops.get_rmsnorm_fused_dynamic_quant_op() |
| 977 | FUSED_OP = rocm_aiter_ops.get_fused_mla_dual_rms_norm_per_token_quant_op() |
| 978 | |
| 979 | def __init__(self, epsilon: float) -> None: |
| 980 | self._epsilon = epsilon |
| 981 | |
| 982 | def get_inputs(self) -> list[torch.Tensor]: |
| 983 | q_dim, kv_c_dim, k_pe_dim = 256, 128, 64 |
| 984 | return [ |
| 985 | self.empty_bf16(5, q_dim + kv_c_dim + k_pe_dim), |
| 986 | self.empty_bf16(q_dim), |
| 987 | self.empty_bf16(kv_c_dim), |
| 988 | ] |
| 989 | |
| 990 | @property |
| 991 | def pattern( |
| 992 | self, |
| 993 | ) -> Callable[ |
| 994 | ..., |
| 995 | tuple[ |
| 996 | torch.Tensor, |
| 997 | torch.Tensor, |
| 998 | torch.Tensor, |
| 999 | torch.Tensor, |
| 1000 | ], |
| 1001 | ]: |
| 1002 | eps = self._epsilon |
| 1003 | dynamic_quant_op = self.DYNAMIC_QUANT_OP |
| 1004 | |
| 1005 | def _pattern( |
| 1006 | projected: torch.Tensor, |
| 1007 | q_weight: torch.Tensor, |
| 1008 | kv_weight: torch.Tensor, |
| 1009 | ) -> tuple[ |
| 1010 | torch.Tensor, |
| 1011 | torch.Tensor, |
| 1012 | torch.Tensor, |
| 1013 | torch.Tensor, |
| 1014 | ]: |
| 1015 | q_dim = q_weight.shape[0] |
| 1016 | kv_dim = projected.shape[-1] - q_dim |
| 1017 | kv_c_dim = kv_weight.shape[0] |
| 1018 | k_pe_dim = kv_dim - kv_c_dim |
| 1019 | q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) |
| 1020 | kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) |
| 1021 | q_quant = dynamic_quant_op( |
| 1022 | x=q_c, |
| 1023 | weight=q_weight, |
| 1024 | epsilon=eps, |
| 1025 | quant_dtype=FP8_DTYPE, |
| 1026 | ) |
| 1027 | kv_normed = vllm.ir.ops.rms_norm(kv_c, kv_weight, eps) |
| 1028 | return q_quant[0], q_quant[1], kv_normed, k_pe |
| 1029 | |
| 1030 | return _pattern |
| 1031 | |
| 1032 | @property |
| 1033 | def replacement( |
| 1034 | self, |
| 1035 | ) -> Callable[ |
| 1036 | ..., |
| 1037 | tuple[ |
| 1038 | torch.Tensor, |
| 1039 | torch.Tensor, |
| 1040 | torch.Tensor, |
| 1041 | torch.Tensor, |
| 1042 | ], |
| 1043 | ]: |
| 1044 | eps = self._epsilon |
| 1045 | fused_op = self.FUSED_OP |
| 1046 | |
| 1047 | def _replacement( |
| 1048 | projected: torch.Tensor, |
| 1049 | q_weight: torch.Tensor, |
| 1050 | kv_weight: torch.Tensor, |
| 1051 | ) -> tuple[ |
| 1052 | torch.Tensor, |
| 1053 | torch.Tensor, |
| 1054 | torch.Tensor, |
| 1055 | torch.Tensor, |
| 1056 | ]: |
| 1057 | q_dim = q_weight.shape[0] |
| 1058 | kv_dim = projected.shape[-1] - q_dim |
| 1059 | kv_c_dim = kv_weight.shape[0] |
| 1060 | k_pe_dim = kv_dim - kv_c_dim |
| 1061 | q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) |
| 1062 | kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) |
| 1063 | at = fused_op( |
| 1064 | q_c, |
| 1065 | q_weight, |
| 1066 | kv_c, |
| 1067 | kv_weight, |
| 1068 | eps, |
| 1069 | eps, |
| 1070 | ) |
| 1071 | # q_fp8, q_scale, kv_normed, k_pe |
| 1072 | return at[0], at[1], at[2], k_pe |
| 1073 | |
| 1074 | return _replacement |
| 1075 | |
| 1076 | |
| 1077 | class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass): |
| 1078 | """Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in |
| 1079 | MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's |
| 1080 | ``fused_qk_rmsnorm`` HIP kernel. |
| 1081 | |
| 1082 | The FP8 attention path is also handled via |
| 1083 | :class:`MLADualRMSPerTokenQuantPattern`, which fuses the q-latent RMSNorm + |
| 1084 | FP8 per-token quant together with the kv-latent RMSNorm into |
| 1085 | ``fused_mla_dual_rms_norm_per_token_quant`` backed by aiter's |
| 1086 | ``fused_qk_rmsnorm_per_token_quant`` HIP kernel. |
| 1087 | """ |
| 1088 | |
| 1089 | def __init__(self, config: VllmConfig) -> None: |
| 1090 | super().__init__(config, "mla_dual_rms_norm_fusion_pass") |
| 1091 | |
| 1092 | for epsilon in [1e-5, 1e-6]: |
| 1093 | self.register(MLADualRMSNormPattern(epsilon)) |
| 1094 | self.register(MLADualRMSPerTokenQuantPattern(epsilon)) |