@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ad8818bb5e
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 06:30 UTC

[Misc][BE] Type coverage for vllm/compilation [3/3] (#31748)

Signed-off-by: Lucas Kabela <lucaskabela@meta.com>

Lucas Kabela committed Jan 12, 2026 at 11:24 UTC ad8818bb5eb8de14ebe18ebf8c8b3b2620092e17
11 files changed +333 -280
vllm/compilation/activation_quant_fusion.py
+33 -28
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 +from typing import Any
6
7 import torch
8 from torch._higher_order_ops.auto_functionalize import auto_functionalized
@@ -52,7 +53,7 @@ class ActivationQuantPattern(ABC):
53 def __init__(
54 self,
55 quant_key: QuantKey,
55 - ):
56 + ) -> None:
57 self.quant_key = quant_key
58 self.quant_dtype = quant_key.dtype
59
@@ -68,12 +69,12 @@ class ActivationQuantPattern(ABC):
69
70 self.silu_and_mul_matcher = MatcherSiluAndMul()
71
71 - def empty_quant(self, *args, **kwargs):
72 + def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor:
73 kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs}
74 return torch.empty(*args, **kwargs)
75
76 @abstractmethod
76 - def register(self, pm_pass: PatternMatcherPass):
77 + def register(self, pm_pass: PatternMatcherPass) -> None:
78 raise NotImplementedError
79
80
@@ -82,15 +83,22 @@ class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
83 Fusion for SiluMul+Fp8StaticQuant Pattern
84 """
85
85 - def __init__(self):
86 + def __init__(self) -> None:
87 super().__init__(kFp8StaticTensorSym)
88 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
89
89 - def register(self, pm_pass: PatternMatcherPass):
90 + def get_inputs(self) -> list[torch.Tensor]:
91 + scale = self.quant_matcher.inputs()[1]
92 + return [
93 + *self.silu_and_mul_matcher.inputs(), # input
94 + scale,
95 + ]
96 +
97 + def register(self, pm_pass: PatternMatcherPass) -> None:
98 def pattern(
99 input: torch.Tensor,
100 scale: torch.Tensor,
93 - ):
101 + ) -> torch.Tensor:
102 result_silu_mul = self.silu_and_mul_matcher(input)
103 result_quant = self.quant_matcher(result_silu_mul, scale)
104 return result_quant[0]
@@ -98,7 +106,7 @@ class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
106 def replacement(
107 input: torch.Tensor,
108 scale: torch.Tensor,
101 - ):
109 + ) -> torch.Tensor:
110 d = input.shape[-1] // 2
111 output_shape = input.shape[:-1] + (d,)
112 result = torch.empty(
@@ -109,13 +117,10 @@ class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
117 )
118 return at[1]
119
112 - inputs = [
113 - *self.silu_and_mul_matcher.inputs(), # input
114 - self.quant_matcher.inputs()[1], # scale
115 - ]
116 - pattern(*inputs)
120 + inps = self.get_inputs()
121 + pattern(*inps)
122
118 - register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
123 + register_replacement(pattern, replacement, inps, fwd_only, pm_pass)
124
125
126 class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
@@ -123,16 +128,23 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
128 Fusion for SiluMul+Nvfp4Quant Pattern
129 """
130
126 - def __init__(self):
131 + def __init__(self) -> None:
132 super().__init__(kNvfp4Quant)
133
129 - def register(self, pm_pass: PatternMatcherPass):
134 + def get_inputs(self) -> list[torch.Tensor]:
135 + result = self.empty_quant(5, 32)
136 + output_scale = empty_i32(128, 4)
137 + input_ = empty_bf16(5, 64)
138 + scale = empty_fp32(1, 1)
139 + return [result, output_scale, input_, scale]
140 +
141 + def register(self, pm_pass: PatternMatcherPass) -> None:
142 def pattern(
143 result: torch.Tensor,
144 output_scale: torch.Tensor,
145 input: torch.Tensor,
146 scale: torch.Tensor,
135 - ):
147 + ) -> tuple[torch.Tensor, torch.Tensor]:
148 result_silu_mul = self.silu_and_mul_matcher(input)
149 at = auto_functionalized(
150 self.QUANT_OP,
@@ -148,7 +160,7 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
160 output_scale: torch.Tensor,
161 input: torch.Tensor,
162 scale: torch.Tensor,
151 - ):
163 + ) -> tuple[torch.Tensor, torch.Tensor]:
164 at = auto_functionalized(
165 self.FUSED_OP,
166 result=result,
@@ -158,14 +170,7 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
170 )
171 return at[1], at[2]
172
161 - inputs = [
162 - self.empty_quant(5, 32), # result
163 - empty_i32(128, 4), # output_scale
164 - empty_bf16(5, 64), # input
165 - empty_fp32(1, 1), # scale
166 - ]
167 -
168 - register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
173 + register_replacement(pattern, replacement, self.get_inputs(), fwd_only, pm_pass)
174
175
176 class ActivationQuantFusionPass(VllmPatternMatcherPass):
@@ -179,7 +184,7 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass):
184 """
185
186 @enable_fake_mode
182 - def __init__(self, config: VllmConfig):
187 + def __init__(self, config: VllmConfig) -> None:
188 super().__init__(config)
189
190 self.patterns: PatternMatcherPass = PatternMatcherPass(
@@ -196,11 +201,11 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass):
201 self.dump_patterns(config, self.patterns)
202
203 @VllmInductorPass.time_and_log
199 - def __call__(self, graph: torch.fx.Graph):
204 + def __call__(self, graph: torch.fx.Graph) -> None:
205 self.matched_count = self.patterns.apply(graph)
206 logger.debug("Replaced %s patterns", self.matched_count)
207
203 - def uuid(self):
208 + def uuid(self) -> str:
209 return VllmInductorPass.hash_source(
210 self,
211 ActivationQuantPattern,
vllm/compilation/collective_fusion.py
+111 -102
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from importlib.util import find_spec
4 +from types import ModuleType
5
6 import torch
7 import torch._inductor.pattern_matcher as pm
@@ -33,15 +34,15 @@ if find_spec("flashinfer"):
34 try:
35 import flashinfer.comm as flashinfer_comm
36
36 - flashinfer_comm = (
37 + flashinfer_comm: ModuleType | None = ( # type: ignore[no-redef]
38 flashinfer_comm
39 if hasattr(flashinfer_comm, "trtllm_allreduce_fusion")
40 else None
41 )
42 except ImportError:
42 - flashinfer_comm = None
43 + flashinfer_comm = None # type: ignore[assignment]
44 else:
44 - flashinfer_comm = None
45 + flashinfer_comm = None # type: ignore[assignment]
46
47 logger = init_logger(__name__)
48
@@ -58,13 +59,13 @@ class BasePattern:
59
60
61 class GEMMReduceScatterPattern(BasePattern):
61 - def get_inputs(self):
62 + def get_inputs(self) -> list[torch.Tensor]:
63 mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
64 mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
65 return [mul, mm_weight]
66
66 - def register(self, pm_pass: PatternMatcherPass):
67 - def pattern(mul: torch.Tensor, mm_weight: torch.Tensor):
67 + def register(self, pm_pass: PatternMatcherPass) -> None:
68 + def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
69 mm = torch.ops.aten.mm.default(mul, mm_weight)
70 reduce_scatter = torch.ops.vllm.reduce_scatter.default(
71 mm,
@@ -74,7 +75,7 @@ class GEMMReduceScatterPattern(BasePattern):
75 )
76 return reduce_scatter
77
77 - def replacement(mul: torch.Tensor, mm_weight: torch.Tensor):
78 + def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
79 gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
80 mul,
81 mm_weight,
@@ -91,17 +92,17 @@ class GEMMReduceScatterPattern(BasePattern):
92
93
94 class AllGatherGEMMPattern(BasePattern):
94 - def get_inputs(self):
95 + def get_inputs(self) -> list[torch.Tensor]:
96 x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
97 weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
98
99 return [x, weight]
100
100 - def register(self, pm_pass: PatternMatcherPass):
101 + def register(self, pm_pass: PatternMatcherPass) -> None:
102 def pattern(
103 x: torch.Tensor,
104 weight: torch.Tensor,
104 - ) -> tuple[torch.Tensor, torch.Tensor]:
105 + ) -> torch.Tensor:
106 all_gather = torch.ops.vllm.all_gather.default(
107 x,
108 dim=0,
@@ -111,9 +112,7 @@ class AllGatherGEMMPattern(BasePattern):
112
113 return torch.ops.aten.mm.default(all_gather, weight)
114
114 - def replacement(
115 - x: torch.Tensor, weight: torch.Tensor
116 - ) -> tuple[torch.Tensor, torch.Tensor]:
115 + def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
116 ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
117 x,
118 [weight],
@@ -128,7 +127,7 @@ class AllGatherGEMMPattern(BasePattern):
127
128
129 class ScaledMMReduceScatterPattern(BasePattern):
131 - def get_inputs(self):
130 + def get_inputs(self) -> list[torch.Tensor]:
131 input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
132 mm_weight = (
133 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
@@ -139,7 +138,7 @@ class ScaledMMReduceScatterPattern(BasePattern):
138 scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
139 return [input, mm_weight, scale_a, scale_b]
140
142 - def register(self, pm_pass: PatternMatcherPass):
141 + def register(self, pm_pass: PatternMatcherPass) -> None:
142 def pattern(
143 input: torch.Tensor,
144 mat2: torch.Tensor,
@@ -196,7 +195,7 @@ class ScaledMMReduceScatterPattern(BasePattern):
195
196
197 class AllGatherScaledMMPattern(BasePattern):
199 - def get_inputs(self):
198 + def get_inputs(self) -> list[torch.Tensor]:
199 x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
200 weight = (
201 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
@@ -211,7 +210,7 @@ class AllGatherScaledMMPattern(BasePattern):
210
211 return [x, weight, scale_a, scale_b]
212
214 - def register(self, pm_pass: PatternMatcherPass):
213 + def register(self, pm_pass: PatternMatcherPass) -> None:
214 def pattern(
215 x: torch.Tensor,
216 weight: torch.Tensor,
@@ -258,7 +257,7 @@ class AllGatherScaledMMPattern(BasePattern):
257
258
259 class CutlassScaledMMReduceScatterPattern(BasePattern):
261 - def get_inputs(self):
260 + def get_inputs(self) -> list[torch.Tensor]:
261 input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
262 mm_weight = (
263 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
@@ -271,7 +270,7 @@ class CutlassScaledMMReduceScatterPattern(BasePattern):
270 cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype)
271 return [input, mm_weight, scale_a, scale_b, cutlass_mm_output]
272
274 - def register(self, pm_pass: PatternMatcherPass):
273 + def register(self, pm_pass: PatternMatcherPass) -> None:
274 def pattern(
275 input: torch.Tensor,
276 weight: torch.Tensor,
@@ -331,7 +330,7 @@ class CutlassScaledMMReduceScatterPattern(BasePattern):
330
331
332 class AllGatherCutlassScaledMMPattern(BasePattern):
334 - def get_inputs(self):
333 + def get_inputs(self) -> list[torch.Tensor]:
334 x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
335 weight = (
336 torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
@@ -349,7 +348,7 @@ class AllGatherCutlassScaledMMPattern(BasePattern):
348
349 return [x, weight, scale_a, scale_b, output]
350
352 - def register(self, pm_pass: PatternMatcherPass):
351 + def register(self, pm_pass: PatternMatcherPass) -> None:
352 def pattern(
353 x: torch.Tensor,
354 weight: torch.Tensor,
@@ -400,7 +399,7 @@ class AllGatherCutlassScaledMMPattern(BasePattern):
399
400 class AsyncTPPass(VllmPatternMatcherPass):
401 @enable_fake_mode
403 - def __init__(self, config: VllmConfig):
402 + def __init__(self, config: VllmConfig) -> None:
403 super().__init__(config)
404
405 # Enable symmetric memory for the TP process group
@@ -445,7 +444,7 @@ class AsyncTPPass(VllmPatternMatcherPass):
444 return compile_range.is_single_size() and compile_range.end % tp_size == 0
445
446 @VllmInductorPass.time_and_log
448 - def __call__(self, graph: fx.Graph):
447 + def __call__(self, graph: fx.Graph) -> None:
448 self.matched_count = self.patterns.apply(graph)
449 logger.debug("Replaced %s patterns", self.matched_count)
450
@@ -512,11 +511,13 @@ if flashinfer_comm is not None:
511 f"max token num {max_token_num} * hidden size {hidden_size} * "
512 f"element size {element_size}"
513 )
515 - device_capability = current_platform.get_device_capability().to_int()
514 + curr_device = current_platform.get_device_capability()
515 + device_capability = curr_device.to_int() if curr_device is not None else None
516 # Get one shot input size limit for the current world size
517 # for the current device capability
518 max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(
519 - device_capability, {}
519 + device_capability, # type: ignore[arg-type]
520 + {},
521 ).get(world_size, None)
522 # Use one shot if no max size is specified
523 use_oneshot = (
@@ -606,7 +607,7 @@ class FlashInferFusedAllReduceParams:
607 world_size: int,
608 use_fp32_lamport: bool = False,
609 max_token_num: int = 1024,
609 - ):
610 + ) -> None:
611 self.rank = rank
612 self.world_size = world_size
613 self.use_fp32_lamport = use_fp32_lamport
@@ -615,7 +616,7 @@ class FlashInferFusedAllReduceParams:
616 self.fp32_acc = True
617 self.max_token_num = max_token_num
618
618 - def get_trtllm_fused_allreduce_kwargs(self):
619 + def get_trtllm_fused_allreduce_kwargs(self) -> dict[str, bool | int]:
620 return {
621 "world_rank": self.rank,
622 "world_size": self.world_size,
@@ -639,26 +640,30 @@ class AllReduceRMSNormPattern(BasePattern):
640 dtype: torch.dtype,
641 device: str | None,
642 allreduce_params: FlashInferFusedAllReduceParams,
642 - ):
643 + ) -> None:
644 super().__init__(dtype, device)
645 self.epsilon = epsilon
646 self.allreduce_params = allreduce_params
647 self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
648
648 - def get_inputs(self):
649 + def get_inputs(self) -> list[torch.Tensor]:
650 input, weight = self.rmsnorm_matcher.inputs()
651
652 # input goes through allreduce first, always 16-bit
653 return [input.to(self.dtype), weight]
654
654 - def register(self, pm_pass: PatternMatcherPass):
655 - def pattern(input: torch.Tensor, weight: torch.Tensor):
655 + def register(self, pm_pass: PatternMatcherPass) -> None:
656 + def pattern(
657 + input: torch.Tensor, weight: torch.Tensor
658 + ) -> tuple[torch.Tensor, torch.Tensor]:
659 allreduce_output = tensor_model_parallel_all_reduce(input)
660 rms = self.rmsnorm_matcher(allreduce_output, weight)
661
662 return rms, allreduce_output
663
661 - def replacement(input: torch.Tensor, weight: torch.Tensor):
664 + def replacement(
665 + input: torch.Tensor, weight: torch.Tensor
666 + ) -> tuple[torch.Tensor, torch.Tensor]:
667 residual = torch.zeros_like(input)
668 rms_result = torch.empty_like(input)
669 allreduce = auto_functionalized(
@@ -694,27 +699,29 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
699 dtype: torch.dtype,
700 device: str | None,
701 allreduce_params: FlashInferFusedAllReduceParams,
697 - ):
702 + ) -> None:
703 super().__init__(dtype, device)
704 self.epsilon = epsilon
705 self.allreduce_params = allreduce_params
706 self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
707
703 - def get_inputs(self):
708 + def get_inputs(self) -> list[torch.Tensor]:
709 input, residual, weight = self.rmsnorm_matcher.inputs()
710
711 # input goes through allreduce first, always 16-bit
712 return [residual, input.to(self.dtype), weight]
713
709 - def register(self, pm_pass: PatternMatcherPass):
710 - def pattern(residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor):
714 + def register(self, pm_pass: PatternMatcherPass) -> None:
715 + def pattern(
716 + residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
717 + ) -> tuple[torch.Tensor, torch.Tensor]:
718 allreduce_output = tensor_model_parallel_all_reduce(input)
719 rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
720 return rms, residual
721
722 def replacement(
723 residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
717 - ):
724 + ) -> tuple[torch.Tensor, torch.Tensor]:
725 allreduce = auto_functionalized(
726 flashinfer_trtllm_fused_allreduce_norm,
727 allreduce_in=input,
@@ -739,8 +746,8 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
746 first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0]
747
748 pm.register_replacement(
742 - first_return_only(pattern),
743 - first_return_only(replacement),
749 + first_return_only(pattern), # type: ignore[no-untyped-call]
750 + first_return_only(replacement), # type: ignore[no-untyped-call]
751 self.get_inputs(),
752 pm.fwd_only,
753 pm_pass,
@@ -761,7 +768,7 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
768 dtype: torch.dtype,
769 device: str | None,
770 allreduce_params: FlashInferFusedAllReduceParams,
764 - ):
771 + ) -> None:
772 super().__init__(dtype, device)
773 self.epsilon = epsilon
774 self.allreduce_params = allreduce_params
@@ -769,25 +776,27 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
776 self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
777 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
778
772 - def register(self, pm_pass: PatternMatcherPass):
773 - def get_inputs():
774 - input, weight = self.rmsnorm_matcher.inputs()
775 - _, scale = self.quant_matcher.inputs()
779 + def get_inputs(self) -> list[torch.Tensor]:
780 + input, weight = self.rmsnorm_matcher.inputs()
781 + _, scale = self.quant_matcher.inputs()
782
777 - # input goes through allreduce first, always 16-bit
778 - return [input.to(self.dtype), weight, scale]
783 + # input goes through allreduce first, always 16-bit
784 + return [input.to(self.dtype), weight, scale]
785
786 + def register(self, pm_pass: PatternMatcherPass) -> None:
787 def pattern(
788 input: torch.Tensor,
789 weight: torch.Tensor,
790 scale: torch.Tensor,
784 - ):
791 + ) -> tuple[torch.Tensor, torch.Tensor]:
792 all_reduce = tensor_model_parallel_all_reduce(input)
793 rms = self.rmsnorm_matcher(all_reduce, weight)
794 quant, _ = self.quant_matcher(rms, scale)
795 return quant, all_reduce
796
790 - def replacement(input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor):
797 + def replacement(
798 + input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
799 + ) -> tuple[torch.Tensor, torch.Tensor]:
800 residual = torch.zeros_like(input)
801 result_rms = torch.empty_like(input)
802 result_quant = torch.empty_like(input, dtype=self.quant_dtype)
@@ -812,7 +821,7 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
821 return allreduce[4], allreduce[1]
822
823 pm.register_replacement(
815 - pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
824 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
825 )
826
827
@@ -830,7 +839,7 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
839 dtype: torch.dtype,
840 device: str | None,
841 allreduce_params: FlashInferFusedAllReduceParams,
833 - ):
842 + ) -> None:
843 super().__init__(dtype, device)
844 self.epsilon = epsilon
845 self.allreduce_params = allreduce_params
@@ -839,20 +848,20 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
848 self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
849 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
850
842 - def register(self, pm_pass: PatternMatcherPass):
843 - def get_inputs():
844 - input, residual, weight = self.rmsnorm_matcher.inputs()
845 - _, scale = self.quant_matcher.inputs()
851 + def get_inputs(self) -> list[torch.Tensor]:
852 + input, residual, weight = self.rmsnorm_matcher.inputs()
853 + _, scale = self.quant_matcher.inputs()
854
847 - # input goes through allreduce first, always 16-bit
848 - return [residual, input.to(self.dtype), weight, scale]
855 + # input goes through allreduce first, always 16-bit
856 + return [residual, input.to(self.dtype), weight, scale]
857
858 + def register(self, pm_pass: PatternMatcherPass) -> None:
859 def pattern(
860 residual: torch.Tensor,
861 input: torch.Tensor,
862 weight: torch.Tensor,
863 scale: torch.Tensor,
855 - ):
864 + ) -> tuple[torch.Tensor, torch.Tensor]:
865 allreduce_output = tensor_model_parallel_all_reduce(input)
866 rms, res = self.rmsnorm_matcher(allreduce_output, weight, residual)
867 quant, _ = self.quant_matcher(rms, scale)
@@ -864,7 +873,7 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
873 input: torch.Tensor,
874 weight: torch.Tensor,
875 scale: torch.Tensor,
867 - ):
876 + ) -> tuple[torch.Tensor, torch.Tensor]:
877 result_quant = torch.empty_like(input, dtype=self.quant_dtype)
878 allreduce = auto_functionalized(
879 flashinfer_trtllm_fused_allreduce_norm,
@@ -886,7 +895,7 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
895 return allreduce[4], allreduce[2]
896
897 pm.register_replacement(
889 - pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
898 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
899 )
900
901
@@ -904,31 +913,31 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
913 dtype: torch.dtype,
914 device: str | None,
915 allreduce_params: FlashInferFusedAllReduceParams,
907 - ):
916 + ) -> None:
917 super().__init__(dtype, device)
918 self.epsilon = epsilon
919 self.allreduce_params = allreduce_params
920 self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
921
913 - def register(self, pm_pass: PatternMatcherPass):
914 - def get_inputs():
915 - input = torch.empty([1, 16, 16], device=self.device, dtype=self.dtype)
916 - quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8)
917 - input_global_scale = torch.empty(
918 - [1, 1], device=self.device, dtype=torch.float32
919 - )
920 - weight = torch.empty([16], device=self.device, dtype=self.dtype)
921 - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
922 + def get_inputs(self) -> list[torch.Tensor]:
923 + input = torch.empty([1, 16, 16], device=self.device, dtype=self.dtype)
924 + quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8)
925 + input_global_scale = torch.empty(
926 + [1, 1], device=self.device, dtype=torch.float32
927 + )
928 + weight = torch.empty([16], device=self.device, dtype=self.dtype)
929 + output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
930
923 - return [input, quant_result, weight, input_global_scale, output_scale]
931 + return [input, quant_result, weight, input_global_scale, output_scale]
932
933 + def register(self, pm_pass: PatternMatcherPass) -> None:
934 def pattern(
935 input: torch.Tensor,
936 quant_result: torch.Tensor,
937 weight: torch.Tensor,
938 input_global_scale: torch.Tensor,
939 output_scale: torch.Tensor,
931 - ):
940 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
941 all_reduce = tensor_model_parallel_all_reduce(input)
942 rms = self.rmsnorm_matcher(all_reduce, weight)
943 quant_out_tuple = auto_functionalized(
@@ -948,7 +957,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
957 weight: torch.Tensor,
958 input_global_scale: torch.Tensor,
959 output_scale: torch.Tensor,
951 - ):
960 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
961 residual = torch.zeros_like(input)
962 result_rms = torch.empty_like(input)
963 allreduce = auto_functionalized(
@@ -972,7 +981,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
981 return allreduce[4], allreduce[1], allreduce[5]
982
983 pm.register_replacement(
975 - pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
984 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
985 )
986
987
@@ -990,33 +999,33 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
999 dtype: torch.dtype,
1000 device: str | None,
1001 allreduce_params: FlashInferFusedAllReduceParams,
993 - ):
1002 + ) -> None:
1003 super().__init__(dtype, device)
1004 self.epsilon = epsilon
1005 self.allreduce_params = allreduce_params
1006 self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
1007
999 - def register(self, pm_pass: PatternMatcherPass):
1000 - def get_inputs():
1001 - input = torch.empty([16, 16], device=self.device, dtype=self.dtype)
1002 -
1003 - residual = torch.empty([16, 16], device=self.device, dtype=self.dtype)
1004 - weight = torch.empty([16, 16], device=self.device, dtype=self.dtype)
1005 - quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8)
1006 - input_global_scale = torch.empty(
1007 - [1, 1], device=self.device, dtype=torch.float32
1008 - )
1009 - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
1010 -
1011 - return [
1012 - quant_result,
1013 - residual,
1014 - input,
1015 - output_scale,
1016 - weight,
1017 - input_global_scale,
1018 - ]
1008 + def get_inputs(self) -> list[torch.Tensor]:
1009 + input = torch.empty([16, 16], device=self.device, dtype=self.dtype)
1010
1011 + residual = torch.empty([16, 16], device=self.device, dtype=self.dtype)
1012 + weight = torch.empty([16, 16], device=self.device, dtype=self.dtype)
1013 + quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8)
1014 + input_global_scale = torch.empty(
1015 + [1, 1], device=self.device, dtype=torch.float32
1016 + )
1017 + output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)
1018 +
1019 + return [
1020 + quant_result,
1021 + residual,
1022 + input,
1023 + output_scale,
1024 + weight,
1025 + input_global_scale,
1026 + ]
1027 +
1028 + def register(self, pm_pass: PatternMatcherPass) -> None:
1029 def pattern(
1030 quant_result: torch.Tensor,
1031 residual: torch.Tensor,
@@ -1024,7 +1033,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
1033 output_scale: torch.Tensor,
1034 weight: torch.Tensor,
1035 input_global_scale: torch.Tensor,
1027 - ):
1036 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1037 allreduce_output = tensor_model_parallel_all_reduce(input)
1038 rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
1039 quant_out_tuple = auto_functionalized(
@@ -1045,7 +1054,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
1054 output_scale: torch.Tensor,
1055 weight: torch.Tensor,
1056 input_global_scale: torch.Tensor,
1048 - ):
1057 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1058 allreduce = auto_functionalized(
1059 flashinfer_trtllm_fused_allreduce_norm,
1060 allreduce_in=input,
@@ -1066,12 +1075,12 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
1075 return allreduce[4], allreduce[2], allreduce[5]
1076
1077 pm.register_replacement(
1069 - pattern, replacement, get_inputs(), pm.fwd_only, pm_pass
1078 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
1079 )
1080
1081
1082 class AllReduceFusionPass(VllmPatternMatcherPass):
1074 - def __init__(self, config: VllmConfig):
1083 + def __init__(self, config: VllmConfig) -> None:
1084 super().__init__(config)
1085 self.disabled = True
1086 self.tp_size = get_tensor_model_parallel_world_size()
@@ -1122,7 +1131,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1131 )
1132
1133 self.ipc_handles, workspace_tensor = (
1125 - flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
1134 + flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion( # type: ignore[misc]
1135 tp_rank=rank,
1136 tp_size=self.tp_size,
1137 max_token_num=self.max_token_num,
@@ -1145,7 +1154,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1154 self.dump_patterns(config, self.patterns)
1155
1156 @enable_fake_mode
1148 - def register_patterns(self):
1157 + def register_patterns(self) -> None:
1158 for epsilon in [1e-5, 1e-6]:
1159 AllReduceFusedRMSNormStaticQuantFP8Pattern(
1160 epsilon,
@@ -1198,7 +1207,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1207 return compile_range.end <= self.max_token_num
1208
1209 @VllmInductorPass.time_and_log
1201 - def __call__(self, graph: fx.Graph):
1210 + def __call__(self, graph: fx.Graph) -> None:
1211 if self.disabled:
1212 logger.debug("AllReduceFusionPass disabled")
1213 return
@@ -1206,7 +1215,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1215 self.matched_count = self.patterns.apply(graph)
1216 logger.debug("Replaced %s patterns", self.matched_count)
1217
1209 - def __del__(self):
1218 + def __del__(self) -> None:
1219 if getattr(self, "disabled", True):
1220 return
1221 if flashinfer_comm is not None:
vllm/compilation/fusion.py
+58 -38
@@ -38,19 +38,19 @@ FP8_DTYPE = current_platform.fp8_dtype()
38 FP4_DTYPE = torch.uint8
39
40
41 -def empty_bf16(*args, **kwargs):
41 +def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor:
42 return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda")
43
44
45 -def empty_fp32(*args, **kwargs):
45 +def empty_fp32(*args: Any, **kwargs: Any) -> torch.Tensor:
46 return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
47
48
49 -def empty_i32(*args, **kwargs):
49 +def empty_i32(*args: Any, **kwargs: Any) -> torch.Tensor:
50 return torch.empty(*args, **kwargs, dtype=torch.int32, device="cuda")
51
52
53 -def empty_i64(*args, **kwargs):
53 +def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor:
54 return torch.empty(*args, **kwargs, dtype=torch.int64, device="cuda")
55
56
@@ -79,7 +79,7 @@ class FusedRMSQuantKey(NamedTuple):
79 quant: QuantKey
80 fused_add: bool
81
82 - def __str__(self):
82 + def __str__(self) -> str:
83 return (
84 f"FusedQuantKey({self.quant}, with"
85 f"{'' if self.fused_add else 'out'} residual)"
@@ -121,7 +121,7 @@ class RMSNormQuantPattern:
121 key: FusedRMSQuantKey,
122 has_col_major_scales: bool = False,
123 is_e8m0: bool = False,
124 - ):
124 + ) -> None:
125 self.epsilon = epsilon
126 self.quant_dtype = key.quant.dtype
127 config = get_current_vllm_config()
@@ -141,7 +141,9 @@ class RMSNormQuantPattern:
141
142
143 class RMSNormStaticQuantPattern(RMSNormQuantPattern):
144 - def __init__(self, epsilon: float, quant_dtype: torch.dtype, symmetric=True):
144 + def __init__(
145 + self, epsilon: float, quant_dtype: torch.dtype, symmetric: bool = True
146 + ) -> None:
147 fused_key = FusedRMSQuantKey(
148 fused_add=False,
149 quant=QuantKey(
@@ -150,13 +152,17 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
152 )
153 super().__init__(epsilon, fused_key)
154
153 - def register(self, pm_pass: PatternMatcherPass):
155 + def register(self, pm_pass: PatternMatcherPass) -> None:
156 # Cannot use methods, as the self argument affects tracing
155 - def pattern(input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor):
157 + def pattern(
158 + input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
159 + ) -> torch.Tensor:
160 result_rms = self.rmsnorm_matcher(input, weight)
161 return self.quant_matcher(result_rms, scale)[0]
162
159 - def replacement(input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor):
163 + def replacement(
164 + input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
165 + ) -> torch.Tensor:
166 # In case we're matching native rms-norm, conversions might be
167 # optimized out. We convert here just to be safe.
168 input = input.to(dtype=self.model_dtype)
@@ -187,7 +193,9 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
193
194
195 class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
190 - def __init__(self, epsilon: float, quant_dtype: torch.dtype, symmetric=True):
196 + def __init__(
197 + self, epsilon: float, quant_dtype: torch.dtype, symmetric: bool = True
198 + ) -> None:
199 key = FusedRMSQuantKey(
200 fused_add=True,
201 quant=QuantKey(
@@ -196,13 +204,13 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
204 )
205 super().__init__(epsilon, key)
206
199 - def register(self, pm_pass: PatternMatcherPass):
207 + def register(self, pm_pass: PatternMatcherPass) -> None:
208 def pattern(
209 input: torch.Tensor,
210 weight: torch.Tensor,
211 residual: torch.Tensor,
212 scale: torch.Tensor,
205 - ):
213 + ) -> tuple[torch.Tensor, torch.Tensor]:
214 result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
215 result, _ = self.quant_matcher(result_rms, scale)
216
@@ -213,7 +221,7 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
221 weight: torch.Tensor,
222 residual: torch.Tensor,
223 scale: torch.Tensor,
216 - ):
224 + ) -> tuple[torch.Tensor, torch.Tensor]:
225 # In case we're matching native rms-norm, conversions might be
226 # optimized out. We convert here just to be safe.
227 input = input.to(dtype=self.model_dtype)
@@ -253,10 +261,10 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
261 epsilon: float,
262 quant_dtype: torch.dtype,
263 group_shape: GroupShape,
256 - symmetric=True,
264 + symmetric: bool = True,
265 has_col_major_scales: bool = False,
266 is_e8m0: bool = False,
259 - ):
267 + ) -> None:
268 scale = ScaleDesc(torch.float32, False, group_shape)
269 key = FusedRMSQuantKey(
270 fused_add=True,
@@ -269,15 +277,17 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
277 epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
278 )
279
272 - def register(self, pm_pass: PatternMatcherPass):
273 - def pattern(input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor):
280 + def register(self, pm_pass: PatternMatcherPass) -> None:
281 + def pattern(
282 + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
283 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
284 result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
285 result, scale = self.quant_matcher(result_rms)
286 return result, residual, scale
287
288 def replacement(
289 input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
280 - ):
290 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
291 # In case we're matching native rms-norm, conversions might be
292 # optimized out. We convert here just to be safe.
293 input = input.to(dtype=self.model_dtype)
@@ -315,10 +325,10 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
325 epsilon: float,
326 quant_dtype: torch.dtype,
327 group_shape: GroupShape,
318 - symmetric=True,
328 + symmetric: bool = True,
329 has_col_major_scales: bool = False,
330 is_e8m0: bool = False,
321 - ):
331 + ) -> None:
332 scale = ScaleDesc(torch.float32, False, group_shape)
333 key = FusedRMSQuantKey(
334 fused_add=False,
@@ -329,13 +339,17 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
339 epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
340 )
341
332 - def register(self, pm_pass: PatternMatcherPass):
333 - def pattern(input: torch.Tensor, weight: torch.Tensor):
342 + def register(self, pm_pass: PatternMatcherPass) -> None:
343 + def pattern(
344 + input: torch.Tensor, weight: torch.Tensor
345 + ) -> tuple[torch.Tensor, torch.Tensor]:
346 result_rms = self.rmsnorm_matcher(input, weight)
347 result, scale = self.quant_matcher(result_rms)
348 return result, scale
349
338 - def replacement(input: torch.Tensor, weight: torch.Tensor):
350 + def replacement(
351 + input: torch.Tensor, weight: torch.Tensor
352 + ) -> tuple[torch.Tensor, torch.Tensor]:
353 # In case we're matching native rms-norm, conversions might be
354 # optimized out. We convert here just to be safe.
355 input = input.to(dtype=self.model_dtype)
@@ -375,8 +389,8 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
389 epsilon: float,
390 quant_dtype: torch.dtype,
391 group_shape: GroupShape = GroupShape.PER_TOKEN,
378 - symmetric=True,
379 - ):
392 + symmetric: bool = True,
393 + ) -> None:
394 scale = ScaleDesc(torch.float32, False, group_shape)
395 key = FusedRMSQuantKey(
396 fused_add=False,
@@ -384,13 +398,17 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
398 )
399 super().__init__(epsilon, key)
400
387 - def register(self, pm_pass: PatternMatcherPass):
388 - def pattern(input: torch.Tensor, weight: torch.Tensor):
401 + def register(self, pm_pass: PatternMatcherPass) -> None:
402 + def pattern(
403 + input: torch.Tensor, weight: torch.Tensor
404 + ) -> tuple[torch.Tensor, torch.Tensor]:
405 result_rms = self.rmsnorm_matcher(input, weight)
406 # result, scale
391 - return self.quant_matcher(result_rms)
407 + return self.quant_matcher(result_rms) # type: ignore[no-any-return]
408
393 - def replacement(input: torch.Tensor, weight: torch.Tensor):
409 + def replacement(
410 + input: torch.Tensor, weight: torch.Tensor
411 + ) -> tuple[torch.Tensor, torch.Tensor]:
412 # In case we're matching native rms-norm, conversions might be
413 # optimized out. We convert here just to be safe.
414 input = input.to(dtype=self.model_dtype)
@@ -426,8 +444,8 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
444 epsilon: float,
445 quant_dtype: torch.dtype,
446 group_shape: GroupShape = GroupShape.PER_TOKEN,
429 - symmetric=True,
430 - ):
447 + symmetric: bool = True,
448 + ) -> None:
449 scale = ScaleDesc(torch.float32, False, group_shape)
450 key = FusedRMSQuantKey(
451 fused_add=True,
@@ -435,8 +453,10 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
453 )
454 super().__init__(epsilon, key)
455
438 - def register(self, pm_pass: PatternMatcherPass):
439 - def pattern(input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor):
456 + def register(self, pm_pass: PatternMatcherPass) -> None:
457 + def pattern(
458 + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
459 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
460 result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
461 result, scale = self.quant_matcher(result_rms)
462
@@ -444,7 +464,7 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
464
465 def replacement(
466 input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
447 - ):
467 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
468 # In case we're matching native rms-norm, conversions might be
469 # optimized out. We convert here just to be safe.
470 input = input.to(dtype=self.model_dtype)
@@ -481,7 +501,7 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
501 """
502
503 @enable_fake_mode
484 - def __init__(self, config: VllmConfig):
504 + def __init__(self, config: VllmConfig) -> None:
505 super().__init__(config)
506
507 self.patterns: PatternMatcherPass = PatternMatcherPass(
@@ -533,11 +553,11 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
553 self.dump_patterns(config, self.patterns)
554
555 @VllmInductorPass.time_and_log
536 - def __call__(self, graph: fx.Graph):
556 + def __call__(self, graph: fx.Graph) -> None:
557 self.matched_count = self.patterns.apply(graph)
558 logger.debug("Replaced %s patterns", self.matched_count)
559
540 - def uuid(self) -> Any:
560 + def uuid(self) -> str:
561 return self.hash_source(
562 self,
563 RMSNormGroupQuantPattern,
vllm/compilation/fusion_attn.py
+24 -20
@@ -3,6 +3,7 @@
3
4 from abc import ABC, abstractmethod
5 from collections.abc import Callable
6 +from typing import Any, ParamSpec
7
8 import torch
9 import torch._inductor.pattern_matcher as pm
@@ -28,7 +29,7 @@ from .matcher_utils import MatcherQuantFP8
29 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
30
31 logger = init_logger(__name__)
31 -
32 +P = ParamSpec("P")
33 FP8_DTYPE = current_platform.fp8_dtype()
34 FP4_DTYPE = torch.uint8
35
@@ -47,7 +48,7 @@ class AttentionQuantPattern(ABC):
48 layer: Attention,
49 quant_key: QuantKey,
50 dtype: torch.dtype,
50 - ):
51 + ) -> None:
52 self.layer = layer
53 self.layer_name = layer.layer_name
54 self.num_heads = layer.num_heads
@@ -61,17 +62,20 @@ class AttentionQuantPattern(ABC):
62 )
63 self.QUANT_OP = QUANT_OPS[self.quant_key]
64
64 - def empty(self, *args, **kwargs):
65 + def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor:
66 kwargs = {"dtype": self.dtype, "device": "cuda", **kwargs}
67 return torch.empty(*args, **kwargs)
68
68 - def empty_quant(self, *args, **kwargs):
69 + def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor:
70 kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs}
71 return torch.empty(*args, **kwargs)
72
73 @staticmethod
73 - def wrap_trace_fn(trace_fn, *process_fx_fns: Callable[[fx.GraphModule], None]):
74 - def wrapped(*args, **kwargs):
74 + def wrap_trace_fn(
75 + trace_fn: Callable[P, fx.GraphModule],
76 + *process_fx_fns: Callable[[fx.GraphModule], None],
77 + ) -> Callable[P, fx.GraphModule]:
78 + def wrapped(*args: P.args, **kwargs: P.kwargs) -> fx.GraphModule:
79 gm = trace_fn(*args, **kwargs)
80 for process_fx in process_fx_fns:
81 process_fx(gm)
@@ -81,13 +85,13 @@ class AttentionQuantPattern(ABC):
85 return wrapped
86
87 @staticmethod
84 - def fx_view_to_reshape(gm: torch.fx.GraphModule):
88 + def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None:
89 from torch._inductor.fx_passes.post_grad import view_to_reshape
90
91 view_to_reshape(gm)
92
93 @staticmethod
90 - def remove_noop_permutes(gm: torch.fx.GraphModule):
94 + def remove_noop_permutes(gm: torch.fx.GraphModule) -> None:
95 for node in gm.graph.nodes:
96 if not is_func(node, torch.ops.aten.permute.default):
97 continue
@@ -100,12 +104,12 @@ class AttentionQuantPattern(ABC):
104 node.replace_all_uses_with(node.args[0])
105 gm.graph.erase_node(node)
106
103 - def register_if_supported(self, pm_pass: PatternMatcherPass):
107 + def register_if_supported(self, pm_pass: PatternMatcherPass) -> None:
108 if self.layer.impl.fused_output_quant_supported(self.quant_key):
109 self._register(pm_pass)
110
111 @abstractmethod
108 - def _register(self, pm_pass: PatternMatcherPass):
112 + def _register(self, pm_pass: PatternMatcherPass) -> None:
113 raise NotImplementedError
114
115
@@ -124,21 +128,21 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
128 layer: Attention,
129 dtype: torch.dtype,
130 symmetric: bool = True,
127 - ):
131 + ) -> None:
132 quant_key = QuantKey(
133 dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=symmetric
134 )
135 super().__init__(layer, quant_key, dtype)
136 self.quant_matcher = MatcherQuantFP8(quant_key)
137
134 - def _register(self, pm_pass: PatternMatcherPass):
138 + def _register(self, pm_pass: PatternMatcherPass) -> None:
139 def pattern(
140 q: torch.Tensor,
141 k: torch.Tensor,
142 v: torch.Tensor,
143 output_attn: torch.Tensor,
144 scale: torch.Tensor,
141 - ):
145 + ) -> torch.Tensor:
146 at1 = auto_functionalized(
147 ATTN_OP,
148 query=q,
@@ -161,7 +165,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
165 v: torch.Tensor,
166 output_attn: torch.Tensor,
167 scale: torch.Tensor,
164 - ):
168 + ) -> torch.Tensor:
169 # attn output in quant_dtype
170 output_attn = torch.ops.aten.full.default(
171 [q.shape[0], self.num_heads, self.head_size],
@@ -212,10 +216,10 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
216 will be passed into Attention op as the `output_scale` argument.
217 """
218
215 - def __init__(self, layer: Attention, dtype: torch.dtype):
219 + def __init__(self, layer: Attention, dtype: torch.dtype) -> None:
220 super().__init__(layer, kNvfp4Quant, dtype)
221
218 - def _register(self, pm_pass: PatternMatcherPass):
222 + def _register(self, pm_pass: PatternMatcherPass) -> None:
223 def pattern(
224 q: torch.Tensor,
225 k: torch.Tensor,
@@ -224,7 +228,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
228 output_quant: torch.Tensor,
229 output_scale: torch.Tensor,
230 input_scale: torch.Tensor,
227 - ):
231 + ) -> tuple[torch.Tensor, torch.Tensor]:
232 at1 = auto_functionalized(
233 ATTN_OP,
234 query=q,
@@ -256,7 +260,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
260 output_quant: torch.Tensor,
261 output_scale: torch.Tensor,
262 input_scale: torch.Tensor,
259 - ):
263 + ) -> tuple[torch.Tensor, torch.Tensor]:
264 # attention output in quant_dtype
265 output_attn = torch.ops.aten.full.default(
266 [q.shape[0], self.num_heads, self.head_size // 2],
@@ -318,7 +322,7 @@ class AttnFusionPass(VllmPatternMatcherPass):
322 """
323
324 @enable_fake_mode
321 - def __init__(self, config: VllmConfig):
325 + def __init__(self, config: VllmConfig) -> None:
326 super().__init__(config)
327
328 self.patterns = PatternMatcherPass(pass_name="attn_fusion_pass")
@@ -350,7 +354,7 @@ class AttnFusionPass(VllmPatternMatcherPass):
354 self.matched_count = self.patterns.apply(graph)
355 logger.debug("Fused quant onto %s attention nodes", self.matched_count)
356
353 - def uuid(self):
357 + def uuid(self) -> str:
358 return VllmInductorPass.hash_source(
359 self,
360 AttentionQuantPattern,
vllm/compilation/inductor_pass.py
+1 -1
@@ -68,7 +68,7 @@ class InductorPass(CustomGraphPass): # type: ignore[misc]
68 This is defined as a convenience and should work in most cases.
69 """
70
71 - def uuid(self) -> Any:
71 + def uuid(self) -> str:
72 """
73 Provide a unique identifier for the pass, used in Inductor code cache.
74 This should depend on the pass implementation, so that changes to the
vllm/compilation/matcher_utils.py
+23 -22
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 +from typing import Any
5
6 import torch
7 from torch._higher_order_ops import auto_functionalized
@@ -47,7 +48,7 @@ SILU_MUL_OP = torch.ops._C.silu_and_mul.default
48
49
50 class MatcherCustomOp(ABC):
50 - def __init__(self, enabled: bool):
51 + def __init__(self, enabled: bool) -> None:
52 config = get_current_vllm_config()
53 self.model_dtype = config.model_config.dtype if config.model_config else None
54 self.device = config.device_config.device if config.device_config else None
@@ -56,24 +57,24 @@ class MatcherCustomOp(ABC):
57 self.forward = self.forward_custom if enabled else self.forward_native
58
59 @abstractmethod
59 - def forward_custom(self, *args, **kws):
60 + def forward_custom(self, *args: Any, **kwargs: Any) -> Any:
61 pass
62
63 @abstractmethod
63 - def forward_native(self, *args, **kws):
64 + def forward_native(self, *args: Any, **kwargs: Any) -> Any:
65 pass
66
66 - def __call__(self, *args, **kws):
67 - return self.forward(*args, **kws)
67 + def __call__(self, *args: Any, **kwargs: Any) -> Any:
68 + return self.forward(*args, **kwargs)
69
69 - def empty(self, *args, **kws):
70 - return torch.empty(*args, dtype=self.model_dtype, device=self.device, **kws)
70 + def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor:
71 + return torch.empty(*args, dtype=self.model_dtype, device=self.device, **kwargs)
72
72 - def empty_int64(self, *args, **kws):
73 - return torch.empty(*args, dtype=torch.int64, device=self.device, **kws)
73 + def empty_int64(self, *args: Any, **kwargs: Any) -> torch.Tensor:
74 + return torch.empty(*args, dtype=torch.int64, device=self.device, **kwargs)
75
75 - def empty_f32(self, *args, **kws):
76 - return torch.empty(*args, dtype=torch.float32, device=self.device, **kws)
76 + def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor:
77 + return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs)
78
79 def inputs(self) -> list[torch.Tensor]:
80 """Utility for inputs to the pattern"""
@@ -157,7 +158,7 @@ class MatcherRMSNorm(MatcherCustomOp):
158 epsilon: float,
159 enabled: bool | None = None,
160 match_rocm_aiter: bool = False,
160 - ):
161 + ) -> None:
162 if enabled is None:
163 enabled = RMSNorm.enabled()
164
@@ -169,7 +170,7 @@ class MatcherRMSNorm(MatcherCustomOp):
170 if match_rocm_aiter:
171 self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_op()
172
172 - def inputs(self):
173 + def inputs(self) -> list[torch.Tensor]:
174 input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)
175 weight = self.empty(16)
176 return [input, weight]
@@ -220,7 +221,7 @@ class MatcherFusedAddRMSNorm(MatcherCustomOp):
221 epsilon: float,
222 enabled: bool | None = None,
223 match_rocm_aiter: bool = False,
223 - ):
224 + ) -> None:
225 if enabled is None:
226 enabled = RMSNorm.enabled()
227
@@ -233,7 +234,7 @@ class MatcherFusedAddRMSNorm(MatcherCustomOp):
234 if match_rocm_aiter:
235 self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_fused_add_op()
236
236 - def inputs(self):
237 + def inputs(self) -> list[torch.Tensor]:
238 input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)
239 weight = self.empty(16)
240 residual = self.empty(5, 16)
@@ -245,7 +246,7 @@ class MatcherFusedAddRMSNorm(MatcherCustomOp):
246 weight: torch.Tensor,
247 residual: torch.Tensor,
248 ) -> tuple[torch.Tensor, torch.Tensor]:
248 - return self._rmsnorm_op(
249 + return self._rmsnorm_op( # type: ignore[no-any-return]
250 x=input, residual=residual, weight=weight, variance_epsilon=self.epsilon
251 )
252
@@ -287,7 +288,7 @@ class MatcherQuantFP8(MatcherCustomOp):
288 has_col_major_scales: bool = False,
289 is_e8m0: bool = False,
290 match_rocm_aiter: bool = False,
290 - ):
291 + ) -> None:
292 if enabled is None:
293 enabled = QuantFP8.enabled()
294
@@ -340,13 +341,13 @@ class MatcherQuantFP8(MatcherCustomOp):
341 ) -> tuple[torch.Tensor, torch.Tensor]:
342 quant_key_group_shape = self.quant_key.scale.group_shape
343 if quant_key_group_shape == GroupShape.PER_TOKEN:
343 - return self.QUANT_OP(
344 + return self.QUANT_OP( # type: ignore[no-any-return]
345 x=input,
346 quant_dtype=self.quant_key.dtype,
347 scale=scale,
348 )
349 else:
349 - return self.QUANT_OP(input, quant_key_group_shape.col)
350 + return self.QUANT_OP(input, quant_key_group_shape.col) # type: ignore[no-any-return]
351
352 def forward_custom(
353 self,
@@ -400,9 +401,9 @@ class MatcherQuantFP8(MatcherCustomOp):
401 input: torch.Tensor,
402 scale: torch.Tensor | None = None,
403 ) -> tuple[torch.Tensor, torch.Tensor]:
403 - return self.quant_fp8(input, scale)
404 + return self.quant_fp8(input, scale) # type: ignore[no-any-return]
405
405 - def make_scale(self, input: torch.Tensor, transposed: bool = False):
406 + def make_scale(self, input: torch.Tensor, transposed: bool = False) -> torch.Tensor:
407 normalized_group_shape = _normalize_quant_group_shape(
408 input, self.quant_key.scale.group_shape
409 )
@@ -427,7 +428,7 @@ class MatcherQuantFP8(MatcherCustomOp):
428
429
430 class MatcherSiluAndMul(MatcherCustomOp):
430 - def __init__(self, enabled: bool | None = None):
431 + def __init__(self, enabled: bool | None = None) -> None:
432 if enabled is None:
433 enabled = SiluAndMul.enabled()
434 super().__init__(enabled)
vllm/compilation/qk_norm_rope_fusion.py
+16 -10
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 +from typing import ParamSpec
6
7 import torch
8 import torch._inductor.pattern_matcher as pm
@@ -23,6 +24,8 @@ logger = init_logger(__name__)
24
25 FUSED_QK_ROPE_OP = torch.ops._C.fused_qk_norm_rope.default
26
27 +P = ParamSpec("P")
28 +
29
30 class QkNormRopePattern:
31 """
@@ -72,7 +75,7 @@ class QkNormRopePattern:
75 use_flashinfer=self.rope_flashinfer,
76 )
77
75 - def get_inputs(self):
78 + def get_inputs(self) -> list[torch.Tensor]:
79 # Sample inputs to help pattern tracing
80 T = 5
81 qkv = empty_bf16(T, self.q_size + 2 * self.kv_size)
@@ -92,8 +95,11 @@ class QkNormRopePattern:
95 ]
96
97 @staticmethod
95 - def wrap_trace_fn(trace_fn, *process_fx_fns: Callable[[fx.GraphModule], None]):
96 - def wrapped(*args, **kwargs):
98 + def wrap_trace_fn(
99 + trace_fn: Callable[P, fx.GraphModule],
100 + *process_fx_fns: Callable[[fx.GraphModule], None],
101 + ) -> Callable[P, fx.GraphModule]:
102 + def wrapped(*args: P.args, **kwargs: P.kwargs) -> fx.GraphModule:
103 gm = trace_fn(*args, **kwargs)
104 for process_fx in process_fx_fns:
105 process_fx(gm)
@@ -103,19 +109,19 @@ class QkNormRopePattern:
109 return wrapped
110
111 @staticmethod
106 - def fx_view_to_reshape(gm: torch.fx.GraphModule):
112 + def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None:
113 from torch._inductor.fx_passes.post_grad import view_to_reshape
114
115 view_to_reshape(gm)
116
111 - def register(self, pm_pass: PatternMatcherPass):
117 + def register(self, pm_pass: PatternMatcherPass) -> None:
118 def pattern(
119 qkv: torch.Tensor,
120 positions: torch.Tensor,
121 q_weight: torch.Tensor,
122 k_weight: torch.Tensor,
123 cos_sin_cache: torch.Tensor,
118 - ):
124 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
125 # split qkv -> q,k,v
126 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
127
@@ -143,7 +149,7 @@ class QkNormRopePattern:
149 q_weight: torch.Tensor,
150 k_weight: torch.Tensor,
151 cos_sin_cache: torch.Tensor,
146 - ):
152 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
153 # Run fused qk_norm_rope op
154 result = auto_functionalized(
155 FUSED_QK_ROPE_OP,
@@ -162,7 +168,7 @@ class QkNormRopePattern:
168 result_qkv = result[1]
169
170 # Split back to q,k,v and return
165 - return result_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
171 + return result_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # type: ignore[no-any-return]
172
173 # NOTE: use fx_view_to_reshape to unify view/reshape to simplify
174 # pattern and increase matching opportunities
@@ -182,7 +188,7 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
188 """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists."""
189
190 @enable_fake_mode
185 - def __init__(self, config: VllmConfig):
191 + def __init__(self, config: VllmConfig) -> None:
192 super().__init__(config)
193 self.patterns: PatternMatcherPass = PatternMatcherPass(
194 pass_name="qk_norm_rope_fusion_pass"
@@ -234,5 +240,5 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
240 self.matched_count = self.patterns.apply(graph)
241 logger.debug("Fused QK Norm+RoPE on %s sites", self.matched_count)
242
237 - def uuid(self):
243 + def uuid(self) -> str:
244 return VllmInductorPass.hash_source(self, QkNormRopePattern)
vllm/compilation/rocm_aiter_fusion.py
+38 -36
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Any
3
4 import torch
5 import torch._inductor.pattern_matcher as pm
@@ -65,8 +64,8 @@ class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
64 quant_dtype: torch.dtype,
65 match_aiter_quant: bool = True,
66 group_shape: GroupShape = GroupShape.PER_TOKEN,
68 - symmetric=True,
69 - ):
67 + symmetric: bool = True,
68 + ) -> None:
69 scale = ScaleDesc(torch.float32, False, group_shape)
70 key = FusedRMSQuantKey(
71 fused_add=False,
@@ -75,11 +74,11 @@ class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
74
75 super().__init__(epsilon, key, match_aiter_quant)
76
78 - def register(self, pm_pass):
77 + def register(self, pm_pass: PatternMatcherPass) -> None:
78 def pattern(
79 input: torch.Tensor,
80 weight: torch.Tensor,
82 - ):
81 + ) -> tuple[torch.Tensor, torch.Tensor]:
82 result_rms = self.rmsnorm_matcher(input, weight)
83 result, scale = self.quant_matcher(result_rms)
84 return result, scale
@@ -87,7 +86,7 @@ class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
86 def replacement(
87 input: torch.Tensor,
88 weight: torch.Tensor,
90 - ):
89 + ) -> tuple[torch.Tensor, torch.Tensor]:
90 result = self.FUSED_OP(
91 x=input,
92 weight=weight,
@@ -117,8 +116,8 @@ class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
116 quant_dtype: torch.dtype,
117 match_aiter_quant: bool = True,
118 group_shape: GroupShape = GroupShape.PER_TOKEN,
120 - symmetric=True,
121 - ):
119 + symmetric: bool = True,
120 + ) -> None:
121 scale = ScaleDesc(torch.float32, False, group_shape)
122 key = FusedRMSQuantKey(
123 fused_add=True,
@@ -127,12 +126,12 @@ class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
126
127 super().__init__(epsilon, key, match_aiter_quant)
128
130 - def register(self, pm_pass):
129 + def register(self, pm_pass: PatternMatcherPass) -> None:
130 def pattern(
131 input: torch.Tensor,
132 weight: torch.Tensor,
133 residual: torch.Tensor,
135 - ):
134 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
135 result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
136 result, scale = self.quant_matcher(result_rms)
137
@@ -140,7 +139,7 @@ class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
139
140 def replacement(
141 input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
143 - ):
142 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
143 result = self.FUSED_OP(
144 x=input,
145 residual=residual,
@@ -174,8 +173,8 @@ class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
173 quant_dtype: torch.dtype,
174 group_shape: GroupShape,
175 match_aiter_quant: bool = True,
177 - symmetric=True,
178 - ):
176 + symmetric: bool = True,
177 + ) -> None:
178 scale = ScaleDesc(torch.float32, False, group_shape)
179 key = FusedRMSQuantKey(
180 fused_add=False,
@@ -184,11 +183,11 @@ class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
183
184 super().__init__(epsilon, key, match_aiter_quant)
185
187 - def register(self, pm_pass: PatternMatcherPass):
186 + def register(self, pm_pass: PatternMatcherPass) -> None:
187 def pattern(
188 input: torch.Tensor,
189 weight: torch.Tensor,
191 - ):
190 + ) -> tuple[torch.Tensor, torch.Tensor]:
191 result_rms = self.rmsnorm_matcher(input, weight)
192 result, scale = self.quant_matcher(result_rms)
193 return result, scale
@@ -196,7 +195,7 @@ class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
195 def replacement(
196 input: torch.Tensor,
197 weight: torch.Tensor,
199 - ):
198 + ) -> tuple[torch.Tensor, torch.Tensor]:
199 at = self.FUSED_OP(
200 x=input,
201 weight=weight,
@@ -225,8 +224,8 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
224 quant_dtype: torch.dtype,
225 group_shape: GroupShape,
226 match_aiter_quant: bool = True,
228 - symmetric=True,
229 - ):
227 + symmetric: bool = True,
228 + ) -> None:
229 scale = ScaleDesc(torch.float32, False, group_shape)
230 key = FusedRMSQuantKey(
231 fused_add=True,
@@ -235,12 +234,12 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
234
235 super().__init__(epsilon, key, match_aiter_quant)
236
238 - def register(self, pm_pass: PatternMatcherPass):
237 + def register(self, pm_pass: PatternMatcherPass) -> None:
238 def pattern(
239 input: torch.Tensor,
240 weight: torch.Tensor,
241 residual: torch.Tensor,
243 - ):
242 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
243 result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
244 result, scale = self.quant_matcher(result_rms)
245
@@ -250,7 +249,7 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
249 input: torch.Tensor,
250 weight: torch.Tensor,
251 residual: torch.Tensor,
253 - ):
252 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
253 at = self.FUSED_OP(
254 x=input,
255 residual=residual,
@@ -275,7 +274,7 @@ class RocmAiterRMSNormFusionPass(VllmPatternMatcherPass):
274 """
275
276 @enable_fake_mode
278 - def __init__(self, config: VllmConfig):
277 + def __init__(self, config: VllmConfig) -> None:
278 super().__init__(config)
279
280 self.patterns: PatternMatcherPass = PatternMatcherPass(
@@ -311,11 +310,11 @@ class RocmAiterRMSNormFusionPass(VllmPatternMatcherPass):
310 self.dump_patterns(config, self.patterns)
311
312 @VllmInductorPass.time_and_log
314 - def __call__(self, graph: fx.Graph):
313 + def __call__(self, graph: fx.Graph) -> None:
314 self.matched_count = self.patterns.apply(graph)
315 logger.debug("Replaced %s patterns", self.matched_count)
316
318 - def uuid(self) -> Any:
317 + def uuid(self) -> str:
318 fusion_patterns = [
319 AiterRMSNormDynamicQuantPattern,
320 AiterFusedAddRMSNormDynamicQuantPattern,
@@ -333,29 +332,32 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
332
333 FUSED_SILU_MUL_QUANT_OP = rocm_aiter_ops.get_act_mul_fused_fp8_group_quant_op()
334
336 - def __init__(self, quant_op: OpOverload):
335 + def __init__(self, quant_op: OpOverload) -> None:
336 self.silu_and_mul_matcher = MatcherSiluAndMul()
337 self.quant_op = quant_op
338
340 - def register(self, pm_pass: PatternMatcherPass):
339 + def get_inputs(self) -> list[torch.Tensor]:
340 + return [
341 + self.silu_and_mul_matcher.inputs()[0],
342 + ]
343 +
344 + def register(self, pm_pass: PatternMatcherPass) -> None:
345 def pattern(
346 input: torch.Tensor,
343 - ):
347 + ) -> tuple[torch.Tensor, torch.Tensor]:
348 at1 = self.silu_and_mul_matcher(input)
349 at2 = self.quant_op(at1, 128)
350 return at2[0], at2[1]
351
352 def replacement(
353 input: torch.Tensor,
350 - ):
354 + ) -> tuple[torch.Tensor, torch.Tensor]:
355 at = self.FUSED_SILU_MUL_QUANT_OP(x=input, group_size=128)
356 return at[0], at[1]
357
354 - inputs = [
355 - self.silu_and_mul_matcher.inputs()[0],
356 - ]
357 -
358 - pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
358 + pm.register_replacement(
359 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
360 + )
361
362
363 class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
@@ -374,7 +376,7 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
376 QUANT_OPS = [AITER_GROUP_FP8_QUANT_OP, TRITON_GROUP_FP8_QUANT_OP]
377
378 @enable_fake_mode
377 - def __init__(self, config: VllmConfig):
379 + def __init__(self, config: VllmConfig) -> None:
380 super().__init__(config)
381
382 self.patterns: PatternMatcherPass = PatternMatcherPass(
@@ -387,11 +389,11 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
389 self.dump_patterns(config, self.patterns)
390
391 @VllmInductorPass.time_and_log
390 - def __call__(self, graph: torch.fx.Graph):
392 + def __call__(self, graph: torch.fx.Graph) -> None:
393 self.matched_count = self.patterns.apply(graph)
394 logger.debug("Replaced %s patterns", self.matched_count)
395
394 - def uuid(self):
396 + def uuid(self) -> str:
397 fusion_patterns = [
398 ActivationQuantPattern,
399 AiterSiluMulFp8GroupQuantPattern,
vllm/compilation/sequence_parallelism.py
+22 -18
@@ -2,6 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import functools
5 +from collections.abc import Callable, Sequence
6 +from typing import Any
7
8 import torch
9 import torch._inductor.pattern_matcher as pm
@@ -26,9 +28,11 @@ from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
28 logger = init_logger(__name__)
29
30
29 -def get_first_out_wrapper(fn):
31 +def get_first_out_wrapper(
32 + fn: Callable[..., Sequence[torch.Tensor]],
33 +) -> Callable[..., torch.Tensor]:
34 @functools.wraps(fn)
31 - def wrapper(*args):
35 + def wrapper(*args: Any) -> torch.Tensor:
36 return fn(*args)[0]
37
38 return wrapper
@@ -68,7 +72,7 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
72 super().__init__(epsilon, dtype, device)
73 self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
74
71 - def get_inputs(self):
75 + def get_inputs(self) -> list[torch.Tensor]:
76 input = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
77 arg3_1 = torch.empty([4], device=self.device, dtype=self.dtype)
78
@@ -78,7 +82,7 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
82 def pattern(
83 input: torch.Tensor,
84 arg3_1: torch.Tensor,
81 - ):
85 + ) -> tuple[torch.Tensor, torch.Tensor]:
86 all_reduce = self._all_reduce(input)
87 rmsnorm = self.rmsnorm_matcher(all_reduce, arg3_1)
88
@@ -87,7 +91,7 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
91 def replacement(
92 input: torch.Tensor,
93 arg3_1: torch.Tensor,
90 - ):
94 + ) -> tuple[torch.Tensor, torch.Tensor]:
95 reduce_scatter = self._reduce_scatter(input)
96
97 rmsnorm = self.rmsnorm_matcher(reduce_scatter, arg3_1)
@@ -100,11 +104,11 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
104
105
106 class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
103 - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None):
107 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:
108 super().__init__(epsilon, dtype, device)
109 self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
110
107 - def get_inputs(self):
111 + def get_inputs(self) -> list[torch.Tensor]:
112 mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
113
114 residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -116,7 +120,7 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
120 rms_norm_weights,
121 ]
122
119 - def register(self, pm_pass: PatternMatcherPass):
123 + def register(self, pm_pass: PatternMatcherPass) -> None:
124 def pattern(
125 residual: torch.Tensor,
126 mm_1: torch.Tensor,
@@ -163,23 +167,23 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
167 epsilon: float,
168 dtype: torch.dtype,
169 device: str | None,
166 - ):
170 + ) -> None:
171 super().__init__(epsilon, dtype, device)
172 self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
173 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
174
171 - def get_inputs(self):
175 + def get_inputs(self) -> list[torch.Tensor]:
176 input = torch.zeros([1, 8, 4], device=self.device, dtype=self.dtype)
177 weight = torch.empty([4], device=self.device, dtype=self.dtype)
178 scale = torch.tensor(1.0, device=self.device, dtype=torch.float32)
179 return [input, weight, scale]
180
177 - def register(self, pm_pass: PatternMatcherPass):
181 + def register(self, pm_pass: PatternMatcherPass) -> None:
182 def pattern(
183 input: torch.Tensor,
184 weight: torch.Tensor,
185 scale: torch.Tensor,
182 - ):
186 + ) -> tuple[torch.Tensor, torch.Tensor]:
187 all_reduce = self._all_reduce(input)
188 rms = self.rmsnorm_matcher(all_reduce, weight)
189 quant, _ = self.quant_matcher(rms, scale)
@@ -189,7 +193,7 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
193 input: torch.Tensor,
194 weight: torch.Tensor,
195 scale: torch.Tensor,
192 - ):
196 + ) -> tuple[torch.Tensor, torch.Tensor]:
197 reduce_scatter = self._reduce_scatter(input)
198 rms = self.rmsnorm_matcher(reduce_scatter, weight)
199 quant, _ = self.quant_matcher(rms, scale)
@@ -203,12 +207,12 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
207
208
209 class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
206 - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None):
210 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:
211 super().__init__(epsilon, dtype, device)
212 self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
213 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
214
211 - def get_inputs(self):
215 + def get_inputs(self) -> list[torch.Tensor]:
216 mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
217 residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)
218 rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -216,7 +220,7 @@ class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
220
221 return [residual, mm_1, rms_norm_weights, scale]
222
219 - def register(self, pm_pass: PatternMatcherPass):
223 + def register(self, pm_pass: PatternMatcherPass) -> None:
224 def pattern(
225 residual: torch.Tensor,
226 mm_1: torch.Tensor,
@@ -302,7 +306,7 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
306 """
307
308 @enable_fake_mode
305 - def __init__(self, config: VllmConfig):
309 + def __init__(self, config: VllmConfig) -> None:
310 super().__init__(config)
311
312 # Used to clean up redundant views created temporarily
@@ -357,7 +361,7 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
361 return (compile_range.is_single_size()) and (compile_range.end % tp_size == 0)
362
363 @VllmInductorPass.time_and_log
360 - def __call__(self, graph: fx.Graph):
364 + def __call__(self, graph: fx.Graph) -> None:
365 self.matched_count = self.patterns.apply(graph)
366 logger.debug("Replaced %s patterns", self.matched_count)
367 # Clean up reshape nodes
vllm/distributed/parallel_state.py
+4 -4
@@ -1529,22 +1529,22 @@ def patch_tensor_parallel_group(tp_group: GroupCoordinator):
1529 _TP = old_tp_group
1530
1531
1532 -def get_tensor_model_parallel_world_size():
1532 +def get_tensor_model_parallel_world_size() -> int:
1533 """Return world size for the tensor model parallel group."""
1534 return get_tp_group().world_size
1535
1536
1537 -def get_tensor_model_parallel_rank():
1537 +def get_tensor_model_parallel_rank() -> int:
1538 """Return my rank for the tensor model parallel group."""
1539 return get_tp_group().rank_in_group
1540
1541
1542 -def get_decode_context_model_parallel_world_size():
1542 +def get_decode_context_model_parallel_world_size() -> int:
1543 """Return world size for the decode context model parallel group."""
1544 return get_dcp_group().world_size
1545
1546
1547 -def get_decode_context_model_parallel_rank():
1547 +def get_decode_context_model_parallel_rank() -> int:
1548 """Return my rank for the decode context model parallel group."""
1549 return get_dcp_group().rank_in_group
1550
vllm/model_executor/layers/rotary_embedding/__init__.py
+3 -1
@@ -20,7 +20,9 @@ from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding
20 from .xdrope import XDRotaryEmbedding
21 from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
22
23 -_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}
23 +_ROPE_DICT: dict[tuple[Any, ...], RotaryEmbedding] = {}
24 +
25 +__all__ = ["RotaryEmbedding"]
26
27
28 def get_rope(