@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 243 lines 7.99 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import patch
5
6 import pytest
7 import torch
8
9 from vllm.platforms import current_platform
10 from vllm.utils.flashinfer import (
11 can_use_trtllm_attention,
12 supports_trtllm_attention,
13 use_trtllm_attention,
14 )
15
16 if not current_platform.is_cuda():
17 pytest.skip(
18 "TRTLLM attention is only supported on CUDA platforms.",
19 allow_module_level=True,
20 )
21
22 MODEL_CONFIGS = {
23 "Llama-3-70B": dict(num_qo_heads=64, num_kv_heads=8),
24 "Llama-3-8B": dict(num_qo_heads=32, num_kv_heads=8),
25 "Qwen2.5-0.5B": dict(num_qo_heads=14, num_kv_heads=2),
26 "Mistral-7B": dict(num_qo_heads=32, num_kv_heads=8),
27 "Gemma-2-9B": dict(num_qo_heads=8, num_kv_heads=4),
28 "Falcon-40B": dict(num_qo_heads=128, num_kv_heads=8),
29 }
30
31
32 def get_config(model: str) -> dict:
33 """Return the attention config for a model."""
34 return MODEL_CONFIGS[model]
35
36
37 DEFAULT_KWARGS = dict(
38 **get_config("Llama-3-70B"),
39 num_tokens=128,
40 max_seq_len=4096,
41 dcp_world_size=1,
42 kv_cache_dtype="auto",
43 q_dtype=torch.bfloat16,
44 is_prefill=False,
45 force_use_trtllm=None,
46 has_sinks=False,
47 has_spec=False,
48 )
49
50
51 def _call(**overrides) -> bool:
52 kwargs = {**DEFAULT_KWARGS, **overrides}
53 return use_trtllm_attention(**kwargs)
54
55
56 @pytest.fixture(autouse=True)
57 def _clear_supports_cache():
58 """Clear functools.cache to ensure each test runs independently."""
59 supports_trtllm_attention.cache_clear()
60
61
62 # supports_trtllm_attention
63
64
65 @patch("vllm.envs.VLLM_BATCH_INVARIANT", True)
66 def test_supports_batch_invariant_disables():
67 assert supports_trtllm_attention() is False
68
69
70 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
71 @patch(
72 "vllm.utils.flashinfer.current_platform.is_device_capability_family",
73 return_value=True,
74 )
75 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
76 def test_supports_sm100_with_artifactory(_art, _cap):
77 assert supports_trtllm_attention() is True
78
79
80 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
81 @patch(
82 "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=False
83 )
84 @patch(
85 "vllm.utils.flashinfer.current_platform.is_device_capability_family",
86 return_value=False,
87 )
88 def test_supports_unsupported_platform(_family, _cap):
89 assert supports_trtllm_attention() is False
90
91
92 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
93 @patch("vllm.utils.flashinfer.current_platform.is_device_capability", return_value=True)
94 @patch(
95 "vllm.utils.flashinfer.current_platform.is_device_capability_family",
96 return_value=False,
97 )
98 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
99 def test_supports_sm90_decode_only(_art, _family, _cap):
100 assert supports_trtllm_attention(is_prefill=False) is True
101 assert supports_trtllm_attention(is_prefill=True) is False
102
103
104 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
105 @patch(
106 "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=False
107 )
108 @patch(
109 "vllm.utils.flashinfer.current_platform.is_device_capability_family",
110 side_effect=lambda capability: capability == 120,
111 )
112 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
113 def test_supports_sm12x_decode_only(_art, _family, _cap):
114 assert supports_trtllm_attention(is_prefill=False) is True
115 assert supports_trtllm_attention(is_prefill=True) is False
116
117
118 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
119 @patch(
120 "vllm.utils.flashinfer.current_platform.is_device_capability_family",
121 return_value=True,
122 )
123 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)
124 def test_supports_sm100_without_artifactory(_art, _cap):
125 assert supports_trtllm_attention() is False
126
127
128 # can_use_trtllm_attention
129
130
131 @patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=False)
132 def test_can_use_force_disabled(_mock):
133 cfg = get_config("Llama-3-70B")
134 assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
135
136
137 @patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
138 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
139 def test_can_use_compatible_heads(_sup, _force):
140 cfg = get_config("Llama-3-70B")
141 assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is True
142
143
144 @patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
145 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
146 def test_can_use_incompatible_heads(_sup, _force):
147 assert can_use_trtllm_attention(40, 6) is False
148
149
150 @pytest.mark.parametrize("model", list(MODEL_CONFIGS.keys()))
151 @patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
152 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
153 def test_can_use_platform_unsupported(_sup, _force, model):
154 cfg = get_config(model)
155 assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
156
157
158 # use_trtllm_attention
159
160
161 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
162 def test_use_force_off(_mock):
163 assert _call(force_use_trtllm=False) is False
164
165
166 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
167 def test_use_dcp_fallback(_mock):
168 assert _call(dcp_world_size=2) is False
169
170
171 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
172 def test_use_dcp_fallback_prefill(_mock):
173 # TRTLLM prefill has no cross-rank LSE combine for DCP-sharded KV.
174 assert _call(dcp_world_size=2, is_prefill=True) is False
175
176
177 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
178 def test_use_dcp_fallback_prefill_force_on_still_false(_mock):
179 assert _call(dcp_world_size=2, is_prefill=True, force_use_trtllm=True) is False
180
181
182 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
183 def test_use_platform_unsupported(_mock):
184 assert _call() is False
185
186
187 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
188 def test_use_platform_unsupported_force_on_still_false(_mock):
189 assert _call(force_use_trtllm=True) is False
190
191
192 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
193 def test_use_incompatible_heads(_mock):
194 assert _call(num_qo_heads=40, num_kv_heads=6) is False
195
196
197 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
198 def test_use_incompatible_heads_force_on_still_false(_mock):
199 assert _call(num_qo_heads=40, num_kv_heads=6, force_use_trtllm=True) is False
200
201
202 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
203 def test_use_spec_decode_enables(_mock):
204 assert _call(has_spec=True, is_prefill=False) is True
205
206
207 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
208 @patch(
209 "vllm.utils.flashinfer.current_platform.fp8_dtype",
210 return_value=torch.float8_e4m3fn,
211 )
212 def test_use_fp8_query_forces_trtllm(_fp8, _sup):
213 assert _call(q_dtype=torch.float8_e4m3fn) is True
214
215
216 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
217 def test_use_sinks_force_trtllm(_mock):
218 assert _call(has_sinks=True) is True
219
220
221 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
222 def test_use_auto_prefill_kv_auto(_mock):
223 assert _call(is_prefill=True, kv_cache_dtype="auto") is True
224
225
226 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
227 def test_use_auto_prefill_kv_fp8(_mock):
228 assert _call(is_prefill=True, kv_cache_dtype="fp8") is False
229
230
231 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
232 def test_use_auto_decode_small_batch(_mock):
233 assert _call(is_prefill=False, num_tokens=128, kv_cache_dtype="auto") is True
234
235
236 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
237 def test_use_auto_decode_large_batch(_mock):
238 assert _call(is_prefill=False, num_tokens=512, kv_cache_dtype="auto") is False
239
240
241 @patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
242 def test_use_force_on(_mock):
243 assert _call(force_use_trtllm=True) is True