b853540388 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/core/test_scheduler.py

+7

@@ -899,6 +899,7 @@ def test_kv_connector_basic():

899 scheduler = create_scheduler(

900 enable_prefix_caching=True,

901 use_kv_connector=True,

902 + disable_hybrid_kv_cache_manager=True,

903 )

904 NUM_TOTAL_BLOCKS = scheduler.kv_cache_manager.block_pool.get_num_free_blocks()

905 BLOCK_SIZE = scheduler.cache_config.block_size

@@ -1024,6 +1025,7 @@ def test_external_prefix_cache_metrics():

1025 scheduler = create_scheduler(

1026 enable_prefix_caching=False,

1027 use_kv_connector=True,

1028 + disable_hybrid_kv_cache_manager=True,

1029 )

1030

1031 # Mock connector to simulate a partial external cache hit

@@ -1088,6 +1090,7 @@ def test_kv_connector_unable_to_allocate():

1090 use_kv_connector=True,

1091 block_size=BLOCK_SIZE,

1092 num_blocks=NUM_BLOCKS,

1093 + disable_hybrid_kv_cache_manager=True,

1094 )

1095 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2

1096 scheduler.connector.get_num_new_matched_tokens = Mock(name="method")

@@ -1171,6 +1174,7 @@ def test_kv_connector_handles_preemption():

1174 use_kv_connector=True,

1175 block_size=BLOCK_SIZE,

1176 num_blocks=NUM_BLOCKS,

1177 + disable_hybrid_kv_cache_manager=True,

1178 )

1179

1180 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE

@@ -1387,6 +1391,7 @@ def create_scheduler_with_priority(

1391 block_size: int = 16,

1392 max_model_len: int | None = None,

1393 num_speculative_tokens: int | None = None,

1394 + disable_hybrid_kv_cache_manager: bool = False,

1395 ) -> Scheduler:

1396 """Create scheduler with priority policy enabled.

1397

@@ -1411,6 +1416,7 @@ def create_scheduler_with_priority(

1416 disable_chunked_mm_input=disable_chunked_mm_input,

1417 enable_chunked_prefill=True,

1418 policy="priority", # Enable priority scheduling

1419 + disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,

1420 )

1421 model_config = ModelConfig(

1422 model=model,

@@ -2018,6 +2024,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv():

2024 num_blocks=5, # Can hold 64 tokens (first block is null)

2025 block_size=16, # Standard block size

2026 use_kv_connector=True,

2027 + disable_hybrid_kv_cache_manager=True,

2028 )

2029

2030 # Create a request and schedule it

tests/v1/core/utils.py

+2

@@ -46,6 +46,7 @@ def create_scheduler(

46 num_speculative_tokens: int | None = None,

47 skip_tokenizer_init: bool = False,

48 async_scheduling: bool = False,

49 + disable_hybrid_kv_cache_manager: bool = False,

50 ) -> Scheduler | AsyncScheduler:

51 """Create scheduler under test.

52

@@ -70,6 +71,7 @@ def create_scheduler(

71 disable_chunked_mm_input=disable_chunked_mm_input,

72 enable_chunked_prefill=True,

73 async_scheduling=async_scheduling,

74 + disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,

75 )

76 model_config = ModelConfig(

77 model=model,

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

+2

@@ -136,6 +136,7 @@ run_tests_for_model() {

136 vllm serve $model_name \

137 --port $PORT \

138 --enforce-eager \

139 + --disable-hybrid-kv-cache-manager \

140 --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \

141 --tensor-parallel-size $PREFILLER_TP_SIZE \

142 --kv-transfer-config '$KV_CONFIG'"

@@ -178,6 +179,7 @@ run_tests_for_model() {

179 --port $PORT \

180 --enforce-eager \

181 --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \

182 + --disable-hybrid-kv-cache-manager \

183 --kv-transfer-config '$KV_CONFIG'"

184

185 # DP-EP attention mode

tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh

+2

@@ -85,6 +85,7 @@ run_tests_for_model() {

85 --port $PREFILL_PORT \

86 --enforce-eager \

87 --gpu-memory-utilization 0.2 \

88 + --disable-hybrid-kv-cache-manager \

89 --kv-transfer-config '$KV_CONFIG'"

90

91 if [ -n "$model_args" ]; then

@@ -103,6 +104,7 @@ run_tests_for_model() {

104 --port $DECODE_PORT \

105 --enforce-eager \

106 --gpu-memory-utilization 0.2 \

107 + --disable-hybrid-kv-cache-manager \

108 --kv-transfer-config '$KV_CONFIG'"

109

110 if [ -n "$model_args" ]; then

tests/v1/kv_connector/unit/test_multi_connector.py

+1

@@ -114,6 +114,7 @@ def test_multi_shared_storage_connector_consistency():

114 enforce_eager=True,

115 gpu_memory_utilization=0.5,

116 kv_transfer_config=kv_transfer_config,

117 + disable_hybrid_kv_cache_manager=True,

118 )

119 # Run generation - this should trigger saving KV cache

120 _ = llm.generate(PROMPTS, SAMPLING_PARAMS)

tests/v1/kv_connector/unit/test_nixl_connector.py

+1

@@ -932,6 +932,7 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):

932 "gpu_memory_utilization": 0.5,

933 "kv_transfer_config": kv_transfer_config,

934 "distributed_executor_backend": distributed_executor_backend,

935 + "disable_hybrid_kv_cache_manager": True,

936 }

937

938 timeout = 6

tests/v1/kv_connector/unit/test_shared_storage_connector.py

+1

@@ -132,6 +132,7 @@ def test_shared_storage_connector_hashes(tmp_path):

132 enforce_eager=True,

133 kv_transfer_config=kv_transfer_config,

134 limit_mm_per_prompt={"image": 2},

135 + disable_hybrid_kv_cache_manager=True,

136 )

137

138 # don't put this import at the top level

tests/v1/kv_connector/unit/utils.py

+3

@@ -91,6 +91,9 @@ def create_vllm_config(

91 max_num_batched_tokens=max_num_batched_tokens,

92 max_model_len=max_model_len,

93 enable_chunked_prefill=enable_chunked_prefill,

94 + # Disable hybrid KV cache manager for testing

95 + # Should be removed after we support hybrid KV cache manager-based testing.

96 + disable_hybrid_kv_cache_manager=True,

97 )

98 model_config = ModelConfig(

99 model=model,

tests/v1/kv_offload/test_cpu_offloading.py

+1

@@ -27,6 +27,7 @@ def test_cpu_offloading(cpu_block_size: int) -> None:

27 model="meta-llama/Llama-3.2-1B-Instruct",

28 gpu_memory_utilization=0.5,

29 kv_transfer_config=kv_transfer_config,

30 + disable_hybrid_kv_cache_manager=True,

31 )

32

33 prompts = ["Hi " * 100]

vllm/config/vllm.py

+3 -3

@@ -40,11 +40,14 @@ if TYPE_CHECKING:

40 from transformers import PretrainedConfig

41

42 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

43 + from vllm.v1.kv_cache_interface import KVCacheConfig

44 else:

45 PretrainedConfig = Any

46

47 QuantizationConfig = Any

48

49 + KVCacheConfig = Any

50 +

51 logger = init_logger(__name__)

52

53

@@ -568,9 +571,6 @@ class VllmConfig:

571 if not current_platform.support_hybrid_kv_cache():

572 # Hybrid KV cache manager is not supported on non-GPU platforms.

573 self.scheduler_config.disable_hybrid_kv_cache_manager = True

571 - if self.kv_transfer_config is not None:

572 - # Hybrid KV cache manager is not compatible with KV transfer.

573 - self.scheduler_config.disable_hybrid_kv_cache_manager = True

574 if self.kv_events_config is not None:

575 # Hybrid KV cache manager is not compatible with KV events.

576 self.scheduler_config.disable_hybrid_kv_cache_manager = True

vllm/distributed/kv_transfer/kv_connector/factory.py

+15 -3

@@ -6,15 +6,18 @@ from collections.abc import Callable

6 from typing import TYPE_CHECKING, cast

7

8 import vllm.envs as envs

9 +from vllm.config import VllmConfig

10 from vllm.distributed.kv_transfer.kv_connector.base import (

11 KVConnectorBase,

12 KVConnectorBaseType,

13 )

13 -from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole

14 +from vllm.distributed.kv_transfer.kv_connector.v1 import (

15 + KVConnectorRole,

16 + supports_hma,

17 +)

18 from vllm.logger import init_logger

19

20 if TYPE_CHECKING:

17 - from vllm.config import VllmConfig

21 from vllm.config.kv_transfer import KVTransferConfig

22

23 logger = init_logger(__name__)

@@ -38,7 +41,7 @@ class KVConnectorFactory:

41 @classmethod

42 def create_connector(

43 cls,

41 - config: "VllmConfig",

44 + config: VllmConfig,

45 role: KVConnectorRole,

46 ) -> KVConnectorBase:

47 if not envs.VLLM_USE_V1:

@@ -51,6 +54,15 @@ class KVConnectorFactory:

54 if kv_transfer_config is None:

55 raise ValueError("kv_transfer_config must be set to create a connector")

56 connector_cls = cls.get_connector_class(kv_transfer_config)

57 +

58 + # check if the connector supports HMA

59 + hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager

60 + if hma_enabled and not supports_hma(connector_cls):

61 + raise ValueError(

62 + f"Connector {connector_cls.__name__} does not support HMA but "

63 + f"HMA is enabled. Please set `--disable-hybrid-kv-cache-manager`."

64 + )

65 +

66 logger.info(

67 "Creating v1 connector with name: %s and engine_id: %s",

68 connector_cls.__name__,

vllm/distributed/kv_transfer/kv_connector/v1/__init__.py

+9 -1

@@ -3,9 +3,17 @@

3 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

4 KVConnectorBase_V1,

5 KVConnectorRole,

6 + SupportsHMA,

7 + supports_hma,

8 )

9 from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import ( # noqa E:501

10 DecodeBenchConnector,

11 )

12

11 -__all__ = ["KVConnectorRole", "KVConnectorBase_V1", "DecodeBenchConnector"]

13 +__all__ = [

14 + "KVConnectorRole",

15 + "KVConnectorBase_V1",

16 + "supports_hma",

17 + "SupportsHMA",

18 + "DecodeBenchConnector",

19 +]

vllm/distributed/kv_transfer/kv_connector/v1/base.py

+40 -1

@@ -70,6 +70,45 @@ CopyBlocksOp = Callable[

70 logger = init_logger(__name__)

71

72

73 +class SupportsHMA(ABC):

74 + """

75 + The class that indicates the corresponding connector supports hybrid memory

76 + allocator (HMA).

77 + This is required to use the connector together with hybrid memory allocator.

78 + """

79 +

80 + @abstractmethod

81 + def request_finished_all_groups(

82 + self,

83 + request: "Request",

84 + block_ids: tuple[list[int], ...],

85 + ) -> tuple[bool, dict[str, Any] | None]:

86 + """

87 + Called exactly once when a request has finished for all kv cache groups,

88 + before its blocks are freed for each group.

89 +

90 + NOTE(Kuntai): This function is only supported by connectors that support HMA.

91 +

92 + The connector may assumes responsibility for freeing the blocks

93 + asynchronously by returning True.

94 +

95 + Returns:

96 + True if the request is being saved/sent asynchronously and blocks

97 + should not be freed until the request_id is returned from

98 + get_finished().

99 + Optional KVTransferParams to be included in the request outputs

100 + returned by the engine.

101 + """

102 + raise NotImplementedError

103 +

104 +

105 +def supports_hma(connector: Any) -> bool:

106 + if isinstance(connector, type):

107 + return issubclass(connector, SupportsHMA)

108 + else:

109 + return isinstance(connector, SupportsHMA)

110 +

111 +

112 class KVConnectorRole(enum.Enum):

113 # Connector running in the scheduler process

114 SCHEDULER = 0

@@ -370,7 +409,7 @@ class KVConnectorBase_V1(ABC):

409 Called exactly once when a request has finished, before its blocks are

410 freed.

411

373 - The connector may assumes responsibility for freeing the the blocks

412 + The connector may assumes responsibility for freeing the blocks

413 asynchronously by returning True.

414

415 Returns:

vllm/v1/core/sched/scheduler.py

+17 -8

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

3 +import copy

4 import itertools

5 import time

6 from collections import defaultdict

@@ -13,6 +13,7 @@ from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory

13 from vllm.distributed.kv_transfer.kv_connector.v1 import (

14 KVConnectorBase_V1,

15 KVConnectorRole,

16 + supports_hma,

17 )

18 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats

19 from vllm.logger import init_logger

@@ -86,15 +87,14 @@ class Scheduler(SchedulerInterface):

87 self.connector = None

88 self.connector_prefix_cache_stats: PrefixCacheStats | None = None

89 if self.vllm_config.kv_transfer_config is not None:

89 - assert len(self.kv_cache_config.kv_cache_groups) == 1, (

90 - "Multiple KV cache groups are not currently supported "

91 - "with KV connectors"

92 - )

90 assert not self.is_encoder_decoder, (

91 "Encoder-decoder models are not currently supported with KV connectors"

92 )

93 +

94 + connector_vllm_config = copy.copy(self.vllm_config)

95 + connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config)

96 self.connector = KVConnectorFactory.create_connector(

97 - config=self.vllm_config, role=KVConnectorRole.SCHEDULER

97 + config=connector_vllm_config, role=KVConnectorRole.SCHEDULER

98 )

99 if self.log_stats:

100 self.connector_prefix_cache_stats = PrefixCacheStats()

@@ -1324,8 +1324,17 @@ class Scheduler(SchedulerInterface):

1324 if self.connector is None:

1325 return False, None

1326

1327 - (block_ids,) = self.kv_cache_manager.get_block_ids(request.request_id)

1328 - return self.connector.request_finished(request, block_ids)

1327 + block_ids = self.kv_cache_manager.get_block_ids(request.request_id)

1328 +

1329 + if not supports_hma(self.connector):

1330 + # NOTE(Kuntai): We should deprecate this code path after we enforce

1331 + # all connectors to support HMA.

1332 + # Hybrid memory allocator should be already turned off for this

1333 + # code path, but let's double-check here.

1334 + assert len(self.kv_cache_config.kv_cache_groups) == 1

1335 + return self.connector.request_finished(request, block_ids[0])

1336 + else:

1337 + return self.connector.request_finished(request, block_ids)

1338

1339 def _update_waiting_for_remote_kv(self, request: Request) -> bool:

1340 """

vllm/v1/worker/gpu_worker.py

+9 -2

@@ -331,6 +331,15 @@ class Worker(WorkerBase):

331 def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:

332 """Allocate GPU KV cache with the specified kv_cache_config."""

333

334 + # Init kv cache connector here, because it requires

335 + # `kv_cache_config`.

336 + # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,

337 + # because `initialize_kv_cache` will inject kv cache groups not

338 + # related to kv cache connector (e.g. kv cache sharing layers).

339 + connector_vllm_config = copy.copy(self.vllm_config)

340 + connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config)

341 + ensure_kv_transfer_initialized(connector_vllm_config)

342 +

343 if self.vllm_config.model_config.enable_sleep_mode:

344 from vllm.device_allocator.cumem import CuMemAllocator

345

@@ -783,5 +792,3 @@ def init_worker_distributed_environment(

792 parallel_config.pipeline_parallel_size,

793 parallel_config.decode_context_parallel_size,

794 )

786 -

787 - ensure_kv_transfer_initialized(vllm_config)

[Core][Hybrid allocator + kv connector 1/n] Enable hybrid allocator + KV cache connector (#25712)