@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8cb24d3aed
Fork 0

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 6, 20:30 UTC

[KV Connector] Support using FlexKV as KV Cache Offloading option. (#34328)

Signed-off-by: phaedonsun <phaedonsun@tencent.com> Co-authored-by: phaedonsun <phaedonsun@tencent.com>

sfeiqiang committed Mar 12, 2026 at 15:46 UTC 8cb24d3aedb9f431fb15a636a3e11a00262f5991
5 files changed +725
docs/features/disagg_prefill.md
+6
@@ -44,6 +44,12 @@ For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
44 --kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
45 ```
46
47 +- **FlexKVConnectorV1**: refer to [examples/offline_inference/prefix_caching_flexkv.py](../../examples/offline_inference/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
48 +
49 + ```bash
50 + --kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
51 + ```
52 +
53 ## Benchmarks
54
55 Please refer to [benchmarks/disagg_benchmarks](../../benchmarks/disagg_benchmarks) for disaggregated prefilling benchmarks.
examples/offline_inference/prefix_caching_flexkv.py new
+221
@@ -0,0 +1,221 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +This example shows how to use FlexKV with vLLM for prefix caching.
5 +
6 +FlexKV is a distributed KV Store and multi-level cache management system for
7 +ultra-large-scale LLM inference.
8 +
9 +Requirements:
10 + - Install FlexKV (https://github.com/taco-project/FlexKV):
11 + 1. git clone git@github.com:taco-project/FlexKV.git
12 + 2. cd FlexKV && bash build.sh
13 + - Ensure FlexKV is compatible with your vLLM version.
14 +
15 +Usage:
16 + 1. Run this script:
17 + python examples/offline_inference/prefix_caching_flexkv.py \
18 + --model /path/to/your/model
19 +
20 + 2. Arguments:
21 + --model Path or name of the model (required)
22 + --tp-size Tensor parallel size (default: 1)
23 + --gpu-memory-util GPU memory utilization (default: 0.4)
24 +
25 + 3. The script will:
26 + - Create a FlexKV configuration file.
27 + - Set the FLEXKV_CONFIG_PATH environment variable.
28 + - Run vLLM with FlexKVConnectorV1 enabled.
29 + - Compare results between regular execution, vLLM's default prefix
30 + caching, and FlexKV.
31 +"""
32 +
33 +import argparse
34 +import json
35 +import os
36 +import time
37 +
38 +from vllm import LLM, SamplingParams
39 +from vllm.distributed import cleanup_dist_env_and_memory
40 +
41 +# NOTE: This is just a running example. For benchmarking purpose,
42 +# please see benchmarks/benchmark_prefix_caching.py
43 +
44 +
45 +def parse_args():
46 + parser = argparse.ArgumentParser(
47 + description="Example of using FlexKV with vLLM for prefix caching."
48 + )
49 + parser.add_argument(
50 + "--model",
51 + type=str,
52 + required=True,
53 + help="Path or name of the model to use.",
54 + )
55 + parser.add_argument(
56 + "--tp-size",
57 + type=int,
58 + default=1,
59 + help="Tensor parallel size (default: 1).",
60 + )
61 + parser.add_argument(
62 + "--gpu-memory-util",
63 + type=float,
64 + default=0.4,
65 + help="GPU memory utilization fraction (default: 0.4).",
66 + )
67 + return parser.parse_args()
68 +
69 +
70 +def main():
71 + args = parse_args()
72 +
73 + flexkv_config = {
74 + "server_recv_port": f"ipc:///tmp/flexkv_test_{os.getpid()}",
75 + "cache_config": {
76 + "enable_cpu": True,
77 + "num_cpu_blocks": 10240,
78 + },
79 + "num_log_interval_requests": 200,
80 + }
81 + flexkv_config_path = f"./flexkv_config_{os.getpid()}.json"
82 + with open(flexkv_config_path, "w") as f:
83 + json.dump(flexkv_config, f)
84 + os.environ["FLEXKV_CONFIG_PATH"] = flexkv_config_path
85 +
86 + try:
87 + _run(args)
88 + finally:
89 + if os.path.exists(flexkv_config_path):
90 + os.remove(flexkv_config_path)
91 +
92 +
93 +def _run(args):
94 + # Common prefix.
95 + prefix = (
96 + "You are an expert school principal, skilled in effectively managing "
97 + "faculty and staff. Draft 10-15 questions for a potential first grade "
98 + "Head Teacher for my K-12, all-girls', independent school that emphasizes "
99 + "community, joyful discovery, and life-long learning. The candidate is "
100 + "coming in for a first-round panel interview for a 8th grade Math "
101 + "teaching role. They have 5 years of previous teaching experience "
102 + "as an assistant teacher at a co-ed, public school with experience "
103 + "in middle school math teaching. Based on these information, fulfill "
104 + "the following paragraph: "
105 + )
106 +
107 + # Sample prompts.
108 + prompts = [
109 + "Hello, my name is",
110 + "The president of the United States is",
111 + "The capital of France is",
112 + "The future of AI is",
113 + ]
114 +
115 + generating_prompts = [prefix + prompt for prompt in prompts]
116 +
117 + # Create a sampling params object.
118 + sampling_params = SamplingParams(temperature=0.0)
119 +
120 + kv_transfer_config = {
121 + "kv_connector": "FlexKVConnectorV1",
122 + "kv_role": "kv_both",
123 + }
124 +
125 + # Create an LLM without prefix caching as a baseline.
126 + regular_llm = LLM(
127 + model=args.model,
128 + enable_prefix_caching=False,
129 + gpu_memory_utilization=args.gpu_memory_util,
130 + tensor_parallel_size=args.tp_size,
131 + )
132 +
133 + print("Results without `enable_prefix_caching`")
134 +
135 + # ruff: noqa: E501
136 + # Generate texts from the prompts. The output is a list of RequestOutput
137 + # objects that contain the prompt, generated text, and other information.
138 + outputs = regular_llm.generate(generating_prompts, sampling_params)
139 +
140 + regular_generated_texts = []
141 + # Print the outputs.
142 + print("-" * 50)
143 + for output in outputs:
144 + prompt = output.prompt
145 + generated_text = output.outputs[0].text
146 + regular_generated_texts.append(generated_text)
147 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
148 + print("-" * 50)
149 +
150 + # Destroy the LLM object and free up the GPU memory.
151 + del regular_llm
152 + cleanup_dist_env_and_memory()
153 +
154 + # Create an LLM with prefix caching enabled.
155 + prefix_cached_llm = LLM(
156 + model=args.model,
157 + enable_prefix_caching=True,
158 + gpu_memory_utilization=args.gpu_memory_util,
159 + tensor_parallel_size=args.tp_size,
160 + kv_transfer_config=kv_transfer_config,
161 + )
162 +
163 + # Warmup so that the shared prompt's KV cache is computed.
164 + prefix_cached_llm.generate(generating_prompts[0], sampling_params)
165 +
166 + # wait for offload kv task finished.
167 + time.sleep(2)
168 +
169 + # Generate with prefix caching.
170 + outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)
171 +
172 + print("Results with `enable_prefix_caching`")
173 +
174 + cached_generated_texts = []
175 + # Print the outputs. You should see the same outputs as before.
176 + print("-" * 50)
177 + for output in outputs:
178 + prompt = output.prompt
179 + generated_text = output.outputs[0].text
180 + cached_generated_texts.append(generated_text)
181 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
182 + print("-" * 50)
183 +
184 + # Compare the results and display the speedup
185 + generated_same = all(
186 + regular_generated_texts[i] == cached_generated_texts[i]
187 + for i in range(len(prompts))
188 + )
189 + print(f"Generated answers are the same: {generated_same}")
190 +
191 + # wait for offload kv task finished.
192 + time.sleep(2)
193 +
194 + # reset prefix cache to use flexkv
195 + prefix_cached_llm.reset_prefix_cache()
196 +
197 + # Generate with prefix caching.
198 + outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)
199 +
200 + print("Results with `flexkv`")
201 +
202 + flexkv_generated_texts = []
203 + # Print the outputs. You should see the same outputs as before.
204 + print("-" * 50)
205 + for output in outputs:
206 + prompt = output.prompt
207 + generated_text = output.outputs[0].text
208 + flexkv_generated_texts.append(generated_text)
209 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
210 + print("-" * 50)
211 +
212 + # Compare the results and display the speedup
213 + generated_same = all(
214 + regular_generated_texts[i] == flexkv_generated_texts[i]
215 + for i in range(len(prompts))
216 + )
217 + print(f"Generated answers are the same: {generated_same}")
218 +
219 +
220 +if __name__ == "__main__":
221 + main()
tests/v1/kv_connector/unit/test_flexkv_connector.py new
+232
@@ -0,0 +1,232 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for FlexKVConnectorV1.
4 +
5 +These tests mock the ``flexkv`` package so they can run without a real FlexKV
6 +installation. They verify:
7 +
8 +1. That ``FlexKVConnectorV1`` raises a helpful ``ImportError`` when FlexKV is
9 + not installed.
10 +2. That all public methods are correctly delegated to the underlying
11 + ``FlexKVConnectorV1Impl``.
12 +"""
13 +
14 +import sys
15 +import types
16 +from unittest.mock import MagicMock, patch
17 +
18 +import pytest
19 +import torch
20 +
21 +from vllm.config import KVTransferConfig, VllmConfig
22 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
23 +from vllm.v1.kv_cache_interface import KVCacheConfig
24 +
25 +from .utils import create_vllm_config
26 +
27 +# ---------------------------------------------------------------------------
28 +# Helpers
29 +# ---------------------------------------------------------------------------
30 +
31 +
32 +def _make_vllm_config(
33 + kv_connector: str = "FlexKVConnectorV1",
34 + kv_role: str = "kv_both",
35 +) -> VllmConfig:
36 + """Return a minimal VllmConfig with a KVTransferConfig attached."""
37 + vllm_config = create_vllm_config(block_size=16, max_num_batched_tokens=512)
38 + vllm_config.kv_transfer_config = KVTransferConfig(
39 + kv_connector=kv_connector,
40 + kv_role=kv_role,
41 + )
42 + return vllm_config
43 +
44 +
45 +def _make_kv_cache_config() -> KVCacheConfig:
46 + return MagicMock(spec=KVCacheConfig)
47 +
48 +
49 +def _make_flexkv_module(
50 + impl_mock: MagicMock,
51 +) -> tuple[types.ModuleType, types.ModuleType]:
52 + """Build a fake ``flexkv`` package hierarchy that returns *impl_mock*
53 + when ``FlexKVConnectorV1Impl`` is instantiated."""
54 + flexkv_mod = types.ModuleType("flexkv")
55 + integration_mod = types.ModuleType("flexkv.integration")
56 + vllm_mod = types.ModuleType("flexkv.integration.vllm")
57 + adapter_mod = types.ModuleType("flexkv.integration.vllm.vllm_v1_adapter")
58 +
59 + # Make FlexKVConnectorV1Impl() return our mock instance.
60 + # The "# type: ignore" markers below are needed because ModuleType does
61 + # not declare these attributes statically; they are set dynamically.
62 + FlexKVConnectorV1ImplCls = MagicMock(return_value=impl_mock)
63 + adapter_mod.FlexKVConnectorV1Impl = FlexKVConnectorV1ImplCls # type: ignore
64 +
65 + flexkv_mod.integration = integration_mod # type: ignore
66 + integration_mod.vllm = vllm_mod # type: ignore
67 + vllm_mod.vllm_v1_adapter = adapter_mod # type: ignore
68 +
69 + return flexkv_mod, adapter_mod
70 +
71 +
72 +def _install_flexkv_mock(impl_mock: MagicMock):
73 + """Insert fake flexkv modules into sys.modules and return a context that
74 + cleans them up afterwards."""
75 + flexkv_mod, adapter_mod = _make_flexkv_module(impl_mock)
76 + mods = {
77 + "flexkv": flexkv_mod,
78 + "flexkv.integration": flexkv_mod.integration,
79 + "flexkv.integration.vllm": flexkv_mod.integration.vllm,
80 + "flexkv.integration.vllm.vllm_v1_adapter": adapter_mod,
81 + }
82 + return patch.dict(sys.modules, mods)
83 +
84 +
85 +def _build_connector(vllm_config: VllmConfig, impl_mock: MagicMock):
86 + """Instantiate FlexKVConnectorV1 with faked flexkv modules."""
87 + from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import (
88 + FlexKVConnectorV1,
89 + )
90 +
91 + with _install_flexkv_mock(impl_mock):
92 + connector = FlexKVConnectorV1(
93 + vllm_config=vllm_config,
94 + role=KVConnectorRole.WORKER,
95 + kv_cache_config=_make_kv_cache_config(),
96 + )
97 + return connector
98 +
99 +
100 +# ---------------------------------------------------------------------------
101 +# Tests
102 +# ---------------------------------------------------------------------------
103 +
104 +
105 +class TestFlexKVConnectorImportError:
106 + """FlexKVConnectorV1 should fail with a helpful message when flexkv is
107 + absent."""
108 +
109 + def test_import_error_message(self):
110 + from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import (
111 + FlexKVConnectorV1,
112 + )
113 +
114 + # Ensure flexkv is NOT in sys.modules
115 + for key in list(sys.modules):
116 + if key.startswith("flexkv"):
117 + del sys.modules[key]
118 +
119 + with pytest.raises(ImportError, match="(?i)flexkv") as exc_info:
120 + FlexKVConnectorV1(
121 + vllm_config=_make_vllm_config(),
122 + role=KVConnectorRole.WORKER,
123 + kv_cache_config=_make_kv_cache_config(),
124 + )
125 +
126 + assert "https://github.com/taco-project/FlexKV" in str(exc_info.value)
127 +
128 +
129 +class TestFlexKVConnectorDelegation:
130 + """All public API methods should be forwarded to the impl."""
131 +
132 + @pytest.fixture()
133 + def connector_and_impl(self):
134 + impl = MagicMock()
135 + cfg = _make_vllm_config()
136 + connector = _build_connector(cfg, impl)
137 + return connector, impl
138 +
139 + def test_shutdown(self, connector_and_impl):
140 + connector, impl = connector_and_impl
141 + connector.shutdown()
142 + impl.shutdown.assert_called_once()
143 +
144 + def test_start_load_kv(self, connector_and_impl):
145 + connector, impl = connector_and_impl
146 + ctx = MagicMock()
147 + connector.start_load_kv(ctx, extra_arg="x")
148 + impl.start_load_kv.assert_called_once_with(ctx, extra_arg="x")
149 +
150 + def test_save_kv_layer(self, connector_and_impl):
151 + connector, impl = connector_and_impl
152 + kv_layer = torch.zeros(4, 4)
153 + attn_meta = MagicMock()
154 + connector.save_kv_layer("layer_0", kv_layer, attn_meta)
155 + impl.save_kv_layer.assert_called_once_with("layer_0", kv_layer, attn_meta)
156 +
157 + def test_wait_for_save(self, connector_and_impl):
158 + connector, impl = connector_and_impl
159 + connector.wait_for_save()
160 + impl.wait_for_save.assert_called_once()
161 +
162 + def test_get_finished(self, connector_and_impl):
163 + connector, impl = connector_and_impl
164 + impl.get_finished.return_value = ({"req1"}, None)
165 + result = connector.get_finished({"req1"})
166 + impl.get_finished.assert_called_once_with({"req1"})
167 + assert result == ({"req1"}, None)
168 +
169 + def test_register_kv_caches(self, connector_and_impl):
170 + connector, impl = connector_and_impl
171 + kv_caches = {"layer_0": torch.zeros(1)}
172 + connector.register_kv_caches(kv_caches)
173 + impl.register_kv_caches.assert_called_once_with(kv_caches)
174 +
175 + def test_get_num_new_matched_tokens(self, connector_and_impl):
176 + connector, impl = connector_and_impl
177 + req = MagicMock()
178 + impl.get_num_new_matched_tokens.return_value = (10, False)
179 + result = connector.get_num_new_matched_tokens(req, 5)
180 + impl.get_num_new_matched_tokens.assert_called_once_with(req, 5)
181 + assert result == (10, False)
182 +
183 + def test_update_state_after_alloc(self, connector_and_impl):
184 + connector, impl = connector_and_impl
185 + req = MagicMock()
186 + blocks = MagicMock()
187 + connector.update_state_after_alloc(req, blocks, 4)
188 + impl.update_state_after_alloc.assert_called_once_with(req, blocks, 4)
189 +
190 + def test_build_connector_meta(self, connector_and_impl):
191 + connector, impl = connector_and_impl
192 + sched_out = MagicMock()
193 + connector.build_connector_meta(sched_out)
194 + impl.build_connector_meta.assert_called_once_with(sched_out)
195 +
196 + def test_update_connector_output(self, connector_and_impl):
197 + connector, impl = connector_and_impl
198 + out = MagicMock()
199 + connector.update_connector_output(out)
200 + impl.update_connector_output.assert_called_once_with(out)
201 +
202 + def test_request_finished(self, connector_and_impl):
203 + connector, impl = connector_and_impl
204 + req = MagicMock()
205 + impl.request_finished.return_value = (True, {"key": "val"})
206 + result = connector.request_finished(req, [1, 2, 3])
207 + impl.request_finished.assert_called_once_with(req, [1, 2, 3])
208 + assert result == (True, {"key": "val"})
209 +
210 + def test_take_events(self, connector_and_impl):
211 + connector, impl = connector_and_impl
212 + impl.take_events.return_value = iter([])
213 + list(connector.take_events())
214 + impl.take_events.assert_called_once()
215 +
216 + def test_get_kv_connector_stats(self, connector_and_impl):
217 + connector, impl = connector_and_impl
218 + impl.get_kv_connector_stats.return_value = None
219 + result = connector.get_kv_connector_stats()
220 + impl.get_kv_connector_stats.assert_called_once()
221 + assert result is None
222 +
223 + def test_get_block_ids_with_load_errors(self, connector_and_impl):
224 + connector, impl = connector_and_impl
225 + impl.get_block_ids_with_load_errors.return_value = {7, 8}
226 + result = connector.get_block_ids_with_load_errors()
227 + assert result == {7, 8}
228 +
229 + def test_wait_for_layer_load(self, connector_and_impl):
230 + connector, impl = connector_and_impl
231 + connector.wait_for_layer_load("layer_0")
232 + impl.wait_for_layer_load.assert_called_once_with("layer_0")
vllm/distributed/kv_transfer/kv_connector/factory.py
+6
@@ -207,3 +207,9 @@ KVConnectorFactory.register_connector(
207 "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
208 "MooncakeConnector",
209 )
210 +
211 +KVConnectorFactory.register_connector(
212 + "FlexKVConnectorV1",
213 + "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",
214 + "FlexKVConnectorV1",
215 +)
vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py new
+260
@@ -0,0 +1,260 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable
4 +from typing import TYPE_CHECKING, Any
5 +
6 +import torch
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
10 + KVConnectorBase_V1,
11 + KVConnectorMetadata,
12 + KVConnectorRole,
13 +)
14 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
15 +from vllm.logger import init_logger
16 +from vllm.v1.core.sched.output import SchedulerOutput
17 +from vllm.v1.outputs import KVConnectorOutput
18 +
19 +if TYPE_CHECKING:
20 + from vllm.distributed.kv_events import KVCacheEvent
21 + from vllm.forward_context import ForwardContext
22 + from vllm.v1.attention.backend import AttentionMetadata
23 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
24 + from vllm.v1.kv_cache_interface import KVCacheConfig
25 + from vllm.v1.request import Request
26 +
27 +logger = init_logger(__name__)
28 +
29 +
30 +# FlexKV is a distributed KV Store and multi-level cache management system for
31 +# ultra-large-scale LLM inference.
32 +# GitHub: https://github.com/taco-project/FlexKV
33 +# Install: git clone git@github.com:taco-project/FlexKV.git \
34 +# && cd FlexKV && bash build.sh
35 +class FlexKVConnectorV1(KVConnectorBase_V1):
36 + """KV Connector that offloads KV cache to FlexKV.
37 +
38 + FlexKV is a distributed KV Store and multi-level cache management system
39 + designed for ultra-large-scale LLM inference. It supports offloading KV
40 + cache to CPU memory, SSD, and remote storage.
41 +
42 + Installation:
43 + See https://github.com/taco-project/FlexKV for installation instructions.
44 + Quick start::
45 +
46 + git clone git@github.com:taco-project/FlexKV.git
47 + cd FlexKV && bash build.sh
48 +
49 + Configuration:
50 + Pass ``kv_connector="FlexKVConnectorV1"`` via ``--kv-transfer-config``::
51 +
52 + --kv-transfer-config \
53 + '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
54 + """
55 +
56 + def __init__(
57 + self,
58 + vllm_config: "VllmConfig",
59 + role: KVConnectorRole,
60 + kv_cache_config: "KVCacheConfig",
61 + ):
62 + super().__init__(
63 + vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
64 + )
65 + try:
66 + from flexkv.integration.vllm.vllm_v1_adapter import FlexKVConnectorV1Impl
67 + except ImportError as e:
68 + raise ImportError(
69 + "FlexKV is not installed. Please install it to use "
70 + "FlexKVConnectorV1. See https://github.com/taco-project/FlexKV "
71 + "for installation instructions."
72 + ) from e
73 +
74 + self._flexkv_connector = FlexKVConnectorV1Impl(vllm_config, role)
75 +
76 + def shutdown(self):
77 + self._flexkv_connector.shutdown()
78 +
79 + # ==============================
80 + # Worker-side methods
81 + # ==============================
82 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
83 + """No-op for FlexKV (currently).
84 +
85 + FlexKV manages all KV transfers on the **scheduler side** via
86 + ``build_connector_meta`` (which calls ``launch_tasks``) and
87 + ``update_connector_output`` (which polls ``query_finished_task``).
88 + KV blocks are transferred directly between the FlexKV server and
89 + vLLM's GPU memory without worker-side intervention during the
90 + forward pass — similar to how NIXL operates.
91 +
92 + These worker-side hooks are kept (rather than omitted) to satisfy
93 + the ``KVConnectorBase_V1`` interface contract and to serve as
94 + extension points for a future worker-side layer-pipelining path.
95 +
96 + Args:
97 + forward_context (ForwardContext): the forward context.
98 + **kwargs (Any): additional arguments (unused).
99 + """
100 + self._flexkv_connector.start_load_kv(forward_context, **kwargs)
101 +
102 + def wait_for_layer_load(self, layer_name: str) -> None:
103 + """No-op for FlexKV (currently).
104 +
105 + FlexKV manages all KV transfers on the scheduler side.
106 + This hook is retained for ``KVConnectorBase_V1`` API compatibility.
107 +
108 + Args:
109 + layer_name: the name of the layer (unused).
110 + """
111 + self._flexkv_connector.wait_for_layer_load(layer_name)
112 +
113 + def save_kv_layer(
114 + self,
115 + layer_name: str,
116 + kv_layer: torch.Tensor,
117 + attn_metadata: "AttentionMetadata",
118 + **kwargs,
119 + ) -> None:
120 + """No-op for FlexKV (currently).
121 +
122 + FlexKV offloads KV cache asynchronously from the scheduler side
123 + after a request finishes (see ``request_finished``). It does not
124 + intercept individual layer tensors during the forward pass.
125 +
126 + This hook is retained to satisfy ``KVConnectorBase_V1`` and as an
127 + extension point for future per-layer async offload support.
128 +
129 + Args:
130 + layer_name (str): the name of the layer (unused).
131 + kv_layer (torch.Tensor): the paged KV buffer (unused).
132 + attn_metadata (AttentionMetadata): the attention metadata (unused).
133 + **kwargs (Any): additional arguments (unused).
134 + """
135 + self._flexkv_connector.save_kv_layer(
136 + layer_name, kv_layer, attn_metadata, **kwargs
137 + )
138 +
139 + def wait_for_save(self):
140 + """No-op for FlexKV (currently).
141 +
142 + KV offload tasks are tracked asynchronously by the scheduler
143 + connector via ``request_finished`` / ``query_finished_task``.
144 + There is no pending worker-side save to wait for at
145 + forward-context exit.
146 +
147 + Retained to satisfy ``KVConnectorBase_V1`` and as an extension
148 + point for future worker-side save-completion signalling.
149 + """
150 + self._flexkv_connector.wait_for_save()
151 +
152 + def get_finished(
153 + self, finished_req_ids: set[str]
154 + ) -> tuple[set[str] | None, set[str] | None]:
155 + """Notify worker-side connector of requests that have finished
156 + generating tokens.
157 +
158 + Returns:
159 + Tuple of (sending/saving ids, recving/loading ids) for requests
160 + that have finished asynchronous transfer. The finished saves/sends
161 + req ids must belong to a set provided in a call to this method
162 + (this call or a prior one).
163 + """
164 + return self._flexkv_connector.get_finished(finished_req_ids)
165 +
166 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
167 + """Initialize with the KV caches. Useful for pre-registering the
168 + KV caches in the KVConnector (e.g. for NIXL).
169 +
170 + Args:
171 + kv_caches: dictionary of layer names to kv cache tensors.
172 + """
173 + self._flexkv_connector.register_kv_caches(kv_caches)
174 +
175 + # ==============================
176 + # Scheduler-side methods
177 + # ==============================
178 + def get_num_new_matched_tokens(
179 + self,
180 + request: "Request",
181 + num_computed_tokens: int,
182 + ) -> tuple[int, bool]:
183 + """Get the number of new tokens that can be loaded from the
184 + external KV cache beyond ``num_computed_tokens``.
185 +
186 + Args:
187 + request (Request): the request object.
188 + num_computed_tokens (int): the number of locally computed
189 + tokens for this request.
190 +
191 + Returns:
192 + Tuple of (num_external_tokens, is_ready) where
193 + num_external_tokens is the number of additional tokens that
194 + can be loaded from the external KV cache.
195 + """
196 + return self._flexkv_connector.get_num_new_matched_tokens(
197 + request, num_computed_tokens
198 + )
199 +
200 + def update_state_after_alloc(
201 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
202 + ):
203 + """Update KVConnector state after block allocation."""
204 + self._flexkv_connector.update_state_after_alloc(
205 + request, blocks, num_external_tokens
206 + )
207 +
208 + def build_connector_meta(
209 + self, scheduler_output: SchedulerOutput
210 + ) -> KVConnectorMetadata:
211 + """Build the connector metadata for this step.
212 +
213 + This function should NOT modify fields in the scheduler_output.
214 + Also, calling this function will reset the state of the connector.
215 +
216 + Args:
217 + scheduler_output (SchedulerOutput): the scheduler output object.
218 + """
219 + return self._flexkv_connector.build_connector_meta(scheduler_output)
220 +
221 + def update_connector_output(self, connector_output: KVConnectorOutput):
222 + """Update KVConnector state from worker-side connectors output.
223 +
224 + Args:
225 + connector_output (KVConnectorOutput): the worker-side
226 + connectors output.
227 + """
228 + self._flexkv_connector.update_connector_output(connector_output)
229 +
230 + def request_finished(
231 + self,
232 + request: "Request",
233 + block_ids: list[int],
234 + ) -> tuple[bool, dict[str, Any] | None]:
235 + """Called when a request has finished, before its blocks are freed.
236 +
237 + Returns:
238 + Tuple of (async_save, kv_transfer_params) where async_save is
239 + True if the request is being saved/sent asynchronously and blocks
240 + should not be freed until the request_id is returned from
241 + :meth:`get_finished`. kv_transfer_params is an optional dict of
242 + KVTransferParams to be included in the request outputs.
243 + """
244 + return self._flexkv_connector.request_finished(request, block_ids)
245 +
246 + def take_events(self) -> Iterable["KVCacheEvent"]:
247 + """Collect buffered KV cache events.
248 +
249 + Returns:
250 + New KV cache events since the last call.
251 + """
252 + return self._flexkv_connector.take_events()
253 +
254 + def get_kv_connector_stats(self) -> KVConnectorStats | None:
255 + """Get the KV connector stats collected during the last interval."""
256 + return self._flexkv_connector.get_kv_connector_stats()
257 +
258 + def get_block_ids_with_load_errors(self) -> set[int]:
259 + """Get the block ids that have failed to load."""
260 + return self._flexkv_connector.get_block_ids_with_load_errors()