[Test] Enable KV cache events for HMA models in CPU offloading test (#47754)
Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Itay Etelis committed
Jul 14, 2026 at 12:22 UTC
f04d3f640e41ff217d4a793c315438aa2878f193
1 file changed
+14
-19
tests/v1/kv_connector/unit/test_offloading_connector.py
+14
-19
@@ -251,23 +251,19 @@ def test_cpu_offloading(
251
kv_connector_extra_config=extra_config,
252
)
253
254
- # KV events are incompatible with HMA (setting kv_events_config
255
- # would force HMA off), so only enable them for non-HMA models.
254
subscriber: MockSubscriber | None = None
257
- kv_events_config: KVEventsConfig | None = None
258
- if not uses_hma:
259
- port: int
260
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
261
- s.bind(("0.0.0.0", 0))
262
- port = s.getsockname()[1]
263
-
264
- events_endpoint = f"tcp://*:{port}"
265
- kv_events_config = KVEventsConfig(
266
- enable_kv_cache_events=True,
267
- publisher="zmq",
268
- endpoint=events_endpoint,
269
- topic="test",
270
- )
255
+ port: int
256
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
257
+ s.bind(("0.0.0.0", 0))
258
+ port = s.getsockname()[1]
259
+
260
+ events_endpoint = f"tcp://*:{port}"
261
+ kv_events_config = KVEventsConfig(
262
+ enable_kv_cache_events=True,
263
+ publisher="zmq",
264
+ endpoint=events_endpoint,
265
+ topic="test",
266
+ )
267
268
# Attention-free / hybrid models disable prefix caching by default
269
# (ModelConfig.is_prefix_caching_supported returns False). Without it,
@@ -289,9 +285,8 @@ def test_cpu_offloading(
285
**({"max_num_seqs": 1} if current_platform.is_rocm() else {}),
286
)
287
292
- if kv_events_config is not None:
293
- events_endpoint = events_endpoint.replace("*", "127.0.0.1")
294
- subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
288
+ events_endpoint = events_endpoint.replace("*", "127.0.0.1")
289
+ subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
290
291
try:
292
_latency_test(llm, subscriber)