@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 180 lines 5.54 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 from vllm.platforms.interface import DeviceCapability
6 from vllm.v1.attention.backend import (
7 AttentionBackend,
8 AttentionImpl,
9 )
10 from vllm.v1.attention.backends.registry import (
11 AttentionBackendEnum,
12 MambaAttentionBackendEnum,
13 register_backend,
14 )
15 from vllm.v1.attention.selector import AttentionSelectorConfig
16
17
18 class CustomAttentionImpl(AttentionImpl):
19 """Mock custom attention implementation for testing."""
20
21 def __init__(self, *args, **kwargs):
22 super().__init__()
23
24 def forward(self, *args, **kwargs):
25 """Mock forward pass."""
26 pass
27
28
29 class CustomAttentionBackend(AttentionBackend):
30 """Mock custom attention backend for testing."""
31
32 @staticmethod
33 def get_name():
34 return "CUSTOM"
35
36 @staticmethod
37 def get_impl_cls():
38 return CustomAttentionImpl
39
40 @staticmethod
41 def get_builder_cls():
42 """Mock builder class."""
43 return None
44
45
46 class CustomMambaAttentionImpl(AttentionImpl):
47 """Mock custom mamba attention implementation for testing."""
48
49 def __init__(self, *args, **kwargs):
50 super().__init__()
51
52 def forward(self, *args, **kwargs):
53 """Mock forward pass."""
54 pass
55
56
57 class CustomMambaAttentionBackend(AttentionBackend):
58 """Mock custom mamba attention backend for testing."""
59
60 @staticmethod
61 def get_name():
62 return "CUSTOM_MAMBA"
63
64 @staticmethod
65 def get_impl_cls():
66 return CustomMambaAttentionImpl
67
68 @staticmethod
69 def get_builder_cls():
70 """Mock builder class."""
71 return None
72
73
74 def test_custom_attention_backend_requires_dcp_opt_in():
75 """A new implementation remains usable without implicitly opting into DCP."""
76 config = AttentionSelectorConfig(
77 head_size=64,
78 dtype=torch.float16,
79 kv_cache_dtype="auto",
80 block_size=16,
81 )
82 kwargs = dict(device_capability=DeviceCapability(9, 0), **config._asdict())
83 assert CustomAttentionBackend.validate_configuration(**kwargs) == []
84
85 kwargs["use_dcp"] = True
86 assert CustomAttentionBackend.validate_configuration(**kwargs) == [
87 "DCP not supported"
88 ]
89
90
91 def test_custom_is_not_alias_of_any_backend():
92 # Get all members of AttentionBackendEnum
93 all_backends = list(AttentionBackendEnum)
94
95 # Find any aliases of CUSTOM
96 aliases = []
97 for backend in all_backends:
98 if backend.name != "CUSTOM" and backend is AttentionBackendEnum.CUSTOM:
99 aliases.append(backend.name)
100
101 # CUSTOM should not be an alias of any other backend
102 assert len(aliases) == 0, (
103 f"BUG! CUSTOM is an alias of: {', '.join(aliases)}!\n"
104 f"CUSTOM.value = {repr(AttentionBackendEnum.CUSTOM.value)}\n"
105 f"This happens when CUSTOM has the same value as another backend.\n"
106 f"When you register to CUSTOM, you're actually registering to {aliases[0]}!\n"
107 f"All backend values:\n"
108 + "\n".join(f" {b.name}: {repr(b.value)}" for b in all_backends)
109 )
110
111 # Verify CUSTOM has its own unique identity
112 assert AttentionBackendEnum.CUSTOM.name == "CUSTOM", (
113 f"CUSTOM.name should be 'CUSTOM', but got '{AttentionBackendEnum.CUSTOM.name}'"
114 )
115
116
117 def test_register_custom_backend_with_class_path():
118 # Register with explicit class path
119 register_backend(
120 backend=AttentionBackendEnum.CUSTOM,
121 class_path="tests.test_attention_backend_registry.CustomAttentionBackend",
122 is_mamba=False,
123 )
124
125 # Check that CUSTOM backend is registered
126 assert AttentionBackendEnum.CUSTOM.is_overridden(), (
127 "CUSTOM should be overridden after registration"
128 )
129
130 # Get the registered class path
131 class_path = AttentionBackendEnum.CUSTOM.get_path()
132 assert class_path == "tests.test_attention_backend_registry.CustomAttentionBackend"
133
134 # Get the backend class
135 backend_cls = AttentionBackendEnum.CUSTOM.get_class()
136 assert backend_cls.get_name() == "CUSTOM"
137 assert backend_cls.get_impl_cls() == CustomAttentionImpl
138
139
140 def test_mamba_custom_is_not_alias_of_any_backend():
141 # Get all mamba backends
142 all_backends = list(MambaAttentionBackendEnum)
143
144 # Find any aliases of CUSTOM
145 aliases = []
146 for backend in all_backends:
147 if backend.name != "CUSTOM" and backend is MambaAttentionBackendEnum.CUSTOM:
148 aliases.append(backend.name)
149
150 # CUSTOM should not be an alias of any other backend
151 assert len(aliases) == 0, (
152 f"BUG! MambaAttentionBackendEnum.CUSTOM is an alias of: {', '.join(aliases)}!\n"
153 f"CUSTOM.value = {repr(MambaAttentionBackendEnum.CUSTOM.value)}\n"
154 f"All mamba backend values:\n"
155 + "\n".join(f" {b.name}: {repr(b.value)}" for b in all_backends)
156 )
157
158
159 def test_register_custom_mamba_backend_with_class_path():
160 # Register with explicit class path
161 register_backend(
162 backend=MambaAttentionBackendEnum.CUSTOM,
163 class_path="tests.test_attention_backend_registry.CustomMambaAttentionBackend",
164 is_mamba=True,
165 )
166
167 # Check that the backend is registered
168 assert MambaAttentionBackendEnum.CUSTOM.is_overridden()
169
170 # Get the registered class path
171 class_path = MambaAttentionBackendEnum.CUSTOM.get_path()
172 assert (
173 class_path
174 == "tests.test_attention_backend_registry.CustomMambaAttentionBackend"
175 )
176
177 # Get the backend class
178 backend_cls = MambaAttentionBackendEnum.CUSTOM.get_class()
179 assert backend_cls.get_name() == "CUSTOM_MAMBA"
180 assert backend_cls.get_impl_cls() == CustomMambaAttentionImpl