A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[BugFix] Fix MultiConnector test after HMA changes (#19291)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Jun 6, 2025 at 13:16 UTC
aad30bd306d30025dc05a7e9eb2aa3ed5c3c1ed0
3 files changed
+14
-9
tests/v1/kv_connector/unit/test_multi_connector.py
+8
-7
@@ -60,7 +60,8 @@ class TestSharedStorageConnector(SharedStorageConnector):
60
if isinstance(arg, int):
61
to_log.append(str(arg))
62
elif isinstance(arg, KVCacheBlocks):
63
- to_log.append(f"num_blocks={len(arg.blocks)}")
63
+ to_log.append(
64
+ f"num_blocks={[len(b) for b in arg.blocks]}")
65
66
# Log the event as a line to the file
67
try:
@@ -176,7 +177,7 @@ def test_multi_shared_storage_connector_consistency():
177
# on each connector in turn.
178
assert events["storage1-SCHEDULER"][:3] == [
179
'get_num_new_matched_tokens 0',
179
- 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
180
+ 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
181
]
182
assert events["storage1-WORKER"][:5] == [
183
'register_kv_caches', 'bind_connector_metadata', 'start_load_kv',
@@ -184,7 +185,7 @@ def test_multi_shared_storage_connector_consistency():
185
]
186
assert events["storage2-SCHEDULER"][:3] == [
187
'get_num_new_matched_tokens 0',
187
- 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
188
+ 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
189
]
190
assert events["storage2-WORKER"][:5] == [
191
'register_kv_caches', 'bind_connector_metadata', 'start_load_kv',
@@ -205,11 +206,11 @@ def test_multi_shared_storage_connector_consistency():
206
# chosen).
207
assert events["storage1-SCHEDULER"][:3] == [
208
'get_num_new_matched_tokens 0',
208
- 'update_state_after_alloc num_blocks=7 96', 'build_connector_meta'
209
+ 'update_state_after_alloc num_blocks=[7] 96', 'build_connector_meta'
210
]
211
assert events["storage2-SCHEDULER"][:3] == [
212
'get_num_new_matched_tokens 0',
212
- 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
213
+ 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
214
]
215
216
# Delete storage1 connector state
@@ -229,11 +230,11 @@ def test_multi_shared_storage_connector_consistency():
230
# blocks for the second connector.
231
assert events["storage1-SCHEDULER"][:3] == [
232
'get_num_new_matched_tokens 0',
232
- 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
233
+ 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
234
]
235
assert events["storage2-SCHEDULER"][:3] == [
236
'get_num_new_matched_tokens 0',
236
- 'update_state_after_alloc num_blocks=7 96', 'build_connector_meta'
237
+ 'update_state_after_alloc num_blocks=[7] 96', 'build_connector_meta'
238
]
239
240
# Clean up
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+2
-2
@@ -153,6 +153,7 @@ class MultiConnector(KVConnectorBase_V1):
153
num_external_tokens: int):
154
chosen_connector = self._requests_to_connector.get(
155
request.request_id, -1)
156
+ empty_blocks = blocks.new_empty()
157
for i, c in enumerate(self._connectors):
158
if i == chosen_connector:
159
# Forward call to the chosen connector (if any).
@@ -160,8 +161,7 @@ class MultiConnector(KVConnectorBase_V1):
161
num_external_tokens)
162
else:
163
# Call with empty blocks for other connectors.
163
- c.update_state_after_alloc(request,
164
- KVCacheBlocks.create_empty(), 0)
164
+ c.update_state_after_alloc(request, empty_blocks, 0)
165
166
def build_connector_meta(
167
self,
vllm/v1/core/kv_cache_manager.py
+4
@@ -61,6 +61,10 @@ class KVCacheBlocks:
61
if block.block_hash is None
62
]
63
64
+ def new_empty(self) -> "KVCacheBlocks":
65
+ """Creates a new KVCacheBlocks instance with no blocks."""
66
+ return KVCacheBlocks([[] for _ in range(len(self.blocks))])
67
+
68
69
class KVCacheManager:
70