@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 296 lines 10.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
4
5 import contextlib
6 import json
7 from unittest.mock import MagicMock, patch
8
9 import pytest
10 import torch
11
12 from vllm.model_executor.layers.attention.mm_encoder_attention import (
13 _FP8_AMAX_HISTORY_LEN,
14 _FP8_MAX,
15 )
16 from vllm.platforms import current_platform
17 from vllm.utils.flashinfer import (
18 is_flashinfer_cudnn_fp8_prefill_attn_supported,
19 )
20
21 pytestmark = pytest.mark.skipif(
22 not is_flashinfer_cudnn_fp8_prefill_attn_supported(),
23 reason="FlashInfer cuDNN FP8 prefill attention not supported",
24 )
25
26 LAYER_0 = "visual.blocks.0.attn.attn"
27 LAYER_1 = "visual.blocks.1.attn.attn"
28 NUM_HEADS = 16
29 HEAD_DIM = 72
30
31
32 def _is_mi3xx() -> bool:
33 if not current_platform.is_rocm():
34 return False
35
36 from vllm.platforms.rocm import on_mi3xx
37
38 return on_mi3xx()
39
40
41 @contextlib.contextmanager
42 def _build_attention(mm_config, attn_backend=None):
43 """Yield an MMEncoderAttention with the given multimodal config.
44
45 The VllmConfig context stays active while the test runs so that
46 ``get_multimodal_config()`` calls during the forward path resolve. Also
47 invokes ``process_weights_after_loading`` to simulate the model loader's
48 auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
49 """
50 from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
51 from vllm.model_executor.layers.attention.mm_encoder_attention import (
52 MMEncoderAttention,
53 )
54 from vllm.v1.attention.backends.registry import AttentionBackendEnum
55
56 if attn_backend is None:
57 attn_backend = AttentionBackendEnum.FLASHINFER
58
59 if (
60 attn_backend == AttentionBackendEnum.FLASHINFER
61 and not is_flashinfer_cudnn_fp8_prefill_attn_supported()
62 ):
63 yield None
64 return
65
66 vllm_config = VllmConfig()
67 vllm_config.model_config = MagicMock(spec=ModelConfig, multimodal_config=mm_config)
68
69 with (
70 set_current_vllm_config(vllm_config),
71 patch(
72 "vllm.model_executor.layers.attention.mm_encoder_attention"
73 ".get_vit_attn_backend",
74 return_value=attn_backend,
75 ),
76 ):
77 attn = MMEncoderAttention(
78 num_heads=NUM_HEADS,
79 head_size=HEAD_DIM,
80 prefix=LAYER_0,
81 )
82 attn.process_weights_after_loading(torch.bfloat16)
83 yield attn
84
85
86 @pytest.fixture
87 def _make_attention():
88 """Create an MMEncoderAttention with dynamic FP8 scaling."""
89 from vllm.config.multimodal import MultiModalConfig
90
91 with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
92 yield attn
93
94
95 @pytest.fixture
96 def _make_static_attention(tmp_path):
97 """Create an MMEncoderAttention with static FP8 scales from a file."""
98 from vllm.config.multimodal import MultiModalConfig
99
100 scale_file = tmp_path / "scales.json"
101 scale_file.write_text(
102 json.dumps(
103 {
104 LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
105 LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
106 }
107 )
108 )
109 with _build_attention(
110 MultiModalConfig(
111 mm_encoder_attn_dtype="fp8",
112 mm_encoder_fp8_scale_path=str(scale_file),
113 )
114 ) as attn:
115 yield attn
116
117
118 def test_dynamic_scaling_updates_scales(_make_attention) -> None:
119 """Verify that _record_amax_and_update_scales updates scale buffers."""
120 attn = _make_attention
121 if attn is None or not attn.fp8_enabled:
122 pytest.skip("FP8 attention not available (FlashInfer backend required)")
123
124 attn = attn.to("cuda")
125
126 S, H, D = 32, NUM_HEADS, HEAD_DIM
127 q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
128 k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
129 v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
130
131 attn._record_amax_and_update_scales(q, k, v)
132
133 expected_q_scale = 2.0 / _FP8_MAX
134 expected_k_scale = 3.0 / _FP8_MAX
135 expected_v_scale = 4.0 / _FP8_MAX
136
137 torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
138 torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
139 torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
140
141
142 def test_circular_buffer_wraps(_make_attention) -> None:
143 """Verify the amax circular buffer wraps at HISTORY_LEN."""
144 attn = _make_attention
145 if attn is None or not attn.fp8_enabled:
146 pytest.skip("FP8 attention not available (FlashInfer backend required)")
147
148 attn = attn.to("cuda")
149 S, H, D = 16, NUM_HEADS, HEAD_DIM
150
151 for i in range(_FP8_AMAX_HISTORY_LEN + 2):
152 mag = float(i + 1)
153 q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
154 k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
155 v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
156 attn._record_amax_and_update_scales(q, k, v)
157
158 assert attn._fp8_amax_pos == 2
159
160 expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
161 expected_scale = expected_max / _FP8_MAX
162 torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
163
164
165 def test_static_scales_loaded(_make_static_attention) -> None:
166 """Verify static scales are loaded from the JSON file."""
167 attn = _make_static_attention
168 if attn is None or not attn.fp8_enabled:
169 pytest.skip("FP8 attention not available (FlashInfer backend required)")
170
171 assert attn.fp8_enabled
172 assert not attn._fp8_dynamic_scale
173
174 # Layer 0 scales (the layer this attention was created with).
175 assert attn._fp8_q_scale.item() == 224.0
176 assert attn._fp8_k_scale.item() == 198.0
177 assert attn._fp8_v_scale.item() == 210.0
178
179 assert not attn.skip_scale_q
180 assert not attn.skip_scale_k
181 assert not attn.skip_scale_v
182
183 # No amax history buffers for static scaling.
184 assert not hasattr(attn, "_fp8_q_amax")
185
186
187 @pytest.mark.skipif(
188 not _is_mi3xx(),
189 reason="AITER FP8 attention requires MI300/MI350",
190 )
191 def test_aiter_static_scales_loaded(tmp_path) -> None:
192 """Verify AITER reuses the existing static FP8 scale loading path."""
193 from vllm.config.multimodal import MultiModalConfig
194 from vllm.v1.attention.backends.registry import AttentionBackendEnum
195
196 scale_file = tmp_path / "aiter_scales.json"
197 scale_file.write_text(json.dumps({LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0}}))
198 mm_config = MultiModalConfig(
199 mm_encoder_attn_dtype="fp8",
200 mm_encoder_fp8_scale_path=str(scale_file),
201 )
202
203 with _build_attention(mm_config, AttentionBackendEnum.ROCM_AITER_FA) as attn:
204 assert attn is not None
205 assert attn.fp8_enabled
206 assert not attn._fp8_dynamic_scale
207 assert attn._fp8_q_scale.item() == 224.0
208 assert attn._fp8_k_scale.item() == 198.0
209 assert attn._fp8_v_scale.item() == 210.0
210
211
212 def test_static_scales_missing_layer(tmp_path) -> None:
213 """Verify error when requested layer is not in the scale file."""
214 from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
215 from vllm.config.multimodal import MultiModalConfig
216 from vllm.v1.attention.backends.registry import AttentionBackendEnum
217
218 if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
219 pytest.skip("FlashInfer cuDNN not available")
220
221 scale_file = tmp_path / "wrong_layer.json"
222 scale_file.write_text(
223 json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
224 )
225 mm_config = MultiModalConfig(
226 mm_encoder_attn_dtype="fp8",
227 mm_encoder_fp8_scale_path=str(scale_file),
228 )
229 vllm_config = VllmConfig()
230 vllm_config.model_config = MagicMock(spec=ModelConfig, multimodal_config=mm_config)
231
232 from vllm.model_executor.layers.attention.mm_encoder_attention import (
233 MMEncoderAttention,
234 )
235
236 with (
237 set_current_vllm_config(vllm_config),
238 patch(
239 "vllm.model_executor.layers.attention.mm_encoder_attention"
240 ".get_vit_attn_backend",
241 return_value=AttentionBackendEnum.FLASHINFER,
242 ),
243 ):
244 attn = MMEncoderAttention(
245 num_heads=NUM_HEADS,
246 head_size=HEAD_DIM,
247 prefix=LAYER_0,
248 )
249 with pytest.raises(ValueError, match="scales not found for layer"):
250 attn.process_weights_after_loading(torch.bfloat16)
251
252
253 def test_dynamic_scales_auto_save(tmp_path) -> None:
254 """Verify scales are saved to disk after the amax buffer fills."""
255 import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
256 from vllm.config.multimodal import MultiModalConfig
257
258 if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
259 pytest.skip("FlashInfer cuDNN not available")
260
261 # Reset module-level state between runs (other tests may have left
262 # state behind after triggering a save).
263 _mod._fp8_scale_save_path = None
264 _mod._fp8_saved_scale_refs.clear()
265
266 save_file = tmp_path / "auto_scales.json"
267 with _build_attention(
268 MultiModalConfig(
269 mm_encoder_attn_dtype="fp8",
270 mm_encoder_fp8_scale_save_path=str(save_file),
271 )
272 ) as attn:
273 if attn is None or not attn.fp8_enabled:
274 pytest.skip("FP8 attention not available")
275
276 attn = attn.to("cuda")
277 S, H, D = 16, NUM_HEADS, HEAD_DIM
278
279 # Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
280 for i in range(_FP8_AMAX_HISTORY_LEN):
281 mag = float(i + 1)
282 q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
283 k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
284 v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
285 attn._record_amax_and_update_scales(q, k, v)
286
287 # File should have been written on the 16th call (buffer wrap).
288 assert save_file.is_file(), "Scale file was not saved"
289 scales = json.loads(save_file.read_text())
290 assert LAYER_0 in scales
291 assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
292 for val in scales[LAYER_0].values():
293 assert isinstance(val, float) and val > 0
294
295 # Path is cleared after the one-shot save fires.
296 assert _mod._fp8_scale_save_path is None