@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 229 lines 7.93 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # NOTE: if your PR has broken one of the tests here (sorry),
5 # kindly patch the corresponding integration in
6 # /vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
7 # or reach out to @aposataC for assistance
8
9 # Assumption vs. Correctness Tests:
10 # these unit tests do *not* test correctness of LMCache-side or vLLM-side logic
11 # it is to ensure that assumptions LMCache makes about vLLM's interface are stable
12
13 import pytest
14
15 from vllm.platforms import current_platform
16
17
18 def assumes(obj, attr, is_callable=False, is_instance_of=None):
19 import inspect
20 from dataclasses import is_dataclass
21
22 assumption_msg = (
23 f"LMCache connector currently assumes that {obj} has a(n) {attr} attribute"
24 )
25 if hasattr(obj, attr):
26 attr_value = getattr(obj, attr)
27 elif is_dataclass(obj) and attr in getattr(obj, "__dataclass_fields__", {}):
28 field = obj.__dataclass_fields__[attr]
29 field_type = field.type
30 origin = getattr(field_type, "__origin__", None)
31 if origin is not None:
32 field_type = origin
33 attr_value = field_type
34 else:
35 raise AssertionError(assumption_msg)
36 if is_callable:
37 assumption_msg += f" and that {obj}.{attr} is a callable"
38 assert callable(attr_value), assumption_msg
39 if is_instance_of:
40 assumption_msg += f" and that {obj}.{attr} is an instance of {is_instance_of}"
41 if isinstance(attr_value, property):
42 fget = attr_value.fget
43 assert fget is not None, f"Property {obj}.{attr} has no fget"
44 sig = inspect.signature(fget)
45 ret_anno = sig.return_annotation
46 assert ret_anno is not inspect._empty, (
47 f"Property {obj}.{attr} has no return annotation"
48 )
49 assert ret_anno == is_instance_of, assumption_msg
50 else:
51 if isinstance(attr_value, type):
52 assert attr_value is is_instance_of, assumption_msg
53 else:
54 assert isinstance(attr_value, is_instance_of), assumption_msg
55
56
57 @pytest.mark.skipif(
58 current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
59 )
60 def test_multimodal_interface():
61 # protect against interface changes
62 from vllm.multimodal.inputs import PlaceholderRange
63
64 assumes(PlaceholderRange, "offset")
65 assumes(PlaceholderRange, "length")
66
67
68 @pytest.mark.skipif(
69 current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
70 )
71 def test_config_interface():
72 # protect against interface changes
73 from vllm.config import VllmConfig
74 from vllm.config.cache import CacheConfig
75 from vllm.config.kv_transfer import KVTransferConfig
76 from vllm.config.model import ModelConfig
77 from vllm.config.parallel import ParallelConfig
78
79 assumes(VllmConfig, "model_config")
80 assumes(VllmConfig, "cache_config")
81 assumes(VllmConfig, "parallel_config")
82 assumes(VllmConfig, "kv_transfer_config")
83
84 assumes(KVTransferConfig, "kv_role")
85 assumes(KVTransferConfig, "kv_connector_extra_config")
86
87 assumes(ModelConfig, "use_mla", is_instance_of=bool)
88 assumes(ModelConfig, "dtype")
89 assumes(ModelConfig, "max_model_len")
90 assumes(ModelConfig, "get_vocab_size", is_callable=True)
91 assumes(ModelConfig, "get_num_attention_heads", is_callable=True)
92 assumes(ModelConfig, "get_num_kv_heads", is_callable=True)
93 assumes(ModelConfig, "get_head_size", is_callable=True)
94 assumes(ModelConfig, "get_num_layers", is_callable=True)
95 assumes(ModelConfig, "get_num_kv_heads", is_callable=True)
96 assumes(ModelConfig, "model")
97
98 assumes(ParallelConfig, "world_size")
99 assumes(ParallelConfig, "rank")
100 assumes(ParallelConfig, "tensor_parallel_size")
101 assumes(ParallelConfig, "pipeline_parallel_size")
102 assumes(ParallelConfig, "data_parallel_size_local")
103 assumes(ParallelConfig, "data_parallel_rank_local")
104
105 assumes(CacheConfig, "cache_dtype")
106 assumes(CacheConfig, "block_size")
107 assumes(CacheConfig, "gpu_memory_utilization")
108
109 # kv metadata minimal case
110 from vllm.utils.torch_utils import get_kv_cache_torch_dtype
111
112 model_config = ModelConfig(dtype="bfloat16")
113 parallel_config = ParallelConfig()
114 cache_config = CacheConfig(cache_dtype="bfloat16")
115 kv_dtype = get_kv_cache_torch_dtype(cache_config.cache_dtype, model_config.dtype)
116 use_mla = False
117 chunk_size = 256
118 num_layer = model_config.get_num_layers(parallel_config)
119 num_kv_head = model_config.get_num_kv_heads(parallel_config)
120 head_size = model_config.get_head_size()
121 kv_shape = (num_layer, 1 if use_mla else 2, chunk_size, num_kv_head, head_size)
122
123 # dummy lmcache metadata creation example
124 _ = (
125 model_config.model,
126 parallel_config.world_size,
127 parallel_config.rank,
128 "vllm",
129 kv_dtype,
130 kv_shape,
131 use_mla,
132 )
133
134
135 @pytest.mark.skipif(
136 current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
137 )
138 def test_request_interface():
139 # protect against interface changes
140 from types import NoneType
141
142 from vllm.sampling_params import SamplingParams
143 from vllm.v1.request import Request
144
145 sampling_params = SamplingParams(max_tokens=10)
146 sampling_params.update_from_generation_config({}, eos_token_id=100)
147
148 req = Request(
149 request_id="test_request",
150 prompt_token_ids=[1, 2, 3],
151 sampling_params=sampling_params,
152 pooling_params=None,
153 lora_request=None,
154 )
155 assumes(req, "mm_features", is_instance_of=(list, NoneType))
156 assumes(req, "request_id")
157 assumes(req, "priority")
158 assumes(req, "prompt_token_ids")
159 assumes(req, "sampling_params")
160 assumes(req, "num_tokens")
161 assumes(req, "kv_transfer_params", is_instance_of=(dict, NoneType))
162
163 from vllm.multimodal.inputs import MultiModalFeatureSpec
164
165 assumes(MultiModalFeatureSpec, "identifier")
166 assumes(MultiModalFeatureSpec, "mm_position")
167
168
169 def test_new_request_interface():
170 # protect against interface changes
171 from vllm.v1.core.sched.output import NewRequestData
172
173 assumes(NewRequestData, "req_id")
174 assumes(NewRequestData, "block_ids")
175 assumes(NewRequestData, "prompt_token_ids")
176 assumes(NewRequestData, "sampling_params")
177
178
179 def test_sampling_params_interface():
180 # protect against interface changes
181 from vllm.sampling_params import SamplingParams
182
183 assumes(SamplingParams, "extra_args")
184
185 # dumb example use case in LMCache
186 kv_transfer_params = {
187 "lmcache.tag.user": "example_user_1",
188 "lmcache.ttl": 60,
189 }
190 sampling_params = SamplingParams(
191 extra_args={"kv_transfer_params": kv_transfer_params}
192 )
193 assert sampling_params.extra_args["kv_transfer_params"] == kv_transfer_params
194
195
196 def test_tp_interface():
197 # protect against interface changes
198 import inspect
199
200 from vllm.distributed.parallel_state import get_tp_group
201
202 sig = inspect.signature(get_tp_group)
203 GroupCoordinator = sig.return_annotation
204
205 assumes(GroupCoordinator, "broadcast", is_callable=True)
206 assumes(GroupCoordinator, "broadcast_object", is_callable=True)
207
208
209 def test_forward_context_interface():
210 # protect against interface changes
211 from vllm.forward_context import ForwardContext
212
213 assumes(ForwardContext, "no_compile_layers", is_instance_of=dict)
214 assumes(ForwardContext, "attn_metadata")
215
216
217 def test_scheduler_output_interface():
218 # protect against interface changes
219 from vllm.v1.core.sched.output import SchedulerOutput
220
221 assumes(SchedulerOutput, "finished_req_ids")
222 assumes(SchedulerOutput, "scheduled_new_reqs", is_instance_of=list)
223 assumes(SchedulerOutput, "num_scheduled_tokens", is_instance_of=dict)
224 assumes(SchedulerOutput, "scheduled_cached_reqs")
225
226 from vllm.v1.core.sched.output import CachedRequestData
227
228 assumes(CachedRequestData, "req_ids", is_instance_of=list)
229 assumes(CachedRequestData, "new_block_ids", is_instance_of=list)