@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 146 lines 4.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for CUDA attention backend selection and probe error handling.
4
5 Environment-shaped probe failures (missing packages, unreadable caches,
6 broken driver installs) must mark the backend unavailable rather than
7 terminating engine init; programming errors must still propagate.
8 See https://github.com/vllm-project/vllm/issues/51658.
9 """
10
11 from unittest.mock import MagicMock, patch
12
13 import pytest
14 import torch
15
16 from vllm.platforms import current_platform
17 from vllm.platforms.cuda import CudaPlatform
18 from vllm.platforms.interface import DeviceCapability
19 from vllm.v1.attention.backends.registry import AttentionBackendEnum
20 from vllm.v1.attention.selector import AttentionSelectorConfig
21
22 pytestmark = pytest.mark.skipif(
23 not current_platform.is_cuda(), reason="CUDA-specific tests"
24 )
25
26 SELECTOR_CONFIG = AttentionSelectorConfig(
27 head_size=64,
28 dtype=torch.float16,
29 kv_cache_dtype=None,
30 block_size=16,
31 )
32
33 SM90 = DeviceCapability(major=9, minor=0)
34
35
36 @pytest.mark.parametrize(
37 "exc",
38 [
39 ImportError("flashinfer is not installed"),
40 PermissionError(13, "Permission denied", "/root/.cache/flashinfer"),
41 OSError("libcuda.so.1: cannot open shared object file"),
42 ],
43 )
44 def test_get_valid_backends_records_environment_failure(exc):
45 with patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=exc):
46 valid, invalid_reasons = CudaPlatform.get_valid_backends(
47 device_capability=SM90,
48 attn_selector_config=SELECTOR_CONFIG,
49 num_heads=32,
50 )
51 assert valid == []
52 assert invalid_reasons
53 for _priority, reasons in invalid_reasons.values():
54 assert reasons == [f"{type(exc).__name__}: {exc}"]
55
56
57 @pytest.mark.parametrize(
58 "exc",
59 [
60 RuntimeError("CUDA error: no kernel image is available"),
61 AttributeError("module has no attribute 'get_builder_cls'"),
62 KeyboardInterrupt(),
63 ],
64 )
65 def test_get_valid_backends_propagates_unexpected_errors(exc):
66 with (
67 patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=exc),
68 pytest.raises(type(exc)),
69 ):
70 CudaPlatform.get_valid_backends(
71 device_capability=SM90,
72 attn_selector_config=SELECTOR_CONFIG,
73 num_heads=32,
74 )
75
76
77 def test_get_valid_backends_keeps_probing_after_failure():
78 healthy = MagicMock()
79 healthy.validate_configuration.return_value = []
80 side_effects = [OSError("probe failed")] + [healthy] * 32
81
82 with patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=side_effects):
83 valid, invalid_reasons = CudaPlatform.get_valid_backends(
84 device_capability=SM90,
85 attn_selector_config=SELECTOR_CONFIG,
86 num_heads=32,
87 )
88 assert len(invalid_reasons) == 1
89 assert valid
90
91
92 def test_sm90_nope_mla_prefers_flashinfer_without_changing_rope_order():
93 backend_cls = MagicMock()
94 backend_cls.validate_configuration.return_value = []
95 sparse_backends = {
96 AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
97 AttentionBackendEnum.FLASHMLA_SPARSE,
98 AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90,
99 }
100
101 def sparse_order(head_size: int) -> list[AttentionBackendEnum]:
102 config = SELECTOR_CONFIG._replace(
103 head_size=head_size,
104 use_mla=True,
105 use_sparse=True,
106 )
107 with patch(
108 "vllm.platforms.cuda._get_attn_backend_class",
109 return_value=backend_cls,
110 ):
111 valid, _ = CudaPlatform.get_valid_backends(
112 device_capability=SM90,
113 attn_selector_config=config,
114 num_heads=32,
115 )
116 return [
117 candidate.backend
118 for candidate in valid
119 if candidate.backend in sparse_backends
120 ]
121
122 assert sparse_order(512) == [
123 AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90,
124 AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
125 AttentionBackendEnum.FLASHMLA_SPARSE,
126 ]
127 assert sparse_order(576) == [
128 AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
129 AttentionBackendEnum.FLASHMLA_SPARSE,
130 AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90,
131 ]
132
133
134 def test_selected_backend_probe_failure_raises_value_error_with_cause():
135 exc = OSError("libcuda.so.1: cannot open shared object file")
136 with (
137 patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=exc),
138 patch.object(CudaPlatform, "get_device_capability", return_value=SM90),
139 pytest.raises(ValueError, match="OSError") as excinfo,
140 ):
141 CudaPlatform.get_attn_backend_cls(
142 selected_backend=AttentionBackendEnum.FLASH_ATTN,
143 attn_selector_config=SELECTOR_CONFIG,
144 num_heads=32,
145 )
146 assert excinfo.value.__cause__ is exc