@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 150 lines 4.73 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import itertools
4 from collections.abc import Callable, Iterable
5 from typing import Any, NamedTuple
6
7 import pytest
8 import regex as re
9
10 from vllm.platforms import current_platform
11 from vllm.v1.attention.backends.registry import AttentionBackendEnum
12
13
14 class Matches(NamedTuple):
15 # simple pointwise
16 aiter_rms_quant_fusion: int = 0
17 rms_quant_fusion: int = 0
18 act_quant_fusion: int = 0
19 norm_rope_fusion: int = 0
20 attn_quant_fusion: int = 0
21 # distributed
22 ar_rms_fusion: int = 0
23 aiter_ar_rms_fusion: int = 0
24 sequence_parallel: int = 0
25 async_tp: int = 0
26
27
28 class ModelFusionInfo(NamedTuple):
29 model_name: str
30 matches: Callable[[int], Matches]
31 """Given number of hidden layers, produces the matches object"""
32 model_kwargs: dict[str, Any] = {}
33 hf_overrides: Callable[[int], dict] = lambda n: {"num_hidden_layers": n}
34
35
36 class AttentionBackendCase(NamedTuple):
37 backend: AttentionBackendEnum
38 model_kwargs: dict[str, Any] = {}
39 """Additional args required for attn+quant fusion"""
40
41
42 is_blackwell = lambda: current_platform.is_device_capability_family(100)
43 """Are we running on Blackwell, a lot of tests depend on it"""
44
45
46 def nvfp4_kernel_exposes_input_quant_key() -> bool:
47 """Check if the NVFP4 kernel selected on this platform exposes input_quant_key.
48
49 FlashInferCuteDslNvFp4LinearKernel does not expose input_quant_key() due to
50 layout incompatibility between the manual fusion kernel (silu_and_mul_nvfp4_quant)
51 output format and the cutedsl backend's expected input format.
52
53 FlashInferCutlassNvFp4LinearKernel does expose input_quant_key() and supports
54 manual fusion.
55 """
56 if not current_platform.is_cuda():
57 return False
58
59 try:
60 from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
61 FlashInferCuteDslNvFp4LinearKernel,
62 )
63
64 is_supported, _ = FlashInferCuteDslNvFp4LinearKernel.is_supported()
65 if is_supported:
66 return False
67 except ImportError:
68 pass
69
70 try:
71 from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
72 FlashInferCutlassNvFp4LinearKernel,
73 )
74
75 is_supported, _ = FlashInferCutlassNvFp4LinearKernel.is_supported()
76 if is_supported:
77 return True
78 except ImportError:
79 pass
80
81 # Fallback: assume no manual fusion support
82 return False
83
84
85 def custom_ops_combos(*custom_ops: str) -> Iterable[str]:
86 """Generate all combinations of custom ops for parametrization."""
87 custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops]
88 for op_list in itertools.product(*custom_ops_lists):
89 yield ",".join(op_list)
90
91
92 # Quick inline validation
93 assert list(custom_ops_combos("silu_and_mul")) == ["-silu_and_mul", "+silu_and_mul"]
94 assert list(custom_ops_combos("quant_fp8", "rms_norm")) == [
95 "-quant_fp8,-rms_norm",
96 "-quant_fp8,+rms_norm",
97 "+quant_fp8,-rms_norm",
98 "+quant_fp8,+rms_norm",
99 ]
100
101
102 def has_cuda_graph_wrapper_metadata() -> bool:
103 from importlib import import_module
104
105 try:
106 module = import_module("torch._inductor.utils")
107 module.CUDAGraphWrapperMetadata # noqa B018
108 except AttributeError:
109 return False
110 return True
111
112
113 INDUCTOR_GRAPH_PARTITION = [
114 pytest.param(
115 True,
116 marks=pytest.mark.skipif(
117 not has_cuda_graph_wrapper_metadata(),
118 reason="torch version does not support Inductor partition",
119 ),
120 id="inductor_partition",
121 ),
122 pytest.param(False, id="dynamo_partition"),
123 ]
124
125 FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
126 "aiter_rms_quant_fusion": re.compile(
127 r"RocmAiterRMSNormQuantFusionPass Replaced (\d+) patterns"
128 ),
129 "rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),
130 "act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),
131 # The combined pass can consume these sites before the standalone pass.
132 "norm_rope_fusion": re.compile(
133 r"(?:qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on|"
134 r"qk_norm_rope_kvcache_fusion.py:\d+] "
135 r"QK-Norm\+RoPE(?:/MRoPE)?\+KVCache fusion: replaced) (\d+)"
136 ),
137 "attn_quant_fusion": re.compile(
138 r"attn_quant_fusion.py:\d+] Fused quant onto (\d+) attention nodes"
139 ),
140 "ar_rms_fusion": re.compile(
141 r"allreduce_rms_fusion.py:\d+] Replaced (\d+) patterns"
142 ),
143 "aiter_ar_rms_fusion": re.compile(
144 r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns"
145 ),
146 "sequence_parallel": re.compile(
147 r"sequence_parallelism.py:\d+] Replaced (\d+) patterns"
148 ),
149 "async_tp": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"),
150 }