| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Regression tests for HMA auto-disable with KV transfer connectors.""" |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.config import DeviceConfig, KVTransferConfig, SchedulerConfig, VllmConfig |
| 9 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 10 | from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 13 | |
| 14 | pytestmark = pytest.mark.cpu_test |
| 15 | |
| 16 | |
| 17 | @pytest.fixture(autouse=True) |
| 18 | def mock_hybrid_kv_cache_supported(monkeypatch): |
| 19 | monkeypatch.setattr(current_platform, "support_hybrid_kv_cache", lambda: True) |
| 20 | |
| 21 | |
| 22 | @pytest.mark.parametrize( |
| 23 | "kv_transfer_config,expect_disabled", |
| 24 | [ |
| 25 | ( # HMA-supporting connector → HMA stays enabled |
| 26 | KVTransferConfig( |
| 27 | kv_connector="SimpleCPUOffloadConnector", |
| 28 | kv_role="kv_both", |
| 29 | kv_connector_extra_config={"cpu_bytes_to_use": 1 << 30}, |
| 30 | ), |
| 31 | False, |
| 32 | ), |
| 33 | ( # Non-HMA connector → HMA is auto-disabled |
| 34 | KVTransferConfig(kv_connector="ExampleConnector", kv_role="kv_both"), |
| 35 | True, |
| 36 | ), |
| 37 | ( # MultiConnector: all HMA children → HMA stays enabled |
| 38 | KVTransferConfig( |
| 39 | kv_connector="MultiConnector", |
| 40 | kv_role="kv_both", |
| 41 | kv_connector_extra_config={ |
| 42 | "connectors": [ |
| 43 | { |
| 44 | "kv_connector": "SimpleCPUOffloadConnector", |
| 45 | "kv_role": "kv_both", |
| 46 | "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, |
| 47 | }, |
| 48 | { |
| 49 | "kv_connector": "OffloadingConnector", |
| 50 | "kv_role": "kv_both", |
| 51 | "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, |
| 52 | }, |
| 53 | ] |
| 54 | }, |
| 55 | ), |
| 56 | False, |
| 57 | ), |
| 58 | ( # MultiConnector: mixed children → HMA is auto-disabled |
| 59 | KVTransferConfig( |
| 60 | kv_connector="MultiConnector", |
| 61 | kv_role="kv_both", |
| 62 | kv_connector_extra_config={ |
| 63 | "connectors": [ |
| 64 | { |
| 65 | "kv_connector": "SimpleCPUOffloadConnector", |
| 66 | "kv_role": "kv_both", |
| 67 | "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, |
| 68 | }, |
| 69 | {"kv_connector": "ExampleConnector", "kv_role": "kv_both"}, |
| 70 | ] |
| 71 | }, |
| 72 | ), |
| 73 | True, |
| 74 | ), |
| 75 | ], |
| 76 | ids=["hma_connector", "non_hma_connector", "multi_all_hma", "multi_mixed"], |
| 77 | ) |
| 78 | def test_hma_auto_config(kv_transfer_config, expect_disabled): |
| 79 | vllm_config = VllmConfig( |
| 80 | device_config=DeviceConfig("cpu"), |
| 81 | kv_transfer_config=kv_transfer_config, |
| 82 | ) |
| 83 | assert ( |
| 84 | vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is expect_disabled |
| 85 | ) |
| 86 | |
| 87 | |
| 88 | def test_explicit_hma_with_non_hma_connector_errors_at_factory(): |
| 89 | vllm_config = VllmConfig( |
| 90 | device_config=DeviceConfig("cpu"), |
| 91 | scheduler_config=SchedulerConfig( |
| 92 | max_model_len=16, |
| 93 | is_encoder_decoder=False, |
| 94 | disable_hybrid_kv_cache_manager=False, |
| 95 | ), |
| 96 | kv_transfer_config=KVTransferConfig( |
| 97 | kv_connector="ExampleConnector", |
| 98 | kv_role="kv_both", |
| 99 | ), |
| 100 | ) |
| 101 | kv_cache_config = KVCacheConfig( |
| 102 | num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[] |
| 103 | ) |
| 104 | with pytest.raises(ValueError, match="does not support HMA but HMA is enabled"): |
| 105 | KVConnectorFactory.create_connector( |
| 106 | vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config |
| 107 | ) |