From d33d65b68e8ac219d459fdc86c242e9163a24579 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:37 +0800 Subject: [PATCH 01/10] [None][fix] Restore is_generation_only_request as a property on LlmRequest The C++ base exposes is_generation_only_request as a read-only property (def_prop_ro, nanobind/batch_manager/bindings.cpp), but the Python LlmRequest subclass redefined it as a plain method with no @property. Reading it as an attribute therefore yields a bound method, which is always truthy. V1 never trips this because the connector path reaches the attribute from C++ with the C++ object, where the property is real. Every Python-side reader gets the method object instead, so `if request.is_generation_only_request` is unconditionally true and the corresponding guard is dead. Restore the property and drop the call parentheses at the six production call sites and in the test mocks. A local shim was rejected: the same attribute is reached from C++ with the C++ object and from Python with the Python one, so a shim would have to sniff callable(). Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/_torch/disaggregation/transceiver.py | 2 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 5 +++++ tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py | 2 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 4 ++-- tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py | 2 +- tests/unittest/_torch/executor/test_mamba_cache_manager.py | 4 ++-- tests/unittest/_torch/executor/test_request_utils.py | 2 +- tests/unittest/disaggregated/test_cache_reuse_adapter.py | 2 +- 8 files changed, 14 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index a85fb1b85796..e2ed3e4d8aab 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -181,7 +181,7 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: assert self._page_table is not None layer_groups = self._page_table.layer_groups - is_gen_only = req.is_generation_only_request() + is_gen_only = req.is_generation_only_request cached_per_lg = ( adapter.get_cached_token_count_per_layer_group(req, layer_groups) if is_gen_only diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index d52e1f086e7e..91fb34606a55 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -849,7 +849,12 @@ def cached_tokens(self, value: int): self._cached_tokens = value self._cached_tokens_set = True + @property def is_generation_only_request(self): + # Must stay a property: the C++ base exposes this as a read-only + # property (nanobind/batch_manager/bindings.cpp), and a plain method + # here shadows it with something always-truthy for any caller that + # reads it as an attribute. return self.py_llm_request_type == LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY def create_response(self, diff --git a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py index cb13757956a0..cd16090c8895 100644 --- a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py @@ -250,7 +250,7 @@ def append_step_metrics(self, request, iter_counter: int, batch_token_time=None) # - py_decoding_iter == 1 and not yet marked complete: last/only chunk # - Gen-only requests (disagg gen server) are never ctx is_ctx = ( - not request.is_generation_only_request() + not request.is_generation_only_request and not perf.ctx_chunks_complete and request.py_decoding_iter <= 1 ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bf91c4ab2ff5..13b30081b423 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5540,7 +5540,7 @@ def _check_gen_cache_transfer_errors_consensus(self) -> None: """Flush generation transfer errors through a TP-uniform path.""" error_requests = [ req for req in self._get_disagg_reqs_in_error_state() - if req.is_generation_only_request() + if req.is_generation_only_request ] local_needs_flush = bool(error_requests) @@ -6797,7 +6797,7 @@ def _handle_responses(self, emit_first_iter: bool = True): new_active_requests.append(request) continue - if request.is_generation_only_request() and not request.is_finished: + if request.is_generation_only_request and not request.is_finished: # If request is in transmission, so we don't need to emit a response # Also, for the first iteration with overlap, we should skip since first # token has already been emitted previously diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 958afc59ac04..0d3afb573538 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -282,7 +282,7 @@ def _schedule_loop(self, active_requests, inflight_request_ids): if self._prioritize_first_token_gen: requests_list.sort( key=lambda req: ( - 0 if (req.is_generation_only_request() and req.py_decoding_iter == 0) else 1 + 0 if (req.is_generation_only_request and req.py_decoding_iter == 0) else 1 ) ) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index e7d61248e7b4..679e1d795022 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -549,7 +549,7 @@ def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) transceiver._page_table = SimpleNamespace(layer_groups=[]) request = SimpleNamespace( - is_generation_only_request=lambda: False, + is_generation_only_request=False, prompt_len=0, py_request_id=123, ) @@ -571,7 +571,7 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) transceiver._page_table = SimpleNamespace(layer_groups=[]) request = SimpleNamespace( - is_generation_only_request=lambda: False, + is_generation_only_request=False, prompt_len=0, py_request_id=123, ) diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index f2d6d6bc605b..ae23af422ca4 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -125,7 +125,7 @@ def test_executor_request_to_llm_request_adopts_context_phase_draft_tokens() -> exclude_last_generation_logits=False, ) - assert llm_request.is_generation_only_request() + assert llm_request.is_generation_only_request assert llm_request.has_draft_tokens() assert llm_request.num_draft_tokens == len(draft_tokens) assert llm_request.draft_tokens == draft_tokens diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index c131183461be..74d1480f6e30 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -311,7 +311,7 @@ def _build_transceiver_for_kv_slice(num_extra_kv_tokens: int, prompt_len: int): prompt_len=prompt_len, py_request_id=0, py_beam_width=1, # beam search disabled; _trim_packed_beam_block_ids is pass-through - is_generation_only_request=lambda: False, + is_generation_only_request=False, ) return transceiver, req From 7d6b8820db4a8470dfab397b94439929cdc82202 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:37 +0800 Subject: [PATCH 02/10] [None][fix] Reject KVCacheManagerV2 with the C++ cache transceiver CacheTransceiverCpp is bound to the V1 BaseKVCacheManager, while KVCacheManagerV2.impl is the Python V2 core's manager. The combination reached BindKvCacheTransceiver and died on a raw nanobind signature mismatch that named neither the manager nor the way out. An equivalent check already existed for MambaHybridCacheManagerV2 but not for its base class, so plain KVCacheManagerV2 fell straight through. The new check sits after the hybrid one so the subclass keeps its more specific message; both messages share the phrase the existing tests match on, so that ordering is pinned by its own test. This makes the failure legible; it does not make V2 work with a default transceiver config. CacheTransceiverConfig.transceiver_runtime defaults to "auto", which is resolved from the model's preference (llm_utils._resolve_transceiver_runtime_auto) and knows nothing about which cache manager will be built, so most models land on the C++ transceiver. Working configuration is backend="NIXL" with transceiver_runtime="PYTHON". Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_transceiver.py | 21 ++++++++ .../test_disagg_inflight_cancel_gate.py | 52 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 9b9ea2556e1e..304224aefb36 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -13,6 +13,7 @@ CacheTransceiverConfig) from tensorrt_llm.mapping import Mapping +from .kv_cache_manager_v2 import KVCacheManagerV2 from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, @@ -174,6 +175,26 @@ def create_kv_cache_transceiver( "MambaHybridCacheManagerV2 requires transceiver_runtime='PYTHON' " "with backend='NIXL'; it cannot use the C++ transceiver.") + # The same applies to KVCacheManagerV2 itself, not only to its hybrid + # subclass: its `impl` is the Python V2 core's KVCacheManager, and + # CacheTransceiverCpp is bound to BaseKVCacheManager. Without this check the + # combination reaches BindKvCacheTransceiver and dies on a nanobind + # signature mismatch that names neither the manager nor the way out. + # + # Note `transceiver_runtime` defaults to "auto", which is resolved from the + # *model's* preference (llm_utils._resolve_transceiver_runtime_auto) and + # knows nothing about which cache manager will be built -- so most models + # land here rather than on the Python transceiver. + if isinstance(kv_cache_manager, + KVCacheManagerV2) and not use_python_transceiver: + raise ValueError( + "KVCacheManagerV2 requires transceiver_runtime='PYTHON' with " + "backend='NIXL' for disaggregated serving; it cannot use the C++ " + "transceiver, which is bound to the V1 BaseKVCacheManager. Either " + "set cache_transceiver_config.transceiver_runtime='PYTHON' and " + "backend='NIXL', or select the V1 manager with " + "kv_cache_config.use_kv_cache_manager_v2=False.") + if use_python_transceiver: if isinstance(mamba_cache_manager, CppMambaHybridCacheManager): raise ValueError( diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 08a6f4701761..f5997197e268 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -21,6 +21,7 @@ from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as transceiver_module from tensorrt_llm._torch.pyexecutor import py_executor as executor_module +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import BindKvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( @@ -550,6 +551,57 @@ def test_cpp_runtime_rejects_v2_mamba_manager(runtime): ) +@pytest.mark.parametrize("runtime", [None, "CPP", "auto"]) +def test_cpp_runtime_rejects_v2_manager(runtime): + """Plain KVCacheManagerV2 cannot drive the C++ transceiver either. + + `CacheTransceiverCpp` is bound to the V1 `BaseKVCacheManager`, while + `KVCacheManagerV2.impl` is the Python V2 core's manager. Without this guard + the combination reaches `BindKvCacheTransceiver` and dies on a nanobind + signature mismatch that names neither the manager nor the way out. + + `auto` is covered because it is the *default*: it resolves from the model's + preferred runtime, which knows nothing about which cache manager will be + built, so most models land on the C++ transceiver here. + """ + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(KVCacheManagerV2) + + with pytest.raises(ValueError, match="KVCacheManagerV2 requires transceiver_runtime='PYTHON'"): + transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), manager, Mock(), config) + + +def test_python_nixl_transceiver_accepts_v2_manager(monkeypatch): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + expected = object() + constructor = Mock(return_value=expected) + fake_module = SimpleNamespace(KvCacheTransceiverV2=constructor) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.disaggregation.transceiver", fake_module) + manager = object.__new__(KVCacheManagerV2) + + result = transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), manager, Mock(), config) + + assert result is expected + constructor.assert_called_once() + + +def test_v2_mamba_manager_keeps_its_specific_rejection(): + """MambaHybridCacheManagerV2 subclasses KVCacheManagerV2. + + So the general guard must sit after the hybrid one, or the more specific + message -- the one that tells a hybrid user which manager they are on -- + would be shadowed. Both messages share the "requires + transceiver_runtime='PYTHON'" phrase, so this matches the manager name. + """ + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="CPP") + manager = object.__new__(MambaHybridCacheManagerV2) + + with pytest.raises(ValueError, match="MambaHybridCacheManagerV2 requires"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + def test_python_runtime_rejects_cpp_mamba_manager(): config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") manager = object.__new__(CppMambaHybridCacheManager) From d9f6ed99642dbdb758361d42bb35877441efbfac Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:38 +0800 Subject: [PATCH 03/10] [None][test] Run the KV connector suite against both KV cache managers Every connector test now runs twice, once per KV cache manager, so the V2 integration that follows has somewhere to land. Selecting V2 must actually reach V2. _fallback_if_unsupported_kv_cache_manager_v2 silently substitutes the V1 manager for combinations it cannot serve, and a connector test that ran on V1 while claiming to test V2 would pass while exercising nothing. test_connector_runs_on_kv_cache_manager_v2 spies on both managers' __init__ to prove positively which one was built, and asserts the fallback warning is absent. assert_kv_caches_registered is the per-test half of the same argument: the managers hand the worker their pools through different entry points, and asserting the V1 one unconditionally would pass on V2 exactly when the connector was never registered at all. Parametrization is spelled out per test as an explicit @pytest.mark.parametrize(..., indirect=True) rather than as a fixture params=, because scripts/check_test_list.py resolves ids from decorators via AST and cannot see fixture-level parametrization. Tests that supply their own KvCacheConfig have the manager forced onto it, otherwise the V2 parametrization degrades into a second V1 run. Also replace the fixed time.sleep(1) barrier -- which carried its author's TODO -- with a poll on the recorded mock call count until the connector goes quiet. That returns as soon as the callbacks settle and stretches automatically when a slower path lengthens the tail. The one end-to-end test evaluated the generated text and discarded it; it now asserts a cold miss, that cache files were written, a warm hit, and token agreement. Comparing two deterministic runs proves nothing on its own -- they agree whether or not the cache is consulted -- so the spy on the connector's matched-token count is what makes it non-tautological. Agreement is a prefix floor rather than equality: skipping prefill changes the attention reduction order, so reuse can legitimately move the last token. Every test id changed, so the CI list is retargeted onto the kv_cache_manager_v1 variants. The kv_cache_manager_v2 variants fail until the connector is implemented on V2, and are registered in no list until then. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../defs/llmapi/test_llm_api_connector.py | 513 ++++++++++++++++-- .../integration/test_lists/test-db/l0_a10.yml | 44 +- 2 files changed, 481 insertions(+), 76 deletions(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 6a0fbd26c7ac..685e2e6ea994 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -13,8 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import math import os +import shutil import sys import tempfile import time @@ -23,15 +25,77 @@ import pytest from tensorrt_llm import LLM, DisaggregatedParams, SamplingParams +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, KvCacheConfig, KvCacheConnectorConfig) from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX from ..conftest import llm_models_root +# Name of the TensorRT-LLM logger. It sets `propagate = False` +# (tensorrt_llm/logger.py:186-187), so pytest's `caplog` only sees its records +# once `caplog.handler` is attached to it directly. +TRTLLM_LOGGER_NAME = "TRT-LLM" + +# Emitted by `_fallback_if_unsupported_kv_cache_manager_v2` +# (tensorrt_llm/_torch/pyexecutor/_util.py:629-631) when a connector run is +# downgraded from V2 to V1. +FALLBACK_WARNING_FRAGMENT = "Falling back to KVCacheManager" + +# V1 `KVCacheManager` methods reached on the KV connector path proper. Each is +# defined in tensorrt_llm/_torch/pyexecutor/resource_manager.py and wraps a +# nanobind method on the C++ manager. `KVCacheManagerV2` implements none of +# them - `test_connector_v1_method_contract_gap` pins that, and the list is the +# Phase 2 worklist. +CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS = ( + # Connector bring-up: registers the KV tensor with the worker. + # py_executor.py:1070. This is the first hard stop under V2. + "get_unique_primary_pool", + # kv_cache_connector.py:322 and py_executor.py:6173. + "get_cache_indices", + # kv_cache_connector.py:328 (sole caller). + "commit_and_get_block_hashes", + # kv_cache_connector.py:353, only when a retention config is set. + "get_priority_by_block_id", +) + +# Reached only when the connector coexists with disaggregated serving +# (`test_connector_disagg_prefill`), via AsyncTransferManager and the V1 cache +# reuse adapter - not from `KvCacheConnectorManager` itself. Tracked separately +# because V2 already has its own answer for some of them (`try_commit_blocks` +# at kv_cache_manager_v2.py:3201, page refcounts instead of explicit pinning), +# so they do not necessarily belong in a connector port. +DISAGG_PATH_KV_CACHE_MANAGER_METHODS = ( + "store_blocks_for_reuse", # py_executor.py:455 + "unpin_blocks_by_id", # py_executor.py:489 + "get_memory_pool_block_indices", # disaggregation/resource/cache_reuse.py:121 + "pin_blocks", # no Python caller today +) + + +@pytest.fixture(scope="function") +def use_kv_cache_manager_v2(request): + """Run each connector test under both KV cache managers. + + Parametrized by an explicit `@pytest.mark.parametrize(..., indirect=True)` + on each test, applied as the innermost decorator so the manager lands first + in the generated test id. It is spelled out per test rather than set as a + fixture `params=` because the test-list validator + (scripts/check_test_list.py) resolves ids from parametrize decorators via + AST and cannot see fixture-level parametrization. + + Selecting V2 must actually reach V2: `_fallback_if_unsupported_kv_cache_manager_v2` + silently substitutes the V1 manager for combinations it cannot serve, and a + connector test that ran on V1 while claiming to test V2 would pass while + exercising nothing. `test_connector_runs_on_kv_cache_manager_v2` guards that. + """ + return request.param + @pytest.fixture(scope="function") -def model_with_connector(): +def model_with_connector(use_kv_cache_manager_v2): with patch("tensorrt_llm._torch.pyexecutor.py_executor_creator.importlib" ) as importlib_mock: mock_scheduler = MagicMock() @@ -56,7 +120,16 @@ def model_fn(*args, **kwargs): "kv_cache_config": KvCacheConfig(free_gpu_memory_fraction=0.1) } - return LLM(*args, **{**default_kwargs, **kwargs}) + merged_kwargs = {**default_kwargs, **kwargs} + + # Tests that supply their own `KvCacheConfig` must still honour the + # manager under test, otherwise the V2 parametrization silently + # degrades into a second V1 run. + kv_cache_config = merged_kwargs.get("kv_cache_config") + if kv_cache_config is not None: + kv_cache_config.use_kv_cache_manager_v2 = use_kv_cache_manager_v2 + + return LLM(*args, **merged_kwargs) yield model_fn, mock_scheduler, mock_worker @@ -68,27 +141,205 @@ def enforce_single_worker(monkeypatch): yield -def generate_and_sleep(model, *args, **kwargs): - # Some KV connector API calls are made after a full response is returned. We want to be able to track these calls. - # However, we don't have any indication of when all the calls are complete. - # To compensate for this, we sleep between the generate call and the return of the outputs. - # TODO(jthomson04): Surely there's a better way to do this? +# Some KV connector API calls are made after a full response is returned +# (`request_finished`, the trailing `get_finished` polls and asynchronous +# saves), and there is no public signal for when they are complete. Instead of +# sleeping a fixed amount, wait until the connector mocks stop recording new +# calls. That returns as soon as the connector goes quiet, and - unlike a fixed +# sleep - stretches automatically when a slower path lengthens the tail. +CONNECTOR_QUIESCE_TIMEOUT_S = 60.0 +CONNECTOR_QUIET_PERIOD_S = 0.5 +CONNECTOR_POLL_INTERVAL_S = 0.01 + +# Fraction of generated tokens a connector-warmed run must reproduce exactly. +# See test_connector_e2e_persistent_cache for why this is not 1.0. +E2E_MIN_TOKEN_AGREEMENT = 0.75 + + +def assert_kv_caches_registered(worker, use_kv_cache_manager_v2): + """The two managers hand the worker its pools through different entry points. + + V1 passes a single pool tensor to `register_kv_caches`; V2 has no such + tensor and passes a `KvCacheLayout` to `register_kv_cache_layout` instead. + Asserting the V1 method unconditionally would silently pass on V2 only if + the connector were never registered at all. + """ + if use_kv_cache_manager_v2: + assert worker.register_kv_cache_layout.call_count == 1 + assert worker.register_kv_caches.call_count == 0 + else: + assert worker.register_kv_caches.call_count == 1 + assert worker.register_kv_cache_layout.call_count == 0 + + +def wait_for_connector_quiescence(scheduler, + worker, + timeout=CONNECTOR_QUIESCE_TIMEOUT_S, + quiet_period=CONNECTOR_QUIET_PERIOD_S): + """Block until no new connector callback lands for `quiet_period` seconds. + + `MagicMock.mock_calls` records every call made on the mock and its children, + so its length is a monotonic progress counter for connector activity. + """ + deadline = time.monotonic() + timeout + + def total_calls(): + return len(scheduler.mock_calls) + len(worker.mock_calls) + + last_seen = total_calls() + quiet_since = time.monotonic() + + while time.monotonic() < deadline: + time.sleep(CONNECTOR_POLL_INTERVAL_S) + + current = total_calls() + if current != last_seen: + last_seen = current + quiet_since = time.monotonic() + elif time.monotonic() - quiet_since >= quiet_period: + return + + raise AssertionError( + f"KV connector callbacks did not go quiet within {timeout}s " + f"({last_seen} calls recorded). The connector is still active or a " + "callback is blocked.") + + +def generate_and_wait(model, scheduler, worker, *args, **kwargs): + """`model.generate`, then block until the connector callbacks settle.""" outputs = model.generate(*args, **kwargs) - time.sleep(1) + wait_for_connector_quiescence(scheduler, worker) return outputs +def test_connector_v1_method_contract_gap(): + """Pin exactly which V1 KV-cache-manager methods `KVCacheManagerV2` lacks. + + This is a static contract check rather than an end-to-end run on purpose: + under V2 the connector dies at the *first* missing method + (`get_unique_primary_pool`, py_executor.py:1070), so no e2e run can ever + report more than one gap. Phase 2 should remove entries from + `CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS` as it implements them. + + Needs no GPU. + """ + stale = [ + name for name in CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS + + DISAGG_PATH_KV_CACHE_MANAGER_METHODS + if not hasattr(KVCacheManager, name) + ] + assert stale == [], ( + f"{stale} are not defined on the V1 KVCacheManager either, so this " + "test is measuring a stale method list rather than a real V2 gap.") + + implemented = [ + name for name in CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS + if hasattr(KVCacheManagerV2, name) + ] + assert implemented == [], ( + f"KVCacheManagerV2 now implements {implemented}. Drop them from " + "CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS and re-check whether the " + "corresponding connector tests can be un-skipped.") + + +@pytest.mark.threadleak(enabled=False) +def test_connector_runs_on_kv_cache_manager_v2(enforce_single_worker, + monkeypatch, caplog): + """Anti-vacuity guard for the `kv_cache_manager_v2` parametrization. + + Without this, every V2-parametrized test below could pass green while + `_fallback_if_unsupported_kv_cache_manager_v2` silently swapped in the V1 + manager. It asserts positively that V2 is constructed, and that the + downgrade warning is absent. + + Any construction failure is recorded rather than asserted on: what must + hold is that the run reached V2 rather than being papered over by a + fallback, which stays true regardless of how far bring-up gets. + """ + constructed = [] + + def record_construction(cls): + original_init = cls.__init__ + + def recording_init(self, *args, **kwargs): + constructed.append(cls.__name__) + return original_init(self, *args, **kwargs) + + monkeypatch.setattr(cls, "__init__", recording_init) + + record_construction(KVCacheManagerV2) + record_construction(KVCacheManager) + + # The TensorRT-LLM logger sets `propagate = False`, so caplog only sees its + # records once its handler is attached to that logger directly. + trtllm_logger = logging.getLogger(TRTLLM_LOGGER_NAME) + trtllm_logger.addHandler(caplog.handler) + + construction_error = None + llm = None + try: + with patch( + "tensorrt_llm._torch.pyexecutor.py_executor_creator.importlib" + ) as importlib_mock: + connector_module = importlib_mock.import_module.return_value + connector_module.KvConnectorScheduler.return_value = MagicMock() + connector_module.KvConnectorWorker.return_value = MagicMock() + + try: + llm = LLM( + model=f"{llm_models_root()}/Qwen2-0.5B", + backend="pytorch", + kv_connector_config=KvCacheConnectorConfig( + connector_module="", + connector_scheduler_class="KvConnectorScheduler", + connector_worker_class="KvConnectorWorker", + ), + cuda_graph_config=None, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1, + use_kv_cache_manager_v2=True), + ) + # The V2 connector path is knowingly incomplete, so any construction + # failure is an acceptable outcome. What must hold is that the run + # reached V2 instead of being papered over by a V1 fallback. + except Exception as exc: # noqa: BLE001 + construction_error = exc + finally: + trtllm_logger.removeHandler(caplog.handler) + if llm is not None: + llm.shutdown() + + # Report the current frontier so the failure mode is visible in CI output + # instead of being silently swallowed by the except above. + print(f"\n[connector+V2 frontier] construction_error=" + f"{type(construction_error).__name__ if construction_error else None}" + f": {construction_error}") + + assert "KVCacheManagerV2" in constructed, ( + "KVCacheManagerV2 was never constructed, so the connector run silently " + f"fell back to V1 (managers constructed: {constructed}; construction " + f"error: {construction_error!r}). Every kv_cache_manager_v2-" + "parametrized connector test in this file is vacuous until this passes." + ) + + assert FALLBACK_WARNING_FRAGMENT not in caplog.text, ( + f"{FALLBACK_WARNING_FRAGMENT!r} was logged, so the connector was " + "downgraded to the V1 KV cache manager.") + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_simple(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -96,7 +347,8 @@ def test_connector_simple(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, ["Hello, world"], sampling_params) + generate_and_wait(model, scheduler, worker, ["Hello, world"], + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 @@ -154,22 +406,26 @@ def test_connector_simple(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_async_onboard(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 16, True worker.get_finished.side_effect = lambda finished_gen, load_async: ( finished_gen, load_async) - generate_and_sleep(model, [ + generate_and_wait(model, scheduler, worker, [ "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua." ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True)) @@ -183,15 +439,18 @@ def test_connector_async_onboard(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_async_save(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -202,7 +461,8 @@ def test_connector_async_save(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, ["Hello, world"], sampling_params) + generate_and_wait(model, scheduler, worker, ["Hello, world"], + sampling_params) assert scheduler.request_finished.call_count == 1 @@ -224,8 +484,12 @@ def test_connector_async_save(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_scheduler_output(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): NUM_INPUT_TOKENS = 48 NUM_TOKENS = 32 BLOCK_SIZE = 32 @@ -234,7 +498,7 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -242,7 +506,8 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) - generate_and_sleep(model, [0] * NUM_INPUT_TOKENS, sampling_params) + generate_and_wait(model, scheduler, worker, [0] * NUM_INPUT_TOKENS, + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 assert len( @@ -292,7 +557,8 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, assert len(scheduler.request_finished.call_args.args[1]) == math.ceil( (NUM_INPUT_TOKENS + NUM_TOKENS) / BLOCK_SIZE) - generate_and_sleep(model, [1] * NUM_INPUT_TOKENS, sampling_params) + generate_and_wait(model, scheduler, worker, [1] * NUM_INPUT_TOKENS, + sampling_params) # The initial computed position should be 0, since we haven't yet onboarded any blocks. assert scheduler.build_connector_meta.call_args_list[0].args[ @@ -301,9 +567,13 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_scheduler_output_chunked_context(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): model_fn, scheduler, worker = model_with_connector CHUNK_SIZE = 128 @@ -313,7 +583,7 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, enable_chunked_prefill=True, max_num_tokens=CHUNK_SIZE) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -321,7 +591,8 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, sampling_params = SamplingParams(max_tokens=BLOCK_SIZE, ignore_eos=True) - generate_and_sleep(model, [0] * (CHUNK_SIZE * 2), sampling_params) + generate_and_wait(model, scheduler, worker, [0] * (CHUNK_SIZE * 2), + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 @@ -365,19 +636,43 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, (CHUNK_SIZE * 2 + BLOCK_SIZE) / BLOCK_SIZE) +def _disagg_transceiver_config(use_kv_cache_manager_v2): + """The transceiver each KV cache manager can actually be driven by. + + `CacheTransceiverCpp` is bound to the V1 `BaseKVCacheManager`, while + `KVCacheManagerV2.impl` is the Python V2 core's manager, so V2 can only use + the Python transceiver -- which in turn only supports NIXL + (kv_cache_transceiver.py, `create_kv_cache_transceiver`). This is spelled + out per manager rather than left at the default because + `transceiver_runtime` defaults to "auto", and "auto" is resolved from the + *model's* preference (llm_utils._resolve_transceiver_runtime_auto), which + knows nothing about which cache manager will be built. Qwen2 declares no + preference, so the default resolves to the C++ transceiver -- which V2 + cannot use, and which is now rejected with an actionable error rather than + a nanobind signature mismatch. + """ + if use_kv_cache_manager_v2: + return CacheTransceiverConfig(backend="NIXL", + transceiver_runtime="PYTHON") + return CacheTransceiverConfig(backend="DEFAULT") + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("save_async", [False, True]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, - save_async): + save_async, use_kv_cache_manager_v2): model_fn, scheduler, worker = model_with_connector - prefill_worker = model_fn( - disable_overlap_scheduler=True, - cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT")) + transceiver_config = _disagg_transceiver_config(use_kv_cache_manager_v2) - decode_worker = model_fn( - cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), - kv_connector_config=None) + prefill_worker = model_fn(disable_overlap_scheduler=True, + cache_transceiver_config=transceiver_config) + + decode_worker = model_fn(cache_transceiver_config=transceiver_config, + kv_connector_config=None) sampling_params = SamplingParams(ignore_eos=True, max_tokens=16) @@ -394,16 +689,20 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, scheduler.request_finished.return_value = False worker.get_finished.return_value = [], [] - result = generate_and_sleep(prefill_worker, [0] * 48, - sampling_params=sampling_params, - disaggregated_params=disaggregated_params) + result = generate_and_wait(prefill_worker, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=disaggregated_params) gen_disagg_params = result.disaggregated_params gen_disagg_params.request_type = "generation_only" - generate_and_sleep(decode_worker, [0] * 48, - sampling_params=sampling_params, - disaggregated_params=gen_disagg_params) + generate_and_wait(decode_worker, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=gen_disagg_params) assert scheduler.build_connector_meta.call_count == 1 @@ -422,6 +721,9 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_multi_request(enforce_single_worker, model_with_connector): model_fn, scheduler, worker = model_with_connector @@ -446,12 +748,16 @@ def test_connector_multi_request(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_priorities(enforce_single_worker, model_with_connector): """Test that retention priorities flow through the connector correctly. This test verifies that when KvCacheRetentionConfig is provided, the RequestData.priorities field is populated with the correct per-block priorities based on the token ranges. + """ BLOCK_SIZE = 32 NUM_INPUT_TOKENS = 64 # 2 blocks @@ -487,9 +793,11 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, [0] * NUM_INPUT_TOKENS, - sampling_params=sampling_params, - kv_cache_retention_config=retention_config) + generate_and_wait(model, + scheduler, + worker, [0] * NUM_INPUT_TOKENS, + sampling_params=sampling_params, + kv_cache_retention_config=retention_config) # Verify that build_connector_meta was called assert scheduler.build_connector_meta.call_count >= 1 @@ -517,6 +825,9 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_priorities_default(enforce_single_worker, model_with_connector): """Test that priorities are None when no retention config is provided.""" @@ -530,7 +841,10 @@ def test_connector_priorities_default(enforce_single_worker, sampling_params = SamplingParams(max_tokens=4, ignore_eos=True) # Generate without retention config - generate_and_sleep(model, [0] * 48, sampling_params=sampling_params) + generate_and_wait(model, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params) first_call = scheduler.build_connector_meta.call_args_list[0] sched_output = first_call.args[0] @@ -564,6 +878,9 @@ def test_connector_priorities_default(enforce_single_worker, ), ], ) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_rejects_unsupported_config(enforce_single_worker, model_with_connector, llm_kwargs, match): @@ -577,12 +894,23 @@ def test_connector_rejects_unsupported_config(enforce_single_worker, @pytest.mark.threadleak(enabled=False) -def test_connector_e2e_persistent_cache(enforce_single_worker): - """Test e2e KV cache connector using PersistentKvCacheConnector from examples. - - Runs generation twice with separate LLM instances sharing a disk-based - connector cache, verifying that outputs are identical (proving cache - save/load works end-to-end). +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_e2e_persistent_cache(enforce_single_worker, + use_kv_cache_manager_v2, monkeypatch): + """End-to-end KV connector test using PersistentKvCacheConnector from examples. + + Runs the same prompt through two separate LLM instances sharing a + disk-backed connector cache and asserts that: + + 1. the first (cold) run matches nothing and writes cache files, + 2. the second (warm) run actually reads blocks back from disk, and + 3. both runs produce identical text and token ids. + + (3) on its own proves nothing - two deterministic runs of the same prompt + agree whether or not the cache is ever consulted - so (2) is what makes + this a real correctness test rather than a tautology. """ examples_dir = os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "examples", "llm-api") @@ -590,9 +918,29 @@ def test_connector_e2e_persistent_cache(enforce_single_worker): sys.path.insert(0, examples_dir) cache_dir = tempfile.mkdtemp() - os.environ["CONNECTOR_CACHE_FOLDER"] = cache_dir + monkeypatch.setenv("CONNECTOR_CACHE_FOLDER", cache_dir) try: + import llm_kv_cache_connector + + # Record how many tokens the connector served from disk on each run. + # The leader logs this, but the TensorRT-LLM logger does not propagate + # to the root logger, so read it from the return value instead. + matched_tokens = [] + leader_cls = llm_kv_cache_connector.PersistentKvCacheConnectorLeader + original_get_num_new_matched_tokens = ( + leader_cls.get_num_new_matched_tokens) + + def recording_get_num_new_matched_tokens(self, request, + num_computed_tokens): + result = original_get_num_new_matched_tokens( + self, request, num_computed_tokens) + matched_tokens.append(result[0]) + return result + + monkeypatch.setattr(leader_cls, "get_num_new_matched_tokens", + recording_get_num_new_matched_tokens) + kv_connector_config = KvCacheConnectorConfig( connector_module="llm_kv_cache_connector", connector_scheduler_class="PersistentKvCacheConnectorLeader", @@ -605,7 +953,9 @@ def test_connector_e2e_persistent_cache(enforce_single_worker): kv_connector_config=kv_connector_config, cuda_graph_config=None, disable_overlap_scheduler=True, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1), + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + use_kv_cache_manager_v2=use_kv_cache_manager_v2), ) prompt = ( @@ -620,21 +970,70 @@ def test_connector_e2e_persistent_cache(enforce_single_worker): sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) llm1 = LLM(**llm_kwargs) - output1 = llm1.generate([prompt], sampling_params) - output1[0].outputs[0].text - del llm1 + try: + output1 = llm1.generate([prompt], sampling_params) + cold_text = output1[0].outputs[0].text + cold_token_ids = list(output1[0].outputs[0].token_ids) + finally: + llm1.shutdown() + + assert matched_tokens and all(count == 0 for count in matched_tokens), ( + "The first run should be a cold miss, but the connector reported " + f"matched token counts {matched_tokens}. The cache directory was " + "not clean, so the comparison below is meaningless.") cache_files = [f for f in os.listdir(cache_dir) if f.endswith(".pt")] assert len(cache_files) > 0, "No cache files written by connector" + matched_tokens.clear() + llm2 = LLM(**llm_kwargs) - llm2.generate([prompt], sampling_params) - del llm2 + try: + output2 = llm2.generate([prompt], sampling_params) + warm_text = output2[0].outputs[0].text + warm_token_ids = list(output2[0].outputs[0].token_ids) + finally: + llm2.shutdown() + + assert matched_tokens and max(matched_tokens) > 0, ( + "The second run read nothing back from the connector cache " + f"(matched token counts {matched_tokens}), so the comparisons " + "below would pass just as well with the connector disabled.") + + assert len(warm_token_ids) == len(cold_token_ids), ( + f"Generation length changed: cold {len(cold_token_ids)} tokens, " + f"warm {len(warm_token_ids)} tokens.") + + # Exact equality is NOT asserted. Reusing cached KV skips prefill for + # the matched blocks, which changes the attention reduction order, so + # the logits differ in the last bits even though the restored K/V are + # bit-identical (the connector round-trips them through torch.save / + # torch.load). Greedy decoding turns a near-tie into a different token. + # Observed on V1: the two runs agreed on 31 of 32 tokens and split on + # the final one ("The company's" vs "The company is"). + # + # A corrupted or misaddressed cache does not look like that - it + # diverges early and degenerates - so requiring a long common prefix + # keeps the test meaningful without making it a coin flip. + common_prefix = 0 + for cold_id, warm_id in zip(cold_token_ids, warm_token_ids): + if cold_id != warm_id: + break + common_prefix += 1 + + min_common_prefix = math.floor( + len(cold_token_ids) * E2E_MIN_TOKEN_AGREEMENT) + assert common_prefix >= min_common_prefix, ( + f"Connector cache reuse diverged at token {common_prefix} of " + f"{len(cold_token_ids)}, below the {min_common_prefix}-token " + "floor. Early divergence indicates the restored KV is wrong, not " + "just numerically different.\n" + f" cold run: {cold_text!r}\n" + f" warm run: {warm_text!r}\n" + f" cold ids: {cold_token_ids}\n" + f" warm ids: {warm_token_ids}") finally: - os.environ.pop("CONNECTOR_CACHE_FOLDER", None) - if examples_dir in sys.path: sys.path.remove(examples_dir) - import shutil shutil.rmtree(cache_dir, ignore_errors=True) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index b5d9a4ba7c18..7b8bec8ea2f1 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -161,25 +161,31 @@ l0_a10: # usage telemetry - unittest/llmapi/test_llm_telemetry.py::TestTelemetryPyTorchBackend - unittest/llmapi/test_llm_telemetry.py::TestTelemetryArchitectureExtraction - - llmapi/test_llm_api_connector.py::test_connector_simple[True] - - llmapi/test_llm_api_connector.py::test_connector_simple[False] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[True] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[False] - - llmapi/test_llm_api_connector.py::test_connector_async_save[True] - - llmapi/test_llm_api_connector.py::test_connector_async_save[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[True] - - llmapi/test_llm_api_connector.py::test_connector_multi_request - - llmapi/test_llm_api_connector.py::test_connector_priorities - - llmapi/test_llm_api_connector.py::test_connector_priorities_default - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[host_offloading] - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[beam_search] - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[attention_dp] - - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache + # KV connector. The suite is parametrized over the KV cache manager; only the + # kv_cache_manager_v1 variants are registered here. The kv_cache_manager_v2 + # variants are expected to fail today (KVCacheManagerV2 implements none of the + # V1 block-ID methods the connector needs) and are run manually, not in the + # merge gate. See test_connector_v1_method_contract_gap for the current gap. + - llmapi/test_llm_api_connector.py::test_connector_v1_method_contract_gap + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-host_offloading] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-beam_search] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-attention_dp] + - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache[kv_cache_manager_v1] # third-party policy checks CPU-only - thirdparty/test_cmake_third_party.py::test_cmake_listfiles - thirdparty/test_git_modules.py::test_gitmodules From c68b566d9bd96eacdc572ef4f2dfc26105710a3a Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:39 +0800 Subject: [PATCH 04/10] [None][feat] Describe KVCacheManagerV2's KV pools to the connector Connector bring-up hands the worker a single tensor from get_unique_primary_pool. KVCacheManagerV2 has no such tensor and cannot grow one: it allocates pool groups of slots, a slot is the set of coalesced buffers belonging to one layer group, and there is one slot address space per pool and one page-index space per layer group. The first hop of bring-up was therefore a hard stop, and the fix is to describe the memory rather than point at it. KvCacheLayout carries, per layer group, the byte ranges its pages live in: slot i of a region is at base + stride * i for size bytes, or equivalently region.as_tensor()[i]. build_kv_cache_layout_v2 is assembly of V2's own public layout API -- layer_grouping, all_buffer_ids, get_aggregated_pages, pool_group_descs -- so coalescing is derived from the allocator rather than assumed by the consumer. For a uniform model a layer group's buffers merge into a single whole-slot region, which is why the example connector's load and save paths need no change. Regions are byte-oriented because one may span roles with different element types, and a connector moving bytes should not have to care. register_kv_cache_layout is non-abstract and raises by default, so existing connectors are untouched on V1 and get an actionable error on V2 rather than a crash. A registered address is only valid while its page is pinned to GPU, and eviction to another tier reassigns the page's slot, so reject any tier below GPU. The resolved tier list is read from the manager rather than from KvCacheConfig.host_cache_size, whose default of None is falsy but still yields a host tier -- a truthiness check there is dead. That automatic tier is also skipped outright when a connector is attached: it exists only to give the MAX_UTILIZATION scheduler's suspend/resume somewhere to spill to, and a connector run is already restricted to GUARANTEED_NO_EVICT, so it is dead weight rather than a capability. The VSWA guard in bring-up is relaxed for V2 in passing. It exists because the single-tensor registration cannot describe one pool per window size, which is the ordinary case for a layout. The creator-level gate above it still rejects, and is dealt with separately. Removing kv_connector_manager from the V2 incompatible-feature list left the hybrid test asserting a reason that is no longer produced, so it now asserts V2 is returned unchanged; the operator-facing message no longer suggests disabling the connector when the connector is not what makes the configuration unsupported. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- examples/llm-api/llm_kv_cache_connector.py | 15 + tensorrt_llm/_torch/pyexecutor/_util.py | 8 +- .../connectors/kv_cache_connector.py | 25 ++ .../pyexecutor/connectors/kv_cache_layout.py | 245 ++++++++++++++ .../_torch/pyexecutor/kv_cache_manager_v2.py | 20 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 66 +++- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../_torch/executor/test_kv_cache_layout.py | 310 ++++++++++++++++++ .../executor/test_mamba_cache_manager.py | 18 +- 9 files changed, 690 insertions(+), 18 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py create mode 100644 tests/unittest/_torch/executor/test_kv_cache_layout.py diff --git a/examples/llm-api/llm_kv_cache_connector.py b/examples/llm-api/llm_kv_cache_connector.py index 882478993e5a..f5d9ba0ba51a 100644 --- a/examples/llm-api/llm_kv_cache_connector.py +++ b/examples/llm-api/llm_kv_cache_connector.py @@ -116,6 +116,21 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor): assert self.kv_cache_tensor is None, "KV cache tensor already registered" self.kv_cache_tensor = kv_cache_tensor + def register_kv_cache_layout(self, layout): + # KVCacheManagerV2 describes its pools rather than handing over one + # tensor. This example targets the uniform case, where a layer group's + # buffers coalesce into a single whole-slot region; indexing that + # region by page slot is the direct analogue of indexing a pool by + # block id, so the load/save paths below need no change. + assert self.kv_cache_tensor is None, "KV cache tensor already registered" + if len(layout.groups) != 1 or len(layout.groups[0].regions) != 1: + raise NotImplementedError( + "This example connector handles a single layer group with a " + "single coalesced region (uniform attention, one pool). Got " + f"{len(layout.groups)} layer group(s) and " + f"{sum(len(g.regions) for g in layout.groups)} region(s).") + self.kv_cache_tensor = layout.groups[0].regions[0].as_tensor() + def start_load_kv(self, stream: torch.cuda.Stream): # Do all loads synchronously, and blockwise. for path, block_id in self._metadata.load: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b8e160cf7cd6..bc9d8b660bf5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -460,9 +460,9 @@ def _fallback_if_unsupported_kv_cache_manager_v2( # (e.g. ``MiniMaxM3KVCacheManagerV2`` from the sparse-attention path) # also go through the V2-incompatible-feature gate below. if issubclass(kv_cache_manager_cls, KVCacheManagerV2): + # The KV connector is supported on V2 via the pool layout + # registration path, so it no longer forces a fallback to V1. incompat: List[str] = [] - if self._kv_connector_manager is not None: - incompat.append("kv_connector_manager") if self._max_beam_width is not None and self._max_beam_width > 1: incompat.append("max_beam_width > 1") if incompat: @@ -493,8 +493,8 @@ def _fallback_if_unsupported_kv_cache_manager_v2( raise NotImplementedError( "Hybrid Mamba cache managers do not support " f"{incompat_str}; CppMambaHybridCacheManager does not " - "provide a compatible fallback. Use max_beam_width=1 " - "and disable the KV connector.") + "provide a compatible fallback. Disable the listed " + "features to run hybrid linear models.") # Plain V2 (explicitly enabled or selected by a model default): # V2 was a preference, not a structural requirement, so we can # safely fall back to V1. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 99da2a42265c..13cff6298b8b 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -54,6 +54,7 @@ if TYPE_CHECKING: from ..resource_manager import KVCacheManager + from .kv_cache_layout import KvCacheLayout # Used to store data for a single inflight request. @@ -135,6 +136,30 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor): kv_cache_tensor: The contiguous KV cache tensor. """ + def register_kv_cache_layout(self, layout: "KvCacheLayout") -> None: + """ + Register the KV cache pools described by ``layout``. + + Called instead of ``register_kv_caches`` when the KV cache manager is + ``KVCacheManagerV2``, whose memory cannot be expressed as a single + tensor: there is one slot address space per pool and one page-index + space per layer group. + + The bytes for page slot ``i`` of a region live at + ``region.base + region.stride * i`` for ``region.size`` bytes, or + equivalently at ``region.as_tensor()[i]``. Page indices arriving in + ``RequestData.new_block_ids`` are scoped to a layer group and index the + regions of that group. + + Args: + layout: Description of the KV cache pools; see ``KvCacheLayout``. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement register_kv_cache_layout, so it " + "cannot run against KVCacheManagerV2. Implement it, or select the V1 KV " + "cache manager with kv_cache_config.use_kv_cache_manager_v2=False." + ) + @abstractmethod def start_load_kv(self, stream: torch.cuda.Stream): """ diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py new file mode 100644 index 000000000000..8c4b1a77198e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""KV cache layout description handed to a KV connector under KVCacheManagerV2. + +V2 allocates *pool groups*, each holding ``num_slots`` slots. A slot holds the +coalesced buffers of one *layer group* (a life cycle), and a buffer is keyed by +``BufferId(layer_id, role)``. A connector therefore cannot be handed a single +pool tensor: there is one slot address space per pool, and one index space per +layer group. + +Instead it is handed :class:`KvCacheLayout` -- a description of the byte ranges +that repeat per page slot. The addressing contract is V2's own, taken verbatim +from ``AggregatedPageDesc``:: + + (base + stride * i + Range(0, size) for i in aggregated_page_indices) + +where ``i`` comes from ``_KVCache.get_aggregated_page_indices(layer_group_id)``. + +Because ranges are described rather than implied, this covers MLA (a pool simply +has no VALUE buffer), sliding-window attention and hybrid models (one layer group +per window size) without any of them being special cases. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +import torch + +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor + +if TYPE_CHECKING: + from ..kv_cache_manager_v2 import KVCacheManagerV2 + +__all__ = [ + "KvCacheBufferRef", + "KvCacheLayerGroupLayout", + "KvCacheLayout", + "KvCacheRegion", + "build_kv_cache_layout_v2", +] + + +@dataclass(frozen=True) +class KvCacheBufferRef: + """One ``(layer, role)`` buffer covered by a region, in memory order.""" + + #: Global model layer index -- the same index space the per-layer connector + #: hooks (``wait_for_layer_load`` / ``save_kv_layer``) receive. + layer_id: int + #: Native role name as the cache manager spells it, e.g. "key" / "value". + #: Deliberately not an enum: roles are an open vocabulary on the manager + #: side, and a connector should not need updating when one is added. + role: str + #: Page expansion factor for heterogeneous tokens-per-block layers. + expansion: int = 1 + + +@dataclass(frozen=True) +class KvCacheRegion: + """A contiguous byte range that repeats once per page slot. + + The data for page slot ``i`` lives at ``base + stride * i``, for ``size`` + bytes. ``size`` is not necessarily equal to ``stride``: a region covers one + run of adjacent buffers within a slot, and a slot may hold several runs. + """ + + base: int + size: int + stride: int + num_slots: int + buffers: Tuple[KvCacheBufferRef, ...] + + def address_of(self, slot_id: int) -> int: + """Device address of this region for ``slot_id``.""" + if not 0 <= slot_id < self.num_slots: + raise IndexError(f"slot_id {slot_id} out of range [0, {self.num_slots})") + return self.base + self.stride * slot_id + + def as_tensor(self, dtype: torch.dtype = torch.uint8) -> torch.Tensor: + """A strided ``[num_slots, size // itemsize]`` view; row ``i`` is slot ``i``. + + Defaults to ``uint8``. A region may span several roles whose element + types differ, and a connector that only moves bytes should not have to + care; callers that want a typed view can pass ``dtype`` explicitly. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + if self.size % itemsize or self.stride % itemsize: + raise ValueError( + f"region size {self.size} and stride {self.stride} must both be " + f"multiples of {dtype} itemsize {itemsize}" + ) + return convert_to_torch_tensor( + TensorWrapper( + self.base, + dtype, + shape=(self.num_slots, self.size // itemsize), + strides=(self.stride // itemsize, 1), + ) + ) + + +@dataclass(frozen=True) +class KvCacheLayerGroupLayout: + """One layer group -- the unit that page indices are scoped to.""" + + layer_group_id: int + #: Global model layer indices belonging to this group. + layer_ids: Tuple[int, ...] + #: Attention window for this group, or None for full attention. + window_size: Optional[int] + regions: Tuple[KvCacheRegion, ...] + + @property + def bytes_per_page(self) -> int: + """Total bytes this group occupies for a single page slot.""" + return sum(region.size for region in self.regions) + + +@dataclass(frozen=True) +class KvCacheLayout: + """What a connector is handed in place of a single KV cache pool tensor.""" + + tokens_per_block: int + groups: Tuple[KvCacheLayerGroupLayout, ...] + + def group(self, layer_group_id: int) -> KvCacheLayerGroupLayout: + for group in self.groups: + if group.layer_group_id == layer_group_id: + return group + raise KeyError(f"no layer group {layer_group_id} in layout") + + def group_of_layer(self, layer_id: int) -> KvCacheLayerGroupLayout: + """The layer group owning a global model layer index.""" + for group in self.groups: + if layer_id in group.layer_ids: + return group + raise KeyError(f"layer {layer_id} is not covered by this layout") + + +def _global_layer_ids(manager: "KVCacheManagerV2", local_layer_ids) -> List[int]: + """Map V2-internal layer ids to global model layer indices. + + ``pp_layers`` is the local-to-global table the manager already keeps. Models + that map several internal layers onto one model layer (the sparse-attention + virtual-layer path) have no single global index per internal layer, so they + are rejected rather than silently mislabelled. + """ + if hasattr(manager, "_layer_attn_to_layer_id"): + raise NotImplementedError( + "KV connector layout is not supported for managers with virtual " + "attention layers (sparse attention): an internal layer does not map " + "to a single model layer index." + ) + pp_layers = manager.pp_layers + return [int(pp_layers[int(lid)]) for lid in local_layer_ids] + + +def _window_size(init_config, local_layer_id: int) -> Optional[int]: + layers = init_config.layers + if local_layer_id >= len(layers): + raise ValueError(f"no layer config for internal layer {local_layer_id}") + window = getattr(layers[local_layer_id], "window_size", None) + return None if window is None else int(window) + + +def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: + """Describe a ``KVCacheManagerV2``'s GPU pools for a KV connector. + + Built only from V2's public layout API -- ``layer_grouping``, + ``all_buffer_ids``, ``get_aggregated_pages`` and ``pool_group_descs``. No + private storage state is touched, and no assumption is made about dimension + order, kv factor, or the number of pools. + """ + impl = manager.impl + init_config = impl.init_config + + # A pool group's slot count applies to every layer group drawn from it. + # Note LayerGroupId and PoolGroupIndex are distinct index spaces; the + # variants of a pool group name the layer groups it backs. + slots_by_group: Dict[int, int] = {} + for pool_group in impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + slots_by_group[int(variant.layer_group_id)] = int(pool_group.num_slots) + + buffers_by_layer: Dict[int, List] = {} + for buffer_id in impl.all_buffer_ids: + buffers_by_layer.setdefault(int(buffer_id.layer_id), []).append(buffer_id) + + groups: List[KvCacheLayerGroupLayout] = [] + for layer_group_id, local_layer_ids in enumerate(impl.layer_grouping): + local_layer_ids = [int(lid) for lid in local_layer_ids] + if not local_layer_ids: + continue + + num_slots = slots_by_group[layer_group_id] + global_by_local = dict(zip(local_layer_ids, _global_layer_ids(manager, local_layer_ids))) + + buffer_ids = [b for lid in local_layer_ids for b in buffers_by_layer.get(lid, ())] + + regions: List[KvCacheRegion] = [] + for desc in impl.get_aggregated_pages(buffer_ids): + if int(desc.layer_group_id) != layer_group_id: + continue + regions.append( + KvCacheRegion( + base=int(desc.base), + size=int(desc.size), + stride=int(desc.stride), + num_slots=num_slots, + buffers=tuple( + KvCacheBufferRef( + layer_id=global_by_local[int(b.id.layer_id)], + role=str(b.id.role), + expansion=int(b.expansion), + ) + for b in desc.buffers + ), + ) + ) + + groups.append( + KvCacheLayerGroupLayout( + layer_group_id=layer_group_id, + layer_ids=tuple(global_by_local[lid] for lid in local_layer_ids), + window_size=_window_size(init_config, local_layer_ids[0]), + regions=tuple(regions), + ) + ) + + return KvCacheLayout( + tokens_per_block=int(manager.tokens_per_block), + groups=tuple(groups), + ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index d7d1484b6fc9..7fe7876fb433 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -772,10 +772,8 @@ def __init__( self.mapping = mapping self.dtype = dtype self.is_disagg = is_disagg + self.kv_connector_manager = kv_connector_manager - assert kv_connector_manager is None, ( - "kv_connector_manager is not supported for KVCacheManagerV2" - ) assert max_beam_width == 1, "max_beam_width must be 1 for KVCacheManagerV2" assert not (mapping.cp_config.get("cp_type") == CpType.STAR), ( "Star attention is not supported for KVCacheManagerV2" @@ -980,7 +978,21 @@ def append_to_kv_heads_per_layer( logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=int(quota))] - if kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: + if kv_connector_manager is not None and kv_cache_config.host_cache_size is None: + # A KV connector registers device addresses for its pages, and a + # page evicted to another tier has its GPU slot reassigned. The + # automatic host tier below exists only to give the MAX_UTILIZATION + # scheduler's suspend/resume somewhere to spill to, and a connector + # run cannot use that policy (py_executor_creator requires + # GUARANTEED_NO_EVICT), so skip it rather than silently migrating + # pages out from under the connector. An explicitly configured + # host_cache_size is left alone and rejected loudly at bring-up. + host_quota = 0 + logger.info( + "KV cache manager v2 host tier disabled: a KV connector is attached " + "and registers GPU page addresses that tier migration would invalidate." + ) + elif kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: host_quota = kv_cache_config.host_cache_size else: # The V2 MAX_UTILIZATION scheduler relies on suspend/resume to diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 13b30081b423..6ed64504842a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -60,6 +60,7 @@ from ..speculative.speculation_gate import SpeculationGate from .adp_iter_stats import ADPIterStatsBuffer from .connectors.kv_cache_connector import KvCacheConnectorManager +from .connectors.kv_cache_layout import build_kv_cache_layout_v2 from .dwdp import DwdpManager from .error_classification import ErrorBudget from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem @@ -1031,8 +1032,16 @@ def _maybe_init_kv_connector_manager(self): "connector scheduler / worker hooks and are not " "distinguished from real requests.") + if self.kv_cache_manager is None: + raise ValueError( + "KV Cache Connector requires a KV Cache Manager.") + + is_kv_cache_manager_v2 = isinstance(self.kv_cache_manager, + KVCacheManagerV2) + kv_cache_config = getattr(self.llm_args, 'kv_cache_config', None) - if kv_cache_config is not None and kv_cache_config.host_cache_size: + if not is_kv_cache_manager_v2 and (kv_cache_config is not None and + kv_cache_config.host_cache_size): raise NotImplementedError( "KV Cache Connector is not supported with KV cache host " "offloading (KvCacheConfig.host_cache_size). The connector " @@ -1042,11 +1051,11 @@ def _maybe_init_kv_connector_manager(self): "streams are not synchronized with the internal " "onboard/offload streams.") - if self.kv_cache_manager is None: - raise ValueError( - "KV Cache Connector requires a KV Cache Manager.") - - if getattr(self.kv_cache_manager, 'is_vswa', False): + # VSWA allocates one pool per window size. That is fatal for the V1 + # single-tensor registration, but is the normal case for V2, whose + # layout describes one region set per layer group. + if not is_kv_cache_manager_v2 and getattr(self.kv_cache_manager, + 'is_vswa', False): raise NotImplementedError( "KV Cache Connector is not supported with variable " "sliding-window attention (per-layer max_attention_window " @@ -1062,8 +1071,20 @@ def _maybe_init_kv_connector_manager(self): "per-layer load/save hooks have nothing meaningful to " "transfer for those layers.") - kv_tensor = self.kv_cache_manager.get_unique_primary_pool() - self.kv_connector_manager.worker.register_kv_caches(kv_tensor) + if is_kv_cache_manager_v2: + # A registered region is only a valid address while its page is + # pinned to GPU. V2 migrates pages between cache tiers, and it + # provisions a host tier by default, so reject any non-GPU tier + # until the connector participates in migration. Read the + # resolved tier list rather than KvCacheConfig.host_cache_size: + # the default of None is falsy but still yields a host tier. + self._reject_non_gpu_cache_tiers(self.kv_cache_manager) + layout = build_kv_cache_layout_v2(self.kv_cache_manager) + self.kv_connector_manager.worker.register_kv_cache_layout( + layout) + else: + kv_tensor = self.kv_cache_manager.get_unique_primary_pool() + self.kv_connector_manager.worker.register_kv_caches(kv_tensor) # For each of our layers, we need to register the pre/post hooks. # These are used for methods like `wait_for_layer_load` and `save_kv_layer`. @@ -1076,6 +1097,35 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() + @staticmethod + def _reject_non_gpu_cache_tiers(kv_cache_manager) -> None: + """Reject KVCacheManagerV2 cache tiers below GPU while a connector runs. + + A registered region is only a valid device address while its page is + pinned to GPU; eviction to another tier reassigns the page's slot. Until + the connector participates in migration, allow GPU-only configurations. + + The resolved tier list is read from the manager rather than from + ``KvCacheConfig.host_cache_size`` on purpose: V2 provisions a host tier + automatically when that field is left at its default of ``None``, which + is falsy and would slip past a truthiness check. + """ + from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier + + cache_tiers = getattr(kv_cache_manager.impl.init_config, "cache_tiers", + None) + if not cache_tiers: + return + extra = [tier for tier in cache_tiers if tier.tier != CacheTier.GPU_MEM] + if extra: + names = ", ".join(str(tier.tier) for tier in extra) + raise NotImplementedError( + "KV Cache Connector is not supported with KVCacheManagerV2 " + f"cache tiers below GPU (found: {names}). Pages evicted to " + "another tier have their GPU slot reassigned, which would " + "invalidate the addresses registered with the connector. Set " + "KvCacheConfig.host_cache_size=0 to disable the host tier.") + def _end_transfer_and_maybe_terminate(self, request: LlmRequest): transfer_failed = request.state == LlmRequestState.DISAGG_TRANS_ERROR if self.kv_cache_transceiver and request in self.active_requests: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 7b8bec8ea2f1..eb87678eca3e 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -40,6 +40,7 @@ l0_a10: - unittest/_torch/executor/test_disagg_index_mapper_early_release.py - unittest/_torch/executor/test_kv_cache_compression_manager.py - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py + - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py diff --git a/tests/unittest/_torch/executor/test_kv_cache_layout.py b/tests/unittest/_torch/executor/test_kv_cache_layout.py new file mode 100644 index 000000000000..83a7d8116ea9 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_layout.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Tests for the KV cache layout description handed to a KV connector under +# KVCacheManagerV2 (``connectors/kv_cache_layout.py``). +# +# The layout replaces the single-pool-tensor registration used by the V1 +# manager, whose memory V2 cannot express: V2 has one slot address space per +# pool and one page-index space per layer group. +# +# Tests in TestKvCacheRegionArithmetic need no GPU. The rest construct a real +# KVCacheManagerV2 and therefore allocate device memory pools. + +import gc +import unittest + +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import ( + KvCacheBufferRef, + KvCacheLayerGroupLayout, + KvCacheLayout, + KvCacheRegion, + build_kv_cache_layout_v2, +) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 +from tensorrt_llm.mapping import Mapping + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + + +def _make_kwargs( + *, + num_layers: int = 4, + num_kv_heads=4, + head_dim=128, + tokens_per_block: int = 8, + max_seq_len: int = 256, + max_batch_size: int = 4, + max_tokens: int = 2048, + dtype=DataType.HALF, + kv_cache_type=CacheType.SELF, + vocab_size: int = 32000, + kv_cache_config=None, +): + return dict( + kv_cache_config=kv_cache_config + or KvCacheConfigV2(max_tokens=max_tokens, enable_block_reuse=False), + kv_cache_type=kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=dtype, + vocab_size=vocab_size, + ) + + +class TestKvCacheRegionArithmetic(unittest.TestCase): + """Address arithmetic and lookup helpers. No GPU required.""" + + def _region(self, base=4096, size=256, stride=1024, num_slots=8): + return KvCacheRegion( + base=base, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + + def test_address_of_follows_stride(self): + region = self._region() + self.assertEqual(region.address_of(0), 4096) + self.assertEqual(region.address_of(1), 4096 + 1024) + self.assertEqual(region.address_of(7), 4096 + 7 * 1024) + + def test_address_of_rejects_out_of_range_slot(self): + region = self._region(num_slots=8) + for bad in (-1, 8, 99): + with self.assertRaises(IndexError): + region.address_of(bad) + + def test_as_tensor_rejects_dtype_not_dividing_extent(self): + # size/stride are byte counts; a dtype whose itemsize does not divide + # them cannot produce a correct view, so it must fail loudly rather + # than silently truncate. + region = self._region(size=6, stride=6) + with self.assertRaises(ValueError): + region.as_tensor(dtype=torch.float32) + + def test_bytes_per_page_sums_regions(self): + group = KvCacheLayerGroupLayout( + layer_group_id=0, + layer_ids=(0, 1), + window_size=None, + regions=(self._region(size=256), self._region(base=8192, size=128)), + ) + self.assertEqual(group.bytes_per_page, 384) + + def test_layout_lookup_by_group_and_layer(self): + group_a = KvCacheLayerGroupLayout(0, (0, 2), None, ()) + group_b = KvCacheLayerGroupLayout(1, (1, 3), 128, ()) + layout = KvCacheLayout(tokens_per_block=8, groups=(group_a, group_b)) + + self.assertIs(layout.group(1), group_b) + self.assertIs(layout.group_of_layer(2), group_a) + self.assertIs(layout.group_of_layer(3), group_b) + with self.assertRaises(KeyError): + layout.group(7) + with self.assertRaises(KeyError): + layout.group_of_layer(99) + + +class TestKvCacheRegionAliasing(unittest.TestCase): + """as_tensor must alias the exact bytes address_of names.""" + + def setUp(self): + torch.cuda.init() + + def test_as_tensor_aliases_strided_slots(self): + # Lay out 4 "slots" of 32 bytes each, and describe the middle 8 bytes + # of every slot as a region. Writing through the view must land at + # base + stride * i, and must not disturb neighbouring bytes. + num_slots, stride, offset, size = 4, 32, 8, 8 + backing = torch.zeros(num_slots * stride, dtype=torch.uint8, device="cuda") + + region = KvCacheRegion( + base=backing.data_ptr() + offset, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + view = region.as_tensor() + self.assertEqual(tuple(view.shape), (num_slots, size)) + + for slot in range(num_slots): + view[slot] = slot + 1 + + flat = backing.cpu() + for slot in range(num_slots): + start = slot * stride + self.assertTrue( + bool((flat[start + offset : start + offset + size] == slot + 1).all()), + f"slot {slot} payload not written at the address address_of() names", + ) + # Bytes outside the described range must be untouched. + self.assertTrue(bool((flat[start : start + offset] == 0).all())) + self.assertTrue(bool((flat[start + offset + size : start + stride] == 0).all())) + + +class TestBuildKvCacheLayoutV2(unittest.TestCase): + """The builder against a real KVCacheManagerV2.""" + + def setUp(self): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + + def tearDown(self): + gc.collect() + torch.cuda.empty_cache() + + def test_layout_covers_every_layer_exactly_once(self): + num_layers = 4 + mgr = KVCacheManagerV2(**_make_kwargs(num_layers=num_layers)) + try: + layout = build_kv_cache_layout_v2(mgr) + + self.assertEqual(layout.tokens_per_block, mgr.tokens_per_block) + self.assertTrue(layout.groups, "layout must describe at least one layer group") + + covered = [lid for group in layout.groups for lid in group.layer_ids] + self.assertCountEqual( + covered, + list(mgr.pp_layers), + "every local layer must appear in exactly one layer group", + ) + + # Every layer that owns storage must be reachable through a region. + in_regions = [ + ref.layer_id + for group in layout.groups + for region in group.regions + for ref in region.buffers + ] + self.assertCountEqual(set(in_regions), set(covered)) + finally: + mgr.shutdown() + del mgr + + def test_regions_are_disjoint_and_inside_the_slot(self): + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + pool_groups = list(mgr.impl.pool_group_descs) + self.assertEqual(len(pool_groups), 1, "test config should yield one pool group") + pool_group = pool_groups[0] + self.assertEqual(len(pool_group.pools), 1, "test config should yield one pool") + pool = pool_group.pools[0] + + for group in layout.groups: + spans = [] + for region in group.regions: + self.assertEqual(region.num_slots, int(pool_group.num_slots)) + self.assertEqual(region.stride, int(pool.slot_bytes)) + + offset = region.base - int(pool.base_address) + self.assertGreaterEqual(offset, 0) + self.assertLessEqual( + offset + region.size, + int(pool.slot_bytes), + "a region must lie inside one slot", + ) + spans.append((offset, offset + region.size)) + + spans.sort() + for (_, prev_end), (next_start, _) in zip(spans, spans[1:]): + self.assertLessEqual(prev_end, next_start, "regions must not overlap") + finally: + mgr.shutdown() + del mgr + + def test_region_addresses_agree_with_pool_descriptor(self): + # Cross-check: region base/stride come from get_aggregated_pages, while + # pool base_address/slot_bytes come from pool_group_descs. These are + # independent public APIs and must agree on where slot i lives. + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + pool = list(mgr.impl.pool_group_descs)[0].pools[0] + pool_base, slot_bytes = int(pool.base_address), int(pool.slot_bytes) + + for group in layout.groups: + for region in group.regions: + offset = region.base - pool_base + for slot in (0, 1, region.num_slots - 1): + self.assertEqual( + region.address_of(slot), + pool_base + slot_bytes * slot + offset, + ) + finally: + mgr.shutdown() + del mgr + + def test_uniform_model_yields_one_full_slot_region(self): + # With uniform K/V sizes every buffer in a layer group is adjacent, so + # coalescing should collapse them into a single region spanning the + # whole slot -- the efficient whole-page transfer, derived rather than + # assumed. + mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4)) + try: + layout = build_kv_cache_layout_v2(mgr) + pool = list(mgr.impl.pool_group_descs)[0].pools[0] + + self.assertEqual(len(layout.groups), 1) + group = layout.groups[0] + self.assertEqual(len(group.regions), 1) + region = group.regions[0] + self.assertEqual(region.size, int(pool.slot_bytes)) + self.assertEqual(region.base, int(pool.base_address)) + # 4 layers * (K, V) + self.assertEqual(len(region.buffers), 8) + self.assertEqual( + [ref.role for ref in region.buffers], + ["key", "value"] * 4, + ) + finally: + mgr.shutdown() + del mgr + + def test_full_attention_reports_no_window(self): + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + for group in layout.groups: + self.assertIsNone(group.window_size) + finally: + mgr.shutdown() + del mgr + + def test_mla_layout_has_no_value_buffers(self): + # SELFKONLY carries a single compressed latent per token. Nothing in the + # layout counts K against V, so this must simply come out as a layout + # with no "value" role rather than needing a kv-factor special case. + mgr = KVCacheManagerV2(**_make_kwargs(kv_cache_type=CacheType.SELFKONLY)) + try: + layout = build_kv_cache_layout_v2(mgr) + roles = { + ref.role + for group in layout.groups + for region in group.regions + for ref in region.buffers + } + self.assertIn("key", roles) + self.assertNotIn("value", roles) + finally: + mgr.shutdown() + del mgr + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 679e1d795022..6b248c1f693b 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -586,8 +586,13 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): "max_beam_width, has_connector, expected", [ (2, False, "max_beam_width > 1"), - (1, True, "kv_connector_manager"), - (2, True, "kv_connector_manager, max_beam_width > 1"), + # A KV connector alone no longer forces a V1 fallback: it is supported on + # KVCacheManagerV2 through the pool-layout registration path, so V2 is + # returned unchanged and nothing is raised. + (1, True, None), + # With beam search still incompatible, the connector must not appear in + # the reason list -- it is not what makes this configuration unsupported. + (2, True, "max_beam_width > 1"), ], ) def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( @@ -606,6 +611,15 @@ def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( creator._kv_connector_manager = object() if has_connector else None creator._max_beam_width = max_beam_width + if expected is None: + assert ( + creator._fallback_if_unsupported_kv_cache_manager_v2( + MambaHybridCacheManagerV2, model_config, KvCacheConfig() + ) + is MambaHybridCacheManagerV2 + ) + return + with pytest.raises(NotImplementedError, match=expected): creator._fallback_if_unsupported_kv_cache_manager_v2( MambaHybridCacheManagerV2, model_config, KvCacheConfig() From 3f369f21a4855acb9a6729866f0a428a785b7e04 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:39 +0800 Subject: [PATCH 05/10] [None][feat] Drive the connector's scheduler hooks from KVCacheManagerV2 Registration alone leaves the connector inert: it receives its pools but no metadata, because its scheduler-side hooks are never invoked. The two managers place allocation differently. V1 allocates in KVCacheManager.prepare_resources and drives the connector from there. V2 allocates in KVCacheV2Scheduler, via prepare_context and resize_context, and its prepare_resources is a no-op for the non-draft path -- so the hooks had no home and build_connector_meta was never called. prepare_resources still runs after scheduling and before the forward pass, and by then every scheduled request has its pages, which is the same position in the iteration where V1 drives the connector. The hooks go there. Page indices come from the manager rather than from the connector reaching into kv_cache_map and impl.layer_grouping. They are reported per layer group and positionally aligned: valid_only=False yields one entry per block ordinal, so a block with no page in a group -- the sliding-window case -- reads back as BAD_PAGE_INDEX in place rather than shortening the list. That is what keeps the ordinal-to-token-range mapping intact and an append-delta over successive calls valid. RequestData carries them as new_block_ids_by_layer_group; with a single group -- every non-VSWA, non-hybrid model -- new_block_ids still carries that group's indices, so connectors that do not reason about layer groups keep working. request_finished is routed through the same accessor. Without it the V1 lookup raises, the surrounding warning path swallows the exception, and a connector on V2 is never told to save anything. Block hashes and retention priorities are reported empty rather than guessed at. V2 has no per-request accessor for the hash chain, and KvCacheRetentionConfig does not reach KVCacheManagerV2 at all -- per-page priority comes from custom_priority_callback, which V2 never overrides, so every page carries the default. Reporting those defaults would be worse than reporting nothing: a connector doing priority-based offload filtering would act on values unrelated to what the user configured. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/kv_cache_connector.py | 58 ++++++++++++++---- .../_torch/pyexecutor/kv_cache_manager_v2.py | 59 +++++++++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 12 +++- .../defs/llmapi/test_llm_api_connector.py | 9 +++ 4 files changed, 127 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 13cff6298b8b..3fa46e6fad1e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -86,6 +86,11 @@ class RequestData: # remote object id) MUST mix cache_salt into their identifiers, # otherwise blocks from a different salt could be incorrectly reused. cache_salt: Optional[str] = None + # New page slot indices keyed by layer group, populated under + # KVCacheManagerV2. Blocks with no page in a group -- the sliding-window + # case -- appear as BAD_PAGE_INDEX in place, keeping ordinals stable. + # Empty under the V1 manager, whose block IDs are a single flat space. + new_block_ids_by_layer_group: Dict[int, List[int]] = field(default_factory=dict) # A class to store some basic data regarding all inflight requests. @@ -341,21 +346,49 @@ def loading_ids(self) -> Set[int]: class KvCacheConnectorSchedulerOutputRequest: def __init__(self): self.block_ids = [] + self.block_ids_by_layer_group: Dict[int, List[int]] = {} self.tokens = [] def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManager"): - block_ids = kv_cache_manager.get_cache_indices(req) + from ..kv_cache_manager_v2 import KVCacheManagerV2 + + is_v2 = isinstance(kv_cache_manager, KVCacheManagerV2) tokens = req.get_tokens(0) - # Commit hashes for any blocks that have become full since the last call - # and read back the full cumulative chain. The C++ side sets each block's - # mBlockKey/mHash on first call, so subsequent calls become pure lookups. - block_hashes = kv_cache_manager.commit_and_get_block_hashes(req) + new_block_ids_by_layer_group: Dict[int, List[int]] = {} + if is_v2: + # Block hashes and retention priorities have no V2 accessor yet, so + # they are reported empty rather than guessed at. + block_hashes = [] + indices_by_group = kv_cache_manager.get_page_indices_by_layer_group(req) + for layer_group_id, indices in indices_by_group.items(): + seen = self.block_ids_by_layer_group.setdefault(layer_group_id, []) + new_ids = indices[len(seen) :] + seen.extend(new_ids) + new_block_ids_by_layer_group[layer_group_id] = new_ids + # With a single layer group -- every non-VSWA, non-hybrid model -- + # ``new_block_ids`` carries that group's indices so connectors that + # do not reason about layer groups keep working. With several groups + # it is left empty and ``new_block_ids_by_layer_group`` is the only + # correct source. + new_block_ids = ( + next(iter(new_block_ids_by_layer_group.values())) + if len(new_block_ids_by_layer_group) == 1 + else [] + ) + self.block_ids.extend(new_block_ids) + else: + block_ids = kv_cache_manager.get_cache_indices(req) + + # Commit hashes for any blocks that have become full since the last call + # and read back the full cumulative chain. The C++ side sets each block's + # mBlockKey/mHash on first call, so subsequent calls become pure lookups. + block_hashes = kv_cache_manager.commit_and_get_block_hashes(req) - new_block_ids = block_ids[len(self.block_ids) :] - new_tokens = tokens[len(self.tokens) :] + new_block_ids = block_ids[len(self.block_ids) :] + self.block_ids.extend(new_block_ids) - self.block_ids.extend(new_block_ids) + new_tokens = tokens[len(self.tokens) :] self.tokens.extend(new_tokens) if req.state in ( @@ -371,9 +404,13 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag ) # Specdec with draft tokens is not supported yet. # Get retention priority for each new block only if retention config is provided - # (for priority-based offload filtering) + # (for priority-based offload filtering). Priorities stay None on + # KVCacheManagerV2: it does not implement `KvCacheRetentionConfig` at + # all -- its per-page priority comes from `custom_priority_callback`, + # which KVCacheManagerV2 never overrides -- so every page carries the + # default and reporting it would misdescribe what the user asked for. priorities = None - if req.kv_cache_retention_config is not None: + if not is_v2 and req.kv_cache_retention_config is not None: priorities = [ kv_cache_manager.get_priority_by_block_id(block_id) for block_id in new_block_ids ] @@ -387,6 +424,7 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag block_hashes=block_hashes, priorities=priorities, cache_salt=req.cache_salt, + new_block_ids_by_layer_group=new_block_ids_by_layer_group, ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 7fe7876fb433..ed90bae2e398 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2441,6 +2441,40 @@ def resume_request(self, req: LlmRequest) -> bool: # ---- prepare_resources ---- @nvtx_range("prepare_resources_kv_cache_manager_v2") + def get_page_indices_by_layer_group(self, request: LlmRequest) -> Dict[int, List[int]]: + """Per-layer-group page slot indices for ``request``, by block ordinal. + + ``valid_only=False`` is deliberate: the positionally-aligned form yields + one entry per block ordinal, which is what preserves the ordinal-to-token + -range mapping a connector needs. A block with no page in a given layer + group -- the sliding-window case -- reads back as ``BAD_PAGE_INDEX`` in + place rather than shortening the list, so ordinals stay stable and an + append-delta over successive calls remains valid. + """ + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return {} + return { + layer_group_id: list( + kv_cache.get_aggregated_page_indices(layer_group_id, valid_only=False) + ) + for layer_group_id in range(len(self.impl.layer_grouping)) + } + + def get_connector_page_indices(self, request: LlmRequest) -> List[int]: + """Flat page slot indices for ``request``, for connectors. + + A page index is scoped to a layer group, so there is no correct way to + flatten indices from several groups into one list. With a single group + -- every non-VSWA, non-hybrid model -- that group's indices are the flat + list; with several, connectors must read + ``new_block_ids_by_layer_group`` off the scheduler output instead. + """ + indices_by_group = self.get_page_indices_by_layer_group(request) + if len(indices_by_group) != 1: + return [] + return next(iter(indices_by_group.values())) + def prepare_resources(self, scheduled_batch: ScheduledRequests): if self.is_draft: # Draft V2 manager: mirror the main manager by creating/resizing @@ -2449,6 +2483,31 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self._prepare_draft_resources(scheduled_batch) return + # V2 allocates in KVCacheV2Scheduler (prepare_context / resize_context) + # rather than here, so by this point every scheduled request already has + # its pages. This is the same point in the iteration at which the V1 + # manager drives the connector's scheduler-side hooks, and page indices + # are available, so the connector is driven from here. + if self.kv_connector_manager is not None: + self._run_kv_connector_hooks(scheduled_batch) + + def _run_kv_connector_hooks(self, scheduled_batch: ScheduledRequests) -> None: + """Report freshly allocated pages, then build the connector metadata.""" + for request in scheduled_batch.context_requests: + # Mirror V1, which reports allocation once per sequence, right after + # the sequence is added -- that is the first context chunk. + if not request.is_first_context_chunk: + continue + # Connectors that do not reason about layer groups see the single + # group's indices; with several groups there is no correct flat + # list, so report none and leave them to the layer-group-aware + # metadata carried on RequestData. + self.kv_connector_manager.update_state_after_alloc( + request, self.get_connector_page_indices(request) + ) + + self.kv_connector_manager.build_scheduler_output(scheduled_batch, self) + def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): """Create/resize KV caches in the draft V2 manager for scheduled requests. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6ed64504842a..a6c2b4e8de38 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6087,7 +6087,17 @@ def _send_kv_async(self, scheduled_requests: List[LlmRequest]): def kv_connector_request_finished(req: LlmRequest): try: - cache_block_ids = self.kv_cache_manager.get_cache_indices(req) + # KVCacheManagerV2 has no primary-pool block list; its page + # indices are scoped to a layer group. Without this the lookup + # below raises, the warning path swallows it, and + # `request_finished` is never called at all on V2 -- so a + # connector is never told to save anything. + if isinstance(self.kv_cache_manager, KVCacheManagerV2): + cache_block_ids = self.kv_cache_manager.get_connector_page_indices( + req) + else: + cache_block_ids = self.kv_cache_manager.get_cache_indices( + req) except Exception as e: logger.warning( f"Unable to get cache blocks for request {req.py_request_id}. Skipping asynchronous saving: {e}" diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 685e2e6ea994..c5e6708e010e 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -758,6 +758,15 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): the RequestData.priorities field is populated with the correct per-block priorities based on the token ranges. + KNOWN GAP -- this fails on `kv_cache_manager_v2`, deliberately left failing. + `KvCacheRetentionConfig` does not reach KVCacheManagerV2 at all: V2's + per-page priority comes from `custom_priority_callback` + (kv_cache_manager_v2/_core/_kv_cache_manager.py), which KVCacheManagerV2 + never overrides, so every page carries the default priority and the + connector reports `priorities=None`. A user who sets a retention config on + V2 silently gets none of it -- not only through the connector. The + assertions below are the correct expectation for both managers and are kept + that way so the gap stays visible rather than being asserted away. """ BLOCK_SIZE = 32 NUM_INPUT_TOKENS = 64 # 2 blocks From 92feaacd7d9875ad20db920132b560276ed844ec Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:40 +0800 Subject: [PATCH 06/10] [None][feat] Let the KV connector reach KVCacheManagerV2 under VSWA py_executor_creator refused any multi-valued max_attention_window before a cache manager was ever chosen, so the V2 carve-out in PyExecutor below it was unreachable and every multi-layer-group path was dead code. VSWA allocates one pool per window size, which the V1 single-tensor registration cannot describe. A KvCacheLayout can: it carries one region set per layer group. So reject here only when V2 is definitively off. kv_cache_config.use_kv_cache_manager_v2 is tri-state (True / False / "auto"), and under "auto" the manager is not chosen yet, so defer -- PyExecutor re-checks against the manager it actually built and rejects there if the selection landed on V1. Layered rather than duplicated. Two tests cover what this unblocks. A uniform sliding window, where every layer shares one window: the registered layout must still collapse to a single layer group and record the window. And VSWA, where V1 still raises and V2 reports two groups with equal-length, positionally aligned index lists -- the flat new_block_ids being empty is the correct answer there, since a page index is scoped to a layer group and there is no correct way to flatten several. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../_torch/pyexecutor/py_executor_creator.py | 17 +- .../defs/llmapi/test_llm_api_connector.py | 147 ++++++++++++++++++ 2 files changed, 161 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index e28096ceaf17..c42ef07fc55a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -836,11 +836,22 @@ def drafting_loop_wrapper(model): "KV connector is only supported with guaranteed no evict scheduler policy." ) + # VSWA allocates one pool per window size, which the V1 connector + # registration cannot describe: it hands the worker a single primary + # pool tensor. KVCacheManagerV2 has no such limitation -- its layout + # describes one region set per layer group -- so only reject here when + # V2 is definitively off. `use_kv_cache_manager_v2` is tri-state + # (True / False / "auto"), and under "auto" the manager is not chosen + # yet, so defer: PyExecutor re-checks against the manager it actually + # built (py_executor.py, `_maybe_init_kv_connector_manager`) and rejects + # there if the selection landed on V1. max_attention_window = kv_cache_config.max_attention_window - if max_attention_window is not None and len( - set(max_attention_window)) > 1: + if (max_attention_window is not None + and len(set(max_attention_window)) > 1 + and kv_cache_config.use_kv_cache_manager_v2 is False): raise NotImplementedError( - "KV connector is not supported with VSWA (Variable Sliding Window Attention)." + "KV connector is not supported with VSWA (Variable Sliding Window Attention) " + "on the V1 KV cache manager. Set kv_cache_config.use_kv_cache_manager_v2=True." ) if mapping.enable_attention_dp: diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index c5e6708e010e..8cbdf5201514 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -636,6 +636,153 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, (CHUNK_SIZE * 2 + BLOCK_SIZE) / BLOCK_SIZE) +# Sliding-window coverage. `KvCacheConfig.max_attention_window` is repeated +# cyclically across layers (llm_args.py:3761-3765), so a one-element list gives +# every layer the same window -- one V2 layer group with a live window -- while +# a two-element list alternates and produces two. A window equal to +# `max_seq_len` is normalised to "no window" (kv_cache_manager_v2.py:856-858), +# which is how the full-attention half of the VSWA pair is spelled. +SWA_WINDOW = 64 +SWA_MAX_SEQ_LEN = 512 +SWA_NUM_INPUT_TOKENS = 256 + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_connector_uniform_sliding_window(enforce_single_worker, + model_with_connector, + use_kv_cache_manager_v2): + """Connector against a KV cache in which every layer slides. + + There was no sliding-window connector coverage before this test. Uniform + SWA is a single layer group, which is the configuration where the + connector's flat `new_block_ids` list is still well defined -- so this pins + the block-reporting contract, and the VSWA test below pins what happens once + that assumption breaks. + + V2 only. The V1 guard rejects *variable* windows + (py_executor_creator.py:845-850), so uniform SWA reaches V1's connector path + and then dies inside it: `commit_and_get_block_hashes` (kv_cache_connector.py:386) + raises "commitAndGetBlockHashesForRequest does not support sliding-window + attention with detached front blocks" (kvCacheManager.cpp:4645) as soon as + the window drops a front block. That is a pre-existing V1 limitation, not + something this work introduces, so it is recorded rather than asserted here. + """ + model_fn, scheduler, worker = model_with_connector + + model = model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW])) + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + sched_output = scheduler.build_connector_meta.call_args_list[0].args[0] + assert len(sched_output.new_requests) == 1 + req = sched_output.new_requests[0] + + assert req.computed_position == 0 + assert req.num_scheduled_tokens == SWA_NUM_INPUT_TOKENS + # A single layer group keeps the flat list meaningful on both managers. + assert req.new_block_ids + + if use_kv_cache_manager_v2: + # Anti-vacuity: prove the window really did collapse to one layer + # group, otherwise the assertion above would hold for the wrong reason. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 1 + assert layout.groups[0].window_size == SWA_WINDOW + assert list(req.new_block_ids_by_layer_group) == [0] + assert req.new_block_ids_by_layer_group[0] == req.new_block_ids + else: + assert req.new_block_ids_by_layer_group == {} + + +# Half the prompt, and well past the window, so the pages the offer does *not* +# need are a large enough fraction to assert on. +SWA_OFFER_TOKENS = 128 + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_vswa_reports_page_indices_per_layer_group( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """VSWA is where the connector's flat block list stops working. + + V1 cannot describe VSWA to a connector at all: it registers a single primary + pool, but VSWA allocates one pool per window size. V2's layout describes one + region set per layer group, so the combination runs there -- but a page index + is scoped to a layer group, and there is no correct way to flatten indices + from several groups into one list. `KVCacheManagerV2` therefore reports an + empty `new_block_ids` and leaves `new_block_ids_by_layer_group` as the only + correct source (kv_cache_manager_v2.py:2519-2526, kv_cache_connector.py:365-377). + + Both halves are asserted here because the rejection is deliberately + conditional on the manager (py_executor_creator.py). Pinning only the V2 half + would let the V1 guard silently disappear. + """ + model_fn, scheduler, worker = model_with_connector + + def build(): + return model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW, SWA_MAX_SEQ_LEN])) + + if not use_kv_cache_manager_v2: + with pytest.raises(NotImplementedError, match="VSWA"): + build() + return + + model = build() + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + # Alternating windows must actually produce two groups, one sliding and one + # full-attention, or the rest of this test is vacuous. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 2 + assert {group.window_size for group in layout.groups} == {SWA_WINDOW, None} + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # The multi-group degradation reaches `update_state_after_alloc` too. A + # connector reading only `new_block_ids` sees nothing and saves nothing, + # which is silent unless pinned here. + assert scheduler.update_state_after_alloc.call_args.args[1] == [] + + sched_output = scheduler.build_connector_meta.call_args_list[0].args[0] + assert len(sched_output.new_requests) == 1 + req = sched_output.new_requests[0] + + assert req.new_block_ids == [] + assert sorted(req.new_block_ids_by_layer_group) == [0, 1] + # Ordinals stay positionally aligned across groups: a block with no page in + # the sliding group reads back as BAD_PAGE_INDEX in place rather than + # shortening the list (kv_cache_manager_v2.py:2481-2490). + lengths = { + len(indices) + for indices in req.new_block_ids_by_layer_group.values() + } + assert len(lengths) == 1 + + def _disagg_transceiver_config(use_kv_cache_manager_v2): """The transceiver each KV cache manager can actually be driven by. From 0512058b7e7c5e42cd81d252e592de04d44b69d3 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:41 +0800 Subject: [PATCH 07/10] [None][feat] Serve a connector-supplied prefix on KVCacheManagerV2 get_num_new_matched_tokens is a protocol, not a query: five of its six clauses are obligations on the runtime. V1 discharges them in one step because its entire connector interaction happens in prepare_resources, on the final batch. V2 cannot. prepare_context runs during scheduling, and a prepared request can still be dropped at the token budget, at resize_context, at multimodal alignment or at cross attention. Rolling computed_position back on an iteration the request never runs in is silent corruption: the next iteration arrives with the position already advanced, the connector loads nothing, and the runtime skips computing tokens whose KV is garbage. V2 already separates match-and-take-ownership -- which needs only the token sequence and cannot fail for lack of memory -- from residency, which claims slots, migrates pages and can fail. The connector interaction is decomposed along the same seam, with transmission as a third phase sharing the forward pass's preconditions: phase 1 _connector_prefix_position ask, before residency exists phase 2 _reserve_connector_prefix cover the offer with pages, after resume phase 3 _deliver_connector_prefix record the external load, on the batch Phase 2 has to follow the resume because _KVCache.resize asserts the cache is ACTIVE and a freshly created or deferred-and-suspended one is not. Capacity and history move in one call: after a reuse match both equal the local match, so raising history alone trips "History length cannot be greater than capacity". Raising history is also what stops a served prefix allocating a page for every block in a sliding-window layer group, since it is the sole input to the stale-range computation. Phase 2 clears enable_swa_scratch_reuse for a served prefix: scratch slots are transient prefill storage and a connector writes real cache content into those blocks, the same reason the disagg generation path opts out. The memoised value is the absolute offer end, never the returned delta. A deferred request re-derives its local match from a tree another request's commit may have grown; adding the delta to the new match would place the position past the union of what is computed and what is loaded, leaving tokens that are neither. The offer end is also clamped below the prompt, because the first generation step consumes the last prompt position's activations -- a connector holding the whole prompt is the steady state of a repeat, not an error. KvCacheConnectorManager.get_num_new_matched_tokens is split into a side-effect-free query and a commit that registers the async hold and the external load. V1 keeps calling the fused entry point, which is correct there: it is invoked from C++ under the block manager's tree mutex, so the match and the query are atomic with respect to the tree and the answer can be committed immediately. cancel_load is additive with a no-op default. It hands an offer back when phase 2 cannot allocate, and when a request is asked and then cancelled, times out or fails before delivery -- otherwise the connector holds remote blocks for the life of the process with nothing left to release them. should_add_sequence stays out of the V2 scheduler. That predicate is false from the moment an asynchronous load completes until request_finished at the end of generation; in the V2 scheduler it means SKIP, so V2 would skip such a request forever and never run the prefill the load was for. What keeps a loading request out of the batch is its DISAGG_GENERATION_TRANS_IN_PROGRESS state, and what stops the connector being asked or told twice is the per-request state machine. Phase 3 asserts that what it records is covered by the position phase 2 advanced. Nothing downstream validates it: the runtime's subtraction is unguarded and connectors divide the result into block ordinals, so a mismatch would silently point a connector at the wrong offset inside its own code. The unit suite models a suspended cache and drives the real _prepare_context_impl rather than restating it, so the ordering is observed rather than declared. Verified by mutation: moving the reserve above the resume turns 17 of its tests red, where an earlier version of the same suite passed all 24 against that same defect. A second suite drives the phases against a real KVCacheManagerV2 with real pools and a deterministic deferral, which is the only place the ask-once path is exercised end to end -- an engine-level deferral test proved to be a race against the executor's request queue and was removed rather than weakened. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/kv_cache_connector.py | 62 +- .../_torch/pyexecutor/kv_cache_manager_v2.py | 177 ++++- tensorrt_llm/_torch/pyexecutor/llm_request.py | 12 + .../pyexecutor/scheduler/scheduler_v2.py | 8 + .../defs/llmapi/test_llm_api_connector.py | 123 +++- .../executor/test_kv_connector_v2_prefix.py | 630 ++++++++++++++++++ ...est_kv_connector_v2_prefix_real_manager.py | 238 +++++++ tests/unittest/_torch/test_connector.py | 77 ++- 8 files changed, 1316 insertions(+), 11 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py create mode 100644 tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 3fa46e6fad1e..49f6f27e9998 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -285,6 +285,30 @@ def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): block_ids: The KV cacheblock IDs that were allocated. """ + def cancel_load(self, request: LlmRequest, start: int, end: int): + """ + Called when the runtime will not consume KV that was offered by + ``get_num_new_matched_tokens`` for prompt tokens ``[start, end)``. + + Offsets are absolute prompt positions, on the same scale as + ``num_computed_tokens``. This is raised when the local cache overtook + part or all of an offer while the request was waiting to be scheduled, + or when the runtime could not allocate pages to cover it. + + Best-effort: release any ownership taken in + ``get_num_new_matched_tokens`` for that range. For a synchronous load + nothing has been transferred yet, so this is exact. For an + asynchronous load the transfer necessarily began inside + ``get_num_new_matched_tokens`` (see ``take_scheduled_requests_pending_load``), + so it may already be in flight and cancelling is genuinely lossy. + + Args: + request: The request whose offer is being handed back. + start: First prompt position that will not be consumed. + end: One past the last prompt position that will not be consumed. + """ + return + def wait_for_initialization(self): """ Some connectors need to wait for some resources to be initialized. @@ -526,7 +550,18 @@ def _run_on_leader(self, f: Callable[[], Any]) -> Any: res = None return mpi_broadcast(res, root=0) - def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> int: + def query_num_new_matched_tokens( + self, request: LlmRequest, num_computed_tokens: int + ) -> Tuple[int, bool]: + """Ask the connector how much of the prompt it can serve. No side effects. + + This is the half of ``get_num_new_matched_tokens`` that must run at most + once per request, because the connector ABC promises exactly one query + per request and connectors take ownership of remote blocks in it. + Callers that cannot commit to consuming the answer in the same iteration + (KVCacheManagerV2, whose scheduling pass is speculative) query here and + commit later via ``commit_new_matched_tokens``. + """ if request.is_generation_only_request: raise RuntimeError("Connector API is not supported for generation-only requests!") @@ -537,6 +572,18 @@ def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: i if num_tokens == 0 and load_kv_async: raise RuntimeError("load_kv_async must be False when num_tokens is 0!") + return num_tokens, load_kv_async + + def commit_new_matched_tokens( + self, request: LlmRequest, num_tokens: int, load_kv_async: bool + ) -> None: + """Register the runtime's side of an answered query. + + Must run in the iteration the request is actually scheduled: + ``external_loads`` is consumed and cleared by every + ``build_scheduler_output``, and ``new_async_requests.loading`` is read + by that same call to suppress the request's ``RequestData``. + """ # TODO(jthomson04): This part is a bit ugly. # When the connector indicates that a request will be loaded # asynchronously, we need to suspend its execution. This is @@ -551,8 +598,21 @@ def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: i request.py_num_connector_matched_tokens = num_tokens + def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> int: + """Query and commit in one step. + + This is the V1 entry point, called from C++ while the block manager + holds the radix-tree mutex, so the local match and the query are atomic + with respect to the tree and the answer can be committed immediately. + """ + num_tokens, load_kv_async = self.query_num_new_matched_tokens(request, num_computed_tokens) + self.commit_new_matched_tokens(request, num_tokens, load_kv_async) return num_tokens + def cancel_load(self, request: LlmRequest, start: int, end: int) -> None: + if self.scheduler is not None: + self.scheduler.cancel_load(request, start, end) + def should_add_sequence(self, request: LlmRequest) -> bool: req_id = request.request_id return req_id not in self.finished_async_loading_requests diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index ed90bae2e398..425aa68eb472 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2311,11 +2311,22 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: kv_cache.num_committed_tokens, self.tokens_per_block ) + # Connector phase 1: ask what the connector can serve past the local + # match and advance the context position over it, before residency + # is decided. The capacity/history bump this implies has to wait for + # the resume below, because `resize` requires an ACTIVE cache and a + # freshly created one is suspended. + position = self._connector_prefix_position(req, kv_cache) + if req.is_disagg_generation_init_state: # Disagg generation receives prompt KV from the context worker; # scratch blocks are only valid for local prefill chunks. kv_cache.enable_swa_scratch_reuse = False - return self._resume_and_restore(req.py_request_id, kv_cache) + if not self._resume_and_restore(req.py_request_id, kv_cache): + return False + if position is not None: + self._reserve_connector_prefix(req, kv_cache, position) + return True else: # Subsequent chunk: cache must exist from first chunk. # It may be suspended (e.g., evicted between chunks), so @@ -2326,6 +2337,161 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: ) return self._resume_and_restore(req.py_request_id, kv_cache) + # ---- KV connector prefix (see docs/source/features/kv-cache-connector.md) ---- + # + # V2 splits "match and take ownership" from "become resident on GPU", and the + # connector gets the same decomposition: + # + # phase 1 _connector_prefix_position ask, and advance the context position + # phase 2 _reserve_connector_prefix cover the offer with pages + # phase 3 _deliver_connector_prefix tell the connector what to transfer + # + # Phases 1 and 2 run in the scheduling pass, which is speculative: a request + # can still be dropped afterwards at the token budget, `resize_context`, + # multimodal alignment or cross attention. Phase 3 runs on the final batch. + + def _connector_may_serve(self, req: LlmRequest) -> bool: + if req.is_dummy: + # Mirrors V1, which skips the query for dummy requests entirely. + return False + if req.is_generation_only_request: + # The connector API rejects these outright; a disagg generation + # server gets its prompt KV from the context worker. + return False + return not req.is_disagg_generation_init_state + + def _connector_prefix_position(self, req: LlmRequest, kv_cache) -> Optional[int]: + """Phase 1. Position through which KV is valid before the forward pass. + + That is ``max(locally committed, connector offer end)``, or None when no + connector prefix applies to this request. The connector is queried at + most once per request: a deferred request re-derives its position here + from the memoised offer end and the *current* local match, without + asking again. + """ + if self.kv_connector_manager is None or self.is_draft: + return None + + if req.py_connector_prefix_end is None: + if not self._connector_may_serve(req): + return None + local_end = kv_cache.num_committed_tokens + num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( + req, local_end + ) + # The last prompt position must be computed locally whatever the + # connector holds, because the first generation step consumes its + # activations. A connector that offers the whole prompt therefore + # loses at most one token, which the forward recomputes. + req.py_connector_prefix_start = local_end + req.py_connector_prefix_end = min(local_end + num_tokens, req.prompt_len - 1) + # Deliberately not conditioned on the clamp above: if the connector + # said it would transfer asynchronously it has already started, and + # the request must be held out of the batch until it reports done. + req.py_connector_load_async = load_async + + # Not `py_connector_prefix_end` on its own: while the request waited, + # another request may have committed blocks that grew the local match + # past the offer. + return max(kv_cache.num_committed_tokens, req.py_connector_prefix_end) + + def _reserve_connector_prefix(self, req: LlmRequest, kv_cache, position: int) -> None: + """Phase 2. Give the offered prefix capacity, history and pages. + + Runs after ``_resume_and_restore`` because ``resize`` asserts the cache + is ACTIVE. Capacity and history move together in one call: after a reuse + match they are both equal to the local match, so raising history alone + would trip "History length cannot be greater than capacity". + + Raising history is what keeps a served prefix from allocating a page per + block in a sliding-window layer group -- ``history_length`` is the sole + input to the stale-range computation. + """ + local_end = kv_cache.num_committed_tokens + if position <= local_end: + # Nothing offered, or the local match has caught up with the offer. + return + + # SWA scratch slots are transient prefill storage; a connector writes + # real cache content into these blocks. Same reason the disagg + # generation path above opts out. + kv_cache.enable_swa_scratch_reuse = False + + if kv_cache.resize(max(kv_cache.capacity, position), position): + req.context_current_position = position + req.set_prepopulated_prompt_len(position, self.tokens_per_block) + return + + # Out of pages. Run local-only rather than leaving the offered range + # unallocated, and hand the offer back instead of silently dropping it. + logger.debug( + "req %s: could not reserve connector prefix up to %d, falling back to " + "the local match at %d", + req.py_request_id, + position, + local_end, + ) + self.kv_connector_manager.cancel_load(req, req.py_connector_prefix_start, position) + req.py_connector_prefix_start = local_end + req.py_connector_prefix_end = local_end + req.py_connector_load_async = False + + def _deliver_connector_prefix(self, req: LlmRequest) -> None: + """Phase 3. Resolve the offer against ownership as it stands now. + + The connector was told it could serve ``[start, end)``. Anything below + the current commit boundary is now locally owned -- shared + ``CommittedPage`` s that this request is not the only writer to -- so it + is handed back rather than transferred, and only ``[committed, end)`` + is reported as an external load. + """ + end = req.py_connector_prefix_end + if end is None: + return + start = req.py_connector_prefix_start + kv_cache = self.kv_cache_map.get(req.py_request_id) + # Never past `end`: the offer bounds what there is to hand back. + committed = min(kv_cache.num_committed_tokens, end) if kv_cache is not None else start + if committed > start: + self.kv_connector_manager.cancel_load(req, start, committed) + recorded = end - committed + # The one place phases 2 and 3 have to agree. The runtime reports + # `context_current_position - recorded` to the connector as the range it + # computed locally, and nothing downstream validates that: the + # subtraction is unguarded (kv_cache_connector.py:480) and connectors + # divide the result into block ordinals rather than checking it (see + # `computed_position // block_size` in + # examples/llm-api/llm_kv_cache_connector.py). So a `recorded` the + # position was never advanced over does not fail here -- it silently + # points the connector at the wrong offset, inside its own code. Fail + # loudly and locally instead. + assert 0 <= recorded <= req.context_current_position, ( + f"req {req.py_request_id}: connector prefix [{start}, {end}) " + f"records {recorded} externally loaded tokens, but the context " + f"position is only {req.context_current_position} -- phase 2 did " + f"not reserve what phase 1 offered" + ) + self.kv_connector_manager.commit_new_matched_tokens( + req, recorded, req.py_connector_load_async + ) + + def _release_undelivered_connector_prefix(self, req: LlmRequest) -> None: + """Hand back an offer for a request that died before it was delivered. + + Phase 1 is speculative, so a request can be asked and then cancelled, + time out, or fail before it ever reaches a batch. The connector took + ownership of remote blocks in the query and would otherwise hold them + for the rest of the process's life. + """ + if self.kv_connector_manager is None or req.py_connector_delivered: + return + end = req.py_connector_prefix_end + start = req.py_connector_prefix_start + if end is None or end <= start: + return + self.kv_connector_manager.cancel_load(req, start, end) + req.py_connector_prefix_end = start + def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: """Resize KV cache to cover context_current_position + num_tokens. @@ -2498,6 +2664,14 @@ def _run_kv_connector_hooks(self, scheduled_batch: ScheduledRequests) -> None: # the sequence is added -- that is the first context chunk. if not request.is_first_context_chunk: continue + # A request that was loaded asynchronously re-enters this batch still + # on its first context chunk, with the same pages and nothing left to + # load. Reporting it again would re-record an external load and fire + # a second `update_state_after_alloc` for one allocation. + if request.py_connector_delivered: + continue + self._deliver_connector_prefix(request) + request.py_connector_delivered = True # Connectors that do not reason about layer groups see the single # group's indices; with several groups there is no correct flat # list, so report none and leave them to the layer-group-aware @@ -3246,6 +3420,7 @@ def release_index_slot(self, request_id: int) -> None: self._early_freed_index_requests.add(request_id) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + self._release_undelivered_connector_prefix(request) if self.conversation_manager is not None: self.conversation_manager.finish_request(request) self._allocated_draft_lens.pop(request.py_request_id, None) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 91fb34606a55..da1d851b5fac 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -812,6 +812,18 @@ def __init__( self.py_num_connector_matched_tokens = 0 + # KV connector query state, used by KVCacheManagerV2 (V1 answers the + # query under the block manager's tree mutex and needs no state). + # UNASKED (`py_connector_prefix_end is None`) -> ASKED -> DELIVERED. + # `py_connector_prefix_end` is the absolute prompt position one past + # the last token the connector offered to serve, deliberately not the + # returned delta: the request may be deferred and re-derive a larger + # locally-matched prefix, against which a delta would overshoot. + self.py_connector_prefix_start: Optional[int] = None + self.py_connector_prefix_end: Optional[int] = None + self.py_connector_load_async = False + self.py_connector_delivered = False + self.py_result = PyResult( prompt_len=self.py_prompt_len, max_new_tokens=self.py_max_new_tokens, diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 0d3afb573538..dd9b67132834 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -512,6 +512,14 @@ def _try_schedule_context( Returns ``(action, tokens, chunking_flag)``. *tokens* and *chunking_flag* are meaningful only when *action* is ``SCHEDULED``. """ + # No `should_add_sequence` gate here, unlike V1. That predicate stays + # false from the moment an asynchronous load completes until + # `request_finished`, which only runs at the end of generation, so V2 + # would SKIP the request forever and never run the prefill the load was + # for. What keeps a loading request out of the batch instead is its + # DISAGG_GENERATION_TRANS_IN_PROGRESS state, and what stops the + # connector being asked or told twice is the per-request query state + # (see KVCacheManagerV2._connector_prefix_position). if self.chunking_enabled: return self._try_schedule_context_chunked(req, budget) return self._try_schedule_context_full(req, budget) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 8cbdf5201514..9bbebc5069a6 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -596,9 +596,17 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, assert scheduler.update_state_after_alloc.call_count == 1 - assert len( - scheduler.update_state_after_alloc.call_args.args[1]) == math.ceil( - CHUNK_SIZE * 2 / BLOCK_SIZE) + # V1 allocates for the whole prompt when the sequence is added, so every + # block exists on the first chunk. V2 allocates per chunk, which is the + # lower-peak-memory behaviour and the one to keep; the remaining blocks + # arrive as append-deltas in `new_block_ids` on the next chunk, so no + # information is lost. The expectation is split rather than V2 changed. + total_blocks = math.ceil(CHUNK_SIZE * 2 / BLOCK_SIZE) + first_chunk_blocks = (math.ceil(CHUNK_SIZE / BLOCK_SIZE) + if use_kv_cache_manager_v2 else total_blocks) + + assert len(scheduler.update_state_after_alloc.call_args.args[1] + ) == first_chunk_blocks for i, call in enumerate(scheduler.build_connector_meta.call_args_list): sched_output = call.args[0] @@ -613,18 +621,18 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, req = sched_output.cached_requests[0] if i == 0: - # The first prefill chunk. - # All of the prefill tokens and all the blocks should be provided upfront. + # The first prefill chunk. All of the prefill tokens are provided + # upfront on both managers; the blocks are whatever has been + # allocated so far. assert req.computed_position == 0 assert len(req.new_tokens) == CHUNK_SIZE * 2 - assert len(req.new_block_ids) == math.ceil(CHUNK_SIZE * 2 / - BLOCK_SIZE) + assert len(req.new_block_ids) == first_chunk_blocks assert req.num_scheduled_tokens == CHUNK_SIZE elif i == 1: # The second prefill chunk. assert req.computed_position == CHUNK_SIZE assert len(req.new_tokens) == 0 - assert len(req.new_block_ids) == 0 + assert len(req.new_block_ids) == total_blocks - first_chunk_blocks assert req.num_scheduled_tokens == CHUNK_SIZE elif i == 2 and use_overlap_scheduler: assert len(req.new_tokens) == 0 @@ -636,6 +644,30 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, (CHUNK_SIZE * 2 + BLOCK_SIZE) / BLOCK_SIZE) +# The mock scheduler answers every query with the same offer; these helpers +# record what it was asked and when, so a test can assert the connector was +# consulted exactly once per request rather than assuming it. +def record_connector_queries(scheduler, num_matched, load_async=False): + """Answer every query with `num_matched`, recording when each one arrived. + + The third element of each record is `build_connector_meta.call_count` at + query time -- the number of iterations whose connector hooks had already + run. Phase 1 runs during scheduling, before those hooks, so a request asked + in the first scheduling pass records 0. That is how the tests below *prove* + a request was asked before the iteration it eventually ran in, rather than + assuming the scheduler deferred anything. + """ + queries = [] + + def side_effect(request, num_computed_tokens): + queries.append((request.request_id, num_computed_tokens, + scheduler.build_connector_meta.call_count)) + return num_matched, load_async + + scheduler.get_num_new_matched_tokens.side_effect = side_effect + return queries + + # Sliding-window coverage. `KvCacheConfig.max_attention_window` is repeated # cyclically across layers (llm_args.py:3761-3765), so a one-element list gives # every layer the same window -- one V2 layer group with a live window -- while @@ -712,6 +744,81 @@ def test_connector_uniform_sliding_window(enforce_single_worker, SWA_OFFER_TOKENS = 128 +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_connector_sliding_window_prefix_is_backed_by_history( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """Phase 2 raises `history_length`, not just capacity -- observably. + + Under full attention the two are indistinguishable from outside: capacity is + `prompt_len` whether or not a prefix was served, so the page count says + nothing. Under a sliding window it does. `history_length` is the sole input + to the stale-range computation, so raising it to the offer end tells V2 that + the blocks the window has already left need no pages -- and raising capacity + alone would allocate one for every block of the served prefix. + + So this is the test that would fail if `_reserve_connector_prefix` called + `resize(capacity, None)`; `test_connector_uniform_sliding_window` above + cannot, because it offers nothing. + """ + model_fn, scheduler, worker = model_with_connector + + model = model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW])) + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + record_connector_queries(scheduler, SWA_OFFER_TOKENS) + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # Anti-vacuity, as in the test above: one layer group with a live window, + # otherwise the page count below is being read off full attention. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 1 + assert layout.groups[0].window_size == SWA_WINDOW + + req = scheduler.build_connector_meta.call_args_list[0].args[0].new_requests[ + 0] + + # The offer was materialized: the position advanced over it, and the + # runtime rolled the reported position back to the local match. + assert req.computed_position == 0 + assert req.num_scheduled_tokens == SWA_NUM_INPUT_TOKENS - SWA_OFFER_TOKENS + + all_blocks = math.ceil(SWA_NUM_INPUT_TOKENS / 32) + page_indices = scheduler.update_state_after_alloc.call_args.args[1] + assert page_indices == req.new_block_ids + # One entry per block ordinal either way -- a block with no page reads back + # as BAD_PAGE_INDEX in place rather than shortening the list, so that block + # ordinals stay aligned to token ranges (kv_cache_manager_v2.py:2628-2636). + # The page count is therefore not the signal; *which* entries are bad is. + assert len(page_indices) == all_blocks + + # With history at the offer end, every block that lies entirely below the + # window got no page. Without the bump history would still be at the local + # match, nothing would be stale, and all 8 blocks would be backed. + stale_blocks = (SWA_OFFER_TOKENS - SWA_WINDOW) // 32 + assert stale_blocks > 0, "test sizes no longer put any block out of window" + assert page_indices[:stale_blocks] == [BAD_PAGE_INDEX] * stale_blocks, ( + f"expected the first {stale_blocks} blocks of a {SWA_OFFER_TOKENS}-token " + f"served prefix to fall outside a {SWA_WINDOW}-token window and hold no " + f"page, got {page_indices}. History was not raised to the offer end.") + + live = page_indices[stale_blocks:] + assert all(index != BAD_PAGE_INDEX for index in live), ( + f"a block inside the window has no page: {page_indices}") + assert len(set(live)) == len(live), ( + f"page slots reported to the connector are not distinct: {live}") + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py new file mode 100644 index 000000000000..25310a617a07 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -0,0 +1,630 @@ +"""Unit tests for the KV connector prefix phases on KVCacheManagerV2. + +KVCacheManagerV2's scheduling pass is speculative: a request can be prepared +and then dropped in the same iteration at the token budget, at +``resize_context``, at multimodal alignment or at cross attention, and retried +later. The connector ABC, meanwhile, promises exactly one +``get_num_new_matched_tokens`` per request and lets connectors take ownership +of remote blocks inside it. + +Reconciling those two is the whole point of the three phases exercised here, +and none of it is observable from the end-to-end connector tests, which never +defer a request. So these drive the phases directly against a stub cache. + +The stub cache models V2's own two-phase split, because that split is what +forces the connector phases apart: ``_create_kv_cache`` returns a **suspended** +cache, ``_KVCache.resize`` asserts the cache is ACTIVE, and +``_resume_and_restore`` is what activates it. So the offer has to be taken -- +and the context position advanced over it -- *before* the resume, while the +capacity and history that back it can only be reserved *after*. ``FakeKvCache`` +therefore starts suspended and refuses ``resize`` until resumed, so a version of +this code that folded the two phases into one would fail here rather than only +on hardware. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + +TOKENS_PER_BLOCK = 32 +PROMPT_LEN = 256 + + +class FakeKvCache: + """The parts of ``_KVCache`` the connector phases touch. + + Starts suspended, as ``_create_kv_cache`` leaves it. ``resize`` reproduces + the real assertions -- ACTIVE status, history never decreasing, history + never above capacity -- since those are exactly what dictate when each + phase may run and how it must call ``resize``. + """ + + def __init__(self, committed=0, resize_ok=True, active=False, trace=None): + self.num_committed_tokens = committed + self.capacity = committed + self.history_length = committed + self.enable_swa_scratch_reuse = True + self.resize_ok = resize_ok + self.is_active = active + self.resize_calls = [] + self.trace = trace if trace is not None else [] + + def resume(self, cuda_stream): + self.is_active = True + self.trace.append("resume") + return True + + def suspend(self): + self.is_active = False + self.trace.append("suspend") + + def resize(self, capacity, history_length=None): + # `_KVCache.resize` asserts ACTIVE. A cache that has just been created, + # or that was suspended when its request was deferred, is not. + assert self.is_active, "resize on a suspended cache" + self.resize_calls.append((capacity, history_length)) + self.trace.append(("resize", capacity, history_length)) + if not self.resize_ok: + return False + if history_length is not None: + if history_length < self.history_length: + raise ValueError("History length cannot be decreased") + if capacity is not None and capacity < history_length: + raise ValueError("History length cannot be greater than capacity") + self.history_length = history_length + if capacity is not None: + self.capacity = capacity + return True + + +class FakeRequest: + def __init__(self, request_id=0, prompt_len=PROMPT_LEN): + self.py_request_id = request_id + self.request_id = request_id + self.prompt_len = prompt_len + self.context_current_position = 0 + self.prepopulated_prompt_len = 0 + self.is_first_context_chunk = True + self.is_dummy = False + self.is_generation_only_request = False + self.is_disagg_generation_init_state = False + self.py_num_connector_matched_tokens = 0 + self.py_connector_prefix_start = None + self.py_connector_prefix_end = None + self.py_connector_load_async = False + self.py_connector_delivered = False + + def set_prepopulated_prompt_len(self, prepopulated_prompt_len, tokens_per_block): + assert prepopulated_prompt_len < self.prompt_len, ( + f"prepopulatedPromptLen ({prepopulated_prompt_len}) >= promptLen ({self.prompt_len})" + ) + self.prepopulated_prompt_len = prepopulated_prompt_len + if prepopulated_prompt_len > 0: + self.context_current_position = prepopulated_prompt_len + + +class FakeConnectorManager: + """Records the calls the phases make, in order.""" + + def __init__(self, num_matched=0, load_async=False, trace=None): + self.num_matched = num_matched + self.load_async = load_async + self.queries = [] + self.commits = [] + self.cancels = [] + self.trace = trace if trace is not None else [] + + def query_num_new_matched_tokens(self, request, num_computed_tokens): + self.queries.append((request.request_id, num_computed_tokens)) + self.trace.append(("query", num_computed_tokens)) + return self.num_matched, self.load_async + + def commit_new_matched_tokens(self, request, num_tokens, load_kv_async): + self.commits.append((request.request_id, num_tokens, load_kv_async)) + self.trace.append(("commit", num_tokens)) + request.py_num_connector_matched_tokens = num_tokens + + def cancel_load(self, request, start, end): + self.cancels.append((request.request_id, start, end)) + self.trace.append(("cancel", start, end)) + + +def make_manager(connector): + """A KVCacheManagerV2 with only the fields the connector phases read. + + Constructing a real one needs a GPU and a pool allocation; the phases under + test are pure request/cache bookkeeping, so bypass __init__ rather than + turning this into an integration test. + """ + manager = object.__new__(KVCacheManagerV2) + manager.kv_connector_manager = connector + manager.is_draft = False + manager.tokens_per_block = TOKENS_PER_BLOCK + manager.kv_cache_map = {} + # Read by `_prepare_context_impl` on the first-chunk path, so that the + # ordering tests below can drive the real thing rather than a re-statement + # of it. + manager.enable_block_reuse = True + manager.conversation_manager = None + manager._stream = SimpleNamespace(cuda_stream=0) + # Page-index buffer plumbing: needs the real IndexMapper and pool tensors, + # and has no bearing on which phase runs when. + manager._restore_page_index_bufs = lambda request_id, kv_cache: None + return manager + + +def prepare(manager, req, kv_cache): + """One scheduling attempt, driving the real ``_prepare_context_impl``. + + Seeding ``kv_cache_map`` skips the ``_create_kv_cache`` branch, which needs + a GPU; everything after it -- the local-match anchor, phase 1, the resume, + phase 2 -- is the production code, so the phase ordering is observed rather + than restated. + """ + manager.kv_cache_map[req.py_request_id] = kv_cache + assert manager._prepare_context_impl(req) + # A memoised re-read of the position the attempt settled on, not a second + # ask -- `py_connector_prefix_end` is set by now, so the connector is not + # consulted again. `TestAskOnce` pins that. + return manager._connector_prefix_position(req, kv_cache) + + +class TestTwoPhaseSplit: + """The connector phases are split because V2's own allocation is. + + V2 separates *match and take ownership*, which needs only the token + sequence, from *become resident on GPU*, which needs slots and can fail. + `_create_kv_cache` returns a suspended cache and `_KVCache.resize` asserts + ACTIVE, so the two connector steps straddle `_resume_and_restore`. Every + test in this class fails if they are folded into one. + """ + + def test_the_offer_is_taken_before_the_cache_is_resident(self): + trace = [] + connector = FakeConnectorManager(num_matched=64, trace=trace) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32, trace=trace) + + prepare(manager, req, kv_cache) + + assert trace == [("query", 32), "resume", ("resize", 96, 96)], ( + "the connector is asked while the cache is still suspended, and the " + "capacity backing its answer is reserved only once it is active" + ) + + def test_phase_one_touches_no_residency(self): + """Asking needs the token sequence and nothing else. + + This is what makes it safe to ask during a speculative scheduling pass: + no slot is claimed, so a request that is then dropped has cost nothing + locally. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + position = manager._connector_prefix_position(req, kv_cache) + + assert position == 96 + assert kv_cache.is_active is False + assert kv_cache.resize_calls == [] + assert kv_cache.capacity == 32 + + def test_phase_two_refuses_a_suspended_cache(self): + """The assertion that forces the split to exist. + + `_KVCache.resize` asserts ACTIVE, so reserving before the resume -- the + shape a single merged step would have -- fails on the first request + that has anything to reserve. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + position = manager._connector_prefix_position(req, kv_cache) + + with pytest.raises(AssertionError, match="suspended"): + manager._reserve_connector_prefix(req, kv_cache, position) + + def test_a_deferred_request_is_resumed_again_before_reserving(self): + """A deferred first chunk is suspended, so the split holds every attempt. + + `resize_context` suspends the cache when it cannot grow it, so the next + attempt starts from the suspended state again -- and must still ask + nothing, resume, and only then reserve. + """ + trace = [] + connector = FakeConnectorManager(num_matched=64, trace=trace) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32, trace=trace) + + prepare(manager, req, kv_cache) + kv_cache.suspend() + trace.clear() + + prepare(manager, req, kv_cache) + + assert trace == ["resume", ("resize", 96, 96)] + assert connector.queries == [(0, 32)] + + +class TestAskOnce: + def test_deferred_request_is_not_asked_again(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + # Iteration 1: prepared, then dropped before it reached the batch. + prepare(manager, req, kv_cache) + # Iteration 2: prepared again. + prepare(manager, req, kv_cache) + + assert connector.queries == [(0, 0)], ( + "the connector takes ownership of remote blocks inside the query, " + "so a second one double-pins them and breaks the ABC's " + "at-most-once promise" + ) + + def test_delivering_records_once_for_the_iteration_that_ran(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + assert connector.commits == [] + + manager._deliver_connector_prefix(req) + assert connector.commits == [(0, 64, False)] + + def test_unasked_request_delivers_nothing(self): + connector = FakeConnectorManager() + manager = make_manager(connector) + req = FakeRequest() + + manager._deliver_connector_prefix(req) + + assert connector.commits == [] + assert connector.cancels == [] + + +class TestOfferEndIsAbsolute: + def test_position_does_not_overshoot_when_the_local_match_grows(self): + """The reason the offer end is memoised rather than the returned delta. + + A deferred request re-derives its local match from the radix tree, and + another request's commit may have grown it in the meantime. Adding the + delta to the new match would set the position past the union of what is + locally computed and what the connector holds, leaving the tokens in + between neither computed nor loaded -- silently garbage KV. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + assert req.context_current_position == 64 + + # Another request committed the first 32 tokens while this one waited. + kv_cache.num_committed_tokens = 32 + position = prepare(manager, req, kv_cache) + + assert position == 64, "0 + 64 offered, so 64 -- not 32 + 64" + assert req.context_current_position == 64 + + def test_position_follows_the_local_match_when_it_overtakes_the_offer(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + + kv_cache.num_committed_tokens = 128 + position = prepare(manager, req, kv_cache) + + assert position == 128 + assert req.context_current_position == 128 + assert connector.queries == [(0, 0)] + + def test_offer_end_is_clamped_below_the_prompt(self): + """A connector holding the whole prompt is the steady state of a repeat. + + The last prompt position must be computed locally regardless, since the + first generation step consumes its activations, so the offer is clamped + rather than rejected -- and the *clamped* delta is what gets recorded, + or the connector is pointed at the wrong offset. + """ + connector = FakeConnectorManager(num_matched=PROMPT_LEN) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert req.py_connector_prefix_end == PROMPT_LEN - 1 + assert req.context_current_position == PROMPT_LEN - 1 + assert connector.commits == [(0, PROMPT_LEN - 1, False)] + + +class TestWriteRangeAtDelivery: + def test_subsumed_head_is_handed_back_and_not_transferred(self): + """``[start, committed)`` is locally owned by the time we deliver. + + Those are committed pages in the radix tree, potentially shared with + other live requests, and V2's rule is that only the owning request + writes into its own padding. So the connector is told to transfer only + what is still privately owned, and the rest is cancelled. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + kv_cache.num_committed_tokens = 32 + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert connector.cancels == [(0, 0, 32)] + assert connector.commits == [(0, 32, False)], "[32, 64), not [0, 64)" + + def test_fully_subsumed_offer_transfers_nothing(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + kv_cache.num_committed_tokens = 128 + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert connector.cancels == [(0, 0, 64)] + assert connector.commits == [(0, 0, False)] + + def test_recorded_delta_restores_the_locally_computed_position(self): + """``computed_position`` is reported as ``end - recorded``. + + That is what a connector uses to find where its blocks start, so the + recorded delta has to be the one anchored at the *current* commit + boundary. + """ + connector = FakeConnectorManager(num_matched=96) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + kv_cache.num_committed_tokens = 64 + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + recorded = connector.commits[0][1] + assert req.context_current_position - recorded == 64 + + +class TestResidency: + def test_capacity_and_history_move_together(self): + """After a reuse match both equal the local match. + + Raising history alone would trip "History length cannot be greater than + capacity"; raising capacity alone would leave a sliding-window group + allocating a page for every block of the served prefix. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + + assert kv_cache.resize_calls == [(96, 96)] + assert kv_cache.capacity == 96 + assert kv_cache.history_length == 96 + + def test_existing_capacity_is_not_shrunk(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + kv_cache.capacity = 256 + + prepare(manager, req, kv_cache) + + assert kv_cache.resize_calls == [(256, 64)] + + def test_swa_scratch_reuse_is_disabled_for_a_served_prefix(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + + assert kv_cache.enable_swa_scratch_reuse is False + + def test_empty_offer_touches_nothing(self): + connector = FakeConnectorManager(num_matched=0) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert kv_cache.resize_calls == [] + assert kv_cache.enable_swa_scratch_reuse is True + assert connector.cancels == [] + assert connector.commits == [(0, 0, False)] + + def test_offer_is_handed_back_when_pages_run_out(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32, resize_ok=False) + + prepare(manager, req, kv_cache) + + assert connector.cancels == [(0, 32, 96)] + assert req.context_current_position == 32 + # The request runs local-only from here, and delivering must not then + # record an external load for a range that was never covered. + manager._deliver_connector_prefix(req) + assert connector.commits == [(0, 0, False)] + + +class TestExclusions: + @pytest.mark.parametrize( + "attribute", + ["is_dummy", "is_generation_only_request", "is_disagg_generation_init_state"], + ) + def test_request_kinds_that_are_never_asked(self, attribute): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + setattr(req, attribute, True) + kv_cache = FakeKvCache(committed=0) + + assert prepare(manager, req, kv_cache) is None + assert connector.queries == [] + assert req.context_current_position == 0 + + def test_no_connector_is_a_no_op(self): + manager = make_manager(None) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + assert manager._connector_prefix_position(req, kv_cache) is None + + def test_draft_manager_never_asks(self): + """The draft manager's prepare_resources skips the connector hooks, so + a query there would never be delivered.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + manager.is_draft = True + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + assert manager._connector_prefix_position(req, kv_cache) is None + assert connector.queries == [] + + +class TestUndeliveredOfferIsReleased: + def test_request_that_dies_before_delivery_hands_the_offer_back(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + # Cancelled, timed out, or failed before it ever reached a batch. + manager._release_undelivered_connector_prefix(req) + + assert connector.cancels == [(0, 0, 64)] + + def test_delivered_request_keeps_its_offer_on_free(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + req.py_connector_delivered = True + manager._release_undelivered_connector_prefix(req) + + assert connector.cancels == [] + + def test_unasked_request_releases_nothing(self): + connector = FakeConnectorManager() + manager = make_manager(connector) + + manager._release_undelivered_connector_prefix(FakeRequest()) + + assert connector.cancels == [] + + +class TestAsyncLoad: + def test_async_flag_is_carried_from_query_to_delivery(self): + connector = FakeConnectorManager(num_matched=64, load_async=True) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + assert req.py_connector_load_async is True + + manager._deliver_connector_prefix(req) + assert connector.commits == [(0, 64, True)] + + def test_async_hold_survives_the_clamp(self): + """An async connector starts its transfer inside the query itself. + + So even when the clamp leaves nothing to load, the request still has to + be held out of the batch until the connector reports the transfer done, + or the forward races the writes. + """ + connector = FakeConnectorManager(num_matched=1, load_async=True) + manager = make_manager(connector) + req = FakeRequest(prompt_len=PROMPT_LEN) + kv_cache = FakeKvCache(committed=PROMPT_LEN - 1) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert req.py_connector_prefix_end == PROMPT_LEN - 1 + assert connector.commits == [(0, 0, True)] + + +class TestPhasesMustAgree: + """Phase 3 refuses to report a load phase 2 never made room for. + + This is the single coupling in the design that nothing downstream checks: + the runtime hands the connector ``context_current_position - recorded`` as + the range it computed locally, the subtraction is unguarded + (``kv_cache_connector.py:480``), and connectors divide the result into block + ordinals rather than validating it. A phase 2 that silently no-ops -- by + regression, or by a future caller reordering the phases -- therefore points + the connector at a negative offset inside its own code rather than failing. + """ + + def test_recording_more_than_the_position_covers_is_an_error(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + # The shape a silently skipped phase 2 leaves behind: the offer is on + # the request, but no capacity, history or position backs it. + req.context_current_position = 0 + + with pytest.raises(AssertionError, match="phase 2 did not reserve"): + manager._deliver_connector_prefix(req) + + assert connector.commits == [], ( + "the load must not be recorded when the assertion fires, or the " + "connector is left holding an offer the runtime denied" + ) + + def test_a_fully_reserved_offer_is_accepted(self): + """Anti-vacuity: the assertion must not fire on the normal path.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert connector.commits == [(0, 64, False)] diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py new file mode 100644 index 000000000000..a8031da32800 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KV connector prefix phases against a *real* ``KVCacheManagerV2``. + +``test_kv_connector_v2_prefix.py`` drives the same three phases against a stub +cache. That is what makes it fast and exhaustive, and it is the right place for +the ordering and state-machine rules -- but a stub cannot show that a real +``_KVCache`` survives the sequence: that ``resize`` finds real pages for the +offered prefix, that ``history_length`` really moves, that suspend/resume across +a deferral leaves the request askable exactly once, and that the page slots +handed to the connector at delivery are distinct and real. + +The engine-level suite cannot show it either, for a different reason. Whether a +request is deferred there depends on whether it reaches the scheduler in the +same pass as the request that outbids it, and that is a race: the same test was +observed asking both requests in the first pass when run alone, and asking the +second one an iteration later when run after the rest of the suite. Deferral is +therefore driven here directly -- ``prepare_context`` without the +``prepare_resources`` that would have followed it in an iteration where the +request actually ran -- which is exactly what the scheduler does to a request +that loses the token budget after being prepared (scheduler_v2.py:554-558). + +These tests allocate device memory pools. +""" + +import gc +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + +# These build a real manager, which allocates device pools. The directory is +# listed in the GPU-less l0_cpu stage, so the requirement is declared rather +# than left to fail at `torch.cuda.init()`. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="allocates real KV cache pools" +) + +TOKENS_PER_BLOCK = 32 +PROMPT_LEN = 96 +OFFER_TOKENS = 32 + + +class FakeConnectorManager: + """Records what the phases tell the connector, in order. + + Same shape as the one in ``test_kv_connector_v2_prefix.py``, plus the + ``build_scheduler_output`` that ``prepare_resources`` calls after delivery. + """ + + def __init__(self, num_matched=OFFER_TOKENS, load_async=False): + self.num_matched = num_matched + self.load_async = load_async + self.queries = [] + self.commits = [] + self.cancels = [] + self.allocs = [] + + def query_num_new_matched_tokens(self, request, num_computed_tokens): + self.queries.append((request.py_request_id, num_computed_tokens)) + return self.num_matched, self.load_async + + def commit_new_matched_tokens(self, request, num_tokens, load_kv_async): + self.commits.append((request.py_request_id, num_tokens, load_kv_async)) + request.py_num_connector_matched_tokens = num_tokens + + def cancel_load(self, request, start, end): + self.cancels.append((request.py_request_id, start, end)) + + def update_state_after_alloc(self, request, block_ids): + self.allocs.append((request.py_request_id, list(block_ids))) + + def build_scheduler_output(self, scheduled_batch, kv_cache_manager): + pass + + +def make_manager(connector, **overrides): + kwargs = dict( + kv_cache_config=KvCacheConfig(max_tokens=2048, enable_block_reuse=True), + kv_cache_type=CacheType.SELF, + num_layers=2, + num_kv_heads=4, + head_dim=64, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=256, + max_batch_size=4, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.HALF, + vocab_size=32000, + kv_connector_manager=connector, + ) + kwargs.update(overrides) + return KVCacheManagerV2(**kwargs) + + +def make_request(request_id=1, prompt_len=PROMPT_LEN): + return LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(range(prompt_len)), + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + + +def deliver(manager, request): + """One `prepare_resources`, i.e. the request reached the final batch.""" + manager.prepare_resources(SimpleNamespace(context_requests=[request], generation_requests=[])) + + +@pytest.fixture +def connector(): + return FakeConnectorManager() + + +@pytest.fixture +def manager(connector): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + mgr = make_manager(connector) + yield mgr + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() + + +def test_offer_is_backed_by_real_pages(manager, connector): + """Phase 2 against a real `_KVCache`: capacity, history, and slots. + + `history_length` is what the stub suite can only observe as an argument to + a fake `resize`. Here it is read back off the cache that the forward pass + would use, which is the difference between "we called resize correctly" and + "the offered prefix is resident". + """ + request = make_request() + + assert manager.prepare_context(request) + + kv_cache = manager.kv_cache_map[request.py_request_id] + assert request.context_current_position == OFFER_TOKENS + assert kv_cache.capacity >= OFFER_TOKENS + assert kv_cache.history_length == OFFER_TOKENS + assert kv_cache.is_active + + deliver(manager, request) + + assert connector.commits == [(request.py_request_id, OFFER_TOKENS, False)] + assert len(connector.allocs) == 1 + + _, page_indices = connector.allocs[0] + assert len(page_indices) >= OFFER_TOKENS // TOKENS_PER_BLOCK + assert all(index != BAD_PAGE_INDEX for index in page_indices) + assert len(set(page_indices)) == len(page_indices) + + +def test_deferred_request_is_asked_once_and_delivered_when_it_runs(manager, connector): + """Prepared, dropped before the batch, prepared again, then run. + + The cache is suspended in between, which is what `_revert_context_resize` + does to a request that loses its slot after being prepared. Both attempts + must reach the same offer end, and the connector must see exactly one query + and exactly one commit -- it took ownership of remote blocks inside the + first query, and a second would double-pin them. + """ + request = make_request() + + # Iteration 1: prepared, then outbid before it reached the batch. + assert manager.prepare_context(request) + kv_cache = manager.kv_cache_map[request.py_request_id] + kv_cache.suspend() + + assert connector.queries == [(request.py_request_id, 0)] + assert connector.commits == [], "nothing may be recorded until it runs" + + # Iteration 2: prepared again, and this time it runs. + assert manager.prepare_context(request) + assert kv_cache.is_active, "phase 2 must have found an active cache" + assert request.context_current_position == OFFER_TOKENS + assert kv_cache.history_length == OFFER_TOKENS + + deliver(manager, request) + + assert connector.queries == [(request.py_request_id, 0)] + assert connector.commits == [(request.py_request_id, OFFER_TOKENS, False)] + assert len(connector.allocs) == 1 + + +def test_delivering_twice_reports_one_allocation(manager, connector): + """An asynchronously loaded request re-enters on its first context chunk. + + `py_connector_delivered` is what stops that second pass re-recording the + load and firing a second `update_state_after_alloc` for one allocation. + """ + request = make_request() + + assert manager.prepare_context(request) + deliver(manager, request) + deliver(manager, request) + + assert len(connector.commits) == 1 + assert len(connector.allocs) == 1 + + +def test_undelivered_offer_is_released_when_the_request_is_freed(manager, connector): + """Phase 1 is speculative, so an offer can outlive the request. + + Without this the connector holds the remote blocks it took ownership of for + the rest of the process's life. + """ + request = make_request() + + assert manager.prepare_context(request) + manager.free_resources(request) + + assert connector.cancels == [(request.py_request_id, 0, OFFER_TOKENS)] + + +def test_delivered_offer_is_not_released_when_the_request_is_freed(manager, connector): + request = make_request() + + assert manager.prepare_context(request) + deliver(manager, request) + manager.free_resources(request) + + assert connector.cancels == [] diff --git a/tests/unittest/_torch/test_connector.py b/tests/unittest/_torch/test_connector.py index 7a1898aa4f36..d0d9ef82f5bc 100644 --- a/tests/unittest/_torch/test_connector.py +++ b/tests/unittest/_torch/test_connector.py @@ -23,7 +23,7 @@ from tensorrt_llm import mpi_rank from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( - AsyncRequests, KvCacheConnectorManager, + AsyncRequests, KvCacheConnectorManager, KvCacheConnectorScheduler, KvCacheConnectorSchedulerOutputManager) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -175,6 +175,81 @@ def test(): run_across_mpi(mpi_pool_executor, test, 2) +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_connector_manager_query_is_side_effect_free(mpi_pool_executor): + """The query and the commit are separable, and the query records nothing. + + KVCacheManagerV2 asks during a speculative scheduling pass and resolves the + answer in whichever iteration the request actually runs. That only works if + asking is inert: `external_loads` is cleared by every + `build_scheduler_output`, and a request registered as loading is dropped + from the batch. Recording at query time would attribute the load to + whichever iteration happened to ask. + """ + + def test(): + worker = MagicMock() + + if mpi_rank() == 0: + scheduler = MagicMock() + scheduler.get_num_new_matched_tokens.return_value = (16, True) + else: + scheduler = None + + manager = KvCacheConnectorManager(worker, scheduler=scheduler) + + req = MagicMock() + req.request_id = 42 + req.is_generation_only_request = False + req.py_num_connector_matched_tokens = 0 + + assert manager.query_num_new_matched_tokens(req, 32) == (16, True) + + assert manager.new_async_requests.loading_ids == set() + assert manager.scheduler_output_manager.external_loads == {} + assert req.py_num_connector_matched_tokens == 0 + + manager.commit_new_matched_tokens(req, 16, True) + + assert manager.new_async_requests.loading_ids == {42} + assert manager.scheduler_output_manager.external_loads == {42: 16} + assert req.py_num_connector_matched_tokens == 16 + + if mpi_rank() == 0: + assert scheduler.get_num_new_matched_tokens.call_count == 1 + + run_across_mpi(mpi_pool_executor, test, 2) + + +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_connector_manager_cancel_load_reaches_the_leader(mpi_pool_executor): + + def test(): + worker = MagicMock() + scheduler = MagicMock() if mpi_rank() == 0 else None + + manager = KvCacheConnectorManager(worker, scheduler=scheduler) + + req = MagicMock() + req.request_id = 42 + + manager.cancel_load(req, 0, 32) + + if mpi_rank() == 0: + assert scheduler.cancel_load.call_args[0] == (req, 0, 32) + + run_across_mpi(mpi_pool_executor, test, 2) + + +def test_cancel_load_is_additive(): + """Existing connectors predate `cancel_load` and must keep working. + + It is only ever raised by KVCacheManagerV2, so it has to stay optional with + a no-op default rather than becoming another abstract method. + """ + assert "cancel_load" not in KvCacheConnectorScheduler.__abstractmethods__ + + def test_scheduler_output_num_scheduled_tokens_with_mtp(): """Test that num_scheduled_tokens is correctly set for MTP (multi-token prediction).""" NUM_DRAFT_TOKENS = 3 From 31235fa25ccaea8d52b1ec3e875d9fae1f4304a6 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:27:41 +0800 Subject: [PATCH 08/10] [None][doc] Document the connector contract on KVCacheManagerV2 BAD_PAGE_INDEX was undocumented. Under a sliding window a connector on V2 receives -1 entries in update_state_after_alloc, RequestData.new_block_ids and request_finished; one that treats them as page slots computes an address from -1. State it, along with the reason the entry is kept in place rather than dropped: ordinals stay aligned to token ranges and append-deltas stay valid. Document cancel_load and the speculative scheduling pass it exists for, and record what does not change -- get_num_new_matched_tokens is still called exactly once per request on both managers, including across a deferral. Correct its trigger list. The local match overtaking an offer cannot arise today: a request's local match is fixed when its cache is created, resume() does not re-match, and only the request's own completed forward passes extend it. The two cases that can occur -- allocation failure, and a request freed before delivery -- were not named at all. Note that update_state_after_alloc covers only the first chunk's blocks under chunked prefill on V2, since V2 allocates per chunk; the rest arrive as append-deltas through build_connector_meta. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index dc16209c31ad..89ae82287b88 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -42,6 +42,19 @@ These methods run on the leader process and drive the connector's behavior. * **`update_state_after_alloc(self, request: LlmRequest, block_ids: list[int])`** * **Description**: a callback to update internal state after KV cache blocks have been allocated for the prefill. + * **Note**: on `KVCacheManagerV2` with chunked prefill, `block_ids` covers only the blocks allocated for the first chunk, because V2 allocates per chunk rather than for the whole prompt. The remaining blocks arrive as append-deltas in `RequestData.new_block_ids` on subsequent chunks. A connector that treats this callback as its only source of block ids will under-plan; drive off `build_connector_meta` instead. + * **Note**: on `KVCacheManagerV2` under sliding-window attention, a block that the window has already passed holds no page, and is reported as `-1` (`BAD_PAGE_INDEX`) **in place** rather than being dropped from the list. This keeps each entry aligned with its block ordinal, so entry `i` always describes prompt tokens `[i * tokens_per_block, (i+1) * tokens_per_block)` and an append-delta over successive calls stays valid. Connectors must skip `-1` entries rather than treating them as page slots. The same applies to `RequestData.new_block_ids` and to `cache_block_ids` in `request_finished`. + +* **`cancel_load(self, request: LlmRequest, start: int, end: int)`** + * **Description**: Optional, with a no-op default. Tells the connector that the runtime will not consume KV it offered from `get_num_new_matched_tokens` for prompt tokens `[start, end)`, so any ownership taken for that range can be released. Offsets are absolute prompt positions, on the same scale as `num_computed_tokens`. + * **When it fires**: only on `KVCacheManagerV2`, which asks during a speculative scheduling pass and resolves the answer later. Two things can happen in between, and both are reported here: the runtime may fail to allocate pages to cover the offer, in which case the request falls back to computing the prefix locally; or the request may be cancelled, time out or fail before it ever reaches a batch, in which case the whole offer is released. A third case -- the local cache overtaking part of the offer because another request committed the same prefix -- is handled by the same callback but cannot arise today, since a request's local match is fixed when its cache is created and only its own completed forward passes extend it. + * **Caveat**: best-effort. For a synchronous load nothing has been transferred yet, so cancelling is exact. For `is_async=True` the transfer necessarily started inside `get_num_new_matched_tokens`, so it may already be in flight. + +##### Serving a prefix on `KVCacheManagerV2` + +V1 answers `get_num_new_matched_tokens` from C++ while the block manager holds its radix-tree mutex, so the local match and the query are atomic and the answer is consumed immediately. V2 has no such mutex, and its scheduling pass is speculative -- a prepared request can still be dropped at the token budget, at resize, at multimodal alignment or at cross attention, and retried in a later iteration. + +The contract for connectors is unchanged, and in particular `get_num_new_matched_tokens` is still called **exactly once per request** on both managers -- a request that is asked and then deferred is not asked again when it comes back. What differs is that on V2 the runtime may resolve the answer in a later iteration than the one it asked in, and may by then be unable to honour part or all of it. That is what `cancel_load` reports. #### 2. Worker Interface (`KvCacheConnectorWorker`) From 41a8efd95f062b062c3ac3ebe7cdb0c149104468 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 19 Aug 2026 16:33:48 +0800 Subject: [PATCH 09/10] [None][test] Gate the KVCacheManagerV2 connector tests in pre-merge CI The connector suite has run under both KV cache managers since it was parametrized, but only the kv_cache_manager_v1 half was registered, behind a comment saying the V2 variants were expected to fail. That is no longer true -- the connector is supported on V2 -- and coverage that runs only by hand decays the first time someone touches _util.py or the bring-up path. Register the 22 kv_cache_manager_v2 ids, plus the two that were never listed: test_connector_runs_on_kv_cache_manager_v2, which is what makes the rest meaningful (the creator silently falls back to the V1 manager for combinations it cannot serve, so without it every V2 id could pass while running V1), and the V1 half of the VSWA test. test_connector_priorities[kv_cache_manager_v2] is marked xfail(strict=True) rather than dropped from the list. KvCacheRetentionConfig does not reach KVCacheManagerV2 at all, so a retention config is silently ignored there -- not only through the connector. Its assertions stay the correct expectation for both managers, so wiring retention into V2 turns the test green rather than needing it rewritten, and strict=True makes that day loud instead of silent. The host_offloading rejection assertion matched the bare string "host", which both managers' messages contain, so it would have passed through exactly the silent V1 fallback the parametrization exists to rule out. Match per manager instead: V1 names the config field, V2 the resolved tier. Retarget the static contract test. It was written as a Phase 2 worklist -- "remove entries as they are implemented" -- and that is not what happened: KVCacheManagerV2 implements none of the V1 block-id methods and is not meant to, because a flat pool-wide block id cannot describe memory whose page indices are scoped to a layer group. What the check is actually worth keeping for is that update_and_build_data reports block_hashes and priorities empty on V2 by branching on the manager type, not on hasattr: those short-circuits are correct only while the accessors are genuinely absent. It now says that, asserts the V2-side contract it does have, and the disagg method list records why none of its entries is a V2 gap rather than describing them as unported. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../defs/llmapi/test_llm_api_connector.py | 126 +++++++++++++----- .../integration/test_lists/test-db/l0_a10.yml | 40 +++++- 2 files changed, 123 insertions(+), 43 deletions(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 9bbebc5069a6..d55dbb9ff421 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -25,6 +25,8 @@ import pytest from tensorrt_llm import LLM, DisaggregatedParams, SamplingParams +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ + KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, KvCacheConfig, @@ -47,30 +49,34 @@ # V1 `KVCacheManager` methods reached on the KV connector path proper. Each is # defined in tensorrt_llm/_torch/pyexecutor/resource_manager.py and wraps a # nanobind method on the C++ manager. `KVCacheManagerV2` implements none of -# them - `test_connector_v1_method_contract_gap` pins that, and the list is the -# Phase 2 worklist. -CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS = ( - # Connector bring-up: registers the KV tensor with the worker. - # py_executor.py:1070. This is the first hard stop under V2. +# them and is not meant to: each assumes a single flat block-id space over one +# primary pool, which cannot describe memory whose page indices are scoped to a +# layer group. V2's connector contract is `register_kv_cache_layout` plus +# `get_page_indices_by_layer_group` / `get_connector_page_indices` instead. +CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS = ( + # Connector bring-up: hands the worker the single primary pool tensor. + # `PyExecutor._maybe_init_kv_connector_manager`. "get_unique_primary_pool", - # kv_cache_connector.py:322 and py_executor.py:6173. + # `KvCacheConnectorSchedulerOutputRequest.update_and_build_data` and + # `PyExecutor.kv_connector_request_finished`. "get_cache_indices", - # kv_cache_connector.py:328 (sole caller). + # `update_and_build_data`, for `RequestData.block_hashes`. "commit_and_get_block_hashes", - # kv_cache_connector.py:353, only when a retention config is set. + # `update_and_build_data`, for `RequestData.priorities`. "get_priority_by_block_id", ) # Reached only when the connector coexists with disaggregated serving # (`test_connector_disagg_prefill`), via AsyncTransferManager and the V1 cache -# reuse adapter - not from `KvCacheConnectorManager` itself. Tracked separately -# because V2 already has its own answer for some of them (`try_commit_blocks` -# at kv_cache_manager_v2.py:3201, page refcounts instead of explicit pinning), -# so they do not necessarily belong in a connector port. +# reuse adapter - not from `KvCacheConnectorManager` itself. None of them is a +# gap for V2, because V2 does not traverse those paths: +# `enable_partial_reuse_for_disagg` excludes V2, so AsyncTransferManager never +# reaches the pin/unpin pair, and the Python transceiver (the only one V2 can +# be driven by) resolves block ids through `_CacheReuseAdapterV2`. DISAGG_PATH_KV_CACHE_MANAGER_METHODS = ( - "store_blocks_for_reuse", # py_executor.py:455 - "unpin_blocks_by_id", # py_executor.py:489 - "get_memory_pool_block_indices", # disaggregation/resource/cache_reuse.py:121 + "store_blocks_for_reuse", + "unpin_blocks_by_id", + "get_memory_pool_block_indices", "pin_blocks", # no Python caller today ) @@ -212,34 +218,54 @@ def generate_and_wait(model, scheduler, worker, *args, **kwargs): return outputs -def test_connector_v1_method_contract_gap(): - """Pin exactly which V1 KV-cache-manager methods `KVCacheManagerV2` lacks. +def test_v2_connector_contract_does_not_reuse_the_v1_methods(): + """The V2 connector path implements none of the V1 accessors, by design. - This is a static contract check rather than an end-to-end run on purpose: - under V2 the connector dies at the *first* missing method - (`get_unique_primary_pool`, py_executor.py:1070), so no e2e run can ever - report more than one gap. Phase 2 should remove entries from - `CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS` as it implements them. + Something depends on that, and it does not ask: `update_and_build_data` + reports `block_hashes` and `priorities` empty on V2 by branching on + `isinstance(manager, KVCacheManagerV2)`, not on `hasattr`. Those + short-circuits are only correct while V2 genuinely has no such accessor - + the day one is added (retention priorities are a known gap; see + `test_connector_priorities`) the branch keeps reporting nothing while the + data exists, and this is what says so. + + A static check rather than an end-to-end run: under V2 the connector would + die at the *first* method it reached, so no run can report more than one at + a time. Needs no GPU. """ stale = [ - name for name in CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS + + name for name in CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS + DISAGG_PATH_KV_CACHE_MANAGER_METHODS if not hasattr(KVCacheManager, name) ] assert stale == [], ( f"{stale} are not defined on the V1 KVCacheManager either, so this " - "test is measuring a stale method list rather than a real V2 gap.") + "test is measuring a stale method list rather than a real difference.") implemented = [ - name for name in CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS + name for name in CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS if hasattr(KVCacheManagerV2, name) ] assert implemented == [], ( - f"KVCacheManagerV2 now implements {implemented}. Drop them from " - "CONNECTOR_REQUIRED_KV_CACHE_MANAGER_METHODS and re-check whether the " - "corresponding connector tests can be un-skipped.") + f"KVCacheManagerV2 now implements {implemented}. A V1-shaped accessor " + "on V2 is not automatically the right answer - a flat block-id list " + "cannot describe more than one layer group - but if it is, revisit the " + "`is_v2` short-circuits in " + "`KvCacheConnectorSchedulerOutputRequest.update_and_build_data`, which " + "report nothing on the strength of these methods being absent.") + + # The other half of the contract: what V2 offers instead. + for name in ("get_page_indices_by_layer_group", + "get_connector_page_indices"): + assert hasattr(KVCacheManagerV2, name), ( + f"KVCacheManagerV2.{name} is the V2 replacement for the V1 " + "block-id accessors and every connector path on V2 goes through " + "it.") + assert hasattr(KvCacheConnectorWorker, "register_kv_cache_layout"), ( + "The worker ABC must keep a default `register_kv_cache_layout`, or " + "every existing connector becomes abstract and fails to instantiate.") @pytest.mark.threadleak(enabled=False) @@ -1002,8 +1028,19 @@ def test_connector_multi_request(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) -@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], - ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], +@pytest.mark.parametrize("use_kv_cache_manager_v2", [ + pytest.param(False, id="kv_cache_manager_v1"), + pytest.param( + True, + id="kv_cache_manager_v2", + marks=pytest.mark.xfail( + strict=True, + reason= + "KvCacheRetentionConfig does not reach KVCacheManagerV2 at all " + "(per-page priority comes from custom_priority_callback, which " + "V2 never overrides), so the connector reports priorities=None."), + ), +], indirect=True) def test_connector_priorities(enforce_single_worker, model_with_connector): """Test that retention priorities flow through the connector correctly. @@ -1012,15 +1049,17 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): the RequestData.priorities field is populated with the correct per-block priorities based on the token ranges. - KNOWN GAP -- this fails on `kv_cache_manager_v2`, deliberately left failing. + KNOWN GAP -- `xfail(strict=True)` on `kv_cache_manager_v2`. `KvCacheRetentionConfig` does not reach KVCacheManagerV2 at all: V2's per-page priority comes from `custom_priority_callback` (kv_cache_manager_v2/_core/_kv_cache_manager.py), which KVCacheManagerV2 never overrides, so every page carries the default priority and the connector reports `priorities=None`. A user who sets a retention config on V2 silently gets none of it -- not only through the connector. The - assertions below are the correct expectation for both managers and are kept - that way so the gap stays visible rather than being asserted away. + assertions below stay the correct expectation for both managers rather than + being relaxed per manager, so wiring retention into V2 turns this green + instead of needing the test rewritten; `strict=True` is what makes it fail + loudly on that day rather than passing silently. """ BLOCK_SIZE = 32 NUM_INPUT_TOKENS = 64 # 2 blocks @@ -1121,22 +1160,33 @@ def test_connector_priorities_default(enforce_single_worker, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize( - "llm_kwargs,match", + "llm_kwargs,match_v1,match_v2", [ + # The two managers refuse offloading for the same reason -- a page that + # leaves GPU has its slot reassigned, invalidating what the connector + # registered -- but say so differently, and V2 names the resolved tier + # rather than the config field so it also catches the disk tier. Match + # the manager-specific wording: both messages contain the bare word + # "host", so matching that would still pass if a silent fallback to V1 + # ever crept back in, which is the one thing this parametrization + # exists to rule out. pytest.param( dict(kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1, host_cache_size=1024**3)), - "host", + "host offloading", + "cache tiers below GPU", id="host_offloading", ), pytest.param( dict(max_beam_width=2), "beam", + "beam", id="beam_search", ), pytest.param( dict(enable_attention_dp=True), "attention data parallelism", + "attention data parallelism", id="attention_dp", ), ], @@ -1145,12 +1195,14 @@ def test_connector_priorities_default(enforce_single_worker, ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], indirect=True) def test_connector_rejects_unsupported_config(enforce_single_worker, - model_with_connector, llm_kwargs, - match): + model_with_connector, + use_kv_cache_manager_v2, + llm_kwargs, match_v1, match_v2): # Configurations the connector cannot handle today must fail loudly at # construction time rather than silently miscompute. This pins the set of # constructor-time exclusions in `_maybe_init_kv_connector_manager`. model_fn, _, _ = model_with_connector + match = match_v2 if use_kv_cache_manager_v2 else match_v1 with pytest.raises(NotImplementedError, match=match): model_fn(**llm_kwargs) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index eb87678eca3e..083895fe6d59 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -162,12 +162,17 @@ l0_a10: # usage telemetry - unittest/llmapi/test_llm_telemetry.py::TestTelemetryPyTorchBackend - unittest/llmapi/test_llm_telemetry.py::TestTelemetryArchitectureExtraction - # KV connector. The suite is parametrized over the KV cache manager; only the - # kv_cache_manager_v1 variants are registered here. The kv_cache_manager_v2 - # variants are expected to fail today (KVCacheManagerV2 implements none of the - # V1 block-ID methods the connector needs) and are run manually, not in the - # merge gate. See test_connector_v1_method_contract_gap for the current gap. - - llmapi/test_llm_api_connector.py::test_connector_v1_method_contract_gap + # KV connector. The suite is parametrized over the KV cache manager and both + # halves are gated: the connector is supported on KVCacheManagerV2, so a V2 + # regression has to be a merge-gate failure rather than something found later + # by hand. test_connector_runs_on_kv_cache_manager_v2 is what makes the V2 + # entries meaningful -- the creator silently falls back to the V1 manager for + # combinations it cannot serve, so without it every V2 id below could pass + # while running V1. test_connector_priorities[kv_cache_manager_v2] is + # xfail(strict=True) in the suite, not omitted here, so the retention gap + # stays counted. + - llmapi/test_llm_api_connector.py::test_v2_connector_contract_does_not_reuse_the_v1_methods + - llmapi/test_llm_api_connector.py::test_connector_runs_on_kv_cache_manager_v2 - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-True] - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-False] - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v1-True] @@ -187,6 +192,29 @@ l0_a10: - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-beam_search] - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-attention_dp] - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-host_offloading] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-beam_search] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-attention_dp] + - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_uniform_sliding_window[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_sliding_window_prefix_is_backed_by_history[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v2] # third-party policy checks CPU-only - thirdparty/test_cmake_third_party.py::test_cmake_listfiles - thirdparty/test_git_modules.py::test_gitmodules From 553506909e10b85125aad560a1a7fefa9b995240 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Fri, 21 Aug 2026 12:11:13 +0800 Subject: [PATCH 10/10] [None][fix] Address review feedback on the KVCacheManagerV2 connector The `prompt_len - 1` clamp in `_connector_prefix_position` shrank the offer without telling the connector. `_deliver_connector_prefix` and `_release_undelivered_connector_prefix` both read the already-clamped end, so the tail past it reached neither, and the connector kept ownership of those remote blocks for the life of the process -- the exact leak those two paths exist to prevent. Cancel the dropped range where it is dropped, and assert it in the two tests that exercise the clamp. `_reject_non_gpu_cache_tiers` rejects every tier below GPU, which includes a disk tier configured through `KvCacheConfig.disk_cache_size`, but told the user to set `host_cache_size=0`. Name both fields. Also add the missing NVIDIA header to `test_kv_connector_v2_prefix.py`. Signed-off-by: Yueh-Ting Chen Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 12 +++++++++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 3 ++- .../_torch/executor/test_kv_connector_v2_prefix.py | 10 ++++++++++ 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 425aa68eb472..944036b6155b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2384,7 +2384,17 @@ def _connector_prefix_position(self, req: LlmRequest, kv_cache) -> Optional[int] # activations. A connector that offers the whole prompt therefore # loses at most one token, which the forward recomputes. req.py_connector_prefix_start = local_end - req.py_connector_prefix_end = min(local_end + num_tokens, req.prompt_len - 1) + offered_end = local_end + num_tokens + req.py_connector_prefix_end = min(offered_end, req.prompt_len - 1) + # Hand back what that clamp just dropped, because nothing + # downstream can: `_deliver_connector_prefix` and + # `_release_undelivered_connector_prefix` both read the + # *clamped* end, so the tail past it would stay owned by the + # connector for the life of the process -- the exact leak + # those two exist to prevent. + dropped_start = max(local_end, req.py_connector_prefix_end) + if offered_end > dropped_start: + self.kv_connector_manager.cancel_load(req, dropped_start, offered_end) # Deliberately not conditioned on the clamp above: if the connector # said it would transfer asynchronously it has already started, and # the request must be held out of the batch until it reports done. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a6c2b4e8de38..d14b993f1e72 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1124,7 +1124,8 @@ def _reject_non_gpu_cache_tiers(kv_cache_manager) -> None: f"cache tiers below GPU (found: {names}). Pages evicted to " "another tier have their GPU slot reassigned, which would " "invalidate the addresses registered with the connector. Set " - "KvCacheConfig.host_cache_size=0 to disable the host tier.") + "KvCacheConfig.host_cache_size=0 and " + "KvCacheConfig.disk_cache_size=0 to run GPU-only.") def _end_transfer_and_maybe_terminate(self, request: LlmRequest): transfer_failed = request.state == LlmRequestState.DISAGG_TRANS_ERROR diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py index 25310a617a07..2515b70548f6 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 """Unit tests for the KV connector prefix phases on KVCacheManagerV2. KVCacheManagerV2's scheduling pass is speculative: a request can be prepared @@ -354,6 +356,11 @@ def test_offer_end_is_clamped_below_the_prompt(self): assert req.py_connector_prefix_end == PROMPT_LEN - 1 assert req.context_current_position == PROMPT_LEN - 1 assert connector.commits == [(0, PROMPT_LEN - 1, False)] + # The clamp is the one place an offer shrinks without passing through + # phase 3 or the release path, both of which read the clamped end. So + # it has to hand the remainder back itself or the connector keeps + # ownership of it for the life of the process. + assert connector.cancels == [(0, PROMPT_LEN - 1, PROMPT_LEN)] class TestWriteRangeAtDelivery: @@ -584,6 +591,9 @@ def test_async_hold_survives_the_clamp(self): assert req.py_connector_prefix_end == PROMPT_LEN - 1 assert connector.commits == [(0, 0, True)] + # Lossy for an async load -- the transfer began inside the query -- but + # still the only signal the connector gets that the tail is dead. + assert connector.cancels == [(0, PROMPT_LEN - 1, PROMPT_LEN)] class TestPhasesMustAgree: