diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index fc4421eb06..0907293fd8 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -36,7 +36,9 @@ from vllm.v1.kv_offload.base import ( ReqContext, make_offload_key, ) +from vllm.v1.kv_offload.tiering.backpressure import EMABackpressureDetector from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob +from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager @@ -180,6 +182,7 @@ def _make_tier( g3: dict[str, object] | None = None, control_ports: list[int] | None = None, data_parallel_rank_local: int | None = None, + backpressure: dict[str, Any] | None = None, ) -> KVCRSecondaryTierManager: def make_control(_bind_host, bind_port, advertise_host): return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") @@ -196,7 +199,25 @@ def _make_tier( monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) - return KVCRSecondaryTierManager( + tier_config = { + "type": "kvcr", + "router_capabilities": ["router_hint"], + "control_host": "127.0.0.1", + "control_ports": control_ports if control_ports is not None else [7777], + "control_advertise_host": "127.0.0.1", + "enable_telemetry": enable_telemetry, + "secondary_g2_slots": secondary_g2_slots, + "kvcr_service_socket_path": kvcr_service_socket_path, + "compatibility_digest": compatibility_digest, + "policy": policy, + "g3": g3, + "local_dram_backend": "UCX", + "remote_fw_dram_backend": "UCX", + } + if backpressure is not None: + tier_config["backpressure"] = backpressure + tier = SecondaryTierFactory.create_secondary_tier( + tier_config=tier_config, offloading_spec=SimpleNamespace( config=SimpleNamespace( parallel=SimpleNamespace( @@ -209,20 +230,9 @@ def _make_tier( ), ), primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), - tier_type="kvcr", - router_capabilities=["router_hint"], - control_host="127.0.0.1", - control_ports=control_ports if control_ports is not None else [7777], - control_advertise_host="127.0.0.1", - enable_telemetry=enable_telemetry, - secondary_g2_slots=secondary_g2_slots, - kvcr_service_socket_path=kvcr_service_socket_path, - compatibility_digest=compatibility_digest, - policy=policy, - g3=g3, - local_dram_backend="UCX", - remote_fw_dram_backend="UCX", ) + assert isinstance(tier, KVCRSecondaryTierManager) + return tier def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): @@ -271,10 +281,19 @@ def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): ) -def test_kvcr_tier_adapts_request_and_load(monkeypatch): - """Check hint forwarding, key conversion, load descriptors, and cleanup.""" +@pytest.mark.parametrize( + "backpressure", + [None, {"backpressure_cls": "EMABackpressureDetector"}], + ids=["default", "backpressure"], +) +def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure): + """Factory-created tiers preserve configuration and complete hinted loads.""" kvcr = RecordingKVCR() - tier = _make_tier(monkeypatch, kvcr) + tier = _make_tier(monkeypatch, kvcr, backpressure=backpressure) + if backpressure is None: + assert tier.bp_detector is None + else: + assert isinstance(tier.bp_detector, EMABackpressureDetector) kv_hint = KvHintsEnvelope(protocol_version="0.1", message_id="msg", actions=[]) ctx = ReqContext( req_id="req", diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 4b6f81b3b7..4804f193e3 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -82,6 +82,7 @@ from vllm.v1.kv_offload.tiering.base import ( if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec + from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector _REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES @@ -295,8 +296,11 @@ class KVCRSecondaryTierManager(SecondaryTierManager): g3: dict[str, Any] | None = None, local_dram_backend: str = "UCX", remote_fw_dram_backend: str = "UCX", + backpressure_detector: "BackpressureDetector | None" = None, ) -> None: - super().__init__(offloading_spec, primary_kv_view, tier_type) + super().__init__( + offloading_spec, primary_kv_view, tier_type, backpressure_detector + ) selected_policy = _resolve_policy(policy) if (kvcr_service_socket_path is None) != (compatibility_digest is None): raise ValueError(