@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 174 lines 5.46 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for MLA prefill backend registry."""
4
5 import pytest
6 import torch
7
8 from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
9 from vllm.v1.attention.backends.mla.prefill.registry import (
10 MLAPrefillBackendEnum,
11 register_mla_prefill_backend,
12 )
13
14
15 class CustomMLAPrefillBackend(MLAPrefillBackend):
16 """Mock custom MLA prefill backend for testing."""
17
18 supported_dtypes = [torch.bfloat16, torch.float16]
19
20 @staticmethod
21 def get_name() -> str:
22 return "CUSTOM"
23
24 def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
25 raise NotImplementedError
26
27 def run_prefill_context_chunk(self, chunk, q, k, v, out=None):
28 raise NotImplementedError
29
30
31 def test_prefill_backend_clone_has_isolated_metadata():
32 backend = CustomMLAPrefillBackend(
33 num_heads=4,
34 scale=0.5,
35 kv_lora_rank=8,
36 qk_nope_head_dim=16,
37 qk_rope_head_dim=8,
38 v_head_dim=32,
39 vllm_config=object(),
40 )
41
42 clone = backend.clone()
43
44 assert isinstance(clone, CustomMLAPrefillBackend)
45 assert clone is not backend
46 assert clone.num_heads == backend.num_heads
47 assert clone.scale == backend.scale
48 backend._prefill_metadata = object()
49 clone._prefill_metadata = object()
50 assert clone._prefill_metadata is not backend._prefill_metadata
51
52
53 @pytest.fixture(autouse=True)
54 def cleanup_overrides():
55 """Clear any overrides after each test."""
56 yield
57 for member in MLAPrefillBackendEnum:
58 member.clear_override()
59
60
61 def test_custom_is_not_alias_of_any_backend():
62 all_backends = list(MLAPrefillBackendEnum)
63
64 aliases = []
65 for backend in all_backends:
66 if backend.name != "CUSTOM" and backend is MLAPrefillBackendEnum.CUSTOM:
67 aliases.append(backend.name)
68
69 assert len(aliases) == 0, (
70 f"BUG! CUSTOM is an alias of: {', '.join(aliases)}!\n"
71 f"CUSTOM.value = {repr(MLAPrefillBackendEnum.CUSTOM.value)}\n"
72 f"All MLA prefill backend values:\n"
73 + "\n".join(f" {b.name}: {repr(b.value)}" for b in all_backends)
74 )
75
76 assert MLAPrefillBackendEnum.CUSTOM.name == "CUSTOM"
77
78
79 def test_custom_unregistered_raises():
80 with pytest.raises(ValueError, match="must be registered before use"):
81 MLAPrefillBackendEnum.CUSTOM.get_path()
82
83
84 def test_register_custom_backend_with_class_path():
85 register_mla_prefill_backend(
86 backend=MLAPrefillBackendEnum.CUSTOM,
87 class_path=(
88 "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
89 ),
90 )
91
92 assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
93
94 class_path = MLAPrefillBackendEnum.CUSTOM.get_path()
95 assert class_path == (
96 "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
97 )
98
99 backend_cls = MLAPrefillBackendEnum.CUSTOM.get_class()
100 assert backend_cls.get_name() == "CUSTOM"
101
102
103 def test_register_custom_backend_as_decorator():
104 @register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM)
105 class DecoratedPrefillBackend(MLAPrefillBackend):
106 supported_dtypes = [torch.bfloat16]
107
108 @staticmethod
109 def get_name() -> str:
110 return "DECORATED"
111
112 def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
113 raise NotImplementedError
114
115 def run_prefill_context_chunk(self, chunk, q, k, v, out=None):
116 raise NotImplementedError
117
118 assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
119 assert "DecoratedPrefillBackend" in MLAPrefillBackendEnum.CUSTOM.get_path()
120
121
122 def test_override_existing_backend():
123 original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
124
125 register_mla_prefill_backend(
126 backend=MLAPrefillBackendEnum.FLASH_ATTN,
127 class_path=(
128 "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
129 ),
130 )
131
132 assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
133 assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() != original_path
134
135 backend_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
136 assert backend_cls.get_name() == "CUSTOM"
137
138
139 def test_clear_override():
140 original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
141
142 register_mla_prefill_backend(
143 backend=MLAPrefillBackendEnum.FLASH_ATTN,
144 class_path=(
145 "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
146 ),
147 )
148 assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
149
150 MLAPrefillBackendEnum.FLASH_ATTN.clear_override()
151 assert not MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
152 assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() == original_path
153
154
155 def test_unknown_backend_name_raises():
156 with pytest.raises(ValueError, match="Unknown MLA prefill backend"):
157 MLAPrefillBackendEnum["NONEXISTENT"]
158
159
160 def test_rocm_aiter_fa_registered():
161 """ROCM_AITER_FA is a known backend pointing at the AITER FA class."""
162 assert "ROCM_AITER_FA" in MLAPrefillBackendEnum.__members__
163
164 path = MLAPrefillBackendEnum.ROCM_AITER_FA.get_path()
165 assert path == (
166 "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn."
167 "AiterFlashAttnPrefillBackend"
168 )
169
170 backend_cls = MLAPrefillBackendEnum.ROCM_AITER_FA.get_class()
171 assert backend_cls.get_name() == "ROCM_AITER_FA"
172 # The AITER FA path is the fp16/bf16 generic-varlen prefill path.
173 assert backend_cls.supports_dtype(torch.bfloat16)
174 assert backend_cls.supports_dtype(torch.float16)