From b95fac653f46b8cf4852e3f108be930d5502f9d9 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 14 Sep 2026 05:45:30 +0000 Subject: [PATCH 1/2] [None][refactor] BREAKING: Remove the Python backend of KVCacheManagerV2 KVCacheManagerV2 shipped two implementations of the same subsystem behind TLLM_KV_CACHE_MANAGER_V2_BACKEND. The C++ port is the default, is what CI exercises, and is the only backend newer features support, so the Python implementation was carrying duplicate block-key hashing, eviction and stats logic that had to stay bit-identical to C++, plus a mypyc and rawref build pipeline that existed only to make it fast enough to matter. tensorrt_llm/runtime/kv_cache_manager_v2/ is now a re-export shim over the nanobind module plus the _introspection dispatcher: 36 tracked files down to 4. Consumers that reached into private submodules move to the package surface. The KV-aware router's V2 hashing is expressed with the existing sequence_to_blockchain_keys, since every caller chains from a reuse-scope root, so v2_sha256_block_hasher is gone and no hashing binding was needed. Only the native disaggregated bounce buffer needed something new: PooledPhysMemAllocator and VirtMem wrap the existing cudaVirtMem, and are registered on the _introspection submodule rather than the package surface because they carry no stability promise. Streaming KV events (kv_cache_config.kv_events_config) are dropped. The sink is duck-typed Python and the C++ radix tree calls its sink natively, so there is no live path; validate_streaming_support now rejects the config and points at the buffered path via event_buffer_max_size. The interface is kept as a stub and its tests are skipped rather than deleted. The --mypyc flag, TRTLLM_ENABLE_MYPYC, setup_mypyc.py, the rawref C extension and the setup.py packaging surgery they required are all removed. Signed-off-by: Yao Yao --- .gitignore | 2 - .../kv_cache_manager_v2/AGENTS.md | 11 +- .../batch_manager/kvCacheManagerV2.cpp | 15 + docs/source/features/kvcache.md | 19 +- docs/source/installation/build-from-source.md | 1 - .../kv_cache_compression/nvfp4_cold_page.md | 5 +- pyproject.toml | 4 - scripts/build_wheel.py | 64 - setup.py | 48 +- .../_torch/attention/backends/flashinfer.py | 2 +- .../sparse/deepseek_v4/cache_manager.py | 2 +- .../sparse/minimax_m3/cache_manager.py | 9 +- .../backends/sparse/qsa/cache_manager.py | 3 +- .../_torch/attention/backends/vanilla.py | 2 +- .../disaggregation/native/bounce/buffer.py | 2 +- tensorrt_llm/_torch/pyexecutor/_util.py | 6 - .../kv_cache/kv_cache_manager_v2.py | 20 +- .../_torch/pyexecutor/kv_cache_events.py | 298 +- tensorrt_llm/runtime/__init__.py | 29 +- .../runtime/kv_cache_manager_v2/AGENTS.md | 122 +- .../runtime/kv_cache_manager_v2/Makefile | 51 - .../runtime/kv_cache_manager_v2/__init__.py | 492 ++-- .../runtime/kv_cache_manager_v2/__init__.pyi | 42 +- .../kv_cache_manager_v2/_block_radix_tree.py | 882 ------ .../runtime/kv_cache_manager_v2/_common.py | 101 - .../runtime/kv_cache_manager_v2/_config.py | 291 -- .../kv_cache_manager_v2/_copy_engine.py | 387 --- .../kv_cache_manager_v2/_core/__init__.py | 42 - .../kv_cache_manager_v2/_core/_kv_cache.py | 2458 ----------------- .../_core/_kv_cache_manager.py | 1165 -------- .../_core/_moving_average.py | 56 - .../_core/_pending_stats.py | 317 --- .../kv_cache_manager_v2/_cuda_virt_mem.py | 228 -- .../kv_cache_manager_v2/_event_manager.py | 701 ----- .../_eviction_controller/__init__.py | 23 - .../_eviction_controller.py | 242 -- .../kv_cache_manager_v2/_exceptions.py | 71 - .../kv_cache_manager_v2/_introspection.py | 351 +-- .../_life_cycle_registry.py | 167 -- .../runtime/kv_cache_manager_v2/_page.py | 561 ---- .../runtime/kv_cache_manager_v2/_stats.py | 142 - .../kv_cache_manager_v2/_storage/__init__.py | 20 - .../kv_cache_manager_v2/_storage/_config.py | 242 -- .../kv_cache_manager_v2/_storage/_core.py | 1000 ------- .../kv_cache_manager_v2/_storage_manager.py | 1175 -------- .../runtime/kv_cache_manager_v2/_utils.py | 1126 -------- .../kv_cache_manager_v2/mypy_mypyc.ini | 51 - .../kv_cache_manager_v2/rawref/README.md | 140 - .../kv_cache_manager_v2/rawref/__init__.py | 35 - .../kv_cache_manager_v2/rawref/__init__.pyi | 82 - .../kv_cache_manager_v2/rawref/rawrefmodule.c | 250 -- .../kv_cache_manager_v2/rawref/setup.py | 28 - .../kv_cache_manager_v2/rawref/test_rawref.py | 269 -- .../kv_cache_manager_v2/setup_mypyc.py | 141 - tensorrt_llm/serve/router.py | 3 +- tensorrt_llm/serve/router_utils.py | 41 +- ..._dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 6 +- .../test_deepseek_v4_cache_manager.py | 2 +- .../attention/test_flashinfer_attention.py | 2 +- .../test_kv_cache_v2_multimodal_runs.py | 4 +- ...anager_v2.py => test_kvcm2_integration.py} | 15 +- .../_torch/modeling/test_modeling_gemma4.py | 53 +- .../executor/test_stats_serializer.py | 127 - .../cuda_test_utils.py | 472 ++++ .../kv_cache_manager_v2_tests/fake_engine.py | 44 +- .../kv_cache_manager_v2_tests/kernels.py | 25 +- .../test_branch_reuse.py | 27 +- .../test_first_new_block_probe.py | 22 +- .../test_kv_cache_concurrency.py | 13 +- .../test_kv_cache_event_manager.py | 485 ++-- .../test_kv_cache_manager_v2.py | 328 +-- .../test_kv_cache_salting.py | 82 +- .../test_kv_cache_stats_api.py | 11 - .../test_kv_cache_stats_life_cycles.py | 167 -- .../test_nvbug_6625710.py | 4 - .../test_streaming_kv_events.py | 81 +- .../test_virt_mem_lifetime.py | 88 + 77 files changed, 1432 insertions(+), 14663 deletions(-) delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/Makefile delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_common.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_config.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_core/_moving_average.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/__init__.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_page.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_storage/__init__.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/mypy_mypyc.ini delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/rawref/README.md delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.pyi delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/rawref/rawrefmodule.c delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/rawref/setup.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/rawref/test_rawref.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py rename tests/unittest/_torch/executor/kv_cache/{test_kv_cache_manager_v2.py => test_kvcm2_integration.py} (99%) create mode 100644 tests/unittest/kv_cache_manager_v2_tests/cuda_test_utils.py delete mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py diff --git a/.gitignore b/.gitignore index b693ef923ccc..a4c1e4accea8 100644 --- a/.gitignore +++ b/.gitignore @@ -59,8 +59,6 @@ tensorrt_llm/flash_mla_cpp_tllm.*.so tensorrt_llm/flash_mla_cpp_tllm.pyi /3rdparty/fmha_sm100/ /3rdparty/nccl_extensions/ -tensorrt_llm/runtime/kv_cache_manager_v2/**/*.so -**/*__mypyc*.so tensorrt_llm/scripts *docs/cpp_docs* *docs/source/_cpp_gen* diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md index bf25eab2056b..b8b6cf612e7f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md @@ -468,11 +468,12 @@ Use extra review and tests for changes involving: `radixBlockTreeTest.cpp`, `kvCacheManagerTest.cpp`, `kvCacheManagerV2DigestPoolTest.cpp`, `kvCacheManagerV2HostMemTest.cpp`, `kvCacheManagerV2StatsTest.cpp`, and `kvCacheManagerV2TypedIndexTest.cpp`. -- Python behavior and backend-parity tests are in - `tests/unittest/kv_cache_manager_v2_tests/`. During development, prefer the - fast path below: set `PYTHONPATH` to `tensorrt_llm/runtime/` and execute the - test file directly with `python`. Do not use `pytest` for this fast path; the - file's test runner avoids importing the full `tensorrt_llm` package. +- Python behavior tests are in `tests/unittest/kv_cache_manager_v2_tests/`, and + drive this C++ implementation through the nanobind bindings. During + development, prefer the fast path below: set `PYTHONPATH` to + `tensorrt_llm/runtime/` and execute the test file directly with `python`. Do + not use `pytest` for this fast path; the file's test runner avoids importing + the full `tensorrt_llm` package. ```bash REPO_ROOT="$(git rev-parse --show-toplevel)" diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index f45743b79290..593be74b8fa5 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -21,6 +21,7 @@ #include "kv_cache_manager_v2/coldPageCodec.h" #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/config.h" +#include "kv_cache_manager_v2/cudaVirtMem.h" #include "kv_cache_manager_v2/eventManager.h" #include "kv_cache_manager_v2/exceptions.h" #include "kv_cache_manager_v2/introspection.h" @@ -2227,6 +2228,20 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("quota"), nb::arg("slot_size_lists"), nb::arg("ratio_list"), nb::arg("granularity"), nb::arg("min_slots"), nb::call_guard()); + // CUDA virtual-memory primitives, reached through _introspection because they carry no + // stability promise: the native disaggregated bounce buffer reserves one contiguous fabric + // region with them and maps physical chunks into it up front. + nb::class_(mIntrospection, "PooledPhysMemAllocator") + .def(nb::init(), nb::arg("phys_mem_size")) + .def_prop_ro("device_id", &kv::PooledPhysMemAllocator::deviceId); + nb::class_(mIntrospection, "VirtMem") + // keep_alive<1, 3>: VirtMem holds PooledPhysMemAllocator by reference, so the allocator + // must outlive it. Argument 3 is the allocator (1 is self, 2 is vm_size). + .def(nb::init(), nb::arg("vm_size"), nb::arg("phys_mem_allocator"), + nb::arg("init_num_phys_mem") = 0, nb::keep_alive<1, 3>()) + .def("destroy", &kv::VirtMem::destroy) + .def_prop_ro("address", &kv::VirtMem::address); + // ---- Cold-page codec -------------------------------------------------- nb::class_(m, "IKvCacheColdPageCodec"); m.def("create_default_kv_cache_cold_page_codec", &kv::createDefaultKvCacheColdPageCodec, diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 7ccea7387c3a..9a84f5d5cf1f 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -278,7 +278,14 @@ Events are buffered per rank, gathered onto rank 0 under attention data parallel pulled per iteration through `LLM.get_kv_cache_events()` / `LLM.get_kv_cache_events_async()`, or over the `/kv_cache_events` endpoint of `trtllm-serve`. -#### Streaming path (prototype) +#### Streaming path (unsupported) + +```{note} +The streaming path has no implementation: `kv_cache_config.kv_events_config` is rejected +at startup. Use the buffered path via `kv_cache_config.event_buffer_max_size` instead. The +wire format and endpoint convention below describe the contract a future native event sink +must satisfy. +``` Configured with ```kv_cache_config.kv_events_config```. Each rank encodes its own events and publishes them directly over a ZeroMQ `PUB` socket from a background thread, so there is no @@ -297,12 +304,10 @@ kv_cache_config = KvCacheConfig( ) ``` -**Constraints.** The streaming path requires KV cache manager V2 running on its Python -backend (`TLLM_KV_CACHE_MANAGER_V2_BACKEND=python`); the default `cpp` backend cannot -consume the Python event sink and raises an error naming this variable. Pipeline -parallelism and context parallelism are rejected. Events are not published for draft -models or during KV-cache-size estimation. When streaming is enabled the buffered pull API -returns an empty list rather than raising. +**Constraints.** Enabling the streaming path raises at startup. A Python event sink cannot +serve it, because the KV cache manager V2 radix tree calls its sink natively rather than +through Python; re-enabling it needs a native sink. Pipeline parallelism and context +parallelism are rejected independently. **Endpoint convention.** Every attention-DP rank binds `base_port + rank` using its **global** rank, so `N` ranks occupy `[base_port, base_port + N - 1]` cluster-wide and diff --git a/docs/source/installation/build-from-source.md b/docs/source/installation/build-from-source.md index ca81b2115468..4b9db6ed789d 100644 --- a/docs/source/installation/build-from-source.md +++ b/docs/source/installation/build-from-source.md @@ -104,7 +104,6 @@ With `--build_root ` set, the following default under `` instead of th | Build virtual environment | `venv-` | run inside an activated venv, or `--no-venv` | | Wheel staging tree and `*.egg-info` | `wheel-staging` | `TRTLLM_WHEEL_STAGING_DIR` | | ccache directory (with `--use_ccache`) | `ccache` | `CCACHE_DIR` | -| Intermediate extension-module objects | `kv_cache_manager_v2-temp` | — | Conan's `cpp/CMakeUserPresets.json` convenience file is also skipped in this mode, since it would reference the (possibly ephemeral) out-of-tree build directory. diff --git a/examples/kv_cache_compression/nvfp4_cold_page.md b/examples/kv_cache_compression/nvfp4_cold_page.md index ee5248008d5a..1d501f3945f7 100644 --- a/examples/kv_cache_compression/nvfp4_cold_page.md +++ b/examples/kv_cache_compression/nvfp4_cold_page.md @@ -71,9 +71,8 @@ and KV-cache block reuse remains supported because Page and token identity are unchanged. One-model MTP-EAGLE and EAGLE3 are supported; HELIX context parallelism is not currently supported. -Set `kv_cache_config.use_kv_cache_manager_v2: true` explicitly, and do not set -`TLLM_KV_CACHE_MANAGER_V2_BACKEND=python`. A nonzero Host or Disk cache is also -required for Pages to cross a compression boundary. +Set `kv_cache_config.use_kv_cache_manager_v2: true` explicitly. A nonzero Host +or Disk cache is also required for Pages to cross a compression boundary. On Linux 6.11 through 6.13, mixed models that need both NVFP4 Attention lifecycles and lossless SSM/GDN fallback lifecycles are not supported. See the diff --git a/pyproject.toml b/pyproject.toml index ca50b55bc9ac..6be1ec5f066b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -924,10 +924,6 @@ disallow_incomplete_defs = false disallow_untyped_defs = false warn_return_any = false -[[tool.mypy.overrides]] -module = ["tensorrt_llm.runtime.kv_cache_manager_v2.*"] -disallow_any_generics = false - [[tool.mypy.overrides]] module = ["tensorrt_llm.bindings.*"] ignore_errors = true diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 33a43424de54..8ff59bbdd3c0 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -531,57 +531,6 @@ def generate_python_stubs_windows(venv_python: Path, pkg_dir: Path, exit(1) -def build_kv_cache_manager_v2(project_dir, - venv_python, - use_mypyc=False, - build_root=None): - print("-- Building kv_cache_manager_v2...") - kv_cache_mgr_dir = project_dir / "tensorrt_llm/runtime/kv_cache_manager_v2" - runtime_dir = project_dir / "tensorrt_llm/runtime" - - # The produced .so files always land in-place (they are final artifacts); - # only the intermediate object files are redirected out of the checkout. - build_temp_arg = "" - if build_root is not None: - build_temp_arg = f' --build-temp "{build_root / "kv_cache_manager_v2-temp"}"' - - # Clean up any existing mypyc artifacts in runtime directory to prevent stale inclusion - # when switching from --mypyc to standard build - if not use_mypyc: - for so_file in runtime_dir.glob("*__mypyc*.so"): - print(f"Removing stale mypyc artifact: {so_file}") - so_file.unlink() - - # Also clean up any .so files inside kv_cache_manager_v2 - for so_file in kv_cache_mgr_dir.rglob("*.so"): - print(f"Removing stale artifact: {so_file}") - so_file.unlink() - - # Build rawref - print("-- Building kv_cache_manager_v2 rawref extension...", end=" ") - rawref_dir = kv_cache_mgr_dir / "rawref" - build_run(f'"{venv_python}" setup.py build_ext --inplace{build_temp_arg}', - cwd=rawref_dir) - print("Done") - - if use_mypyc: - # Build mypyc - print("-- Building kv_cache_manager_v2 mypyc extensions...", end=" ") - # setup_mypyc.py is in kv_cache_manager_v2 but executed from runtime dir - setup_mypyc = kv_cache_mgr_dir / "setup_mypyc.py" - build_run( - f'"{venv_python}" "{setup_mypyc}" build_ext --inplace{build_temp_arg}', - cwd=runtime_dir) - - # Verify that the shared library was generated - if not list(runtime_dir.glob("*__mypyc*.so")): - raise RuntimeError( - "Failed to build kv_cache_manager_v2: no shared library generated." - ) - print("Done") - print("-- Done building kv_cache_manager_v2.") - - def _tar_pipe_copy(src: Path, dst: Path) -> bool: """Populate dst from src as one streamed tar pipeline. @@ -841,7 +790,6 @@ def main(*, generate_fmha: bool = False, no_venv: bool = False, nvrtc_dynamic_linking: bool = False, - mypyc: bool = False, require_dynamic_attributions: bool = False, plat_name: Optional[str] = None, yes: bool = False, @@ -1504,11 +1452,6 @@ def get_binding_lib(subdirectory, name): nixl_root is not None or mooncake_root is not None, binding_lib_file_name) - build_kv_cache_manager_v2(wheel_project_dir, - venv_python, - use_mypyc=mypyc, - build_root=build_root) - if not skip_building_wheel: if dist_dir is None: dist_dir = build_root / "dist" if out_of_tree else project_dir / "build" @@ -1576,10 +1519,6 @@ def get_binding_lib(subdirectory, name): ) env = os.environ.copy() - if mypyc: - env["TRTLLM_ENABLE_MYPYC"] = "1" - else: - env["TRTLLM_ENABLE_MYPYC"] = "0" build_run( f'\"{venv_python}\" -m build {wheel_project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"', @@ -1751,9 +1690,6 @@ def add_arguments(parser: ArgumentParser): "--nvrtc_dynamic_linking", action="store_true", help="Link against dynamic NVRTC libraries instead of static ones") - parser.add_argument("--mypyc", - action="store_true", - help="Compile kv_cache_manager_v2 with mypyc") parser.add_argument("--require_dynamic_attributions", action="store_true", help="Fail the build if attribution generation fails") diff --git a/setup.py b/setup.py index 5ad206aa4325..9c730747344b 100644 --- a/setup.py +++ b/setup.py @@ -195,13 +195,6 @@ def has_ext_modules(self): 'flash_mla/LICENSE', 'flash_mla/*.py', 'flash_mla_cpp_tllm.*.so', - 'runtime/kv_cache_manager_v2/*.so', - 'runtime/kv_cache_manager_v2/**/*.so', - 'runtime/kv_cache_manager_v2/*.pyi', - 'runtime/kv_cache_manager_v2/**/*.pyi', - 'runtime/kv_cache_manager_v2/rawref/*.py', - 'runtime/kv_cache_manager_v2/rawref/*.pyi', - 'runtime/*__mypyc*.so', ] package_data += [ @@ -568,15 +561,11 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], # (deep_gemm, deep_ep, flash_mla Python files are generated during build) if file.filename.endswith(".py"): allowed_dirs = ( - "tensorrt_llm/deep_gemm/", "tensorrt_llm/deep_ep/", + "tensorrt_llm/deep_gemm/", + "tensorrt_llm/deep_ep/", "tensorrt_llm/flash_mla/", - "tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.py" ) if not any(file.filename.startswith(d) for d in allowed_dirs): - # Exclude all .py files in kv_cache_manager_v2 except rawref/__init__.py - if file.filename.startswith("tensorrt_llm/runtime/kv_cache_manager_v2/") and \ - not file.filename.endswith("rawref/__init__.py"): - continue continue for filename_pattern in package_data: @@ -610,37 +599,8 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], with open("README.md", "r", encoding="utf-8") as fh: long_description = fh.read() - # We use find_packages with a custom exclude filter to handle the mypyc compiled modules. - # We want to exclude the .py source files for modules that are compiled to .so. - # We exclude the kv_cache_manager_v2 package entirely from the source list, - # but explicitly add back the rawref subpackage (which is not compiled by mypyc). - # The .so and .pyi files for kv_cache_manager_v2 are added via package_data. -enable_mypyc = os.getenv("TRTLLM_ENABLE_MYPYC", "0") == "1" -if enable_mypyc: - packages = find_packages(exclude=[ - "tensorrt_llm.runtime.kv_cache_manager_v2", - "tensorrt_llm.runtime.kv_cache_manager_v2.*", - ]) + ["tensorrt_llm.runtime.kv_cache_manager_v2.rawref"] - exclude_package_data = { - "tensorrt_llm": [ - "runtime/kv_cache_manager_v2/*.py", - "runtime/kv_cache_manager_v2/**/*.py" - ], - "tensorrt_llm.runtime.kv_cache_manager_v2": ["*.py", "**/*.py"], - } -else: - packages = find_packages() - exclude_package_data = {} - - # Remove mypyc shared objects from package_data to avoid packaging stale files - package_data = [ - p for p in package_data if p not in [ - 'runtime/kv_cache_manager_v2/*.so', - 'runtime/kv_cache_manager_v2/**/*.so', 'runtime/*__mypyc*.so' - ] - ] - # Ensure rawref is included - package_data.append('runtime/kv_cache_manager_v2/rawref/*.so') +packages = find_packages() +exclude_package_data = {} # Add vendored triton_kernels as an explicit top-level package. # This is vendored from the Triton project and kept at repo root so its diff --git a/tensorrt_llm/_torch/attention/backends/flashinfer.py b/tensorrt_llm/_torch/attention/backends/flashinfer.py index b800b0cd453f..e3373666692c 100644 --- a/tensorrt_llm/_torch/attention/backends/flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/flashinfer.py @@ -37,7 +37,7 @@ from tensorrt_llm.functional import AttentionMaskType from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX from ...metadata import KVCacheParams from ...utils import get_global_attrs, get_model_extra_attrs, torch_multi_arange diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index a367efdd011b..1e28573b7cab 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -43,6 +43,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime import ModelConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, AttentionLayerConfig, BufferConfig, DataRole, @@ -51,7 +52,6 @@ ScratchDesc, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX from .compressor import NVFP4_COMPRESS_RESIDUAL_DIM, KVCacheDtype from .params import ( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index e450281c98fb..876c6d01e2a9 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -44,9 +44,12 @@ copy_batch_block_offsets_to_device, ) from tensorrt_llm.logger import logger -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX -from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, + BufferConfig, + DataRole, + PageIndexMode, +) class MiniMaxM3SparseIndexCache: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py index 394cd9e697fd..a5ddcaa70614 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -12,8 +12,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode -from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, DataRole, PageIndexMode from .constants import ( QSA_INDEX_K_CACHE_DTYPE, diff --git a/tensorrt_llm/_torch/attention/backends/vanilla.py b/tensorrt_llm/_torch/attention/backends/vanilla.py index 52da0a7d33bf..2477c271b1e2 100644 --- a/tensorrt_llm/_torch/attention/backends/vanilla.py +++ b/tensorrt_llm/_torch/attention/backends/vanilla.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX try: from transformers.modeling_attn_mask_utils import AttentionMaskConverter diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py b/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py index bb77eab32908..c45ff75297e8 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py @@ -21,7 +21,7 @@ from tensorrt_llm import logger from tensorrt_llm._torch.disaggregation.base.agent import RegMemoryDescs -from tensorrt_llm.runtime.kv_cache_manager_v2._cuda_virt_mem import PooledPhysMemAllocator, VirtMem +from tensorrt_llm.runtime.kv_cache_manager_v2._introspection import PooledPhysMemAllocator, VirtMem _MIB = 1024 * 1024 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b5e17fbeae15..e4d243c7fdf1 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3172,12 +3172,6 @@ def validate_kv_cache_compression_compatibility( ) -> None: """Reject unsupported KV-cache compression feature combinations.""" if config.algorithm == "quantization_for_cold_page": - from tensorrt_llm.runtime.kv_cache_manager_v2 import _BACKEND - - if _BACKEND == "python": - raise ValueError( - "Cold-page quantization requires the C++ KVCacheManagerV2 backend" - ) if config.quant == "nvfp4" and not is_sm_100f(): raise RuntimeError( "NVFP4 cold-page quantization requires an SM100-family device " diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 0774d0affc9f..6cc5a54a7f88 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -84,7 +84,6 @@ gen_multimodal_cache_key_tokens, typed_range, ) -from tensorrt_llm.runtime.kv_cache_manager_v2 import BACKEND as KV_CACHE_MANAGER_V2_BACKEND from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as KVCacheManagerPy from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfMemoryError as KVCacheOutOfMemoryError @@ -1283,16 +1282,9 @@ def __init__( kv_events_config is not None and kv_events_config.enable_kv_cache_events ) if streaming_events_enabled: - if self.event_buffer_max_size > 0: - logger.warning( - "Both kv_cache_config.event_buffer_max_size and streaming " - "kv_events_config are enabled; streaming publishing takes " - "precedence and the buffered get_kv_cache_events() poll path " - "will return no events." - ) assert kv_events_config is not None - # Rejects unsupported parallelism, a non-Python V2 backend and colliding - # publish/replay port ranges, all before any socket is bound. + # Rejects unsupported parallelism and streaming itself, before any socket is + # bound and before any claim is made about which event path is in use. validate_streaming_support( kv_events_config, pp_size=mapping.pp_size, @@ -1300,8 +1292,14 @@ def __init__( # Ranks bind by global rank; only those sharing a host can collide. ranks_per_host=min(mapping.dp_size, mapping.gpus_per_node), data_parallel_size=mapping.dp_size, - backend=KV_CACHE_MANAGER_V2_BACKEND, ) + if self.event_buffer_max_size > 0: + logger.warning( + "Both kv_cache_config.event_buffer_max_size and streaming " + "kv_events_config are enabled; streaming publishing takes " + "precedence and the buffered get_kv_cache_events() poll path " + "will return no events." + ) if mapping.enable_attention_dp or mpi_rank() == 0: # Constructing it is side-effect free; start() below binds the socket # and starts the publisher thread once every other check has passed. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index f2d4dab1f7b7..c8459838d92c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -40,7 +40,7 @@ from tensorrt_llm.llmapi.llm_args import KVEventsConfig from tensorrt_llm.logger import logger from tensorrt_llm.runtime.kv_cache_hash import truncate_sha256_hash_to_int64 -from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import KVCacheEvent, KVCacheEventDiff +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheEvent, KVCacheEventDiff # Subscribers decode block hashes as 64-bit ints, so a bytes value would fail the # decode for the entire batch. @@ -401,29 +401,30 @@ def validate_streaming_support( cp_size: int, ranks_per_host: int, data_parallel_size: int, - backend: str, ) -> None: """Reject streaming-KV-event configurations the engine cannot honour. Split out of ``KVCacheManagerV2.__init__`` so the preconditions are testable without building a manager, which needs a GPU. + + Streaming is currently unsupported, so this always raises; ``config``, + ``ranks_per_host`` and ``data_parallel_size`` are retained for the caller's + signature and are consumed again once a native event sink exists. """ + del config, ranks_per_host, data_parallel_size if pp_size > 1: raise ValueError("Streaming KV events do not support pipeline parallelism") if cp_size > 1: raise ValueError("Streaming KV events do not support context parallelism") - if backend != "python": - # StreamingKVCacheEventManager is a duck-typed Python event sink, which cannot - # satisfy the nanobind constructor's nb::cast> - # (and the C++ radix tree calls the sink natively, not through Python). Fail - # with an actionable message instead of an opaque TypeError from the cast. - raise ValueError( - "Streaming KV events (kv_cache_config.kv_events_config) are only supported " - f"by the Python KV cache manager V2 backend, but '{backend}' is active. Set " - "TLLM_KV_CACHE_MANAGER_V2_BACKEND=python to enable streaming KV events, or " - "use the buffered path via kv_cache_config.event_buffer_max_size." - ) - validate_endpoint_ranges(config, ranks_per_host, data_parallel_size) + # StreamingKVCacheEventManager is a duck-typed Python event sink. It cannot satisfy the + # nanobind constructor's nb::cast>, and the C++ radix + # tree calls its sink natively rather than through Python, so there is no live path that + # can produce streaming events. Fail with an actionable message pointing at the buffered + # alternative instead of an opaque TypeError from the cast. + raise ValueError( + "Streaming KV events (kv_cache_config.kv_events_config) are not supported. Use the " + "buffered path via kv_cache_config.event_buffer_max_size instead." + ) def validate_endpoint_ranges( @@ -515,13 +516,17 @@ class _MultimodalBlockError(ValueError): class StreamingKVCacheEventManager: - """Scheduler-local fast path that produces KV cache event wire messages directly. + """Event-sink hook interface for out-of-band KV cache event publishing. + + The interface is duck typed rather than derived from ``KVCacheEventManager``: a sink + fully replaces event production (reusing the radix block hashes) and shares none of + the base manager's state. - Implements the V2 KV-cache-manager event-sink hook interface by duck - typing rather than inheriting ``KVCacheEventManager``: it fully replaces - event production (reusing the radix block hashes) and shares none of the - base manager's state, so subclassing would only risk partially initialised - base attributes. + No implementation is currently wired up. The C++ radix tree invokes its sink natively + rather than through Python, so a Python object cannot receive these callbacks; the + signatures below record the contract a native sink must satisfy. Constructing this + class raises -- ``validate_streaming_support`` rejects the configuration earlier, so + reaching here means that check was bypassed. """ def __init__( @@ -533,208 +538,36 @@ def __init__( max_window_size: int, max_entries: int = 50_000, ) -> None: - self._rank = data_parallel_rank - self._publisher = create_event_publisher(config, data_parallel_rank) - self._block_size = block_size - self._max_window_size = max_window_size - self._max_entries = max_entries - self._target_life_cycle_id: int | None = None - self._stored_blocks: dict[bytes, int] = {} - self._pending_events: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] - self._pending_entries = 0 - self._closed = False - self.stored_blocks = 0 - self.removed_blocks = 0 - self.partial_blocks_suppressed = 0 - self.multimodal_blocks_suppressed = 0 - self.non_target_life_cycles_ignored = 0 - self.dropped_events = 0 - self.enqueued_batches = 0 - self.enqueued_events = 0 - self.dropped_batches = 0 + raise NotImplementedError( + "Streaming KV events are not supported. Use the buffered path via " + "kv_cache_config.event_buffer_max_size instead." + ) def needs_token_digest_context(self) -> bool: # Streaming events do not emit multimodal keys. return False def start(self) -> None: - """Bind the publisher's sockets and start its background thread. - - Construction is side-effect free, so the owner calls this only once every - other initialization check has passed. A failure before this point therefore - leaves no socket bound and no thread running. - """ - self._publisher.start() + """Bind the publisher's sockets and start its background thread.""" def set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: - target_ids = [ - int(life_cycle_id) - for life_cycle_id, window_size in window_sizes.items() - if int(window_size) == self._max_window_size - ] - if not target_ids and window_sizes: - largest_window = max(window_sizes.values()) - target_ids = [ - int(life_cycle_id) - for life_cycle_id, window_size in window_sizes.items() - if window_size == largest_window - ] - if not target_ids: - raise ValueError("Streaming KV events require an attention KV cache life cycle") - self._target_life_cycle_id = min(target_ids) - logger.info( - "Streaming KV event fast path selected " - f"lifecycle_id={self._target_life_cycle_id} " - f"window_size={self._max_window_size}" - ) + """Select the attention life cycle whose blocks are published.""" def add_created_event( self, num_blocks_per_cache_level: Any, layer_group_ids: Any = None, - ) -> None: - return + ) -> None: ... - def add_stored_event(self, *args: Any, **kwargs: Any) -> None: - # Streaming publishing derives stored events from the per-block hooks - # below; the aggregate stored-event hook is intentionally unused. - return + def add_stored_event(self, *args: Any, **kwargs: Any) -> None: ... - def add_stored_block_event_from_block(self, block: Any) -> None: - if self._closed or self._target_life_cycle_id is None: - return - life_cycle_id = self._target_life_cycle_id - if life_cycle_id >= len(block.storage): - return - page_ref = block.storage[life_cycle_id] - page = None if page_ref is None else page_ref() - if page is None: - return - # A non-null page does not imply it covers the whole radix block: V2 can attach - # a page adopted from a shorter sibling. Publishing that as a BlockStored would - # tell the router the engine holds a prefix it cannot fully reuse. The buffered - # manager applies the same rule in _life_cycle_ids_from_radix_block(). - if page.num_tokens_in_block < len(block.tokens): - self.partial_blocks_suppressed += 1 - return - self._add_full_block(block) + def add_stored_block_event_from_block(self, block: Any) -> None: ... - def add_stored_life_cycle_event_from_block(self, block: Any, life_cycle_id: int) -> None: - if life_cycle_id is None or self._target_life_cycle_id is None: - return - if int(life_cycle_id) != self._target_life_cycle_id: - self.non_target_life_cycles_ignored += 1 - return - self.add_stored_block_event_from_block(block) + def add_stored_life_cycle_event_from_block(self, block: Any, life_cycle_id: int) -> None: ... - def _add_full_block(self, block: Any) -> None: - key = bytes(block.key) - if key in self._stored_blocks: - return - if len(block.tokens) != self._block_size: - self.partial_blocks_suppressed += 1 - return - if not self._reserve_entries(1): - return - try: - token_ids = self._token_ids(block.tokens) - block_hash, parent_hash = self._block_hashes(block) - except _MultimodalBlockError: - # Expected for multimodal cache-key blocks; skip without the - # malformed-data traceback that would otherwise flood the log. - self.multimodal_blocks_suppressed += 1 - self._pending_entries -= 1 - return - except ValueError: - self.dropped_events += 1 - self._pending_entries -= 1 - logger.error( - "Dropping streaming KV store event with unsupported token data\n" - f"{traceback.format_exc()}" - ) - return - self._stored_blocks[key] = block_hash - if self._pending_events and isinstance(self._pending_events[-1], BlockStored): - previous = self._pending_events[-1] - if previous.block_hashes and previous.block_hashes[-1] == parent_hash: - previous.block_hashes.append(block_hash) - previous.token_ids.extend(token_ids) - self.stored_blocks += 1 - return - self._pending_events.append( - BlockStored( - block_hashes=[block_hash], - parent_block_hash=parent_hash, - token_ids=token_ids, - block_size=self._block_size, - lora_id=None, - medium="GPU", - lora_name=None, - ) - ) - self.stored_blocks += 1 - - @staticmethod - def _token_ids(tokens: Any) -> list[int]: - token_ids: list[int] = [] - for token in tokens: - if type(token) is bytes: - # Multimodal cache-key digest; not representable as a wire int. - raise _MultimodalBlockError - if type(token) is not int: - raise ValueError("KV cache event wire format requires integer token IDs") - token_ids.append(token) - return token_ids - - def _block_hashes( - self, - block: Any, - ) -> tuple[int, int | None]: - parent = block.prev - is_root_child = getattr(parent, "ordinal", -1) == -1 - block_hash = _kv_event_wire_hash_from_radix_key(bytes(block.key)) - parent_hash = ( - None if is_root_child else _kv_event_wire_hash_from_radix_key(bytes(parent.key)) - ) - return block_hash, parent_hash - - def add_removed_event(self, block_hashes: Any) -> None: - if self._closed: - return - if isinstance(block_hashes, (bytes, str, int)): - block_hashes = (block_hashes,) - removed_hashes: list[ExternalBlockHash] = [] - for block_key in block_hashes: - if not isinstance(block_key, bytes): - continue - stored_hash = self._stored_blocks.pop(block_key, None) - if stored_hash is not None: - removed_hashes.append(stored_hash) - self._add_removed_hashes(removed_hashes) - - def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: - if self._closed or life_cycle_id is None or self._target_life_cycle_id is None: - return - if int(life_cycle_id) != self._target_life_cycle_id: - self.non_target_life_cycles_ignored += 1 - return - stored_hash = self._stored_blocks.pop(block_hash, None) - if stored_hash is not None: - self._add_removed_hashes([stored_hash]) + def add_removed_event(self, block_hashes: Any) -> None: ... - def _add_removed_hashes(self, block_hashes: list[ExternalBlockHash]) -> None: - if not block_hashes: - return - # Removals are never dropped by the per-iteration cap and, unlike stores, - # do not consume the _pending_entries budget: each hash was already - # reported as stored (so removals are bounded by the stored set), and - # counting them against the store budget would starve legitimate - # BlockStored events in a removal-heavy iteration. - if self._pending_events and isinstance(self._pending_events[-1], BlockRemoved): - self._pending_events[-1].block_hashes.extend(block_hashes) - else: - self._pending_events.append(BlockRemoved(block_hashes=block_hashes, medium="GPU")) - self.removed_blocks += len(block_hashes) + def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: ... def add_updated_event( self, @@ -743,46 +576,10 @@ def add_updated_event( cache_level: KVCacheEventDiff | None = None, priority: KVCacheEventDiff | None = None, layer_group_id: int | None = None, - ) -> None: - return - - def _reserve_entries(self, num_entries: int) -> bool: - if self._pending_entries + num_entries <= self._max_entries: - self._pending_entries += num_entries - return True - self.dropped_events += num_entries - if self.dropped_events == num_entries or ( - self.dropped_events & (self.dropped_events - 1) == 0 - ): - logger.warning( - "Dropping streaming KV events because the per-iteration safety " - f"cap was exceeded; dropped_events={self.dropped_events}" - ) - return False + ) -> None: ... def flush_iteration_events(self) -> None: - if self._closed or not self._pending_events: - return - events = self._pending_events - self._pending_events = [] - self._pending_entries = 0 - batch = KVEventBatch( - ts=time.time(), - events=events, - data_parallel_rank=self._rank, - ) - try: - if self._publisher.publish(batch): - self.enqueued_batches += 1 - self.enqueued_events += len(events) - else: - self.dropped_batches += 1 - except Exception: - self.dropped_batches += 1 - logger.error( - f"Dropping streaming KV event iteration batch on rank={self._rank}\n" - f"{traceback.format_exc()}" - ) + """Publish the events accumulated during this iteration.""" def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: # Streaming publishing pushes events out-of-band, so the pull API has @@ -791,19 +588,4 @@ def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEven return [] def shutdown(self) -> None: - if self._closed: - return - self.flush_iteration_events() - self._closed = True - self._publisher.shutdown() - logger.info( - "Streaming KV event fast path " - f"rank={self._rank} " - f"stored_blocks={self.stored_blocks} " - f"removed_blocks={self.removed_blocks} " - f"partial_blocks_suppressed={self.partial_blocks_suppressed} " - f"non_target_life_cycles_ignored={self.non_target_life_cycles_ignored} " - f"dropped_events={self.dropped_events} " - f"enqueued_batches={self.enqueued_batches} " - f"dropped_batches={self.dropped_batches}" - ) + """Flush pending events and stop the publisher.""" diff --git a/tensorrt_llm/runtime/__init__.py b/tensorrt_llm/runtime/__init__.py index eaad2f1d3ab9..0c34876952d1 100644 --- a/tensorrt_llm/runtime/__init__.py +++ b/tensorrt_llm/runtime/__init__.py @@ -13,34 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import sys -from contextlib import contextmanager -from typing import Iterator - - -# Duplicated from kv_cache_manager_v2._utils. We need this both inside and outside of -# kv_cache_manager_v2 due to restriction of mypyc build process. -@contextmanager -def temporary_sys_path(path: str) -> Iterator[None]: - already_in_path = path in sys.path - if not already_in_path: - sys.path.insert(0, path) - try: - yield - finally: - if not already_in_path: - sys.path.remove(path) - - -# Add current directory to sys.path so kv_cache_manager_v2 can be imported as top-level package. -# This is required because when kv_cache_manager_v2 is compiled with mypyc, it is compiled as -# a top-level package (to avoid complex build paths), but at runtime it is used as a submodule. -# The compiled extension might try to import its submodules using absolute imports based on its -# compiled name. -with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): - import kv_cache_manager_v2 - +from . import kv_cache_manager_v2 from .model_config import ModelConfig try: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md b/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md index e1f904ceb063..481648c2e9f4 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md @@ -4,9 +4,24 @@ This file provides guidance to coding agents when working with code in this dire ## What This Is -KVCacheManagerV2 is the KV cache management subsystem for TensorRT-LLM. It manages GPU/host/disk memory for key-value caches used during LLM inference, handling page allocation, eviction, multi-tier caching, radix-tree-based prefix sharing, and disaggregated serving. +The Python surface of KVCacheManagerV2, the KV cache management subsystem for TensorRT-LLM. +It manages GPU/host/disk memory for key-value caches used during LLM inference, handling page +allocation, eviction, multi-tier caching, radix-tree-based prefix sharing, and disaggregated +serving. -This is a **pure Python implementation** designed to be compilable with **mypyc** for production performance. There is also a `rawref` C extension for mutable object references. +**The implementation is C++**, in `cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/`, exposed +through nanobind as `tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2`. Read +`cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md` for the architecture, the +concurrency contract, and the C++ test suite. + +This directory contains only: + +- `__init__.py` — re-exports the nanobind surface, plus the plain-Python aliases, constants and + two `__dataclass_fields__` grafts the bindings do not carry. Adding a type to the bindings + means adding a rebind here and an entry in `__all__`. +- `_introspection.py` — white-box hooks for tests and accuracy harnesses. Every function + forwards to the bindings' native `_introspection` submodule; the indirection exists so callers + import a stable Python path and get plain Python containers back. ## Commands @@ -34,90 +49,37 @@ PYTHONPATH="$REPO_ROOT/" \ python "$REPO_ROOT/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py" -v ``` -### Building Extensions +### Rebuilding after a C++ change ```bash -# Build rawref C extension -cd rawref && python setup.py build_ext --inplace - -# Build mypyc-compiled version (run from runtime/ parent dir) -cd .. && python kv_cache_manager_v2/setup_mypyc.py build_ext --inplace - -# Or use the Makefile (builds both) -make all +cd cpp/build +cmake --build . --target bindings -j$(nproc) +cp tensorrt_llm/libtensorrt_llm.so ../../tensorrt_llm/libs/ +cp tensorrt_llm/thop/libth_common.so ../../tensorrt_llm/libs/ +cp tensorrt_llm/nanobind/bindings.cpython-*.so ../../tensorrt_llm/ ``` ### Debug Mode -Set `TLLM_DEBUG_MODE=1` to enable debug assertions (`NDEBUG=False`). Default is release mode (`NDEBUG=True`). - -## Architecture - -### Dual Import Trick - -The test file uses `find_spec("kv_cache_manager_v2")` to detect whether the package is importable as a top-level module (fast mode via `PYTHONPATH=.../runtime/`) or must be imported via the full path `tensorrt_llm.runtime.kv_cache_manager_v2` (production mode). This allows the same test file to work in both contexts. - -### Core Layers (bottom-up) - -1. **`_common.py`** — Fundamental types: `TokenId`, `TokenIdExt`, `CacheLevel`, `CacheTier`, `PageStatus`, `CudaStream`, `BeamIndex`, `BlockOrdinal`. All use `NewType` for type safety. - -2. **`rawref/`** — C extension providing `ref[T]`, a substitute for `weakref.ref` which is not compatible with mypyc. Uses raw object IDs instead of Python's weak reference machinery. Used throughout for parent/back-references that must not prevent GC. Objects must define `__rawref__ = NULL` class attribute and call `invalidate()` in `__del__`. - -3. **`_storage/`** — Low-level memory pool management. `_config.py` defines buffer/pool configurations. `_core.py` provides `CacheLevelStorage` with slot-based allocation across `PoolGroup`s and `Pool`s. - -4. **`_storage_manager.py`** — Coordinates storage across cache tiers (GPU → Host → Disk). Manages the `PerLevelEvictionController` and `batched_copy` for cross-tier data movement. - -5. **`_page.py`** — Page abstraction: `CommittedPage` (finalized, prefix-shareable), `UncommittedPage` (being filled), `BlockPage` (within a radix tree block). Includes `_SharedPageLock` and `batched_lock_to_gpu` for multi-level page locking. - -6. **`_life_cycle_registry.py`** — Maps `LayerGroupId`→`LifeCycleId`. Each layer group has either `AttnLifeCycle` (with optional sliding window + sink tokens) or `SsmLifeCycle`. Controls which blocks are "stale" and eligible for eviction. - -7. **`_block_radix_tree.py`** — Radix tree for prefix sharing across sequences. Blocks store pages and token IDs. Supports multi-modal tokens via `gen_multimodal_cache_key_tokens`. - -8. **`_eviction_controller/`** — Decides which pages to evict when memory is low, per cache level. - -9. **`_copy_engine.py`** — `CopyTask` and `batched_copy` for efficient GPU↔Host↔Disk data transfers. - -10. **`_core/_kv_cache.py`** (`_KVCache`) — Per-sequence cache state. Manages the block chain, commit/uncommit lifecycle, beam search forks, page locking, and cross-level migration. This is the most complex module. - -11. **`_core/_kv_cache_manager.py`** (`KVCacheManager`) — Top-level manager. Owns all `_KVCache` instances, handles batched operations (`prepareStep`, `acceptStep`, `cleanStep`), quota management, and disaggregated serving coordination. - -### Key Type Aliases - -- `LayerGroupId` = public alias of `LifeCycleId` (same value, different semantic) -- `PoolGroupIndex` ≠ `LayerGroupId` — pool group index is the storage-level index -- `PageIndex` = index into a pool's slot array -- `BlockOrdinal` = position of a block in a sequence's block chain - -### Page Lifecycle - -`UncommittedPage` → (commit) → `CommittedPage` → (evict to host) → `CommittedPage` at lower level → (recall to GPU) → `CommittedPage` at GPU level - -Pages inside radix tree blocks are `BlockPage` wrappers that delegate to the underlying `CommittedPage`/`UncommittedPage`. - -## Concurrency - -**This pure-Python backend is not thread-safe.** There is no lock around the -manager API; the only `threading.Lock` in the package guards copy-engine state -in `_copy_engine.py`. The GIL makes individual bytecode operations atomic but -gives no atomicity across a multi-step operation like `resize()`, which walks -the radix tree, evicts, and migrates pages. - -The **C++ backend is** thread-safe, via a manager-wide reader-writer lock. The -two backends are selected by `TLLM_KV_CACHE_MANAGER_V2_BACKEND` and are -otherwise interchangeable, so this is an easy asymmetry to trip over: code that -drives `resize()` from a background thread works against the C++ backend and -races against this one. - -See `cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md` -("Concurrency model") for the guarantee, its scope, and the supported -helper-thread handoff pattern. If this backend ever needs the same property, it -needs its own design -- the C++ lock is not reachable from here. +Set `TLLM_DEBUG_MODE=1` to enable debug assertions (`NDEBUG=False`). Default is release mode +(`NDEBUG=True`). ## Gotchas -- **`rawref` must be built first** — the package imports `rawref` at the top level. Run `make rawref` before anything else. -- **mypyc compilation is from the `runtime/` directory** — `setup_mypyc.py` expects `kv_cache_manager_v2/` prefixes in module paths. -- **`_exceptions.py` excluded from mypyc** — mypyc can't compile classes inheriting from builtin `Exception`. -- **`NDEBUG` controls assertions** — many hot-path assertions are gated behind `if not NDEBUG:`. Don't remove these guards. -- **`stopCommitting()` must NOT call `commit()`** — it would double-append tokens to the block. -- **Cross-stream sync on `cuda_stream` setter** — changing the CUDA stream records an event on the old stream and waits on the new one. This is intentional, and it is how a cache built on one thread is adopted by another without a blocking host-side sync (C++ backend; see its AGENTS.md). +- **Dual import trick.** `_load_cpp_module()` reaches the bindings two ways: through + `tensorrt_llm.bindings...` when `tensorrt_llm` is already imported, otherwise by walking up + from `find_spec("kv_cache_manager_v2")` to the `tensorrt_llm` root and importing + `bindings.internal.batch_manager.kv_cache_manager_v2` directly. The second path is what makes + fast mode work; the test files mirror the same branch. Keep both in sync. +- **`__dataclass_fields__` grafts.** `BatchDesc` and `KVCacheManagerConfig` are nanobind classes + that callers pass to `dataclasses.replace()`. `replace()` is keyed on `__dataclass_fields__`, + so `__init__.py` attaches a field spec to each. A new constructor field on either binding must + be added to the matching `_*FieldSpec` or `replace()` silently drops it. +- **Block-key hashing is a cross-language contract.** `root_block_key` and `block_key` must stay + byte-identical to the C++ `RootBlock::makeKey` / `Block::makeKey` they wrap, because + `tensorrt_llm/serve/router_utils.py` computes routing hashes with them and compares against + hashes the engine produced. `TestBlockKeyHashing` in the unit tests guards the format. +- **Streaming KV events have no implementation.** `kv_cache_config.kv_events_config` is rejected + by `validate_streaming_support`; the supported route is the buffered path via + `kv_cache_config.event_buffer_max_size`. A native event sink is what would re-enable it — a + Python sink cannot work, because the C++ radix tree calls its sink natively. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/Makefile b/tensorrt_llm/runtime/kv_cache_manager_v2/Makefile deleted file mode 100644 index 853ab20d5d2d..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/Makefile +++ /dev/null @@ -1,51 +0,0 @@ -# ################################################################################################## -# 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. -# ################################################################################################## - -# Makefile for kv_cache_manager_v2 -# Replaces build_mypyc.sh and rawref/build.sh - -PYTHON ?= python3 -RUNTIME_DIR ?= .. - -.PHONY: all mypyc rawref clean clean_mypyc clean_rawref - -# Default target -all: rawref mypyc - -# Build the rawref C extension -rawref: - cd rawref && $(PYTHON) setup.py build_ext --inplace - -# Build the mypyc extension -# Must be run from the parent directory (runtime) as setup_mypyc.py expects -# kv_cache_manager_v2/ prefixes in module names. -mypyc: - cd $(RUNTIME_DIR) && $(PYTHON) kv_cache_manager_v2/setup_mypyc.py build_ext --inplace - -# Clean everything -clean: clean_rawref clean_mypyc - -# Clean rawref build artifacts -clean_rawref: - cd rawref && rm -rf build - find rawref -name "*.so" -type f -delete - -# Clean mypyc build artifacts -# Cleans build/ directory in runtime (parent) and .so files in this directory -clean_mypyc: - rm -rf $(RUNTIME_DIR)/build - find . -name "*.so" -type f ! -path "./rawref/*" -delete diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 5a5eee4ecda0..a8a71706f7ca 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -13,320 +13,190 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Python surface of KVCacheManagerV2. + +The implementation lives in C++ under ``cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2`` +and is reached through the nanobind module +``tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2``. This package only +re-exports that surface, adds the few plain-Python aliases and constants the bindings do +not carry, and hosts the backend-agnostic ``_introspection`` helpers. +""" + import os import sys from importlib.util import find_spec from pathlib import Path -from types import ModuleType -from typing import NamedTuple, Optional, Union - -_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() - -#: Name of the active backend ("cpp" or "python"). Exposed so callers can gate -#: Python-only extension points, such as duck-typed event sinks, on the selection. -BACKEND = _BACKEND - -if _BACKEND == "python": - from . import rawref # noqa: F401 - from ._block_radix_tree import ( # noqa: F401 - ReuseScope, - gen_multimodal_cache_key_tokens, - sequence_to_blockchain_keys, - ) - from ._common import ( # noqa: F401 - BAD_PAGE_INDEX, - CACHE_LEVEL1, - GPU_LEVEL, - NDEBUG, - CacheLevel, - CacheTier, - CudaStream, - LayerId, - MemAddress, - PageIndexMode, - PageStatus, - Priority, - SlidingWindowSize, - TokenId, - TokenIdExt, - ) - from ._config import ( # noqa: F401 - AttentionLayerConfig, - BatchDesc, - BufferConfig, - CacheTierConfig, - DataRole, - DiskCacheTierConfig, - GpuCacheTierConfig, - HostCacheTierConfig, - KVCacheDesc, - KVCacheManagerConfig, - SsmLayerConfig, - SwaScratchReuseConfig, - ) - from ._core import ( # noqa: F401 - DEFAULT_BEAM_INDEX, - AggregatedPageDesc, - BeamIndex, - ExpandedBuffer, - KVCacheManager, - PageIndexConverter, - PlannedDropHandle, - PoolDesc, - PoolGroupDesc, - PoolGroupPeakBlockStats, - ScratchDesc, - _KVCache, - ) - from ._core._kv_cache import _Status as KvCacheStatus # noqa: F401 - from ._event_manager import ( # noqa: F401 - KVCacheCreatedData, - KVCacheEvent, - KVCacheEventDiff, - KVCacheEventManager, - KVCacheRemovedData, - KVCacheStoredBlockData, - KVCacheStoredData, - KVCacheUpdatedData, - UniqueToken, - ) - from ._exceptions import ( # noqa: F401 - CorruptedError, - CuError, - OutOfMemoryError, - OutOfPagesError, - ) - from ._life_cycle_registry import AttnLifeCycle, LayerGroupId, LifeCycleId # noqa: F401 - from ._stats import ( # noqa: F401 - _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, - CountsByLevel, - KVCacheIterationStatsDelta, - KVCacheStatsDelta, - ReusedBlocksByLevel, - SsmSnapshotIterationStatsDelta, - ) - from ._storage import BufferId # noqa: F401 - from ._storage._config import CoalescedBuffer, SlotDesc, SlotDescVariant # noqa: F401 - from ._storage._core import PoolGroupIndex, PoolIndex # noqa: F401 - from ._storage_manager import StorageStatistics # noqa: F401 - from ._utils import HalfOpenRange, exact_div, typed_range # noqa: F401 - - def poison_reason() -> str | None: - """First recorded KVCM2 invariant violation, or None. - - The pure-Python backend has no poison latch, so this is always None. - """ - return None - - def take_poison() -> str | None: - """Report the recorded violation and clear it. Always None on this backend.""" - return None - - def num_live_managers() -> int: - """Number of constructed, not-yet-destroyed managers. Not tracked on this backend.""" - return 0 - - _cpp_introspection = None -else: - - class ReuseScope(NamedTuple): - lora_id: int | None = None - salt: int | None = None - - def to_bytes(self) -> bytes: - ret = sum((value is not None) << i for i, value in enumerate(self)).to_bytes( - 1, "little", signed=False - ) - for value in self: - if value is not None: - ret += value.to_bytes(8, "little", signed=False) - return ret - - def _load_cpp_module(): - if "tensorrt_llm" in sys.modules: - from tensorrt_llm.bindings.internal.batch_manager import kv_cache_manager_v2 - - return kv_cache_manager_v2 - - spec = find_spec("kv_cache_manager_v2") - assert spec is not None and spec.origin is not None - trtllm_root = str(Path(spec.origin).parent.parent.parent) - sys.path.insert(0, trtllm_root) - try: - from bindings.internal.batch_manager import kv_cache_manager_v2 - - return kv_cache_manager_v2 - finally: - sys.path.remove(trtllm_root) - - _cpp = _load_cpp_module() - - AggregatedPageDesc = _cpp.AggregatedPageDesc - AttentionLayerConfig = _cpp.AttentionLayerConfig - BatchDesc = _cpp.BatchDesc - # BatchDesc is also consumed via dataclasses.replace(): MambaCacheManager's - # _build_cache_config appends dummy KVCacheDesc slots to each constraint with - # replace(batch, kv_caches=[...]). Like KVCacheManagerConfig below, the C++ - # binding replaces the Python @dataclass, so advertise the dataclass field - # set (replace() is keyed on __dataclass_fields__: reads fields via getattr, - # rebuilds via BatchDesc(**fields)). The binding already has a keyword - # __init__ and readable kv_caches / system_prompt_length fields. - import dataclasses as _dataclasses_bd - - @_dataclasses_bd.dataclass - class _BatchDescFieldSpec: - kv_caches: object = None - system_prompt_length: int = 0 - - BatchDesc.__dataclass_fields__ = _BatchDescFieldSpec.__dataclass_fields__ - del _BatchDescFieldSpec, _dataclasses_bd - BufferConfig = _cpp.BufferConfig - BufferId = _cpp.BufferId - CoalescedBuffer = _cpp.CoalescedBuffer - CacheTier = _cpp.CacheTier - DiskCacheTierConfig = _cpp.DiskCacheTierConfig - GpuCacheTierConfig = _cpp.GpuCacheTierConfig - ExpandedBuffer = _cpp.ExpandedBuffer - HostCacheTierConfig = _cpp.HostCacheTierConfig - KVCacheDesc = _cpp.KVCacheDesc - KVCacheCreatedData = _cpp.KVCacheCreatedData - KVCacheEvent = _cpp.KVCacheEvent - KVCacheEventDiff = _cpp.KVCacheEventDiff - KVCacheEventManager = _cpp.KVCacheEventManager - KVCacheIterationStatsDelta = _cpp.KVCacheIterationStatsDelta - KVCacheManager = _cpp.KVCacheManager - KVCacheManagerConfig = _cpp.KVCacheManagerConfig - IKvCacheColdPageCodec = _cpp.IKvCacheColdPageCodec - create_default_kv_cache_cold_page_codec = _cpp.create_default_kv_cache_cold_page_codec - # The C++ KVCacheManagerConfig binding replaces the Python @dataclass, but - # callers (the DeepSeek-V4 cache manager's _build_cache_config and our own - # host-tier fallback) use dataclasses.replace() on it. dataclasses.replace() - # is a free function keyed on __dataclass_fields__: it reads each field via - # getattr and rebuilds via cls(**fields). The binding already has a full - # keyword __init__ and readable fields, so we only need to advertise the - # dataclass field set. Field defaults/types are irrelevant here — replace() - # only uses the field names + init flag. The read-only - # enable_swa_scratch_reuse property is intentionally excluded (not a ctor - # field), matching the Python dataclass. - import dataclasses as _dataclasses - - @_dataclasses.dataclass - class _KVCacheManagerConfigFieldSpec: - tokens_per_block: int = 0 - cache_tiers: object = None - layers: object = None - max_util_for_resume: float = 0.97 - enable_partial_reuse: bool = True - reuse_match_backoff: int = 0 - constraints: object = None - typical_step: object = None - initial_pool_ratio: object = None - swa_scratch_reuse: object = None - commit_min_snapshot: bool = False - enable_stats: bool = True - text_only: bool = False - - KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ - del _KVCacheManagerConfigFieldSpec, _dataclasses - KVCacheRemovedData = _cpp.KVCacheRemovedData - KVCacheStatsDelta = _cpp.KVCacheStatsDelta - KVCacheStoredBlockData = _cpp.KVCacheStoredBlockData - KVCacheStoredData = _cpp.KVCacheStoredData - KVCacheUpdatedData = _cpp.KVCacheUpdatedData - KvCacheStatus = _cpp.KvCacheStatus - OutOfPagesError = _cpp.OutOfPagesError - PageStatus = _cpp.PageStatus - PoolDesc = _cpp.PoolDesc - PoolGroupDesc = _cpp.PoolGroupDesc - PoolGroupPeakBlockStats = _cpp.PoolGroupPeakBlockStats - SlotDesc = _cpp.SlotDesc - SlotDescVariant = _cpp.SlotDescVariant - SsmLayerConfig = _cpp.SsmLayerConfig - StorageStatistics = _cpp.StorageStatistics - _KVCache = _cpp._KVCache - poison_reason = _cpp.poison_reason - take_poison = _cpp.take_poison - num_live_managers = _cpp.num_live_managers - _cpp_introspection = getattr(_cpp, "_introspection", None) - _KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple(KVCacheIterationStatsDelta._field_names) - PlannedDropHandle = _cpp.PlannedDropHandle - CuError = _cpp.CuError - CorruptedError = _cpp.CorruptedError - - # Symbols added on main that are not yet ported to the C++ backend. - # TODO(kvCacheManagerV2-cpp): port these and replace the fallbacks. - AttnLifeCycle = getattr(_cpp, "AttnLifeCycle", None) - OutOfMemoryError = getattr(_cpp, "OutOfMemoryError", MemoryError) - PageIndexConverter = getattr(_cpp, "PageIndexConverter", None) - ReuseScope = getattr(_cpp, "ReuseScope", ReuseScope) - ScratchDesc = getattr(_cpp, "ScratchDesc", None) - SsmSnapshotIterationStatsDelta = _cpp.SsmSnapshotIterationStatsDelta - ReusedBlocksByLevel = _cpp.ReusedBlocksByLevel - SwaScratchReuseConfig = getattr(_cpp, "SwaScratchReuseConfig", None) - UniqueToken = _cpp.UniqueToken - - BeamIndex = int - CacheLevel = int - CacheTierConfig = Union[GpuCacheTierConfig, HostCacheTierConfig, DiskCacheTierConfig] - CudaStream = int - DataRole = str - HalfOpenRange = getattr(_cpp, "HalfOpenRange", tuple) - LayerGroupId = int - LayerId = int - LifeCycleId = int - MemAddress = int - PoolGroupIndex = int - PoolIndex = int - Priority = int - SlidingWindowSize = Optional[int] - TokenId = int - TokenIdExt = Union[int, bytes] - - BAD_PAGE_INDEX = -1 - DEFAULT_BEAM_INDEX = 0 - GPU_LEVEL = 0 - CACHE_LEVEL1 = 1 - NDEBUG = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] != "1" - - class _RawRef: - def __init__(self, obj=None): - self._obj = obj - - def __call__(self): - return self._obj - - def invalidate(self) -> None: - self._obj = None - - @classmethod - def __class_getitem__(cls, _item): - return cls - - rawref = ModuleType(f"{__name__}.rawref") - rawref.ReferenceType = _RawRef - rawref.ref = _RawRef - rawref.NULL = _RawRef() - sys.modules.setdefault(f"{__name__}.rawref", rawref) - - class PageIndexMode(int): - SHARED = 0 - PER_LAYER = 1 - - gen_multimodal_cache_key_tokens = _cpp.gen_multimodal_cache_key_tokens - sequence_to_blockchain_keys = _cpp.sequence_to_blockchain_keys - - def exact_div(x: int, y: int) -> int: - assert x % y == 0 - return x // y - - def typed_range(*args: int) -> range: - return range(*args) +from typing import Optional, Union + + +def _load_cpp_module(): + if "tensorrt_llm" in sys.modules: + from tensorrt_llm.bindings.internal.batch_manager import kv_cache_manager_v2 + + return kv_cache_manager_v2 + + # Dev mode: the package is importable as a top-level ``kv_cache_manager_v2`` (via + # PYTHONPATH=.../tensorrt_llm/runtime/), so the bindings are reached by walking up to + # the tensorrt_llm root rather than importing the whole package. + spec = find_spec("kv_cache_manager_v2") + assert spec is not None and spec.origin is not None + trtllm_root = str(Path(spec.origin).parent.parent.parent) + sys.path.insert(0, trtllm_root) + try: + from bindings.internal.batch_manager import kv_cache_manager_v2 + + return kv_cache_manager_v2 + finally: + sys.path.remove(trtllm_root) + + +_cpp = _load_cpp_module() + +AggregatedPageDesc = _cpp.AggregatedPageDesc +AttentionLayerConfig = _cpp.AttentionLayerConfig +AttnLifeCycle = _cpp.AttnLifeCycle +BatchDesc = _cpp.BatchDesc +# BatchDesc is also consumed via dataclasses.replace(): MambaCacheManager's +# _build_cache_config appends dummy KVCacheDesc slots to each constraint with +# replace(batch, kv_caches=[...]). Like KVCacheManagerConfig below, the C++ +# binding replaces the Python @dataclass, so advertise the dataclass field +# set (replace() is keyed on __dataclass_fields__: reads fields via getattr, +# rebuilds via BatchDesc(**fields)). The binding already has a keyword +# __init__ and readable kv_caches / system_prompt_length fields. +import dataclasses as _dataclasses_bd # noqa: E402 + + +@_dataclasses_bd.dataclass +class _BatchDescFieldSpec: + kv_caches: object = None + system_prompt_length: int = 0 + + +BatchDesc.__dataclass_fields__ = _BatchDescFieldSpec.__dataclass_fields__ +del _BatchDescFieldSpec, _dataclasses_bd +BufferConfig = _cpp.BufferConfig +BufferId = _cpp.BufferId +CoalescedBuffer = _cpp.CoalescedBuffer +CacheTier = _cpp.CacheTier +CorruptedError = _cpp.CorruptedError +CuError = _cpp.CuError +DiskCacheTierConfig = _cpp.DiskCacheTierConfig +GpuCacheTierConfig = _cpp.GpuCacheTierConfig +ExpandedBuffer = _cpp.ExpandedBuffer +HalfOpenRange = _cpp.HalfOpenRange +HostCacheTierConfig = _cpp.HostCacheTierConfig +KVCacheDesc = _cpp.KVCacheDesc +KVCacheCreatedData = _cpp.KVCacheCreatedData +KVCacheEvent = _cpp.KVCacheEvent +KVCacheEventDiff = _cpp.KVCacheEventDiff +KVCacheEventManager = _cpp.KVCacheEventManager +KVCacheIterationStatsDelta = _cpp.KVCacheIterationStatsDelta +KVCacheManager = _cpp.KVCacheManager +KVCacheManagerConfig = _cpp.KVCacheManagerConfig +IKvCacheColdPageCodec = _cpp.IKvCacheColdPageCodec +create_default_kv_cache_cold_page_codec = _cpp.create_default_kv_cache_cold_page_codec +# The C++ KVCacheManagerConfig binding replaces the Python @dataclass, but +# callers (the DeepSeek-V4 cache manager's _build_cache_config and our own +# host-tier fallback) use dataclasses.replace() on it. dataclasses.replace() +# is a free function keyed on __dataclass_fields__: it reads each field via +# getattr and rebuilds via cls(**fields). The binding already has a full +# keyword __init__ and readable fields, so we only need to advertise the +# dataclass field set. Field defaults/types are irrelevant here — replace() +# only uses the field names + init flag. The read-only +# enable_swa_scratch_reuse property is intentionally excluded (not a ctor +# field), matching the binding's constructor. +import dataclasses as _dataclasses # noqa: E402 + + +@_dataclasses.dataclass +class _KVCacheManagerConfigFieldSpec: + tokens_per_block: int = 0 + cache_tiers: object = None + layers: object = None + max_util_for_resume: float = 0.97 + enable_partial_reuse: bool = True + reuse_match_backoff: int = 0 + constraints: object = None + typical_step: object = None + initial_pool_ratio: object = None + swa_scratch_reuse: object = None + commit_min_snapshot: bool = False + enable_stats: bool = True + text_only: bool = False + + +KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ +del _KVCacheManagerConfigFieldSpec, _dataclasses +KVCacheRemovedData = _cpp.KVCacheRemovedData +KVCacheStatsDelta = _cpp.KVCacheStatsDelta +KVCacheStoredBlockData = _cpp.KVCacheStoredBlockData +KVCacheStoredData = _cpp.KVCacheStoredData +KVCacheUpdatedData = _cpp.KVCacheUpdatedData +KvCacheStatus = _cpp.KvCacheStatus +OutOfMemoryError = _cpp.OutOfMemoryError +OutOfPagesError = _cpp.OutOfPagesError +PageIndexConverter = _cpp.PageIndexConverter +PageIndexMode = _cpp.PageIndexMode +PageStatus = _cpp.PageStatus +PlannedDropHandle = _cpp.PlannedDropHandle +PoolDesc = _cpp.PoolDesc +PoolGroupDesc = _cpp.PoolGroupDesc +PoolGroupPeakBlockStats = _cpp.PoolGroupPeakBlockStats +ReuseScope = _cpp.ReuseScope +ReusedBlocksByLevel = _cpp.ReusedBlocksByLevel +ScratchDesc = _cpp.ScratchDesc +SlotDesc = _cpp.SlotDesc +SlotDescVariant = _cpp.SlotDescVariant +SsmLayerConfig = _cpp.SsmLayerConfig +StorageStatistics = _cpp.StorageStatistics +SsmSnapshotIterationStatsDelta = _cpp.SsmSnapshotIterationStatsDelta +SwaScratchReuseConfig = _cpp.SwaScratchReuseConfig +UniqueToken = _cpp.UniqueToken +_KVCache = _cpp._KVCache +_cpp_introspection = getattr(_cpp, "_introspection", None) +_KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple(KVCacheIterationStatsDelta._field_names) + +gen_multimodal_cache_key_tokens = _cpp.gen_multimodal_cache_key_tokens +num_live_managers = _cpp.num_live_managers +poison_reason = _cpp.poison_reason +sequence_to_blockchain_keys = _cpp.sequence_to_blockchain_keys +take_poison = _cpp.take_poison + +BeamIndex = int +CacheLevel = int +CacheTierConfig = Union[GpuCacheTierConfig, HostCacheTierConfig, DiskCacheTierConfig] +CudaStream = int +DataRole = str +LayerGroupId = int +LayerId = int +LifeCycleId = int +MemAddress = int +PoolGroupIndex = int +PoolIndex = int +Priority = int +SlidingWindowSize = Optional[int] +TokenId = int +TokenIdExt = Union[int, bytes] + +BAD_PAGE_INDEX = -1 +DEFAULT_BEAM_INDEX = 0 +GPU_LEVEL = 0 +CACHE_LEVEL1 = 1 +NDEBUG = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] != "1" + + +def exact_div(x: int, y: int) -> int: + assert x % y == 0 + return x // y + + +def typed_range(*args: int) -> range: + return range(*args) __all__ = [ "AggregatedPageDesc", "AttentionLayerConfig", - "BACKEND", "BAD_PAGE_INDEX", "CACHE_LEVEL1", "BatchDesc", @@ -346,6 +216,7 @@ def typed_range(*args: int) -> range: "GpuCacheTierConfig", "HalfOpenRange", "HostCacheTierConfig", + "IKvCacheColdPageCodec", "KVCacheDesc", "KVCacheCreatedData", "KVCacheEvent", @@ -393,15 +264,12 @@ def typed_range(*args: int) -> range: "CuError", "OutOfMemoryError", "_KVCache", + "create_default_kv_cache_cold_page_codec", "exact_div", "gen_multimodal_cache_key_tokens", - "sequence_to_blockchain_keys", - "rawref", - "typed_range", + "num_live_managers", "poison_reason", + "sequence_to_blockchain_keys", "take_poison", - "num_live_managers", + "typed_range", ] - -if _BACKEND != "python": - __all__.extend(["IKvCacheColdPageCodec", "create_default_kv_cache_cold_page_codec"]) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 91ccd9142528..d8273bdaff18 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -35,12 +35,20 @@ from typing import ( # From _common.py NDEBUG: Final[int] DEFAULT_BEAM_INDEX: Final[BeamIndex] +BAD_PAGE_INDEX: Final[int] +GPU_LEVEL: Final[CacheLevel] +CACHE_LEVEL1: Final[CacheLevel] class CorruptedError(Exception): - """Raised by every public entry point once a broken invariant has been recorded. + """Raised by every public entry point once a broken invariant has been recorded.""" - Only the C++ backend has the latch that raises this; the pure-Python backend never does. - """ +class CuError(Exception): + """A CUDA driver call failed; carries the driver's own status code.""" + + error_code: Any + +class OutOfMemoryError(Exception): ... +class OutOfPagesError(OutOfMemoryError): ... def poison_reason() -> str | None: """First recorded invariant violation, or None. Never clears, so it is safe to poll.""" @@ -56,12 +64,31 @@ class CacheTier(enum.IntEnum): HOST_MEM = 1 DISK = 2 +class PageStatus(enum.Enum): + LOCKED = enum.auto() + HELD = enum.auto() + DROPPABLE = enum.auto() + class PageIndexMode(enum.IntEnum): SHARED = 0 PER_LAYER = 1 LifeCycleId = NewType("LifeCycleId", int) LayerGroupId: TypeAlias = LifeCycleId + +class AttnLifeCycle: + """The attention life cycle, keyed by its sliding-window and sink-token shape.""" + + @staticmethod + def make( + window_size: int | None, num_sink_tokens: int | None, tokens_per_block: int + ) -> "AttnLifeCycle": ... + @property + def window_size(self) -> int | None: ... + @property + def num_sink_blocks(self) -> int: ... + def get_stale_range(self, history_length: int, tokens_per_block: int) -> HalfOpenRange: ... + CacheLevel = NewType("CacheLevel", int) TokenId = NewType("TokenId", int) TokenIdExt = Union[TokenId, bytes] @@ -73,6 +100,7 @@ class ReuseScope(NamedTuple): lora_id: int | None = None salt: int | None = None +SlidingWindowSize: TypeAlias = int | None LayerId = NewType("LayerId", int) CudaStream = NewType("CudaStream", int) BeamIndex = NewType("BeamIndex", int) @@ -334,7 +362,7 @@ class KVCacheEventManager: def flush_iteration_events(self) -> None: ... def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: ... -# Backend-neutral key builders (native C++ under the C++ backend, pure-Python otherwise). +# Native key builders, shared with the radix tree so routing hashes match the engine's. def gen_multimodal_cache_key_tokens( id_offset: int, multi_modal_data_digest: bytes, @@ -353,6 +381,8 @@ class _Status(enum.Enum): SUSPENDED = enum.auto() CLOSED = enum.auto() +KvCacheStatus: TypeAlias = _Status + IndexSeq = array.array[int] | memoryview[int] class _KVCache: @@ -548,7 +578,6 @@ class KVCacheManager: self, config: KVCacheManagerConfig, event_manager: KVCacheEventManager | None = None, - # C++ backend only; the pure-Python backend does not accept this parameter. cold_page_codec: IKvCacheColdPageCodec | None = None, ) -> None: ... def __del__(self) -> None: ... @@ -645,3 +674,6 @@ class KVCacheManager: def need_adjustment(self) -> bool: ... @property def commit_min_snapshot(self) -> bool: ... + +def exact_div(x: int, y: int) -> int: ... +def typed_range(*args: int) -> range: ... diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py deleted file mode 100644 index d1e3ec662cc7..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ /dev/null @@ -1,882 +0,0 @@ -# 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. - -import hashlib -import itertools -from array import array -from itertools import chain -from typing import TYPE_CHECKING, Iterable, Iterator, NamedTuple, Sequence, TypeVar, cast - -from . import rawref -from ._common import NDEBUG, BlockOrdinal, PageStatus, TokenId, TokenIdExt -from ._life_cycle_registry import AttnLifeCycle, LifeCycle, LifeCycleId, LifeCycleRegistry -from ._utils import TypedIndexList, filled_list, map_optional, typed_range, unwrap_rawref - -if TYPE_CHECKING: - from ._event_manager import KVCacheEventManager - from ._page import CommittedPage - - -BlockKey = bytes -TokenBlock = list[TokenIdExt] - -_SHA256_DIGEST_SIZE = hashlib.sha256().digest_size -_UINT_ITEM_SIZE = array("I").itemsize -if _UINT_ITEM_SIZE != 4: - raise RuntimeError("Hasher requires a platform with 4-byte unsigned ints") - - -# id_offset is usually vocab_size. Backend-neutral (depends only on _common); the -# C++ backend exposes a native gen_multimodal_cache_key_tokens via nanobind instead. -def gen_multimodal_cache_key_tokens( - id_offset: int, multi_modal_data_digest: bytes, num_tokens: int, token_offset: int = 0 -) -> list[TokenIdExt]: - """Create synthetic tokens used only when building multimodal KV-cache keys. - - Item-local token 0 carries the content digest; later offsets use deterministic IDs above the vocab. - - Args: - id_offset: First synthetic id, usually ``vocab_size``, so generated ids cannot - collide with real token ids. - multi_modal_data_digest: Content digest of the multimodal item; must be exactly - ``_SHA256_DIGEST_SIZE`` bytes. - num_tokens: Number of synthetic tokens to generate. Must be positive. - token_offset: Item-local index of the first generated token. Must be non-negative; - only offset 0 carries the digest. - - Returns: - The generated tokens, digest first when ``token_offset`` is 0. - - Raises: - ValueError: If the digest length is wrong, ``num_tokens`` is not positive, or - ``token_offset`` is negative. - """ - if len(multi_modal_data_digest) != _SHA256_DIGEST_SIZE: - raise ValueError(f"multi_modal_data_digest must have length {_SHA256_DIGEST_SIZE}") - if num_tokens <= 0: - raise ValueError("num_tokens must be positive") - if token_offset < 0: - raise ValueError("token_offset must be non-negative") - return [ - multi_modal_data_digest if token_offset + i == 0 else TokenId(id_offset + token_offset + i) - for i in range(num_tokens) - ] - - -class Hasher: - """Incremental SHA-256 hasher used to derive block keys for the radix tree. - - Accepts ints (encoded as 4 little-endian bytes each, matching the C++ backend's - 4-byte ``TokenIdExt`` layout), raw ``bytes`` (multimodal content digests and - reuse-scope fields), or a sequence mixing the two. Both backends must produce - identical digests for the same logical input, so the encoding is part of the - on-disk/cross-process contract and cannot change unilaterally. - - Args: - data: Optional initial value, hashed immediately as if passed to ``update``. - """ - - # SECURITY INVARIANT: the block-key hash MUST stay cryptographically - # collision-resistant and >= 256-bit. The radix tree is a globally shared, - # cross-request/cross-tenant cache index; prefix matches are decided purely by - # digest equality with NO re-check of the underlying tokens; and the hashed - # input (tokens, the user-supplied cache_salt, multimodal content bytes) is - # attacker-influenceable. A collision therefore silently reuses another - # request's KV blocks (cross-request corruption / data leak), and cache_salt - # tenant isolation relies entirely on this hash's collision resistance. Do NOT - # swap in a non-cryptographic hash (xxHash, HighwayHash, ...) or truncate below - # 256 bits without first adding a token-content equality check on match. The - # C++ backend (blockRadixTree) mirrors this with SHA-256 (CSHA256). - __slots__ = "_hasher" - _hasher: "hashlib._Hash" - - def __init__(self, data: int | bytes | Sequence[int | bytes] | None = None) -> None: - self._hasher = hashlib.sha256() - if data is not None: - self.update(data) - - def update(self, data: int | bytes | Sequence[int | bytes]) -> "Hasher": - """Fold ``data`` into the running digest. - - Args: - data: An int token id (0 <= id < 2**31), raw ``bytes``, or a sequence of - either. An all-int sequence takes a single-call fast path; a sequence - containing ``bytes`` (multimodal blocks) falls back to per-item hashing. - - Returns: - This ``Hasher``, to allow chaining. - """ - # This function is perf-critical. Expect compromised code quality. - if type(data) is int: - assert NDEBUG or (data >= 0 and data < (1 << 31)) - self._hasher.update(data.to_bytes(4, "little")) - elif type(data) is bytes: - self._hasher.update(data) - else: - # Hash the whole token block in one C call instead of one per token. - # array("I", data).tobytes() packs each int as 4 native-endian bytes - # (unsigned int); all NVIDIA GPU host platforms (x86_64, aarch64/Grace) - # are little-endian so this is byte-identical to the per-token - # to_bytes(4, "little") loop AND to the C++ backend's 4-byte TokenIdExt - # layout (normal token = little-endian id, high tag bit clear). Falls - # back to that loop for multimodal blocks (which contain bytes items). - try: - self._hasher.update(array("I", data).tobytes()) # type: ignore - except (TypeError, OverflowError): - for item in data: # type: ignore - assert ( - NDEBUG - or (type(item) is int and (0 <= item < (1 << 31))) - or type(item) is bytes - ) - self._hasher.update(item.to_bytes(4, "little") if (type(item) is int) else item) # type: ignore - return self - - @property - def digest(self) -> bytes: - return self._hasher.digest() - - -def reuse_scope_to_bytes(reuse_scope: Iterable[int | None]) -> bytes: - """Serialize a reuse scope to its reuse-namespace bytes. - - Backend-neutral: reads the scope's fields by iteration, so it works for both - the pure-Python ``ReuseScope`` NamedTuple and the C++ binding without relying - on a ``to_bytes()`` method. The layout mirrors the C++ ``emitReuseScopeBytes``: - a mask byte (one bit per field, set when the field is present) followed by one - little-endian ``uint64`` per present field (``signed=False``). - """ - values = list(reuse_scope) - mask = sum((value is not None) << i for i, value in enumerate(values)) - ret = mask.to_bytes((len(values) + 7) // 8, "little", signed=False) - for value in values: - if value is not None: - ret += int(value).to_bytes(8, "little", signed=False) - return ret - - -def sequence_to_blockchain_keys( - tokens_per_block: int, reuse_scope: Iterable[int | None], tokens: Sequence[TokenIdExt] -) -> Iterator[tuple[TokenBlock, BlockKey]]: - """Yield ``(token_block, key)`` pairs seeding a blockchain of KV-cache keys. - - The first pair is the root (``[]``, reuse-scope digest); each subsequent pair - hashes one ``tokens_per_block`` chunk on top of the previous digest. - """ - digest = Hasher(reuse_scope_to_bytes(reuse_scope)).digest - yield [], digest - iterator = iter(tokens) - while True: - token_block = list(itertools.islice(iterator, tokens_per_block)) - if not token_block: - break - digest = Hasher(digest).update(token_block).digest - yield token_block, digest - - -class ReuseScope(NamedTuple): - """Per-request namespace for prefix reuse.""" - - lora_id: int | None = None - salt: int | None = None - - def to_bytes(self) -> bytes: - return reuse_scope_to_bytes(self) - - -class ReuseMatch(NamedTuple): - """Volatile result of a KV cache prefix match. - - ``num_reusable_tokens_before_hybrid_pruning`` is retained for internal - diagnostics. It is the prefix the attention pages alone would support, - before recurrent snapshot availability shortens it. - - ``num_reusable_tokens_before_pruning`` is the raw token-path walk depth, - before any pruning at all. It locates where this request's content diverges - from the tree, independent of which pages happen to still be resident, so - ``num_reusable_tokens_before_pruning == num_lookup_tokens`` means the whole - lookup range matched and there is no fork here. - """ - - blocks: list["Block"] - num_tokens: int - num_lookup_tokens: int - num_reusable_tokens_before_hybrid_pruning: int - num_reusable_tokens_before_pruning: int - - -Child = TypeVar("Child", bound="Block | RootBlock") -Children = dict[BlockKey, Child] - - -def try_get_tree(block: "RootBlock | Block") -> "BlockRadixTree | None": - node = block - while not isinstance(node, BlockRadixTree): - node = node._prev() - if node is None: - return None - return node - - -def get_tree(block: "RootBlock | Block") -> "BlockRadixTree": - tree = try_get_tree(block) - if tree is None: - raise ValueError("Dereferencing a dangling rawref") - return tree - - -def detach_next(parent: "Block | RootBlock", key: BlockKey) -> "Block | None": - child = parent.next.pop(key, None) - if child is None: - return None - - child._prev = rawref.NULL - if isinstance(parent, RootBlock) and not parent.next: - tree = parent._prev() - if tree is not None and parent.key in tree.next: - detached_root = tree.next.pop(parent.key) - parent._prev = rawref.NULL - assert detached_root is parent - return child - - -def remove_subtree(root: "Block") -> None: - # taking O(1) space - # remove leaf blocks one by one, in post-order - removed_block_hashes: list[BlockKey] = [] - tree = try_get_tree(root) - event_manager = tree.event_manager if tree is not None else None - block: Block = root - while True: - if block.next: - block = next(iter(block.next.values())) - else: - removed_block_hashes.append(block.key) - if block._prev() is None: - assert block is root - break - prev_block: Block | RootBlock = block.prev - detached = detach_next(prev_block, block.key) - assert detached is block - if block is root: - break - assert isinstance(prev_block, Block) - block = prev_block - if event_manager is not None: - event_manager.add_removed_event(removed_block_hashes) - - -def traverse_post_order(root: "Block") -> Iterator["Block"]: - "post-order traversal of the subtree rooted at root" - stack: list[Iterator[Block]] = [] - block: Block | None = root - while True: - assert block is not None - if block.next: - child_iter = iter(block.next.values()) - stack.append(child_iter) - block = next(child_iter) - else: - yield (last_yielded := block) - while stack and (block := next(stack[-1], None)) is None: - yield (last_yielded := cast(Block, last_yielded.prev)) - stack.pop() - if not stack: - break - - -def find_best_partial_match_in_next_nodes( - block: "Block | RootBlock", tokens: TokenBlock -) -> tuple["Block | None", int]: - """ - Among all child nodes (self.next), finds the one whose tokens have the longest leading match with the given tokens. - Returns a tuple of (best_block, num_matched_tokens). - If no child matches any tokens, returns (None, 0). - """ - if len(block.next) >= 32: - # TODO: build a database to accelerate partial matching. (TRTLLM-7784) - # For now, it might be too slow to iterate over all children, so let's just skip. - return None, 0 - best_block = None - best_match_len = 0 - for b in block.next.values(): - match_len = b._partial_match_this_node(tokens) - if match_len > best_match_len: - best_match_len = match_len - best_block = b - return best_block, best_match_len - - -class DuplicateKeyError(Exception): - "Another block with the same key already exists" - - key: BlockKey - - def __init__(self, key: BlockKey) -> None: - super().__init__(f"Block with key {key.hex()} already exists") - self.key = key - - -class UselessBlockError(Exception): - block: "Block" - - def __init__(self, block: "Block") -> None: - super().__init__( - f"Block is useless because all its tokens are covered by another block with key = {block.key.hex()}" - ) - self.block = block - - -def _add_or_get_existing( - parent: "RootBlock | Block", tokens: Sequence[TokenIdExt] -) -> "Block | None": - try: - return Block(tokens, parent) - except DuplicateKeyError as e: - return parent.next[e.key] - except UselessBlockError: - return None - - -class RootBlock: - __slots__ = ("__rawref__", "_needs_token_digest_context", "_prev", "key", "next", "reuse_scope") - key: BlockKey - reuse_scope: ReuseScope - _prev: rawref.ref["BlockRadixTree"] - next: Children["Block"] - _needs_token_digest_context: bool - __rawref__: rawref.ref["RootBlock"] - - def __init__(self, reuse_scope: ReuseScope, prev: "BlockRadixTree") -> None: - self.key = self.make_key(reuse_scope) - assert self.key not in prev.next, "Root block already exists" - self.reuse_scope = reuse_scope - self._prev = rawref.ref(prev) - self.next = {} - self.__rawref__ = rawref.NULL - event_manager = prev.event_manager - self._needs_token_digest_context = ( - event_manager is not None and event_manager.needs_token_digest_context() - ) - prev.next[self.key] = self - - def __del__(self) -> None: - self.__rawref__.invalidate() - - @property - def ordinal(self) -> BlockOrdinal: - return BlockOrdinal(-1) - - @property - def prev(self) -> "BlockRadixTree": - return unwrap_rawref(self._prev) - - @property - def num_life_cycles(self) -> LifeCycleId: - return self.prev.num_life_cycles - - @property - def tokens_per_block(self) -> int: - return self.prev.tokens_per_block - - @staticmethod - def make_key(reuse_scope: ReuseScope) -> BlockKey: - return Hasher(reuse_scope.to_bytes()).digest - - -class Block: - """ - A block of tokens. Manages data for all layers. - """ - - __slots__ = ( - "__rawref__", - "_needs_token_digest_context", - "_prev", - "key", - "last_token_digest", - "next", - "ordinal", - "storage", - "tokens", - ) - key: BlockKey - tokens: Sequence[TokenIdExt] - last_token_digest: bytes | None - ordinal: BlockOrdinal - _needs_token_digest_context: bool - _prev: rawref.ref["Block | RootBlock"] - next: Children["Block"] - __rawref__: rawref.ref["Block"] - - # indexed with LifeCycleId - storage: TypedIndexList[LifeCycleId, rawref.ref["CommittedPage"] | None] - - @staticmethod - def make_key(prev_key: BlockKey, tokens: Sequence[TokenIdExt]) -> BlockKey: - return Hasher(prev_key).update(tokens).digest - - def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> None: - assert prev.tokens_per_block == prev.prev.tokens_per_block, "prev must be a full block" - self.key = self.make_key(prev.key, tokens) - self.tokens = tokens - self.ordinal = BlockOrdinal(prev.ordinal + 1) - self._prev = rawref.ref(prev) - self.next = {} - self.storage = filled_list(None, prev.num_life_cycles) - self.__rawref__ = rawref.NULL - self._needs_token_digest_context = prev._needs_token_digest_context - self.last_token_digest = None - # a Block is useless if all its tokens are covered by a sibling block. Raise UselessBlockError if so. - if self.key in prev.next: - raise UselessBlockError(prev.next[self.key]) - if len(tokens) < self.tokens_per_block: - # @TODO: when we have the database for find_best_partial_match_in_next_nodes, we may use - # that for faster check. - for b in prev.next.values(): - if b.tokens[: len(tokens)] == tokens: - raise UselessBlockError(b) - if self._needs_token_digest_context: - # Share the last digest through text-only descendants, including ancestors - # without committable pages that never publish a stored event themselves. - self.last_token_digest = prev.last_token_digest if isinstance(prev, Block) else None - for token in reversed(tokens): - if isinstance(token, bytes): - self.last_token_digest = token - break - # A later turn may extend a partial endpoint to this longer block, replacing the - # partial sibling. That turn may not have a committable SWA page for this block: - # commit_min_snapshot releases out-of-window pages, while SWA scratch reuse uses - # temporary shared storage that is not preserved. Adopt the partial sibling's - # pages to keep the shorter endpoint reusable, retaining each page's recorded token - # count (see CommittedPage.num_tokens_in_block). - to_remove = [] - for k, b in prev.next.items(): - if len(b.tokens) < len(tokens) and tokens[: len(b.tokens)] == b.tokens: - assert NDEBUG or (not b.is_full and b is not self and b.key == k and not b.next) - to_remove.append(k) - # Two covered siblings would be prefixes of each other; the insertion logic - # would already have replaced the shorter one. - assert NDEBUG or len(to_remove) <= 1 - event_manager = get_tree(prev).event_manager if to_remove else None - # Keep RootBlock attached while covered children are replaced. Adding - # the replacement first prevents detach_next() from pruning an - # otherwise-empty root before this block becomes its new child. - prev.next[self.key] = self - for k in to_remove: - b = detach_next(prev, k) - assert isinstance(b, Block) - self._adopt_pages_from(b) - if event_manager is not None: - event_manager.add_removed_event(b.key) - assert b.is_orphan # _KVCache may still hold it. - # prev.next keeps a strong ref to this _Block, so no need to remove self from prev.next in __del__(). - - def page_coverage(self, lc_idx: LifeCycleId) -> int: - """Return the page's recorded token count, or zero if the slot is empty. - - For attention this is prefix coverage; for SSM it is an exact checkpoint position. - """ - page = self.get_page(lc_idx) - return page.num_tokens_in_block if page is not None else 0 - - def holds_page(self, page: "CommittedPage") -> bool: - return self.get_page(page.life_cycle) is page - - def can_replace_page(self, lc_idx: LifeCycleId, num_tokens_in_block: int) -> bool: - """Whether a page recording `num_tokens_in_block` may take over slot `lc_idx`. - - A slot keeps only the page with the largest recorded token count. For attention, - greater coverage strictly dominates lesser coverage. For SSM, this deliberately - keeps only the latest checkpoint -- two conversation turns rarely end inside the - same block, and if they do, a reuse miss is acceptable. - - Pure; use replace_page() to install. - """ - existing = self.get_page(lc_idx) - return existing is None or existing.num_tokens_in_block < num_tokens_in_block - - def replace_page(self, lc_idx: LifeCycleId, page: "CommittedPage") -> None: - """Install `page` in slot `lc_idx`, detaching whatever it supersedes. - - The superseded page may outlive this call while a request still holds it, so - unlink_page() must clear its back-pointer: _release_pages() walks `storage`, so - nothing would clear it later and it would dangle once this block dies. - """ - assert NDEBUG or self.can_replace_page(lc_idx, page.num_tokens_in_block) - existing = self.unlink_page(lc_idx) - if existing is not None and existing.scheduled_for_eviction: - existing.manager.exclude_from_eviction(existing) - page.block = rawref.ref(self) - self.storage[lc_idx] = rawref.ref(page) - - def _adopt_pages_from(self, other: "Block") -> None: - """Move `other`'s pages into self without changing their recorded token counts.""" - assert other.ordinal == self.ordinal - for lc_idx in typed_range(self.num_life_cycles): - page = other.get_page(lc_idx) - if page is None or not self.can_replace_page(lc_idx, page.num_tokens_in_block): - continue - # Clear the source slot directly rather than via unlink_page(), which would - # null the back-pointer replace_page() is about to overwrite. - other.storage[lc_idx] = None - self.replace_page(lc_idx, page) - - def _release_pages(self) -> None: - """Reclaim every page held by this block. - - Nulls each page's back-pointer and, for pages still scheduled for eviction, - removes them from the eviction controller (releasing their storage slots). - Idempotent: afterwards ``storage`` holds no pages, so it is safe to call again - from ``__del__``. - - Cleanup is normally deferred to ``__del__``. An orphan block may remain - referenced by a live ``_KVCache`` and retain its pages until that cache closes; - every cache must close before ``StorageManager`` teardown. - """ - for lc_idx in typed_range(self.num_life_cycles): - page = self.get_page(lc_idx) - if page is not None: - self.unlink_page(lc_idx) - if page.status == PageStatus.DROPPABLE: - if page.scheduled_for_eviction: - page.manager.exclude_from_eviction(page) - - def __del__(self) -> None: - self._release_pages() - self.__rawref__.invalidate() - - def _partial_match_this_node(self, tokens: TokenBlock) -> int: - """ - Returns the number of leading tokens that match between the given tokens and this block's tokens. - """ - for i, (a, b) in enumerate(zip(tokens, self.tokens)): - if a != b: - return i - return min(len(tokens), len(self.tokens)) - - @property - def num_life_cycles(self) -> LifeCycleId: - return LifeCycleId(len(self.storage)) - - @property - def prev(self) -> "Block | RootBlock": - return unwrap_rawref(self._prev) - - def get_page(self, lc_idx: LifeCycleId) -> "CommittedPage | None": - """Return the page in slot `lc_idx`, or None when the slot is empty. - - A non-empty slot always resolves: CommittedPage.__del__ unlinks the page from its - block before invalidating its rawref, so `storage` never retains a dangling ref. - """ - return map_optional(self.storage[lc_idx], lambda f: f()) - - def unlink_page( - self, lc_idx: LifeCycleId, expected_page: "CommittedPage | None" = None - ) -> "CommittedPage | None": - """Detach slot `lc_idx`, returning the page that was there, or None. - - The sole place a block-page link is severed. - """ - # Called from CommittedPage.__del__, which invalidates the page's rawref only - # afterwards, so the dying page is still reachable here. - page = self.get_page(lc_idx) - if page is None: - return None - # Only unlink when the slot still holds the expected page. During rebase - # another block with the same key may have replaced the stored page, and - # unlinking then would clobber the newer page's back-pointer. - if expected_page is not None and page is not expected_page: - return None - page.block = rawref.NULL - self.storage[lc_idx] = None - return page - - @staticmethod - def clear_stale_blocks_after_page_unlink( - start: "Block", lc_idx: LifeCycleId, lc: LifeCycle - ) -> None: - assert start.get_page(lc_idx) is None - ordinal = start.ordinal - tree = try_get_tree(start) - event_manager = tree.event_manager if tree is not None else None - if type(lc) is AttnLifeCycle and (lc.window_size is None or ordinal < lc.num_sink_blocks): - remove_subtree(start) - elif event_manager is not None: - event_manager.add_removed_life_cycle_event(start.key, int(lc_idx)) - # It's possible to implement more sophisticated logic to remove useless blocks for SWA, e.g. - # check if consecutive available blocks is sufficient for window_size. (TRTLLM-8802) - # But for simplicity, we leave it for now. - curr = start - while ( - ( - isinstance(curr, Block) - and all( - curr.get_page(life_cycle) is None - for life_cycle in typed_range(curr.num_life_cycles) - ) - ) - and not curr.next - and curr._prev() is not None - ): - prev = curr.prev - detached = detach_next(prev, curr.key) - assert detached is curr - if event_manager is not None: - event_manager.add_removed_event(curr.key) - curr = prev - - @property - def tokens_per_block(self) -> int: - # we assume non-leaf blocks are always full. - prev = self.prev - return prev.tokens_per_block if isinstance(prev, RootBlock) else len(prev.tokens) - - @property - def is_full(self) -> bool: - return len(self.tokens) == self.tokens_per_block - - @property - def is_orphan(self) -> bool: - prev = self._prev() - assert prev is None or (self.key in prev.next and prev.next[self.key] is self) - return prev is None - - -class BlockRadixTree: - __slots__ = ( - "__rawref__", - "_event_manager", - "_life_cycles", - "_tokens_per_block", - "next", - ) - _life_cycles: LifeCycleRegistry - _tokens_per_block: int - _event_manager: "KVCacheEventManager | None" - next: Children[RootBlock] - __rawref__: rawref.ref["BlockRadixTree"] - - def __init__( - self, - life_cycles: LifeCycleRegistry, - tokens_per_block: int, - event_manager: "KVCacheEventManager | None" = None, - ) -> None: - self._life_cycles = life_cycles - self._tokens_per_block = tokens_per_block - self._event_manager = event_manager - self.next = {} - self.__rawref__ = rawref.NULL - - def __del__(self) -> None: - self.__rawref__.invalidate() - - def add_or_get_existing(self, reuse_scope: ReuseScope) -> RootBlock: - key = RootBlock.make_key(reuse_scope) - if key in self.next: - return self.next[key] - return RootBlock(reuse_scope, self) - - @property - def tokens_per_block(self) -> int: - return self._tokens_per_block - - @property - def life_cycles(self) -> LifeCycleRegistry: - return self._life_cycles - - @property - def event_manager(self) -> "KVCacheEventManager | None": - return self._event_manager - - @property - def num_life_cycles(self) -> LifeCycleId: - return self.life_cycles.size - - def clear(self) -> None: - # taking O(1) space - # remove leaf blocks one by one, in post-order - # Block.__del__() handles page cleanup when the last owner releases each block. - # detach_next() auto-prunes empty RootBlocks from the tree. - while self.next: - root = next(iter(self.next.values())) - while root.next: - remove_subtree(next(iter(root.next.values()))) - assert not self.next - - def _num_matched_tokens(self, matched: list[tuple[Block, int]]) -> int: - if not matched: - return 0 - return self._tokens_per_block * (len(matched) - 1) + matched[-1][1] - - # yields tuples of (block, num_matched_tokens). num_matched_tokens should be equal to - # tokens_per_block except the last one. - def _match_token_path( - self, - reuse_scope: ReuseScope, - tokens: Sequence[TokenIdExt], - enable_partial_match: bool = False, - ) -> Iterator[tuple[Block, int]]: - block: Block | RootBlock | BlockRadixTree = self - mismatched_token_block: TokenBlock = [] - for token_block, key in sequence_to_blockchain_keys( - self._tokens_per_block, reuse_scope, tokens - ): - if key in block.next: - block = block.next[key] - if token_block: - assert isinstance(block, Block) - yield block, len(token_block) - else: - mismatched_token_block = token_block - break - if mismatched_token_block and enable_partial_match: - partial_block, match_len = find_best_partial_match_in_next_nodes( - cast(Block | RootBlock, block), mismatched_token_block - ) - if partial_block is not None: - block = partial_block - yield block, match_len - - def _prune_match( - self, matched: list[tuple[Block, int]], ssm_lc_id: LifeCycleId | None - ) -> list[tuple[Block, int]]: - """Shorten `matched` to the prefix that is actually reusable. - - Passing ssm_lc_id=None skips the recurrent-snapshot constraint and yields - the attention-only prefix (used for - num_reusable_tokens_before_hybrid_pruning). - """ - tokens_per_block = self._tokens_per_block - assert all(b[1] == tokens_per_block for b in matched[:-1]) - - attn_life_cycles = list(self._life_cycles.attention_life_cycles()) - - # Fixed-point loop: SSM may select an earlier exact snapshot, while attention may - # shorten the match to the coverage of a required page. Every retry strictly - # shortens the match, so the loop terminates. - while matched: - # Check SSM snapshot availability first: truncating to the last reusable SSM - # snapshot changes the matched length that all the attention checks use. - if ssm_lc_id is not None: - ssm_trunc = 0 - ssm_match_len = 0 - for i in reversed(range(len(matched))): - # An SSM page holds the recurrent state after exactly this many tokens, - # so reuse must stop there instead of anywhere inside the block. - snapshot_len = matched[i][0].page_coverage(ssm_lc_id) - if snapshot_len > 0 and matched[i][1] >= snapshot_len: - ssm_trunc = i + 1 - ssm_match_len = snapshot_len - break - matched = matched[:ssm_trunc] - if not matched: - break - matched[-1] = (matched[-1][0], ssm_match_len) - - # Only pages that are active at this candidate endpoint constrain attention - # reuse. Full attention requires every block. SWA requires sink blocks and the - # trailing window, but not the stale blocks between them. In particular, at an - # exact block boundary with window_size=1, every historical block is stale. - num_tokens = self._num_matched_tokens(matched) - shortened = False - for lc_idx, lc in attn_life_cycles: - stale = lc.get_stale_range(num_tokens, tokens_per_block) - for i in chain(range(stale.beg), range(stale.end, len(matched))): - block, num_matched = matched[i] - coverage = block.page_coverage(lc_idx) - if coverage >= num_matched: - continue - if coverage > 0: - matched = matched[: i + 1] - matched[-1] = (block, coverage) - else: - matched = matched[:i] - shortened = True - break - if shortened: - break - if not shortened: - break - return matched - - @staticmethod - def _back_off_match( - matched: list[tuple["Block", int]], backoff: int - ) -> list[tuple["Block", int]]: - """Drop `backoff` tokens from the tail of a match. - - Shortens the last entry, dropping whole blocks while the backoff outruns - them. Leading entries stay full blocks, so _prune_match's invariant holds. - """ - while backoff > 0 and matched: - block, num_matched = matched[-1] - if num_matched > backoff: - matched[-1] = (block, num_matched - backoff) - break - backoff -= num_matched - matched.pop() - return matched - - def match( - self, - reuse_scope: ReuseScope, - tokens: Sequence[TokenIdExt], - enable_partial_match: bool = False, - backoff: int = 0, - ) -> ReuseMatch: - """ - Return the currently reusable prefix match without holding pages. - - The result is volatile: callers that need to reuse the returned blocks must - acquire ownership of the pages before depending on them. - - `backoff` trims that many tokens off the tail (see - KVCacheManagerConfig.reuse_match_backoff). - """ - raw_matched = list(self._match_token_path(reuse_scope, tokens, enable_partial_match)) - num_reusable_tokens_before_pruning = self._num_matched_tokens(raw_matched) - ssm_lc_id = self._life_cycles.ssm_life_cycle_id - # Page requirements depend on the final endpoint. Back off before pruning - # so SWA coverage and recurrent snapshots are validated at that endpoint. - if backoff > 0: - raw_matched = self._back_off_match(raw_matched, backoff) - # Diagnostic only: re-prune ignoring recurrent-snapshot availability to get - # the prefix the attention pages alone support. Only hybrid models pay for - # the second pass; without an SSM life cycle the two results are identical. - num_reusable_tokens_before_hybrid_pruning = ( - self._num_matched_tokens(self._prune_match(list(raw_matched), None)) - if ssm_lc_id is not None - else None - ) - matched = self._prune_match(raw_matched, ssm_lc_id) - num_tokens = self._num_matched_tokens(matched) - return ReuseMatch( - [block for block, _ in matched], - num_tokens, - len(tokens), - ( - num_tokens - if num_reusable_tokens_before_hybrid_pruning is None - else num_reusable_tokens_before_hybrid_pruning - ), - num_reusable_tokens_before_pruning, - ) - - def _check_sanity(self) -> bool: - raise NotImplementedError( - "[KVCacheManager] Check if there are any unusable blocks that should have been removed." - ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_common.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_common.py deleted file mode 100644 index 5782e59df03b..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_common.py +++ /dev/null @@ -1,101 +0,0 @@ -# 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. - -import enum -import os -from dataclasses import dataclass -from typing import Final, NewType - -NDEBUG: Final[bool] = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] != "1" - - -class PageStatus(enum.IntEnum): - LOCKED = 0 # Required in GPU. Eviction/dropping not allowed - HELD = 1 # Allow eviction but not dropping - DROPPABLE = 2 # Allow eviction and dropping - - -# Can extend to more tiers in the future, e.g. object storage like AWS S3. -class CacheTier(enum.IntEnum): - GPU_MEM = 0 - HOST_MEM = 1 - DISK = 2 - - -class PageIndexMode(enum.IntEnum): - # Converted index list is shared across layers in the same LayerGroup. - # Base pointer is per-layer (includes attr.offset). - SHARED = 0 - # Converted index list is per-layer. - # Base pointer is shared (pool group base, no attr.offset). - PER_LAYER = 1 - - -CacheLevel = NewType("CacheLevel", int) - - -GPU_LEVEL: Final[CacheLevel] = CacheLevel(0) -# First cache level below GPU. Its semantic tier depends on the configured -# cache_tiers: host when a host tier exists, otherwise disk. -CACHE_LEVEL1: Final[CacheLevel] = CacheLevel(1) - -# Normal token id that falls in the tokenizer vocabulary. -TokenId = NewType("TokenId", int) - -# For multi-modal tokens, we can handle it in either of the following ways: -# 1. Hash combine image digest and local_token_id, then use digest for every multi-modal token. -# 2. Use digest only for the first multi-modal token, and use int(vocab_size + local_token_id) for the rest. -# 3. Hash the multi-modal token embedding data and use the digest as TokenIdExt for every multi-modal token. -# If we do this, we can't skip the encoder. -TokenIdExt = TokenId | bytes - - -BlockOrdinal = NewType("BlockOrdinal", int) -BlockOrdinalT = type(BlockOrdinal(0)) -BAD_BLOCK_ORDINAL: Final[BlockOrdinal] = BlockOrdinal(-1) - -LayerId = NewType("LayerId", int) - -CudaStream = NewType("CudaStream", int) - -BeamIndex = NewType("BeamIndex", int) -DEFAULT_BEAM_INDEX: Final[BeamIndex] = BeamIndex(0) - -UserId = NewType("UserId", int) - -MemAddress = NewType("MemAddress", int) - -FileDescriptor = NewType("FileDescriptor", int) - -BAD_FILE_DESCRIPTOR: Final[FileDescriptor] = FileDescriptor(-1) - -PageIndex = NewType("PageIndex", int) -BAD_PAGE_INDEX: Final[PageIndex] = PageIndex(-1) - - -@dataclass(slots=True, frozen=True) -class DiskAddress: - fd: FileDescriptor - pos: int - - -Address = MemAddress | DiskAddress - -SlidingWindowSize = int | None - -Priority = NewType("Priority", int) -PRIORITY_MIN: Final[Priority] = Priority(0) -PRIORITY_MAX: Final[Priority] = Priority(100) -PRIORITY_DEFAULT: Final[Priority] = Priority(35) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py deleted file mode 100644 index c47c59c5434c..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py +++ /dev/null @@ -1,291 +0,0 @@ -# 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. - -# Currently, our nvfp4 kernels require that KV data and its corresponding KV block scale use the same -# block index, but different base address. -# As the ratio between KV data size and KV block scale size is fixed, we can simply use a pool with -# smaller block size and the same number of blocks for block scale. -import os -from dataclasses import dataclass, field -from enum import IntEnum -from typing import ClassVar, NewType, Protocol - -from ._common import CacheTier, LayerId - -# The data role of a buffer inside one layer. -# Must be unique for each buffer inside a layer. -# Examples: "key", "value", "key_block_quant", "value_block_quant". -DataRole = NewType("DataRole", str) - - -class CacheTierConfig(Protocol): - """Protocol for cache tier configuration.""" - - quota: int # in bytes - - @property - def tier(self) -> CacheTier: ... - - def assert_valid(self) -> None: ... - - -@dataclass(slots=True) -class GpuCacheTierConfig: - quota: int # in bytes - - @property - def tier(self) -> CacheTier: - return CacheTier.GPU_MEM - - def assert_valid(self) -> None: - assert self.quota > 0, "Quota must be positive" - - -@dataclass(slots=True) -class HostCacheTierConfig: - quota: int # in bytes - - @property - def tier(self) -> CacheTier: - return CacheTier.HOST_MEM - - def assert_valid(self) -> None: - assert self.quota > 0, "Quota must be positive" - - -@dataclass(slots=True) -class DiskCacheTierConfig: - quota: int # in bytes - path: str # a folder where we will store data as files - - @property - def tier(self) -> CacheTier: - return CacheTier.DISK - - def assert_valid(self) -> None: - assert self.quota > 0, "Quota must be positive" - assert os.path.isdir(self.path), ( - f"Disk path {self.path} does not exist or is not a directory" - ) - - -@dataclass(slots=True) -class BufferConfig: - role: DataRole - size: int - - tokens_per_block_override: int | None = None - """ - If not None, overrides the tokens_per_block in KVCacheManagerConfig. Must be a factor of tokens_per_block in - KVCacheManagerConfig and size should be based on tokens_per_block_override. - """ - - -class LayerType(IntEnum): - ATTENTION = 0 - SSM = 1 - - -@dataclass(slots=True) -class AttentionLayerConfig: - type: ClassVar[LayerType] = LayerType.ATTENTION - - layer_id: LayerId - # Each page can have multiple sub-pages, e.g. separate K and V data, block quantization scales for K and/or V, etc. - # KV cache manager will automatically group sub-pages of the same size, and redirect pages of different sizes to - # different memory pools - - # BufferConfig.role should not duplicate - buffers: list[BufferConfig] - # Note that we use None to represent "no sliding window". Sink tokens are excluded. - sliding_window_size: int | None = None - num_sink_tokens: int | None = None - - @property - def window_size(self) -> int | None: - return self.sliding_window_size - - def __post_init__(self) -> None: - assert len(set(buffer.role for buffer in self.buffers)) == len(self.buffers), ( - "duplicate buffer role" - ) - - -@dataclass(slots=True) -class SsmLayerConfig: - type: ClassVar[LayerType] = LayerType.SSM - - layer_id: LayerId - - buffers: list[BufferConfig] - - def __post_init__(self) -> None: - assert len(set(buffer.role for buffer in self.buffers)) == len(self.buffers), ( - "duplicate buffer role" - ) - assert all(buf.tokens_per_block_override is None for buf in self.buffers) - - -LayerConfig = AttentionLayerConfig | SsmLayerConfig - - -@dataclass(slots=True, frozen=True) -class KVCacheDesc: - capacity: int - history_length: int - - def __post_init__(self) -> None: - assert 0 <= self.history_length <= self.capacity - - -# A batch of requests, working as a use case the KVCacheManager must always support. -@dataclass(slots=True, frozen=True) -class BatchDesc: - kv_caches: list[KVCacheDesc] - # Tokens shared by all requests. Set to 0 if no kv cache reuse. - system_prompt_length: int = 0 - - def __post_init__(self) -> None: - assert self.system_prompt_length >= 0 - - -@dataclass(slots=True) -class SwaScratchReuseConfig: - """ - Configuration for SWA scratch reuse. - - Args: - max_rewind_len: Maximum number of tail tokens that can be rewound after - scratch-enabled allocation. Scratch reuse will not cover blocks that - may be needed to preserve those tokens. - """ - - max_rewind_len: int = 0 - - def __post_init__(self) -> None: - assert self.max_rewind_len >= 0, "max_rewind_len must be non-negative" - - -@dataclass(slots=True) -class KVCacheManagerConfig: - """ - Configuration for the KV cache manager. - """ - - tokens_per_block: int - # cache tiers are sorted from warm to cold. The first one must be GPU memory. - cache_tiers: list[CacheTierConfig] - - # AttentionLayerConfig.layer_id should not duplicate - layers: list[LayerConfig] - - # When memory utilization is above this threshold, KV cache resuming will fail. This helps - # reserving some memory for KVCache growth and avoids frequent suspend/resume for dynamic batch size. - max_util_for_resume: float = 0.97 - - enable_partial_reuse: bool = True - """ - If True, we will try to reuse tokens from partially matched blocks. - """ - - reuse_match_backoff: int = 0 - """ - Tokens dropped from the tail of every prefix match. - - For a pool whose KV at position i is a function of tokens [0, i] this is 0: a - match of m tokens proves all m are reusable. Set it to D when the pool also - holds state that reads D tokens ahead -- one-model speculative decoding draft - layers -- where a match of m only describes the first m - D positions. - - Applied inside the match so a single tree walk yields the usable depth. - """ - - constraints: list[BatchDesc] = field(default_factory=list) - """ - A list of step configurations that must always be supported. - """ - - typical_step: BatchDesc | None = None - """ - A typical step configuration used to decide initial memory partitioning between - layer groups. - """ - - initial_pool_ratio: list[float] | None = None - """ - One positive, normalized hot-tier byte-quota weight per layer group. Cold-tier - initialization preserves the implied layer-group slot-count proportions while - accounting for cold page sizes. When set, this takes precedence over typical_step - and constraints for initial ratio selection; constraints remain hot-level feasibility - floors. - """ - - swa_scratch_reuse: SwaScratchReuseConfig | None = None - """ - When set, SWA layers reuse physical pages for out-of-window blocks during prefill. - Scratch blocks share coalesced slot sub-pages across blocks for the currently executing - layer, reducing peak memory. Trade-off: KV cache reuse is degraded because scratch blocks - have no preserved data after the step. - - If max_rewind_len is non-zero, the rewindable tail is excluded from scratch reuse so - draft/target shared KV cache can preserve tokens that may survive speculative rewind. - - Most useful for disaggregated prefill servers handling long prompts or long prompt chunks, - where the number of out-of-window blocks dominates memory usage. - """ - - commit_min_snapshot: bool = False - """ - If True, commit() records only the minimum cache snapshot reusable at the post-call - num_committed_tokens. Only the minimum amount of pages required for such reuse will - be preserved. - - Required when SSM layers are present. - """ - - enable_stats: bool = True - """ - Collect V2 KV cache allocation, reuse, and transfer statistics. - """ - - text_only: bool = False - """ - Deployment-level guarantee that no request carries multi-modal content, so token - sequences never contain digests. A per-_KVCache text_only override may only tighten - this (a text-only deployment forbids a request claiming otherwise). Default False. - - (In this pure-Python backend the block hasher has no digest-free fast path, so this - flag is carried for API/behavior parity with the C++ backend but changes no hashing.) - """ - - @property - def enable_swa_scratch_reuse(self) -> bool: - return self.swa_scratch_reuse is not None - - def __post_init__(self) -> None: - assert self.cache_tiers and self.cache_tiers[0].tier == CacheTier.GPU_MEM - assert len(set(layer.layer_id for layer in self.layers)) == len(self.layers), ( - "duplicate layer id" - ) - assert all( - buffer.tokens_per_block_override is None - or self.tokens_per_block % buffer.tokens_per_block_override == 0 - for layer in self.layers - for buffer in layer.buffers - ) - if any(layer.type == LayerType.SSM for layer in self.layers): - assert self.commit_min_snapshot, ( - "commit_min_snapshot must be True when SSM layers are present" - ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py deleted file mode 100644 index 1b9a1a59ed0a..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py +++ /dev/null @@ -1,387 +0,0 @@ -# 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. - -import atexit -import sys -import threading -from _thread import LockType -from collections.abc import Callable, Iterator -from dataclasses import dataclass - -# avoid importing the whole tensorrt_llm module, which takes time during debugging. -from importlib.util import find_spec -from pathlib import Path -from typing import ClassVar, NamedTuple, Sequence, cast - -import cuda.bindings.driver as drv - -from ._common import Address, CacheTier, CudaStream, MemAddress -from ._utils import ( - CachedCudaEvent, - HomoTuple, - HostMem, - _unwrap, - div_up, - stream_wait_events, - temporary_sys_path, -) - -if "tensorrt_llm" in sys.modules: - from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( # noqa # type: ignore - DiskAddress, - DiskToDiskTask, - DiskToHostTask, - HostToDiskTask, - MemToMemTask, - copy_device_to_device, - copy_device_to_host, - copy_disk_to_disk, - copy_disk_to_host, - copy_host_to_device, - copy_host_to_disk, - copy_host_to_host, - ) -else: - # fast path for dev, avoids importing the whole tensorrt_llm module - spec = find_spec("kv_cache_manager_v2") - assert spec is not None and spec.origin is not None - with temporary_sys_path(str(Path(spec.origin).parent.parent.parent)): - from bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( # noqa - DiskAddress, - DiskToDiskTask, - DiskToHostTask, - HostToDiskTask, - MemToMemTask, - copy_device_to_device, - copy_device_to_host, - copy_disk_to_disk, - copy_disk_to_host, - copy_host_to_device, - copy_host_to_disk, - copy_host_to_host, - ) - - -class CopyTask(NamedTuple): - dst: Address - src: Address - - -def _copy_gpu_to_gpu(tasks: Sequence[CopyTask], num_bytes: int, stream: CudaStream): - _unwrap( - drv.CUresult( - copy_device_to_device([MemToMemTask(dst, src) for dst, src in tasks], num_bytes, stream) - ) - ) - - -def _copy_host_to_host(tasks: Sequence[CopyTask], num_bytes: int, stream: CudaStream): - _unwrap( - drv.CUresult( - copy_host_to_host([MemToMemTask(dst, src) for dst, src in tasks], num_bytes, stream) - ) - ) - - -def _copy_disk_to_disk(tasks: Sequence[CopyTask], num_bytes: int, stream: CudaStream): - _unwrap( - drv.CUresult( - copy_disk_to_disk( - [ - DiskToDiskTask( - DiskAddress( - cast(DiskAddress, dst).fd, - cast(DiskAddress, dst).pos, - ), - DiskAddress( - cast(DiskAddress, src).fd, - cast(DiskAddress, src).pos, - ), - ) - for dst, src in tasks - ], - num_bytes, - stream, - ) - ) - ) - - -def _copy_gpu_to_host(tasks: Sequence[CopyTask], num_bytes: int, stream: CudaStream): - _unwrap( - drv.CUresult( - copy_device_to_host([MemToMemTask(dst, src) for dst, src in tasks], num_bytes, stream) - ) - ) - - -def _copy_host_to_gpu(tasks: Sequence[CopyTask], num_bytes: int, stream: CudaStream): - _unwrap( - drv.CUresult( - copy_host_to_device([MemToMemTask(dst, src) for dst, src in tasks], num_bytes, stream) - ) - ) - - -def _copy_disk_to_host(tasks: Sequence[CopyTask], num_bytes: int, stream: CudaStream): - _unwrap( - drv.CUresult( - copy_disk_to_host( - [ - DiskToHostTask( - cast(MemAddress, dst), - DiskAddress(cast(DiskAddress, src).fd, cast(DiskAddress, src).pos), - ) - for dst, src in tasks - ], - num_bytes, - stream, - ) - ) - ) - - -def _copy_host_to_disk(tasks: Sequence[CopyTask], num_bytes: int, stream: CudaStream): - _unwrap( - drv.CUresult( - copy_host_to_disk( - [ - HostToDiskTask( - DiskAddress( - cast(DiskAddress, dst).fd, - cast(DiskAddress, dst).pos, - ), - cast(MemAddress, src), - ) - for dst, src in tasks - ], - num_bytes, - stream, - ) - ) - ) - - -Copier = Callable[[Sequence[CopyTask], int, CudaStream], None] - - -def get_copier(dst: CacheTier, src: CacheTier) -> Copier | HomoTuple[Copier]: - copiers: HomoTuple[HomoTuple[Copier | HomoTuple[Copier]]] = ( - # dst = GPU_MEM - ( - _copy_gpu_to_gpu, # src = GPU_MEM - _copy_host_to_gpu, # src = HOST_MEM - (_copy_disk_to_host, _copy_host_to_gpu), # src = DISK - ), - # dst = HOST_MEM - ( - _copy_gpu_to_host, # src = GPU_MEM - _copy_host_to_host, # src = HOST_MEM - _copy_disk_to_host, # src = DISK - ), - # dst = DISK - ( - (_copy_gpu_to_host, _copy_host_to_disk), # src = GPU_MEM - _copy_host_to_disk, # src = HOST_MEM - _copy_disk_to_disk, # src = DISK - ), - ) - return copiers[dst][src] - - -@dataclass(slots=True) -class GrainMetadata: - mutex: LockType - ready_event: CachedCudaEvent # protects the buffer grain. - - -class StagingBuffer: - __slots__ = ("manager", "min_size", "max_size", "_size", "start_grain", "stream") - manager: "StagingBufferManager" - min_size: int - max_size: int - _size: int - start_grain: int - stream: CudaStream - - def __init__( - self, manager: "StagingBufferManager", min_size: int, max_size: int, stream: CudaStream - ): - self.manager = manager - self.min_size = min_size - self.max_size = max_size - self.stream = stream - - @property - def address(self) -> MemAddress: - return MemAddress(self.manager.buffer.address + self.start_grain * self.manager.GRANULARITY) - - @property - def size(self) -> int: - return self._size - - @property - def num_grains(self) -> int: - return div_up(self._size, self.manager.GRANULARITY) - - @property - def grains(self) -> list[GrainMetadata]: - return self.manager.grains[self.start_grain : self.start_grain + self.num_grains] - - def __enter__(self) -> "StagingBuffer": - manager = self.manager - if self.min_size > manager.size: - raise ValueError(f"Requested min_size {self.min_size} is too large for the manager") - with manager.mutex: - # If the tail cannot satisfy min_size, wrap to the front before allocating. - available = manager._suggest_next_max_size_unsafe() - if self.min_size > available: - manager.next = 0 - available = manager._suggest_next_max_size_unsafe() - assert self.min_size <= available - self._size = max(min(self.max_size, available), self.min_size) - self.start_grain = manager.next - manager.next += self.num_grains - assert manager.next <= manager.num_grains - if manager.next == manager.num_grains: - manager.next = 0 - - def lock_and_consume_events() -> Iterator[CachedCudaEvent]: - for grain in self.grains: - grain.mutex.acquire() - yield grain.ready_event - grain.ready_event = CachedCudaEvent.NULL - - stream_wait_events(self.stream, lock_and_consume_events()) - return self - - def __exit__(self, exc_type, exc_value, traceback) -> None: - event = CachedCudaEvent(self.stream) - for grain in reversed(self.grains): - grain.ready_event = event - grain.mutex.release() - - -class StagingBufferManager: - __slots__ = ("mutex", "buffer", "grains", "next") - GRANULARITY: ClassVar[int] = 1 << 20 - - mutex: LockType - buffer: HostMem - grains: list[GrainMetadata] - next: int - - def __init__(self, size: int) -> None: - assert size % self.GRANULARITY == 0 - self.mutex = threading.Lock() - num_grains = size // self.GRANULARITY - self.buffer = HostMem(size) - self.grains = [ - GrainMetadata(threading.Lock(), CachedCudaEvent.NULL) for _ in range(num_grains) - ] - self.next = 0 - - @property - def size(self) -> int: - "Requesting more than this will fail." - assert len(self.grains) * self.GRANULARITY == self.buffer.size - return self.buffer.size - - @property - def num_grains(self) -> int: - return len(self.grains) - - def _suggest_next_max_size_unsafe(self) -> int: - "Requesting more than this may degrade performance. Must be called with self.mutex held." - return self.GRANULARITY * (self.num_grains - self.next) - - # max_size is just a hint, the actual size may be smaller. - def new(self, min_size: int, max_size: int, stream: CudaStream) -> StagingBuffer: - """ - min_size is the min required size. max_size is for best efforts. Your should query the actual - size after entering the context. - """ - return StagingBuffer(self, min_size, max_size, stream) - - -class CopyEngine: - __slots__ = ("_staging_buffer_manager",) - _staging_buffer_manager: StagingBufferManager | None - - def __init__(self) -> None: - self._staging_buffer_manager = None - - def close(self) -> None: - self._staging_buffer_manager = None - - @property - def staging_buffer_manager(self) -> StagingBufferManager: - if self._staging_buffer_manager is None: - self._staging_buffer_manager = StagingBufferManager(64 << 20) - return self._staging_buffer_manager - - # @TODO: Use a dedicated stream for each different Copier, take set[CachedCudaEvent] instead of - # stream, and return a new CachedCudaEvent. - def transfer( - self, - dst_cache_tier: CacheTier, - src_cache_tier: CacheTier, - num_bytes: int, - tasks: Sequence[CopyTask], - stream: CudaStream, - ) -> None: - copier = get_copier(dst_cache_tier, src_cache_tier) - if not isinstance(copier, tuple): - return copier(tasks, num_bytes, stream) - assert len(copier) == 2, "for now, we only support 2 copiers via host memory" - manager = self.staging_buffer_manager - remaining = tasks - while remaining: - with manager.new(num_bytes, num_bytes * len(remaining), stream) as buf: - addr = buf.address - n = buf.size // num_bytes - assert n <= len(remaining) - batch = remaining[:n] - copier[0]( - [ - CopyTask(MemAddress(addr + num_bytes * i), t.src) - for i, t in enumerate(batch) - ], - num_bytes, - buf.stream, - ) - copier[1]( - [ - CopyTask(t.dst, MemAddress(addr + num_bytes * i)) - for i, t in enumerate(batch) - ], - num_bytes, - buf.stream, - ) - remaining = remaining[n:] - - -_copy_engine = CopyEngine() -atexit.register(_copy_engine.close) - - -def batched_copy( - dst_cache_tier: CacheTier, - src_cache_tier: CacheTier, - num_bytes: int, - tasks: Sequence[CopyTask], - stream: CudaStream, -) -> None: - _copy_engine.transfer(dst_cache_tier, src_cache_tier, num_bytes, tasks, stream) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py deleted file mode 100644 index 4e52e154c47b..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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. - -from .._common import DEFAULT_BEAM_INDEX, BeamIndex -from ._kv_cache import PlannedDropHandle, _KVCache -from ._kv_cache_manager import ( - AggregatedPageDesc, - ExpandedBuffer, - KVCacheManager, - PageIndexConverter, - PoolDesc, - PoolGroupDesc, - PoolGroupPeakBlockStats, - ScratchDesc, -) - -__all__ = [ - "KVCacheManager", - "_KVCache", - "PlannedDropHandle", - "BeamIndex", - "DEFAULT_BEAM_INDEX", - "AggregatedPageDesc", - "ExpandedBuffer", - "PageIndexConverter", - "PoolDesc", - "PoolGroupDesc", - "PoolGroupPeakBlockStats", - "ScratchDesc", -] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py deleted file mode 100644 index 9eb3d20b15f0..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ /dev/null @@ -1,2458 +0,0 @@ -# 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. - -import array -import enum -import math -from collections.abc import Iterable, Sequence -from contextlib import contextmanager -from dataclasses import dataclass -from itertools import chain -from typing import TYPE_CHECKING, Callable, ClassVar, Iterator, NamedTuple, Type, cast - -from .. import rawref -from .._block_radix_tree import Block, ReuseMatch, ReuseScope, RootBlock, UselessBlockError -from .._common import ( - BAD_BLOCK_ORDINAL, - BAD_PAGE_INDEX, - DEFAULT_BEAM_INDEX, - GPU_LEVEL, - NDEBUG, - BeamIndex, - BlockOrdinal, - BlockOrdinalT, - CacheLevel, - CudaStream, - PageIndex, - PageIndexMode, - PageStatus, - Priority, - TokenIdExt, -) -from .._copy_engine import CopyTask, batched_copy -from .._exceptions import LogicError, OutOfPagesError -from .._life_cycle_registry import ( - AttnLifeCycle, - LayerGroupId, - LifeCycle, - LifeCycleId, - SsmLifeCycle, - compute_scratch_range, -) -from .._page import ( - BatchedLockTarget, - BlockPage, - CommittedPage, - Page, - ScratchSlotLock, - UncommittedPage, - _PageHolder, - _SharedPageLock, - batched_lock_to_gpu, -) -from .._stats import ( - CountsByLevel, - KVCacheIterationStatsDelta, - KVCacheStatsDelta, - ReusedBlocksByLevel, -) -from .._storage._core import Slot -from .._storage_manager import StorageManager -from .._utils import ( - CachedCudaEvent, - HalfOpenRange, - TypedIndexList, - div_up, - expect_type, - filled_list, - intersect, - make_typed, - stream_wait_events, - to_typed, - typed_enumerate, - typed_len, - typed_map, - typed_range, - unwrap_optional, - value_or, -) -from ._moving_average import Average -from ._pending_stats import _PendingStats - -if TYPE_CHECKING: - from ._kv_cache_manager import KVCacheManager, ScratchDesc - - -@dataclass(slots=True) -class SeqBlock: - pages: TypedIndexList[BeamIndex, TypedIndexList[LifeCycleId, BlockPage]] - # In rare cases, this may be the only strong reference to this block. Assume it's the last block we - # committed on stop_committing(), and it's partial. At the same time, we have another _KVCache - # generating same tokens plus some additional tokens. The block committed by the other _KVCache will - # fully cover tokens of this block. In that case, we will remove this block from the radix tree. - # Which means `tree_block not in tree_block.prev.next` will be True. - tree_block: Block | None - - @property - def is_committed(self) -> bool: - ret = self.tree_block is not None - assert NDEBUG or not ret or len(self.pages) == 1 - assert ( - NDEBUG - or not ret - or all( - p is None or isinstance(p.page, CommittedPage) - for p in chain.from_iterable(self.pages) - ) - ) - assert ( - NDEBUG - or ret - or all( - p is None or isinstance(p.page, UncommittedPage) - for p in chain.from_iterable(self.pages) - ) - ) - return ret - - def __del__(self) -> None: - self.tree_block = None - self.pages.clear() - - -class PlannedDropHandle: - """Track committed pages planned for dropping without owning them. - - The handle stores weak references and does not keep pages alive. Dropping it - decrements each live page's planned-drop count and removes an already-droppable - page from eviction tracking when no plans remain. - """ - - __slots__ = ("_page_refs",) - - _page_refs: tuple[rawref.ref[CommittedPage], ...] | None - - def __init__(self, pages: Iterable[CommittedPage]) -> None: - planned_pages = tuple({id(page): page for page in pages}.values()) - self._page_refs = tuple(rawref.ref(page) for page in planned_pages) - for page in planned_pages: - page.planned_drop_count += 1 - - def drop(self) -> None: - """Apply this drop plan and invalidate the handle. - - A live page is removed from eviction tracking only when this is its final - plan and it is already droppable and queued for eviction. Calling this - method twice is invalid. - """ - page_refs = self._page_refs - if page_refs is None: - raise RuntimeError("Planned drop handle has already been dropped") - - pages = list[CommittedPage]() - for page_ref in page_refs: - page = page_ref() - if page is not None: - if page.planned_drop_count <= 0: - raise RuntimeError("Committed page has no planned drop") - pages.append(page) - - self._page_refs = None - for page in pages: - page.planned_drop_count -= 1 - if ( - page.planned_drop_count == 0 - and page.status == PageStatus.DROPPABLE - and page.scheduled_for_eviction - ): - page.manager.exclude_from_eviction(page) - - def __del__(self) -> None: - if self._page_refs is not None: - self.drop() - - -class _Status(enum.Enum): - ACTIVE = enum.auto() - SUSPENDED = enum.auto() - CLOSED = enum.auto() - - -class _CommitState(enum.Enum): - ALLOWED = enum.auto() - # user did not stop but we can't commit any more due to conflict with other blocks - VIRTUAL_STOP = enum.auto() - # user called stop_committing() or close() - USER_STOP = enum.auto() - - -IndexSeq = array.array | memoryview - - -# The _KVCache holds unique/shared ownership of memory blocks. On deletion, the ownership if destroys -# and KVCacheManager takes control of them. A KV cache maintains three lengths: -# 1. num_committed_tokens: the number of tokens that are finalized, immutable and ready for reuse. -# 2. history_length: a cursor separating history and the space for next input tokens. History tokens -# are defined as tokens without query data for the next inference step. For SWA layers, it decides -# which blocks are out-of-window and can be evicted/dropped. In most cases, you don't need to touch -# history_length as it's automatically bumped by the increase of num_committed_tokens, except a few -# cases: -# a. Beam search where we can't commit tokens generated by the last step. But it still makes sense -# to evict uncommitted pages for SWA layers to save memory. -# b. Disaggregated serving with SWA and the reusable tokens are in the other server. We need to -# reserve space for history. Knowing history_length helps us accurately decide which blocks -# needs to be allocated. Then users only transfer data for what is needed. -# c. Multi-round conversation with chain of thoughts (CoT) and excluding CoT tokens for the next -# round. In this case, users should not commit tokens starting from CoT. Then history_length -# needs to be explicitly bumped. -# 3. capacity: the number of tokens that can be stored in the KV cache. It should include the number -# of both historical tokens and input tokens for the next inference step, no matter if it's prefill, -# chunked prefill or generation without/without speculative decoding. For tree-based speculative -# decoding, the number of input tokens here should be the flatten draft length. For beam search, -# multiple candidate tokens at the same position are counted as one. -# num_committed_tokens <= history_length <= capacity always holds. A newly created KV cache has all -# three lengths equal to the number of reused tokens. -# TODO: in __del__, we should check if committed pages are usable for SWA cases. e.g. all pages are -# dropped except the last one. The last one is not usable. -class _KVCache: - __slots__ = ( - "id", - "_manager", - "_reuse_scope", - "_get_priority", - "_cuda_stream", - "_status", - "_beam_width", - "_expected_prompt_length", - "_generation_alloc_ready", - "_capacity", - "_history_length", - "_commit_state", - "_blocks", - "_base_page_indices", - "_committed_tokens", - "_cached_tokens_by_level", - "_last_cached_token_level", - "_num_reusable_tokens_before_hybrid_pruning", - "_num_reusable_tokens_before_pruning", - "_num_committed_blocks", - "_finish_event", - "_tokens_per_block", - "_avg_history_length", - "_avg_capacity", - "_ssm_blocks", - "_never_resumed", - "_enable_swa_scratch_reuse", - "_text_only", - "_scratch_slots", - "_enable_request_stats", - "_pending_stats", - "__rawref__", - ) - - Status: ClassVar[Type[_Status]] = _Status - CommitState: ClassVar[Type[_CommitState]] = _CommitState - - id: int | None - _manager: "KVCacheManager" - _reuse_scope: ReuseScope - _get_priority: Callable[[BlockOrdinal, LifeCycle], Priority] - _cuda_stream: CudaStream | None - _status: _Status - _beam_width: BeamIndex - _expected_prompt_length: int | None - _generation_alloc_ready: bool - _capacity: int - _history_length: int - _commit_state: _CommitState - - _blocks: TypedIndexList[BlockOrdinal, SeqBlock] - # we maintain _base_page_indices to accelerate the get_base_page_indices() API. In principle it can - # be computed on the fly, but that would be slow due to python. - _base_page_indices: TypedIndexList[BeamIndex, TypedIndexList[LifeCycleId, IndexSeq]] - _committed_tokens: list[TokenIdExt] - # Initial current-residency provenance, observed while the reused pages are held. - _cached_tokens_by_level: CountsByLevel - _last_cached_token_level: CacheLevel | None - # Internal diagnostic captured from the reuse match: see ReuseMatch. - _num_reusable_tokens_before_hybrid_pruning: int - _num_reusable_tokens_before_pruning: int - # Sometimes we can't commit a block because all its tokens are already covered by another block in - # the radix tree. But it's unsafe to just use the other block because: 1. the data may have numeric - # difference, 2. if our block is a partial block, we can't write to memory of the other blocks. - # Internally, we stop committing from such a block, but still give user an illusion that the block is - # committed. In such cases, _committed_tokens contains what users have fed with commit(), while - # _num_committed_blocks contains the number of blocks that are actually committed. - _num_committed_blocks: BlockOrdinal - # set when switch away from ACTIVE, cleared when switching to ACTIVE. - _finish_event: CachedCudaEvent | None - - _tokens_per_block: int - _avg_history_length: Average - _avg_capacity: Average - - _ssm_blocks: TypedIndexList[BeamIndex, TypedIndexList[LifeCycleId, BlockPage]] - _never_resumed: bool - _enable_swa_scratch_reuse: bool - _enable_request_stats: bool - # Scratch slots for SWA prefill memory reuse, per life cycle. These hold coalesced slots - # whose sub-pages are reinterpreted as per-block storage for the currently executing layer. - # Number of scratch blocks depends on diff between history_length and capacity. - # Managed via delta in resize(): existing slots are reused across resize calls, - # only the additional needed slots are allocated. Freed on teardown/suspend. - _scratch_slots: TypedIndexList[LifeCycleId, list[ScratchSlotLock]] - _pending_stats: _PendingStats - - def __init__( - self, - manager: "KVCacheManager", - reuse_scope: ReuseScope, - reuse_match: ReuseMatch | None, - id: int | None, - custom_priority_callback: Callable[[BlockOrdinal, LifeCycle], Priority], - expected_prompt_length: int | None = None, - text_only: bool | None = None, - enable_request_stats: bool = False, - ) -> None: - # Keep a partially constructed cache inert if validation fails. - self.__rawref__ = rawref.NULL - self._status = self.Status.CLOSED - self.id = id - self._manager = manager - self._reuse_scope = reuse_scope - self._get_priority = custom_priority_callback - self._cuda_stream = None - self._beam_width = BeamIndex(1) - self._expected_prompt_length = ( - max(expected_prompt_length, 0) if expected_prompt_length is not None else None - ) - self._generation_alloc_ready = False - self._capacity = 0 - self._history_length = 0 - self._commit_state = self.CommitState.ALLOWED - self._blocks = cast(TypedIndexList, []) - self._base_page_indices = make_typed( - lambda _: make_typed(lambda _: array.array("i"), self.manager._storage.num_life_cycles), - self.beam_width, - ) - self._committed_tokens = [] - # Filled by _setup_for_reuse(), which observes the source levels in the same walk that - # holds the matched pages. Stays all-zero when there is no reuse match. - self._cached_tokens_by_level = filled_list(0, manager._storage.num_cache_levels) - self._last_cached_token_level = None - self._num_reusable_tokens_before_hybrid_pruning = ( - reuse_match.num_reusable_tokens_before_hybrid_pruning if reuse_match is not None else 0 - ) - self._num_reusable_tokens_before_pruning = ( - reuse_match.num_reusable_tokens_before_pruning if reuse_match is not None else 0 - ) - self._num_committed_blocks = BlockOrdinal(0) - self._finish_event = None - self._tokens_per_block = manager.tokens_per_block - self._ssm_blocks = make_typed( - lambda _: filled_list(cast(BlockPage, None), manager._storage.num_life_cycles), - self.beam_width, - ) - self._never_resumed = True - self._enable_swa_scratch_reuse = manager.enable_swa_scratch_reuse - if text_only is False and manager.text_only: - raise ValueError( - "text_only=False is not allowed when the manager is configured text_only=True" - ) - self._text_only = manager.text_only if text_only is None else text_only - self._scratch_slots = make_typed( - lambda _: list[ScratchSlotLock](), manager._storage.num_life_cycles - ) - self._enable_request_stats = enable_request_stats - self._pending_stats = _PendingStats() - self._status = self.Status.SUSPENDED - if reuse_match is not None: - self._setup_for_reuse(reuse_match) - self._refresh_generation_alloc_ready() - self._avg_history_length = Average() - self._avg_capacity = Average() - self._avg_history_length.update(self.history_length) - manager._living_kv_caches.add(rawref.ref(self)) - manager._avg_reused_length.update(self.history_length) - manager._num_created_kv_caches += 1 - assert NDEBUG or self._check_sanity() - - def set_base_page_index_buf( - self, beam_idx: BeamIndex, layer_group_id: LayerGroupId, buf: memoryview | None - ) -> None: - """ - Set the buffer for base page indices, so we directly update indices in user buffer to - avoid user-side copy. This is the zero-copy alternative of get_base_page_indices(). - - Note that base page indices are not meant for direct use in the kernels. They need to - be scaled by kv_cache_manager.page_index_scale(). - """ - length = self.num_blocks - old_indices = self._base_page_indices[beam_idx][layer_group_id] - new_indices: IndexSeq - if buf is None: - new_indices = array.array("i", old_indices[:length]) - else: - assert buf.ndim == 1 and buf.format == "i" and len(buf) >= length - buf[:length] = old_indices[:length] - buf[length:] = array.array("i", [BAD_PAGE_INDEX]) * (len(buf) - length) - new_indices = buf - self._base_page_indices[beam_idx][layer_group_id] = new_indices - - @property - def manager(self) -> "KVCacheManager": - return self._manager - - @property - def cuda_stream(self) -> CudaStream: - return unwrap_optional(self._cuda_stream) - - @cuda_stream.setter - def cuda_stream(self, cuda_stream: CudaStream) -> None: - if self._cuda_stream is not None: - if self.is_active: - CachedCudaEvent(self._cuda_stream).wait_in_stream(cuda_stream) - else: - assert self.status == self.Status.SUSPENDED and self._finish_event is None - self._cuda_stream = cuda_stream - - @property - def finish_event(self) -> CachedCudaEvent: - "Event recorded when switching from active to suspended/closed state. Unavailable when active." - return unwrap_optional(self._finish_event) - - @property - def num_blocks(self) -> int: - return len(self._blocks) - - def _stats_excluded(self) -> bool: - return self.manager.is_stats_excluded(self.id) - - def _should_record_manager_stats(self) -> bool: - return self.manager._stats_enabled and not self._stats_excluded() - - def _should_record_request_stats(self) -> bool: - # Manager stats historically also produced the request delta returned by - # commit_pending_stats(). Preserve that contract while allowing callers - # to opt in to request-only accounting through _enable_request_stats. - return ( - self.manager._stats_enabled or self._enable_request_stats - ) and not self._stats_excluded() - - def commit_pending_stats(self) -> KVCacheStatsDelta: - record_manager_stats = self._should_record_manager_stats() - record_request_stats = self._should_record_request_stats() - if not (record_manager_stats or record_request_stats): - self.discard_pending_stats() - return KVCacheStatsDelta() - if record_manager_stats: - self.manager.commit_stats( - self._pending_stats.global_stats, - self._pending_stats.iteration_stats_by_life_cycle, - ) - self.manager._commit_ssm_snapshot_iteration_stats( - self._pending_stats.ssm_snapshot_iteration_stats_by_life_cycle - ) - self.manager._commit_reused_blocks_by_level( - self._pending_stats.reused_blocks_by_level_by_life_cycle - ) - self.manager._commit_cached_tokens_by_level(self._pending_stats.cached_tokens_by_level) - request_stats = ( - self._pending_stats.request_stats.copy() - if record_request_stats - else KVCacheStatsDelta() - ) - self._pending_stats.clear() - self.manager.clear_stats_dirty(self.id) - return request_stats - - def discard_pending_stats(self) -> None: - self._pending_stats.clear() - self.manager.clear_stats_dirty(self.id) - - def _refresh_stats_dirty_state(self) -> None: - if not self._pending_stats.empty: - self.manager.mark_stats_dirty(self.id) - else: - self.manager.clear_stats_dirty(self.id) - - def _is_attention_life_cycle(self, life_cycle: LifeCycleId) -> bool: - return isinstance(self.manager._life_cycles.get_life_cycle(life_cycle), AttnLifeCycle) - - def _stats_life_cycle_key(self, life_cycle: LifeCycleId) -> LifeCycleId | None: - """Key for the attention-only block-reuse (hit/miss range) accounting.""" - return life_cycle if self._is_attention_life_cycle(life_cycle) else None - - def _refresh_generation_alloc_ready(self) -> None: - expected_prompt_length = self._expected_prompt_length - if expected_prompt_length is not None and self._history_length >= expected_prompt_length: - self._generation_alloc_ready = True - - def _should_record_generation_alloc_stats(self, capacity: int) -> bool: - return self._generation_alloc_ready and capacity > self._capacity - - @staticmethod - def _block_ranges_excluding( - block_begin: BlockOrdinal, - block_end: BlockOrdinal, - excluded: HalfOpenRange[BlockOrdinal], - ) -> Iterator[HalfOpenRange[BlockOrdinal]]: - first_end = min(block_end, excluded.beg) - if block_begin < first_end: - yield HalfOpenRange(block_begin, first_end) - second_begin = max(block_begin, excluded.end) - if second_begin < block_end: - yield HalfOpenRange(second_begin, block_end) - - def _record_resize_pending_allocations( - self, - block_begin: BlockOrdinal, - block_end: BlockOrdinal, - beam_width: BeamIndex, - excluded_ranges: TypedIndexList[LifeCycleId, HalfOpenRange[BlockOrdinal]], - count_as_generation: bool, - ) -> None: - record_manager_stats = self._should_record_manager_stats() - record_request_stats = self._should_record_request_stats() - if not (record_manager_stats or record_request_stats) or block_begin >= block_end: - return - # V2 includes generation allocations in per-request alloc_total/new - # metrics. This intentionally differs from the legacy V1 C++ manager, - # where addToken() only updates manager-level generation counters. - changed = False - for lc_idx, _ in self.manager._life_cycles.attention_life_cycles(): - for block_range in self._block_ranges_excluding( - block_begin, block_end, excluded_ranges[lc_idx] - ): - changed |= self._pending_stats.record_allocation_range( - lc_idx, - block_range.beg, - block_range.end, - beam_width=int(beam_width), - count_as_missed=not count_as_generation, - count_as_generation=count_as_generation, - record_manager_stats=record_manager_stats, - record_request_stats=record_request_stats, - ) - if changed: - self.manager.mark_stats_dirty(self.id) - - @staticmethod - def _has_reuse_source(page: BlockPage) -> bool: - if page is None or not isinstance(page.page, CommittedPage): - return False - return page.page.block() is not None - - def _subtract_pending_allocation_range( - self, block_begin: BlockOrdinal, block_end: BlockOrdinal - ) -> None: - if self._pending_stats.subtract_allocation_range(block_begin, block_end): - self._refresh_stats_dirty_state() - - def _record_direct_iteration_stats( - self, life_cycle: LifeCycleId, iteration_stats: KVCacheIterationStatsDelta - ) -> None: - # Every life cycle is reported, including SSM / recurrent ones: iteration - # statistics are keyed by life cycle, so recurrent page movement stays - # distinguishable from attention movement downstream. - if iteration_stats.empty or not self._should_record_manager_stats(): - return - self.manager.commit_stats(KVCacheStatsDelta(), {life_cycle: iteration_stats}) - - def _record_migrated_slots( - self, - pages: Sequence[Page], - slots: Sequence[Slot], - src_level: CacheLevel, - dst_level: CacheLevel, - ) -> None: - if not self._should_record_manager_stats(): - return - assert len(pages) == len(slots) - # One migration batch moves every page between the same pair of cache levels and does - # not mutate the storage while we record, so the per-page contributions differ only by - # life cycle. Aggregate them and commit once: commit_stats() re-samples the peak block - # statistics on every call, which is O(cache levels x pool groups), so committing per - # page turns a long-sequence eviction into thousands of redundant full scans. - stats = KVCacheStatsDelta() - iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] = {} - recorded = False - for page in pages: - is_attention = self._is_attention_life_cycle(page.life_cycle) - pg_idx = self.manager._storage.get_pool_group_index(page.life_cycle) - page_size = sum(self.manager._storage.slot_size(pg_idx)) - if src_level == GPU_LEVEL and dst_level > GPU_LEVEL: - iteration_stats = iteration_stats_by_life_cycle.setdefault( - page.life_cycle, KVCacheIterationStatsDelta() - ) - iteration_stats.iter_offload_blocks += 1 - iteration_stats.iter_offload_bytes += page_size - recorded = True - elif dst_level == GPU_LEVEL: - # Global cache-hit accounting is attention-only. SSM movement is - # reported by life-cycle/pool-group iteration statistics instead. - if is_attention: - stats.alloc_total_blocks += 1 - stats.alloc_new_blocks += 1 - iteration_stats = iteration_stats_by_life_cycle.setdefault( - page.life_cycle, KVCacheIterationStatsDelta() - ) - iteration_stats.iter_alloc_total_blocks += 1 - iteration_stats.iter_alloc_new_blocks += 1 - if src_level > GPU_LEVEL: - iteration_stats.iter_onboard_blocks += 1 - iteration_stats.iter_onboard_bytes += page_size - elif src_level == GPU_LEVEL: - iteration_stats.iter_intra_device_copy_blocks += 1 - iteration_stats.iter_intra_device_copy_bytes += page_size - recorded = True - if recorded: - self.manager.commit_stats(stats, iteration_stats_by_life_cycle) - - def _record_dropped_pages( - self, - pages: Sequence[Page], - cache_level: CacheLevel, - ) -> None: - """Record host-tier LRU drops (pages released without onboarding back to GPU). - - Mirrors _record_migrated_slots in structure: per-life-cycle attribution, - gated on manager statistics, per-page bytes computed from slot_size. - cache_level is unused for now (we only have a 2-tier setup in practice; - all last-level drops are host drops) but kept in the signature for future - per-tier disambiguation. - """ - if not self._should_record_manager_stats() or not pages: - return - # Aggregate per life cycle and commit once -- see _record_migrated_slots for why. - iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] = {} - for page in pages: - pg_idx = self.manager._storage.get_pool_group_index(page.life_cycle) - page_size = sum(self.manager._storage.slot_size(pg_idx)) - iteration_stats = iteration_stats_by_life_cycle.setdefault( - page.life_cycle, KVCacheIterationStatsDelta() - ) - iteration_stats.iter_host_dropped_blocks += 1 - iteration_stats.iter_host_dropped_bytes += page_size - self.manager.commit_stats(KVCacheStatsDelta(), iteration_stats_by_life_cycle) - - # destroy ownership of memory blocks, so KV cache manager can decide to evict or drop them. After - # close, uncommitted data in blocks for (beam_index >= beam_width) will be lost. - def close(self) -> None: - assert NDEBUG or self._check_sanity() - if self.status == self.Status.CLOSED: - return - self.discard_pending_stats() - self.stop_committing() - assert NDEBUG or self._check_sanity() - manager = self.manager - # Dummy/warmup caches are reserved at the model's full declared context, - # not at a realistic sequence length, and _avg_sqr_capacity is an RMS -- - # so a handful of them dominates the statistic outright and the tuner - # sizes pools for sequences that never arrive. They are already tracked - # as stats-excluded at creation; honour that here too. - if self.capacity > 0 and not manager.is_stats_excluded(self.id): - self._avg_capacity.update(self.capacity) - manager._avg_sqr_capacity.update(self._avg_capacity.value**2) - manager._avg_sqr_history_length.update(self._avg_history_length.value**2) - manager._num_sampled_kv_caches += 1 - manager._try_update_target_ratios() - with self._record_event(): - self._clear_blocks() - self._status = self.Status.CLOSED - manager._living_kv_caches.remove(self.__rawref__) - - def __del__(self) -> None: - self.close() - self.__rawref__.invalidate() - - @property - def beam_width(self) -> BeamIndex: - return self._beam_width - - # beam_width > 1 is only for generation. If decreasing beam_width, uncommitted data in blocks for - # (beam_index >= beam_width) will be lost. - @beam_width.setter - def beam_width(self, beam_width: BeamIndex) -> None: - raise NotImplementedError("Not implemented yet for beam search") - - # Get the indices of memory blocks for each beam. - def get_base_page_indices( - self, layer_group_id: LayerGroupId, beam_id: BeamIndex = DEFAULT_BEAM_INDEX - ) -> IndexSeq: - indices = self._base_page_indices[beam_id][layer_group_id] - assert NDEBUG or all( - v == value_or(r, BAD_PAGE_INDEX) - for v, r in zip(indices, self._get_base_page_indices_ref(layer_group_id, beam_id)) - ) - return indices - - def get_ssm_block_base_index( - self, layer_group_id: LayerGroupId, beam_id: BeamIndex = DEFAULT_BEAM_INDEX - ) -> int: - entry = self._ssm_blocks[beam_id][layer_group_id] - if entry is None: - return BAD_PAGE_INDEX - return expect_type(_SharedPageLock, entry).page.slot_id - - def get_aggregated_page_indices( - self, - layer_group_id: LayerGroupId, - beam_id: BeamIndex = DEFAULT_BEAM_INDEX, - valid_only: bool = False, - ) -> Iterator[int]: - """ - Get the internal slot indices for the given layer group and beam. - Each slot is a group of coalesced buffers in one memory pool group. - This API exposes internal slot indices, mainly for efficient data transfer. - For computation, use get_page_indices() instead. - - Args: - layer_group_id: Layer group to inspect. - beam_id: Beam index to read. Defaults to DEFAULT_BEAM_INDEX. - - Returns: - Aggregated page index for each block, or BAD_PAGE_INDEX for invalid blocks. - """ - for b in self._blocks: - if (holder := b.pages[beam_id][layer_group_id]) is None: - if not valid_only: - yield BAD_PAGE_INDEX - else: - yield holder.page.slot_id - - def get_scratch_desc(self, layer_group_id: LayerGroupId) -> "ScratchDesc | None": - """ - Get scratch metadata for the given layer group, or None if scratch is not active. - - The returned ScratchDesc contains the scratch block ordinal range and the - slot IDs for the scratch coalesced slots. Pass this to PageIndexConverter - together with get_base_page_indices() to produce per-layer page indices. - - The returned ScratchDesc is invalidated by the next capacity/history_length update. - """ - lc = self.manager._life_cycles[layer_group_id] - sr = self._get_scratch_range(lc) - if not sr: - return None - from ._kv_cache_manager import ScratchDesc - - return ScratchDesc( - range=sr, - slot_ids=[s.slot.slot_id for s in self._scratch_slots[layer_group_id]], - ) - - @property - def has_scratch_slots(self) -> bool: - """True if this KV cache currently has scratch slots allocated.""" - return any(len(s) > 0 for s in self._scratch_slots) - - @property - def enable_swa_scratch_reuse(self) -> bool: - return self._enable_swa_scratch_reuse - - @enable_swa_scratch_reuse.setter - def enable_swa_scratch_reuse(self, enable: bool) -> None: - if enable == self._enable_swa_scratch_reuse: - return - if enable: - if not self.manager.enable_swa_scratch_reuse: - raise ValueError( - "Cannot enable SWA scratch reuse for a request when it is disabled in " - "KV cache manager config" - ) - if self._would_use_swa_scratch_blocks(): - raise ValueError( - "Cannot enable SWA scratch reuse while the current request state would " - "need scratch blocks" - ) - self._enable_swa_scratch_reuse = True - return - - if self._would_use_swa_scratch_blocks(): - raise ValueError("Cannot disable SWA scratch reuse while scratch blocks are needed") - assert not self.has_scratch_slots - self._enable_swa_scratch_reuse = False - - @property - def text_only(self) -> bool: - return self._text_only - - @text_only.setter - def text_only(self, text_only: bool) -> None: - # A text-only deployment is a hard guarantee: a request may not opt out. - if not text_only and self.manager.text_only: - raise ValueError( - "Cannot set text_only=False for a request when the KV cache manager is " - "configured text_only=True" - ) - # Claiming text-only is a fast-path claim; verify committed tokens are digest-free. - if text_only and any(isinstance(t, bytes) for t in self._committed_tokens): - raise ValueError( - "Cannot set text_only=True: this sequence has already committed digest tokens" - ) - self._text_only = text_only - - def supports_index_mode(self, mode: PageIndexMode) -> bool: - match mode: - case PageIndexMode.PER_LAYER: - return True - case PageIndexMode.SHARED: - return not self.has_scratch_slots - - # reserve space for next inference. Request new blocks from KVCacheManager if necessary. - # if capacity is increased and beam_width > 1, blocks containing new tokens should be allocated for each beam. - # Decrease of capacity may destroy stale blocks (if not used by other requests). - # Decrease of capacity cannot remove historical or committed tokens. - # History length cannot be decreased. - # Increase of history length may trigger out-of-window block eviction/dropping for SWA layers. - # If we use two separate APIs for capacity and history length, sometimes we will need to increase - # capacity first to maintain capacity >= history_length. But then we may have a middle state (between - # two APIs) where we use more pages than necessary for SWA layers. So we use a single API to avoid - # this. Usually this is a concern only for prefill phase where we create many tokens in one step. For - # other cases, we can just set the capacity and history_length properties instead. - def resize(self, capacity: int | None, history_length: int | None = None) -> bool: - assert self.status == self.Status.ACTIVE - tokens_per_block = self.tokens_per_block - assert div_up(self._capacity, tokens_per_block) == len(self._blocks) - if capacity is None: - capacity = self._capacity - else: - self._avg_capacity.update(capacity) - if history_length is None: - history_length = self._history_length - else: - self._avg_history_length.update(history_length) - if history_length < self._history_length: - raise ValueError("History length cannot be decreased") - if capacity < history_length: - raise ValueError("History length cannot be greater than capacity") - manager = self.manager - # Scratch reuse: compute scratch ranges and slot delta - enable_scratch = self.enable_swa_scratch_reuse - if enable_scratch and capacity != self._capacity: - max_rewind_len = self._swa_scratch_max_rewind_len() - min_history_length = max(0, self._capacity - max_rewind_len) - assert min_history_length <= history_length <= self._capacity, ( - "SWA scratch requires " - f"old_capacity - max_rewind_len ({min_history_length}) <= " - f"history_length ({history_length}) <= " - f"old_capacity ({self._capacity})" - ) - record_generation_alloc_stats = self._should_record_generation_alloc_stats(capacity) - if ( - not enable_scratch - and self._shortcut_set_capacity(capacity) - and self._shortcut_set_history_length(history_length) - ): - self._refresh_generation_alloc_ready() - return True - ssm_lc_id = manager._life_cycles.ssm_life_cycle_id - beam_width = self.beam_width - backup_holders = self._unlock_stale_blocks(history_length) - old_num_blocks = BlockOrdinal(div_up(self._capacity, tokens_per_block)) - new_num_blocks = BlockOrdinal(div_up(capacity, tokens_per_block)) - num_life_cycles = manager._life_cycles.size - if new_num_blocks < old_num_blocks: - assert not self.has_scratch_slots, "Cannot shrink while scratch slots exist" - self._subtract_pending_allocation_range(new_num_blocks, old_num_blocks) - with self._record_event(): - del self._blocks[new_num_blocks:] - for beam_indices in self._base_page_indices: - for indices in beam_indices: - assert all(i == BAD_PAGE_INDEX for i in indices[new_num_blocks:]) - if type(indices) is array.array: - del indices[new_num_blocks:] - else: - indices[new_num_blocks:] = array.array("i", [BAD_PAGE_INDEX]) * ( - len(indices) - new_num_blocks - ) - - excess_scratch_slots, delta_scratch_slots, scratch_ranges = self._take_excess_scratch_slots( - capacity, history_length - ) - - if new_num_blocks >= old_num_blocks: - num_new_slots = filled_list(0, num_life_cycles) - stale_ranges = [ - _KVCache._get_stale_range(tokens_per_block, history_length, lc) - for _, lc in manager._life_cycles.items() - ] - for lc in typed_range(num_life_cycles): - if lc == ssm_lc_id: - continue - stale_beg, stale_end = stale_ranges[lc] - if enable_scratch: - # Only newly added blocks consume slots below; scratch range may - # extend before old_num_blocks when history_length < old_capacity. - new_block_range = HalfOpenRange(old_num_blocks, new_num_blocks) - num_new_blocks_using_scratch = len( - intersect(scratch_ranges[lc], new_block_range) - ) - num_new_normal_blocks = len(new_block_range) - num_new_blocks_using_scratch - num_new_slots[lc] = num_new_normal_blocks * beam_width - else: - if old_num_blocks < stale_beg: - assert new_num_blocks >= stale_end - num_new_blocks_to_add = (stale_beg - old_num_blocks) + ( - new_num_blocks - stale_end - ) - else: - num_new_blocks_to_add = new_num_blocks - max(stale_end, old_num_blocks) - num_new_slots[lc] = num_new_blocks_to_add * beam_width - - net_alloc_counts = make_typed( - lambda lc: num_new_slots[lc] + delta_scratch_slots[lc], num_life_cycles - ) - storage = self._storage - if any(c > 0 for c in net_alloc_counts): - try: - new_slots = storage.new_gpu_slots( - make_typed(lambda lc: max(0, net_alloc_counts[lc]), num_life_cycles), - self._record_migrated_slots, - self._record_dropped_pages, - ) - except OutOfPagesError: - self._recover_excess_scratch_slots(excess_scratch_slots) - self._lock_held_blocks(backup_holders) - return False - else: - new_slots = make_typed(lambda _: list[Slot](), num_life_cycles) - - # Wait on newly allocated slots - stream_wait_events( - self.cuda_stream, (s.ready_event for s in chain.from_iterable(new_slots)) - ) - - # Combine slots and distribute - slots = make_typed(lambda _: list[Slot](), num_life_cycles) - for lc in typed_range(num_life_cycles): - slots[lc] = new_slots[lc] + [ - lock.detach_slot() for lock in excess_scratch_slots[lc] - ] - new_slots[lc].clear() - excess_scratch_slots[lc].clear() - - if any(cnt < 0 for cnt in net_alloc_counts): - with self._record_event(): - for lc in typed_range(num_life_cycles): - for _ in range(-net_alloc_counts[lc]): - slot = slots[lc].pop() - slot.ready_event = self.finish_event - storage.release_slot(lc, GPU_LEVEL, slot) - - assert all( - len(slots[lc]) == num_new_slots[lc] + max(0, delta_scratch_slots[lc]) - for lc in typed_range(num_life_cycles) - ) - - # Fulfill additional scratch slots - for lc in typed_range(num_life_cycles): - for _ in range(delta_scratch_slots[lc]): - slot = slots[lc].pop() - self._scratch_slots[lc].append(ScratchSlotLock(slot, self, lc, skip_wait=True)) - - for beam_indices in self._base_page_indices: - for indices in beam_indices: - if type(indices) is array.array: - assert len(indices) == old_num_blocks - indices.extend([BAD_PAGE_INDEX] * (new_num_blocks - old_num_blocks)) - else: - if len(indices) < new_num_blocks: - raise ValueError("User-provided base page indices is too short") - - stream_wait_events( - self.cuda_stream, (s.ready_event for s in chain.from_iterable(slots)) - ) - - # Scratch blocks use temporary shared SWA slots instead of normal - # per-request KV pages, so they are excluded from alloc/miss stats. - excluded_ranges = ( - scratch_ranges if enable_scratch else to_typed(LifeCycleId, stale_ranges) - ) - self._record_resize_pending_allocations( - old_num_blocks, - new_num_blocks, - beam_width, - excluded_ranges, - record_generation_alloc_stats, - ) - for ordinal in typed_range(old_num_blocks, new_num_blocks): - block = make_typed( - lambda _: filled_list(cast(BlockPage, None), num_life_cycles), beam_width - ) - for beam_index in typed_range(beam_width): - for lc in typed_range(num_life_cycles): - if lc == ssm_lc_id: - continue # SSM pages live in _ssm_blocks, not in _blocks - if enable_scratch: - # Assertion guarantees no new block is stale. - if ordinal in scratch_ranges[lc]: - continue # Scratch block — no per-block page allocation - else: - stale_beg, stale_end = stale_ranges[lc] - if stale_beg <= ordinal < stale_end: - continue - slot = slots[lc].pop() - # We have already waited for ready_event of the slots. - block[beam_index][lc] = UncommittedPage( - self, ordinal, lc, GPU_LEVEL, slot, beam_index - ).lock(self, beam_index, ordinal, lc, skip_wait=True) - self._blocks.append(SeqBlock(block, None)) - assert all(len(slots[lc]) == 0 for lc in typed_range(num_life_cycles)) - self._capacity = capacity - self._history_length = history_length - self._refresh_generation_alloc_ready() - assert NDEBUG or self._check_sanity() - return True - - @property - def capacity(self) -> int: - "Get the current capacity in number of tokens." - return self._capacity - - @capacity.setter - def capacity(self, capacity: int) -> None: - """ - Reserve space for next inference. Capacity cannot be smaller than history length. - Use resize() instead if you need to change both capacity and history length. If you use two - separate APIs, you may have a middle state (between two APIs) where we use more pages than - necessary for SWA layers. - Expect OutOfPagesError exception if there are not enough pages in GPU memory. - """ - if self.enable_swa_scratch_reuse: - raise ValueError( - "Cannot use capacity setter when SWA scratch reuse is enabled. " - "Use resize(capacity, history_length) instead." - ) - success = self.resize(capacity, None) - if not success: - raise OutOfPagesError("Not enough pages in GPU memory") - - @property - def history_length(self) -> int: - """ - Get the current history length in number of tokens. history_length decides how many blocks - needs to be in GPU memory for SWA layers. - """ - return self._history_length - - @history_length.setter - def history_length(self, history_length: int) -> None: - "History length cannot be decreased. Increase may trigger out-of-window block eviction/dropping for SWA layers." - success = self.resize(None, history_length) - assert success - - # notify KV cache manager that we have some finalized/accepted tokens. If a block becomes full, - # also commit the block for reuse. - # In case of beam search, this should be called only with finalized (converged) tokens, and the - # token data must be in the 0th beam. - # We'll destroy memory blocks for other beams if the whole block is full and committed. - # Committed tokens are always history, so history_length will be automatically updated to maintain - # (num_committed_tokens <= history_length). Note that history_length increase may trigger out-of-window - # block eviction/dropping for SWA layers. - # beam_search_indices: indices indicating which candidate to choose for each token. A block with all - # tokens committed will be unified to one memory page and the other memory pages are dropped. Only for - # beam search. - # is_end: if True, this call records a final reusable snapshot and stops committing. - # This is a terminal-memory contract: callers must not perform later writes to this - # _KVCache memory. The final live pages may be moved into the radix tree instead - # of copied, including SSM state and the last partial block for commit_min_snapshot. - def commit( - self, - accepted_input_tokens: Sequence[TokenIdExt], - beam_search_indices: Sequence[int] | None = None, - is_end: bool = False, - ): - if self.beam_width != 1: - raise NotImplementedError("Not implemented yet for beam search") - if not accepted_input_tokens: - if is_end: - self.stop_committing() - return - assert beam_search_indices is None - assert self.status == self.Status.ACTIVE - if self._commit_state == self.CommitState.USER_STOP: - raise LogicError("Cannot commit tokens after stop_committing()") - commit_min_snapshot = self.manager.commit_min_snapshot - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - if commit_min_snapshot: - new_num_committed_tokens = self.num_committed_tokens + len(accepted_input_tokens) - assert self.history_length in ( - self.num_committed_tokens, - new_num_committed_tokens, - ), "commit_min_snapshot requires commit() to start or end at history_length" - self._committed_tokens.extend(accepted_input_tokens) - if self._commit_state == self.CommitState.VIRTUAL_STOP: - if is_end: - self._commit_state = self.CommitState.USER_STOP - return - if self.history_length < self.num_committed_tokens: - self.history_length = self.num_committed_tokens - num_committed_blocks = self._num_committed_blocks - new_num_full_blocks = BlockOrdinal(self.num_committed_tokens // self.tokens_per_block) - has_partial_snapshot = ( - commit_min_snapshot - and self.num_committed_tokens % self.tokens_per_block != 0 - and self.num_committed_tokens > 0 - ) - has_new_full_blocks = new_num_full_blocks > num_committed_blocks - if has_new_full_blocks or has_partial_snapshot: - ssm_snapshot_ordinal = _KVCache._to_block_ordinal( - self.tokens_per_block, self.num_committed_tokens - 1 - ) - with self._record_event(): - for ordinal in typed_range(num_committed_blocks, new_num_full_blocks): - self._commit_block( - ordinal, - False, - commit_ssm=( - commit_min_snapshot - and ssm_lc_id is not None - and ordinal == ssm_snapshot_ordinal - ), - move_ssm=is_end, - ) - # _commit_block transitions out of ALLOWED (to USER_STOP) when a - # block cannot be committed (VIRTUAL_STOP). Stop here so we don't - # re-enter _commit_block on an already-stopped cache. - if self._commit_state != self.CommitState.ALLOWED: - break - if has_partial_snapshot and self._commit_state == self.CommitState.ALLOWED: - partial_ordinal = BlockOrdinal(new_num_full_blocks) - if is_end: - self._commit_block( - partial_ordinal, - True, - commit_ssm=ssm_lc_id is not None, - move_ssm=ssm_lc_id is not None, - ) - else: - self._snapshot_partial_block_to_tree( - partial_ordinal, commit_ssm=ssm_lc_id is not None - ) - if is_end and self._commit_state != self.CommitState.USER_STOP: - self.stop_committing() - - # Note that the tokens may not be ready yet, if the event passed to the past commit() calls are not yet signaled. - @property - def num_committed_tokens(self) -> int: - return len(self._committed_tokens) - - @property - def cached_tokens_by_level(self) -> CountsByLevel: - """Reused-token counts indexed by each token's coldest required source cache level.""" - return list(self._cached_tokens_by_level) - - def _get_last_cached_token_level(self) -> int | None: - """Return the source cache level of the last initially reused logical block.""" - return ( - int(self._last_cached_token_level) - if self._last_cached_token_level is not None - else None - ) - - def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: - """Return the pre-hybrid-pruning prefix for internal diagnostics.""" - return self._num_reusable_tokens_before_hybrid_pruning - - def _get_num_reusable_tokens_before_pruning(self) -> int: - """Return the raw token-path walk depth, before any pruning.""" - return self._num_reusable_tokens_before_pruning - - @property - def committed_tokens(self) -> list[TokenIdExt]: - return list(self._committed_tokens) - - @property - def reuse_scope(self) -> ReuseScope: - return self._reuse_scope - - def plan_committed_block_drop(self) -> PlannedDropHandle | None: - """Plan dropping pages needed only by the next conversation turn. - - The plan covers committed pages in each SWA life cycle's current - attention window and the exact SSM snapshot for the committed prefix. - Full-attention and attention-sink blocks are excluded because later - turns may still need them. This must be called after stop_committing(). - Returns None without creating a plan if any required page is unavailable. - """ - if self._commit_state != self.CommitState.USER_STOP: - raise LogicError("plan_committed_block_drop() requires stop_committing()") - - if self.num_committed_tokens == 0: - return None - - # Locate pages through the radix tree rather than - # SeqBlock.tree_block: the latter is not guaranteed to identify a - # partial snapshot after reuse. Requiring an exact match keeps the - # preceding conversation plan intact if this turn no longer has a - # complete reusable endpoint. All PP ranks use the same lookup so - # attention-only ranks include the final partial SWA block too. - match = self.manager._radix_tree.match( - self.reuse_scope, - self._committed_tokens, - self.manager.enable_partial_match, - ) - if match.num_tokens != self.num_committed_tokens or not match.blocks: - return None - - end = BlockOrdinal(len(match.blocks)) - pages_to_drop: list[CommittedPage] = [] - for lc_idx, lc in self.manager._life_cycles.items(): - if isinstance(lc, AttnLifeCycle): - if lc.window_size is None: - continue - stale_range = _KVCache._get_stale_range( - self.tokens_per_block, self.num_committed_tokens, lc - ) - window_start = min(stale_range.end, end) - else: - window_start = BlockOrdinal(end - 1) - for ordinal in typed_range(window_start, end): - page = match.blocks[ordinal].get_page(lc_idx) - if page is None: - return None - pages_to_drop.append(page) - return PlannedDropHandle(pages_to_drop) - - # Users promise to not commit any more tokens. For cases where we shouldn't reuse generated tokens - # (eg. CoT), this helps us drop (instead of evict) out-of-window blocks for SWA layers. - # If there is a uncommitted block containing committed tokens, we will commit the block immediately. - def stop_committing(self) -> None: - assert self.status != self.Status.CLOSED - if self._commit_state == self.CommitState.USER_STOP: - return - assert NDEBUG or self._check_sanity() - if self._commit_state == self.CommitState.VIRTUAL_STOP: - self._commit_state = self.CommitState.USER_STOP - return - assert self._commit_state == self.CommitState.ALLOWED - if self.num_committed_tokens % self.tokens_per_block != 0: - ordinal = _KVCache._to_block_ordinal(self.tokens_per_block, self.num_committed_tokens) - with self._record_event(): - self._commit_block(ordinal, True) - else: - self._commit_state = self.CommitState.USER_STOP - self._on_stop_committing() - # TODO: check if the last committed pages are usable, in case some prior pages are already - # dropped. For SWA, this can be done only when we stop committing. (TRTLLM-8802) - assert self._commit_state == self.CommitState.USER_STOP - - # Suspend, allow the KV cache manager to evict buffers from GPU, but don't drop them. - # suspend+resume allows us to implement dynamic batch size. May also be used to support HSTU model. - def suspend(self) -> None: - assert self.status == self.Status.ACTIVE - assert self._check_sanity() - assert self._finish_event is None - for beam_idx, beam_indices in typed_enumerate(self._base_page_indices): - for lc, indices in typed_enumerate(beam_indices): - if type(indices) is memoryview: - self.set_base_page_index_buf(beam_idx, lc, None) - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - with self._record_event(): # used by _SharedPageLock.__del__ - for ordinal, beam_idx, lc_idx in self._active_pages(): - beam_block = ( - self._block(ordinal, beam_idx) - if lc_idx != ssm_lc_id - else self._ssm_blocks[beam_idx] - ) - holder = expect_type(_SharedPageLock, beam_block[lc_idx]).holder - # after this assignment, __del__ of the original _SharedPageLock will use self.finish_event - # to indicate end of usage for the page. - beam_block[lc_idx] = holder - # Free scratch slots on suspend since the data is ephemeral - self._free_scratch_slots() - self._status = self.Status.SUSPENDED - # Manager-level counter, so gate on the manager predicate: it also honours - # the per-cache stats exclusion (dummy / CUDA-graph caches). C++ spells the - # gate _shouldRecordStats() (manager OR request); the two are equivalent - # here because KVCacheManager.record_request_suspended already returns - # early when stats are disabled. - if self._should_record_manager_stats(): - self.manager.record_request_suspended() - - # Resume, migrate buffers to GPU memory. - def resume(self, cuda_stream: CudaStream | None = None) -> bool: - assert self.status == self.Status.SUSPENDED - if cuda_stream is not None: - self.cuda_stream = cuda_stream - utilization = max(self._storage.get_utilization(GPU_LEVEL)) - if utilization > self.manager._init_config.max_util_for_resume: - return False - assert self._cuda_stream is not None, "cuda_stream is never set" - assert self._finish_event is None - storage = self._storage - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - life_cycles = self.manager._life_cycles - num_life_cycles = life_cycles.size - - # Pre-allocate GPU slots for deferred copies (partial blocks + SSM) before locking, - # so we never end up in a state where pages are locked but we can't allocate for the copy. - deferred_slots: TypedIndexList[LifeCycleId, Slot | None] = filled_list( - None, storage.num_life_cycles - ) - - excess_scratch_slots, delta_scratch_slots, _ = self._take_excess_scratch_slots( - self.capacity, self.history_length - ) - assert all(len(s) == 0 for s in excess_scratch_slots) - - num_slots = filled_list(0, num_life_cycles) - has_partial = False - if self._never_resumed: - assert self.beam_width == 1 - has_partial = self.num_committed_tokens % self.tokens_per_block != 0 - for lc_idx, lc in life_cycles.items(): - if type(lc) is SsmLifeCycle or has_partial: - num_slots[lc_idx] += 1 - - for lc_idx in typed_range(num_life_cycles): - num_slots[lc_idx] += delta_scratch_slots[lc_idx] - - if any(c > 0 for c in num_slots): - try: - tmp_slots = storage.new_gpu_slots( - num_slots, self._record_migrated_slots, self._record_dropped_pages - ) - except OutOfPagesError: - return False - - # Wait for scratch slots to be ready - scratch_slots_to_add = make_typed(lambda _: list[Slot](), num_life_cycles) - for lc_idx, slot_lst in zip(typed_range(num_life_cycles), tmp_slots, strict=True): - if self._never_resumed and ( - type(life_cycles[lc_idx]) is SsmLifeCycle or has_partial - ): - deferred_slots[lc_idx] = slot_lst.pop() - scratch_slots_to_add[lc_idx] = slot_lst - - stream_wait_events( - self.cuda_stream, (s.ready_event for s in chain.from_iterable(scratch_slots_to_add)) - ) - - for lc_idx in typed_range(num_life_cycles): - for slot in scratch_slots_to_add[lc_idx]: - self._scratch_slots[lc_idx].append( - ScratchSlotLock(slot, self, lc_idx, skip_wait=True) - ) - - tasks = list[BatchedLockTarget]() - for ordinal, beam_idx, lc_idx in self._active_pages(): - beam_block = ( - self._block(ordinal, beam_idx) - if lc_idx != ssm_lc_id - else self._ssm_blocks[beam_idx] - ) - page = expect_type(_PageHolder, beam_block[lc_idx]).page - tasks.append(BatchedLockTarget(page, beam_idx, ordinal, lc_idx)) - try: - locks = batched_lock_to_gpu( - self, tasks, self._record_migrated_slots, self._record_dropped_pages - ) - except OutOfPagesError: - for lc_idx, slot in typed_enumerate(deferred_slots): - if slot is not None: - storage.release_slot(lc_idx, GPU_LEVEL, slot) - return False - - # Replace all holders with locks. - for (ordinal, beam_idx, lc_idx), lock in zip(self._active_pages(), locks): - beam_block = ( - self._block(ordinal, beam_idx) - if lc_idx != ssm_lc_id - else self._ssm_blocks[beam_idx] - ) - page = expect_type(_PageHolder, beam_block[lc_idx]).page - assert page is lock.page - beam_block[lc_idx] = lock - - # Deferred copy: for partial blocks and SSM, copy from now-locked source pages - # to pre-allocated GPU slots, then unlock sources and replace with new pages. - if self._never_resumed: - beam_idx = DEFAULT_BEAM_INDEX - last_ordinal = self._to_block_ordinal( - self.tokens_per_block, self.num_committed_tokens - 1 - ) - # Phase 1: Copy GPU→GPU from locked source pages to pre-allocated slots. - src_locks: list[_SharedPageLock] = [] - gpu_tier = storage.cache_tiers[GPU_LEVEL] - # wait for all new slots to be ready - stream_wait_events( - self.cuda_stream, (slot.ready_event for slot in deferred_slots if slot is not None) - ) - for lc_idx, new_slot in typed_enumerate(deferred_slots): - if new_slot is None: - continue - if lc_idx == ssm_lc_id: - if self.num_committed_tokens == 0: - continue # fresh SSM — no source to copy from - lock = self._ssm_blocks[beam_idx][lc_idx] - else: - lock = self._block(last_ordinal, beam_idx)[lc_idx] - assert type(lock) is _SharedPageLock - # V2 still copies a partial reuse into a private slot before writing to it. - # The copy allocates a block, but it is a miss only without a reusable source. - has_partial_reuse_source = self._has_reuse_source(lock) - src_locks.append(lock) - pg_idx = storage._life_cycle_grouping[lc_idx] - slot_size = storage.slot_size(pg_idx) - for p in typed_range(storage.num_pools(pg_idx)): - dst = storage.slot_address(GPU_LEVEL, pg_idx, new_slot.slot_id, p) - src = storage.slot_address(GPU_LEVEL, pg_idx, lock.page.slot_id, p) - # todo: add another batched copy supporting non-uniform size. - batched_copy( - gpu_tier, - gpu_tier, - slot_size[p], - [CopyTask(dst, src)], - self.cuda_stream, - ) - if lc_idx != ssm_lc_id: - life_cycle_key = self._stats_life_cycle_key(lc_idx) - record_manager_stats = self._should_record_manager_stats() - record_request_stats = self._should_record_request_stats() - if life_cycle_key is not None and ( - record_manager_stats or record_request_stats - ): - changed = self._pending_stats.record_allocation_range( - life_cycle_key, - last_ordinal, - BlockOrdinal(last_ordinal + 1), - beam_width=1, - count_as_missed=not has_partial_reuse_source, - record_manager_stats=record_manager_stats, - record_request_stats=record_request_stats, - ) - if changed: - self.manager.mark_stats_dirty(self.id) - # Block-reuse accounting above is attention-only, but the copy - # itself is reported for every life cycle, SSM included — - # matching the C++ backend's deferred-copy loop (kvCache.cpp). - self._record_direct_iteration_stats( - lc_idx, - KVCacheIterationStatsDelta( - iter_intra_device_copy_blocks=1, - iter_intra_device_copy_bytes=sum(storage.slot_size(pg_idx)), - ), - ) - # Unlock source pages — _record_event captures all prior cuda work - # so the original pages know when we're done reading from them. - if src_locks: - with self._record_event(): - for lock in src_locks: - lock.unlock() - # Phase 2: Replace with new UncommittedPages (both copied and fresh SSM). - for lc_idx, new_slot in typed_enumerate(deferred_slots): - if new_slot is None: - continue - if lc_idx == ssm_lc_id: - beam_block = self._ssm_blocks[beam_idx] - block_ordinal = BAD_BLOCK_ORDINAL - else: - beam_block = self._block(last_ordinal, beam_idx) - block_ordinal = last_ordinal - new_page = UncommittedPage( - self, block_ordinal, lc_idx, GPU_LEVEL, new_slot, beam_idx - ) - new_lock = new_page.lock(self, beam_idx, block_ordinal, lc_idx, skip_wait=True) - beam_block[lc_idx] = new_lock - # Clear tree_block for the partial block — it's now uncommitted. - if self.num_committed_tokens % self.tokens_per_block != 0: - self._blocks[last_ordinal].tree_block = None - # A freshly-created cache starts SUSPENDED and is activated by this same - # resume() call, so gate the counter on _never_resumed: only a cache that - # was previously ACTIVE and got suspended counts as a preemption recovery. - # Without this, the counter would track request admissions, not preemption. - first_activation = self._never_resumed - self._never_resumed = False - self._status = self.Status.ACTIVE - if not first_activation and self._should_record_manager_stats(): - self.manager.record_request_resumed() - return True - - def prefetch(self, target: CacheLevel) -> bool: - """Best-effort prefetch active pages to the target cache level. - - The cache must be suspended. Prefetch is only a performance hint: a False - return value means the requested pages could not be recalled due to cache - pressure, but the cache remains functionally valid. - - Args: - target: Destination cache level for active pages in lower tiers. - - Returns: - True if the prefetch was dispatched, False if storage could not reserve enough pages. - """ - assert self.status == self.Status.SUSPENDED - manager = self.manager - storage = manager._storage - num_tiers = storage.num_cache_levels - assert CacheLevel(0) <= target < num_tiers - - num_pool_groups = storage.num_pool_groups - lc2pg = storage.get_pool_group_index - all_pages = make_typed( - lambda _: make_typed(lambda _: list[Page](), num_tiers), num_pool_groups - ) - - for ordinal, beam_idx, lc_idx in self._active_pages(): - holder = self._page(ordinal, beam_idx, lc_idx) - if holder is None: - continue - page = expect_type(_PageHolder, holder).page - lvl = page.cache_level - if lvl < target: - continue - pg_idx = lc2pg(lc_idx) - all_pages[pg_idx][lvl].append(page) - - try: - # StorageManager reports what it actually migrated. Blocks, the unit - # iter_offload_blocks and iter_onboard_blocks use: one page per block per life cycle. - # Unrelated to _cached_tokens_by_level, which answers where matched tokens lived. - disk_blocks_migrated = storage.prefetch(target, all_pages) - if disk_blocks_migrated > 0 and self._should_record_manager_stats(): - manager.record_disk_prefetch_blocks(disk_blocks_migrated) - except OutOfPagesError: - return False - return True - - def _active_pages(self) -> Iterator[tuple[BlockOrdinal, BeamIndex, LifeCycleId]]: - """Yields (ordinal, beam_idx, lc_idx) for all active pages. - - For attention life cycles, yields non-stale blocks from _blocks (excluding scratch blocks). - For SSM, yields entries from _ssm_blocks with ordinal=BAD_BLOCK_ORDINAL. - """ - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - for lc_idx, lc in self.manager._life_cycles.items(): - if lc_idx == ssm_lc_id: - assert ssm_lc_id is not None - for beam_idx, beam_block in typed_enumerate(self._ssm_blocks): - if beam_block[ssm_lc_id] is not None: - yield BAD_BLOCK_ORDINAL, beam_idx, lc_idx - continue - stale_start, stale_end = _KVCache._get_stale_range( - self.tokens_per_block, self.history_length, lc - ) - scratch_range = self._get_scratch_range(lc) - sink_blocks = typed_range(stale_start) - window_blocks = typed_range(stale_end, typed_len(self._blocks)) - for ordinal in chain(sink_blocks, window_blocks): - block = self._blocks[ordinal] - for beam_idx, beam_block in typed_enumerate(block.pages): - is_scratch = ordinal in scratch_range - assert is_scratch == (beam_block[lc_idx] is None) - if not is_scratch: - yield ordinal, beam_idx, lc_idx - - @property - def status(self) -> _Status: - return self._status - - @property - def is_active(self) -> bool: - return self.status == self.Status.ACTIVE - - @property - def tokens_per_block(self) -> int: - return self._tokens_per_block - - def _page( - self, block_ordinal: BlockOrdinal, beam_index: BeamIndex, life_cycle: LifeCycleId - ) -> BlockPage: - """Return the page holder for an attention block or the SSM block.""" - is_ssm = block_ordinal == BAD_BLOCK_ORDINAL - assert (life_cycle == self.manager._life_cycles.ssm_life_cycle_id) == is_ssm - return ( - self._ssm_blocks[beam_index][life_cycle] - if is_ssm - else self._blocks[block_ordinal].pages[beam_index][life_cycle] - ) - - def _block( - self, block_ordinal: BlockOrdinal, beam_index: BeamIndex - ) -> TypedIndexList[LifeCycleId, BlockPage]: - """Return the life-cycle page list for an attention block or the SSM block.""" - is_ssm = block_ordinal == BAD_BLOCK_ORDINAL - return ( - self._ssm_blocks[beam_index] - if is_ssm - else self._blocks[block_ordinal].pages[beam_index] - ) - - def _copy_page_to_tree_block( - self, - tree_block: Block, - lc_idx: LifeCycleId, - src_page: Page, - num_tokens_in_block: int, - ) -> CommittedPage | None: - if not tree_block.can_replace_page(lc_idx, num_tokens_in_block): - return tree_block.get_page(lc_idx) - - storage = self.manager._storage - pg_idx = storage.get_pool_group_index(lc_idx) - for lvl in typed_range(src_page.cache_level, storage.num_cache_levels): - try: - new_slot = storage.new_slots_for_pool_group(lvl, pg_idx, 1)[0] - except OutOfPagesError: - continue - cuda_stream = self.cuda_stream - new_slot.ready_event.wait_in_stream(cuda_stream) - slot_size = storage.slot_size(pg_idx) - for p in typed_range(storage.num_pools(pg_idx)): - dst = storage.slot_address(lvl, pg_idx, new_slot.slot_id, p) - src = storage.slot_address(src_page.cache_level, pg_idx, src_page.slot_id, p) - batched_copy( - storage.cache_tiers[lvl], - storage.cache_tiers[src_page.cache_level], - slot_size[p], - [CopyTask(dst, src)], - cuda_stream, - ) - new_slot.ready_event = CachedCudaEvent(cuda_stream) - priority = self._get_priority(tree_block.ordinal, self.manager._life_cycles[lc_idx]) - committed = CommittedPage( - storage, tree_block, lc_idx, lvl, new_slot, num_tokens_in_block, priority - ) - # Drops the superseded page, deferred until the copy is issued: an - # OutOfPagesError above must not destroy a usable shorter snapshot. - tree_block.replace_page(lc_idx, committed) - storage.schedule_for_eviction(committed) - return committed - return None - - def _snapshot_ssm_to_tree_block( - self, tree_block: Block, ssm_lc_id: LifeCycleId, num_tokens: int, move: bool = False - ) -> None: - """Snapshot live SSM state to tree_block for the given committed token count.""" - tokens_per_block = self.tokens_per_block - num_tokens_in_block = num_tokens - tree_block.ordinal * tokens_per_block - assert 0 < num_tokens_in_block <= tokens_per_block - if tree_block.page_coverage(ssm_lc_id) >= num_tokens_in_block: - return - - ssm_block = self._ssm_blocks[DEFAULT_BEAM_INDEX] - ssm_lock = expect_type(_SharedPageLock, ssm_block[ssm_lc_id]) - src_page = ssm_lock.page - if move: - src_page = expect_type(UncommittedPage, ssm_lock.unlock()) - ssm_block[ssm_lc_id] = None - # convert_to_committed() reserves the slot, dropping any shorter snapshot. - committed = src_page.convert_to_committed( - tree_block, self.finish_event, num_tokens_in_block - ) - storage = self.manager._storage - storage.schedule_for_eviction(committed) - return - - self._copy_page_to_tree_block(tree_block, ssm_lc_id, src_page, num_tokens_in_block) - - def _snapshot_partial_block_to_tree(self, ordinal: BlockOrdinal, commit_ssm: bool) -> None: - tokens_per_block = self.tokens_per_block - start = ordinal * tokens_per_block - tokens = self._committed_tokens[start : start + tokens_per_block] - num_tokens = len(tokens) - assert 0 < num_tokens < tokens_per_block - prev: RootBlock | Block - if ordinal == 0: - prev = self.manager._radix_tree.add_or_get_existing(self._reuse_scope) - else: - prev = self._get_tree_block(BlockOrdinal(ordinal - 1)) - try: - tree_block = Block(tokens, prev) - is_new = True - except UselessBlockError as e: - tree_block = e.block - assert tree_block.tokens[:num_tokens] == tokens - is_new = False - - beam_idx = DEFAULT_BEAM_INDEX - beam_block = self._blocks[ordinal].pages[beam_idx] - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - # `tree_block` may be a longer existing sibling that covers these tokens. Attaching - # the partial attention pages to it is still correct because each page records the - # token span it covers, and prefix matching honours that span. - attached_lcs = list[LifeCycleId]() - for lc_idx, _ in self.manager._life_cycles.attention_life_cycles(): - holder = beam_block[lc_idx] - if holder is None or tree_block.page_coverage(lc_idx) >= num_tokens: - continue - if ( - self._copy_page_to_tree_block(tree_block, lc_idx, holder.page, num_tokens) - is not None - ): - attached_lcs.append(lc_idx) - if commit_ssm: - assert ssm_lc_id is not None - self._snapshot_ssm_to_tree_block(tree_block, ssm_lc_id, start + num_tokens) - event_manager = self.manager.event_manager - if event_manager is not None: - if is_new: - event_manager.add_stored_block_event_from_block(tree_block) - else: - # The block was already announced, so report just the life cycles this - # snapshot added. The event manager itself drops the ones whose page does - # not span the whole block: the payload carries the block's full token - # list and cannot express a shorter valid prefix. - for lc_idx in attached_lcs: - event_manager.add_stored_life_cycle_event_from_block(tree_block, int(lc_idx)) - - def _commit_block( - self, - ordinal: BlockOrdinal, - is_last: bool, - commit_ssm: bool = False, - move_ssm: bool = False, - ) -> None: - """ - Commit one sequence block into the radix tree. - - `is_last` controls the commit-state transition. It is set when this is - the last block committed before a virtual or user stop-committing - transition; it permits a partial final block and runs stop cleanup. - - `commit_ssm` snapshots the current SSM state for this block. `move_ssm` - controls SSM page ownership for that snapshot: it requires the caller to - know there will be no later data writes to this _KVCache's memory pages, - because the live SSM page may be moved into the radix tree instead of - copied. `move_ssm` is independent from `is_last`. - """ - assert self._commit_state == self.CommitState.ALLOWED - assert ( - ordinal == self._num_committed_blocks or self._commit_state != self.CommitState.ALLOWED - ) - seq_block = self._blocks[ordinal] - assert typed_len(seq_block.pages) == 1, "Must have 1 beam only" - beam_idx = DEFAULT_BEAM_INDEX - beam_block = seq_block.pages[beam_idx] - tokens_per_block = self.tokens_per_block - start = ordinal * tokens_per_block - tokens = self._committed_tokens[start : start + tokens_per_block] - num_tokens = len(tokens) - is_full = num_tokens == tokens_per_block - if not is_last and not is_full: - raise LogicError("Cannot commit block that is not full except last block") - prev: RootBlock | Block - if ordinal == 0: - prev = self.manager._radix_tree.add_or_get_existing(self._reuse_scope) - else: - prev = self._get_tree_block(BlockOrdinal(ordinal - 1)) - try: - tree_block = Block(tokens, prev) - is_new = True - except UselessBlockError as e: - tree_block = e.block - assert tree_block.tokens[:num_tokens] == tokens - is_new = False - - assert tree_block.tokens_per_block == tokens_per_block - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - did_commit = False - if is_new: - # We are the only writer to padding. Other _KVCache reusing it should make copies. - skip_lcs = {ssm_lc_id} if ssm_lc_id is not None else None - uncommitted_pages = self._take_uncommitted_page(ordinal, beam_idx, skip_lcs) - # convert uncommitted pages to committed pages and create a new block in the radix tree. - for lc, (page, locked) in typed_enumerate(uncommitted_pages): - if page is None: - continue - p = page.convert_to_committed(tree_block, self.finish_event, num_tokens) - # The page comes from uncommitted page of self, so safe to skip wait. - beam_block[lc] = ( - p.lock(self, beam_idx, ordinal, lc, skip_wait=True) if locked else p.hold() - ) - seq_block.tree_block = tree_block - assert self._get_tree_block(ordinal) is tree_block - self._num_committed_blocks = BlockOrdinal(ordinal + 1) - event_manager = self.manager.event_manager - if event_manager is not None: - event_manager.add_stored_block_event_from_block(tree_block) - did_commit = True - elif tree_block.is_full and self.manager.allow_seq_rebasing and is_full: - # Happens when a concurrent request committed the same tokens before us. - # Try to replace our pages with pages from the existing block to save memory. - reuse_list = list[tuple[LifeCycleId, CommittedPage]]() - for lc in typed_range(typed_len(beam_block)): - if lc == ssm_lc_id: - continue # SSM pages are not rebased - if beam_block[lc] is None: - continue - # A page covering fewer tokens than this block spans (moved in from a - # shorter sibling) is NOT a substitute for our own full page: adopting it - # would feed uninitialized KV for the uncovered tail into a live request. - existing_page = tree_block.get_page(lc) - if existing_page is not None and existing_page.num_tokens_in_block < num_tokens: - existing_page = None - locked = isinstance(beam_block[lc], _SharedPageLock) - if existing_page is None: - # The reusable page is gone. We put our own page into the tree block. - # Keep this a single expression: a local holding the lock/holder would - # keep it alive past `beam_block[lc] = None`, and convert_to_committed() - # requires the page to be droppable by then. - page = cast( - UncommittedPage, - cast("_SharedPageLock | _PageHolder", beam_block[lc]).page, - ) - beam_block[lc] = None - p = page.convert_to_committed(tree_block, self.finish_event, num_tokens) - event_manager = self.manager.event_manager - if event_manager is not None: - event_manager.add_stored_life_cycle_event_from_block(tree_block, int(lc)) - # The page comes from uncommitted page of self, so safe to skip wait. - beam_block[lc] = ( - p.lock(self, beam_idx, ordinal, lc, skip_wait=True) if locked else p.hold() - ) - else: - if locked: - beam_block[lc] = cast(_SharedPageLock, beam_block[lc]).holder - reuse_list.append((lc, existing_page)) - locks = batched_lock_to_gpu( - self, - [BatchedLockTarget(p, beam_idx, ordinal, lc) for lc, p in reuse_list], - self._record_migrated_slots, - self._record_dropped_pages, - ) - for (lc, _), lock in zip(reuse_list, locks): - beam_block[lc] = lock - seq_block.tree_block = tree_block - assert self._get_tree_block(ordinal) is tree_block - self._num_committed_blocks = BlockOrdinal(ordinal + 1) - did_commit = True - else: - # We can't commit and can't reuse existing block. Just stop committing. - self._commit_state = self.CommitState.VIRTUAL_STOP - - if did_commit and commit_ssm: - assert ssm_lc_id is not None - self._snapshot_ssm_to_tree_block( - tree_block, ssm_lc_id, start + num_tokens, move=move_ssm - ) - - if seq_block.is_committed: - for lc_idx, lc in self.manager._life_cycles.attention_life_cycles(): - stale_range = _KVCache._get_stale_range(tokens_per_block, self.history_length, lc) - if ordinal in stale_range: - for beam_block in seq_block.pages: - beam_block[lc_idx] = None - - if is_last or self._commit_state == self.CommitState.VIRTUAL_STOP: - self._commit_state = self.CommitState.USER_STOP - self._on_stop_committing() - - def _on_stop_committing(self) -> None: - # If there are stale held uncommitted pages, release them. - # @TODO: add test for this. - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - for lc_idx, lc in self.manager._life_cycles.items(): - if lc_idx == ssm_lc_id: - continue # SSM pages live in _ssm_blocks, not in _blocks - start, end = _KVCache._get_stale_range(self.tokens_per_block, self.history_length, lc) - start = max(start, self._num_committed_blocks) - for ordinal in typed_range(start, end): - block = self._blocks[ordinal] - assert not block.is_committed - for beam_block in block.pages: - if beam_block[lc_idx] is None: - # Nothing to release: scratch block, commit_min_snapshot early - # release, or a stale block created unallocated by resize(). - continue - assert isinstance(beam_block[lc_idx], _PageHolder) - beam_block[lc_idx] = None - assert NDEBUG or self._check_sanity() - - def _unlock_stale_blocks( - self, new_history_length: int - ) -> list[tuple[BlockOrdinal, BeamIndex, LifeCycleId, _PageHolder]]: - "For SWA layers, unlock out-of-window blocks." - if new_history_length == self.history_length: - return [] - with self._record_event(): - ret = list[tuple[BlockOrdinal, BeamIndex, LifeCycleId, _PageHolder]]() - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - for lc_idx, lc in self.manager._life_cycles.items(): - if lc_idx == ssm_lc_id: - continue # SSM pages live in _ssm_blocks, not in _blocks - if isinstance(lc, AttnLifeCycle) and lc.window_size is None: - continue - _, old_end = _KVCache._get_stale_range( - self.tokens_per_block, self.history_length, lc - ) - new_beg, new_end = _KVCache._get_stale_range( - self.tokens_per_block, new_history_length, lc - ) - for ordinal in typed_range( - max(old_end, new_beg), min(typed_len(self._blocks), new_end) - ): - block = self._blocks[ordinal] - is_committed = block.is_committed - hold_for_commit = ( - not self.manager.commit_min_snapshot - and not is_committed - and self._commit_state == self.CommitState.ALLOWED - ) - for beam_idx, beam_block in typed_enumerate(block.pages): - if beam_block[lc_idx] is None: - # No page to unlock: scratch block, commit_min_snapshot early - # release, or a stale block created unallocated by resize(). - continue - holder = expect_type(_SharedPageLock, beam_block[lc_idx]).holder - ret.append((ordinal, beam_idx, lc_idx, holder)) - beam_block[lc_idx] = holder if hold_for_commit else None - # Scratch slot lifetime is handled by resize() after target scratch ranges are recomputed. - return ret - - def _lock_held_blocks( - self, backup_holders: list[tuple[BlockOrdinal, BeamIndex, LifeCycleId, _PageHolder]] - ): - "Revert _unlock_unused_blocks() by locking the held blocks." - locks = batched_lock_to_gpu( - self, - [ - BatchedLockTarget(holder.page, beam_idx, ordinal, lc) - for ordinal, beam_idx, lc, holder in backup_holders - ], - self._record_migrated_slots, - self._record_dropped_pages, - ) - for lock in locks: - user = lock._user - self._block(user.ordinal, user.beam_index)[user.life_cycle] = lock - - class DeltaScratchSlots(NamedTuple): - excess: TypedIndexList[LifeCycleId, list[ScratchSlotLock]] - delta_cnt: TypedIndexList[LifeCycleId, int] - scratch_ranges: TypedIndexList[LifeCycleId, HalfOpenRange[BlockOrdinal]] - - def _take_excess_scratch_slots(self, capacity: int, history_length: int) -> DeltaScratchSlots: - """ - Calculate scratch slot requirements and extract excess scratch slots. - - Returns: - excess_scratch_slots: List of ScratchSlotLocks taken from `self._scratch_slots`. - additional_scratch_slots: Number of extra slots needed per lifecycle (we have deficit). - scratch_ranges: The scratch ranges per lifecycle for the new capacity/history_length. - """ - num_life_cycles = self.manager._life_cycles.size - excess = make_typed(lambda _: list[ScratchSlotLock](), num_life_cycles) - delta_cnt = filled_list(0, num_life_cycles) - scratch_ranges = make_typed( - lambda _: HalfOpenRange[BlockOrdinal](BlockOrdinal(0), BlockOrdinal(0)), num_life_cycles - ) - - for lc_idx, lc in self.manager._life_cycles.items(): - scratch_range = self._get_scratch_range(lc, history_length, capacity) - scratch_ranges[lc_idx] = scratch_range - num_scratch_blocks = len(scratch_ranges[lc_idx]) - frac_max = self._storage._slot_util_frac_max[lc_idx] - needed_slots = math.ceil(num_scratch_blocks * frac_max) - existing_slots = len(self._scratch_slots[lc_idx]) - delta = needed_slots - existing_slots - delta_cnt[lc_idx] = delta - - if delta < 0: - for _ in range(-delta): - lock = self._scratch_slots[lc_idx].pop() - excess[lc_idx].append(lock) - - return self.DeltaScratchSlots(excess, delta_cnt, scratch_ranges) - - def _recover_excess_scratch_slots( - self, excess_scratch_slots: TypedIndexList[LifeCycleId, list[ScratchSlotLock]] - ) -> None: - for lc_idx, locks in typed_enumerate(excess_scratch_slots): - self._scratch_slots[lc_idx].extend(locks) - locks.clear() - - @property - def _storage(self) -> StorageManager: - return self.manager._storage - - @staticmethod - def _to_block_ordinal(tokens_per_block: int, token_ordinal: int) -> BlockOrdinal: - return BlockOrdinal(token_ordinal // tokens_per_block) - - def _get_tree_block(self, ordinal: BlockOrdinal) -> Block: - assert self._blocks[ordinal].is_committed - ret = unwrap_optional(self._blocks[ordinal].tree_block) - if not NDEBUG: - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - for lc, b in typed_enumerate(self._block(ordinal, DEFAULT_BEAM_INDEX)): - if lc == ssm_lc_id: - assert b is None # SSM pages live in _ssm_blocks - elif b is not None: - # b.page.block() may differ from `ret`: a page can be moved to a longer - # sibling block, or replaced there by one with larger token coverage. - assert isinstance(b.page, CommittedPage) - return ret - - def _take_uncommitted_page( - self, - ordinal: BlockOrdinal, - beam_idx: BeamIndex, - skip_lcs: set[LifeCycleId] | None = None, - ) -> TypedIndexList[LifeCycleId, tuple[UncommittedPage | None, bool]]: - """ - Take ownership of the uncommitted pages, together with bool flag indicating if it was locked. - And reset holders to None. SSM life cycles in skip_lcs are left in place. - """ - holders = self._block(ordinal, beam_idx) - num_life_cycles = self.manager._life_cycles.size - ret: TypedIndexList[LifeCycleId, tuple[UncommittedPage | None, bool]] = filled_list( - (None, False), num_life_cycles - ) - for lc, holder in typed_enumerate(holders): - if holder is None: - continue - if skip_lcs and lc in skip_lcs: - continue - assert isinstance(holder.page, UncommittedPage) - locked = isinstance(holder, _SharedPageLock) - ret[lc] = (holder.page, locked) - # When using debugpy with breakpoints on exceptions enabled, the lock/holder is not GC'ed even - # after return from this function. That will likely lead to assertion failures later. - holders[lc] = None - return ret - - def _check_sanity(self) -> bool: - is_closed = self.status == self.Status.CLOSED - if is_closed: - return self.num_blocks == 0 - assert self.num_committed_tokens <= self.history_length <= self.capacity - assert self.num_blocks == div_up(self.capacity, self.tokens_per_block) - - def get_range(lc: LifeCycle): - return _KVCache._get_stale_range(self.tokens_per_block, self.history_length, lc) - - stale_ranges = typed_map(self.manager._life_cycles.get(), get_range) - num_life_cycles = self.manager._life_cycles.size - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - for ordinal, block in typed_enumerate(self._blocks): - is_committed = self._never_resumed or ordinal < self._num_committed_blocks - assert is_committed == block.is_committed - for beam_block in block.pages: - assert typed_len(beam_block) == num_life_cycles - for lc in typed_range(num_life_cycles): - holder = beam_block[lc] - if lc == ssm_lc_id: - # SSM pages live in _ssm_blocks, not in _blocks - assert holder is None - continue - start, end = stale_ranges[lc] - lc_obj = self.manager._life_cycles[lc] - if start <= ordinal < end: - if is_committed or self._commit_state != self.CommitState.ALLOWED: - assert holder is None - elif ordinal in self._get_scratch_range(lc_obj, 0): - # It is uncertain whether this block should hold a page: it may - # have been scratch-allocated by an earlier chunk, but the prior - # history length and capacity are not retained. - pass - else: - # For the decoder-side disagg case, for the first step, we will skip the - # out-of-window blocks. - assert isinstance(holder, _PageHolder) or ( - holder is None - and (not self._committed_tokens or self.manager.commit_min_snapshot) - ) - else: - # Scratch blocks have None pages but valid base_page_indices - sr = self._get_scratch_range(lc_obj) - is_scratch = ordinal in sr - if is_scratch: - assert holder is None - else: - assert isinstance( - holder, (_SharedPageLock if self.is_active else _PageHolder) - ) - if holder is not None: - assert is_committed == isinstance(holder.page, CommittedPage) - # Check SSM blocks - if ssm_lc_id is not None: - for beam_block in self._ssm_blocks: - holder = beam_block[ssm_lc_id] - if holder is not None: - if self._never_resumed: - # Deferred copy: SSM holds CommittedPage from matched snapshot - assert isinstance(holder, _PageHolder) - assert isinstance(holder.page, CommittedPage) - else: - assert isinstance(holder, _SharedPageLock) - assert isinstance(holder.page, UncommittedPage) - return True - - def _get_scratch_range( - self, - life_cycle: LifeCycle, - history_length_override: int | None = None, - capacity_override: int | None = None, - ) -> HalfOpenRange[BlockOrdinal]: - """ - Range of blocks that should use scratch (shared) slots during SWA prefill. - - Scratch = stale_at_capacity ∩ input_blocks, where: - - stale_at_capacity: blocks out-of-window when all non-rewindable capacity tokens - become history. - - input_blocks: [div_up(history_length, tpb), div_up(capacity, tpb)) — new blocks - for the current chunk. Blocks before this range already contain real KV data - from previous chunks and must not be overwritten. - - The configured max_rewind_len excludes a speculative tail from scratch reuse. - """ - if not self.enable_swa_scratch_reuse: - return HalfOpenRange(BlockOrdinal(0), BlockOrdinal(0)) - history_length = value_or(history_length_override, self.history_length) - capacity = value_or(capacity_override, self.capacity) - max_rewind_len = self._swa_scratch_max_rewind_len() - return compute_scratch_range( - life_cycle, - history_length, - capacity, - self.tokens_per_block, - max_rewind_len, - ) - - def _would_use_swa_scratch_blocks(self) -> bool: - max_rewind_len = self._swa_scratch_max_rewind_len() - return any( - compute_scratch_range( - lc, - self.history_length, - self.capacity, - self.tokens_per_block, - max_rewind_len, - ) - for lc in self.manager._life_cycles - ) - - def _swa_scratch_max_rewind_len(self) -> int: - return unwrap_optional(self.manager.init_config.swa_scratch_reuse).max_rewind_len - - @staticmethod - def _get_stale_range( - tokens_per_block: int, - history_length: int, - life_cycle: LifeCycle, - ) -> HalfOpenRange[BlockOrdinal]: - """ - Range of the stale blocks. Stale blocks are no longer needed for inference. Stale pages should be - held if we may commit them later, or droppable otherwise. - """ - beg, end = life_cycle.get_stale_range(history_length, tokens_per_block) - return HalfOpenRange(BlockOrdinal(beg), BlockOrdinal(end)) - - def _get_matched_tokens(self, match: ReuseMatch) -> list[TokenIdExt]: - ret: list[TokenIdExt] = [] - remaining = match.num_tokens - for block in match.blocks: - assert remaining > 0 - num_block_tokens = min(remaining, len(block.tokens)) - ret.extend(block.tokens[:num_block_tokens]) - remaining -= num_block_tokens - assert remaining == 0 - return ret - - def _finalize_cached_tokens_by_level( - self, - num_tokens: int, - attention_levels: TypedIndexList[BlockOrdinal, CacheLevel], - ssm_level: CacheLevel | None, - ) -> None: - """Turn the per-block cache levels collected while holding the matched pages into logical - token counts. - - Cache levels are ordered hottest first, so the coldest level backing a block is simply the - largest one. Attention pages cover their block's matched token span; the final SSM - checkpoint summarizes the entire recurrent prefix, so it is merged into every block. - """ - counts = filled_list(0, self.manager._storage.num_cache_levels) - self._cached_tokens_by_level = counts - self._last_cached_token_level = None - if num_tokens == 0: - return - - tokens_per_block = self.manager.tokens_per_block - last_cached_token_level: CacheLevel | None = None - for ordinal, attention_level in enumerate(attention_levels): - block_start = ordinal * tokens_per_block - block_end = min(num_tokens, block_start + tokens_per_block) - if block_start >= block_end: - break - source_level = ( - max(attention_level, ssm_level) if ssm_level is not None else attention_level - ) - last_cached_token_level = source_level - counts[source_level] += block_end - block_start - - self._last_cached_token_level = last_cached_token_level - assert NDEBUG or sum(counts) == num_tokens - assert NDEBUG or last_cached_token_level is not None - - # Stage the attribution alongside the reuse counters derived from the same walk, so it is - # committed (or discarded) with them instead of being pushed back in through a public setter. - if ( - self._should_record_manager_stats() - and self._pending_stats.record_cached_tokens_by_level(counts) - ): - self.manager.mark_stats_dirty(self.id) - - def drop_cached_token_attribution(self) -> None: - """Drop this sequence's staged cached-token attribution so its reuse match is not reported - as a local cache hit. - - For traffic that is not user-visible reuse at all (KV cache size estimation), and for a - disaggregated-serving generation request whose whole matched prefix the incoming transfer - overwrites (a recurrent state summarizes and replaces the entire local slot). - """ - self._pending_stats.cached_tokens_by_level = [] - self._refresh_stats_dirty_state() - - def drop_partial_block_cached_token_attribution(self) -> None: - """Drop only the trailing partial block from the staged attribution. - - A disaggregated-serving transfer overwrites the incomplete tail block of the local match - while complete blocks survive, so only the tail stops counting as a local hit. - """ - partial_tokens = sum(self._cached_tokens_by_level) % self.manager.tokens_per_block - if partial_tokens == 0: - return - # A partial cached prefix always has a final block, so it always has a level to discount. - assert NDEBUG or self._last_cached_token_level is not None - if self._last_cached_token_level is None: - return - # Admission can retry the same cache after resume/resize fails. Cap against the initial - # match so repeated calls preserve full blocks on the same source level. - full_tokens = self._cached_tokens_by_level[self._last_cached_token_level] - partial_tokens - self._pending_stats.limit_cached_tokens_by_level(self._last_cached_token_level, full_tokens) - self._refresh_stats_dirty_state() - - def _setup_for_reuse(self, match: ReuseMatch) -> None: - manager = self.manager - matched = match.blocks - tokens_per_block = manager.tokens_per_block - num_tokens = match.num_tokens - life_cycles = manager._life_cycles - ssm_lc_id = life_cycles.ssm_life_cycle_id - self._committed_tokens = self._get_matched_tokens(match) - self._history_length = num_tokens - self._capacity = num_tokens - full_reused_end = BlockOrdinal(num_tokens // tokens_per_block) - has_partial_match = num_tokens % tokens_per_block != 0 - # fill self._blocks - self._blocks = to_typed( - BlockOrdinalT, - [ - SeqBlock( - make_typed( - lambda _: filled_list(cast(BlockPage, None), life_cycles.size), - self.beam_width, - ), - block, - ) - for block in matched - ], - ) - - beam_idx = DEFAULT_BEAM_INDEX - - record_manager_stats = self._should_record_manager_stats() - record_request_stats = self._should_record_request_stats() - record_shared_stats = record_manager_stats or record_request_stats - storage = manager._storage - num_cache_levels = storage.num_cache_levels - # Source-level attribution for the reused tokens, merged across attention life cycles at - # block granularity. It is observed in this same walk rather than in a separate pre-pass: - # hold() only takes a page holder and unschedules eviction, it never migrates data, so a - # page still reports the level it was matched on. Levels are ordered hottest first, so - # merging takes the max. - attention_levels = filled_list(CacheLevel(0), BlockOrdinal(len(matched))) - for lc_idx, lc in life_cycles.items(): - if lc_idx == ssm_lc_id: - continue # SSM is handled separately below - stale_start, stale_end = _KVCache._get_stale_range(tokens_per_block, num_tokens, lc) - is_attention = isinstance(lc, AttnLifeCycle) - full_reused_blocks = 0 - partial_reused_blocks = 0 - # Indexed by CacheLevel, so entry i is the i-th configured tier rather than a fixed - # gpu/host/disk bucket; a deployment with two GPU levels gets two distinct entries. - by_level = ( - ReusedBlocksByLevel( - full=filled_list(0, num_cache_levels), - partial=filled_list(0, num_cache_levels), - ) - if record_shared_stats and is_attention - else None - ) - life_cycle_levels: TypedIndexList[BlockOrdinal, CacheLevel | None] = ( - filled_list(None, BlockOrdinal(len(matched))) if is_attention else [] - ) - for ordinal in chain( - typed_range(stale_start), typed_range(stale_end, BlockOrdinal(len(matched))) - ): - block = self._block(ordinal, beam_idx) - page = unwrap_optional(matched[ordinal].get_page(lc_idx)) - level = page.cache_level - holder = page.hold() - # For partial blocks (last block, not full), we defer the copy to first resume(). - # Just store the holder of the original committed page for now. - block[lc_idx] = holder - if not is_attention: - continue - life_cycle_levels[ordinal] = level - if record_shared_stats: - assert by_level is not None - if ordinal < full_reused_end: - full_reused_blocks += 1 - by_level.full[level] += 1 - elif ( - has_partial_match - and ordinal == full_reused_end - and self._has_reuse_source(holder) - ): - partial_reused_blocks = 1 - by_level.partial[level] += 1 - if is_attention: - # For SWA, only sink and live-window pages at the matched endpoint are - # materialized. Stale spans inherit the next live page's level, because that later - # anchor is what enables those logical tokens to be skipped. - next_anchor_level: CacheLevel | None = None - for ordinal in reversed(range(len(matched))): - level = life_cycle_levels[ordinal] - if level is not None: - next_anchor_level = level - else: - level = next_anchor_level - if level is not None and level > attention_levels[ordinal]: - attention_levels[ordinal] = level - if record_shared_stats and is_attention: - changed = self._pending_stats.record_reuse( - lc_idx, - full_reused_blocks=full_reused_blocks, - partial_reused_blocks=partial_reused_blocks, - by_level=by_level, - record_manager_stats=record_manager_stats, - record_request_stats=record_request_stats, - ) - if changed: - self.manager.mark_stats_dirty(self.id) - # SSM reuse: hold the snapshot from the last matched block. Copy is deferred to first - # resume(). Reuse onboards only this final recurrent checkpoint, and it summarizes the - # entire matched prefix, so its source tier applies to every logical reused token; older - # checkpoints are traversal history, not inputs. - ssm_level: CacheLevel | None = None - if ssm_lc_id is not None and matched: - snapshot_block = matched[-1] - snapshot_page = snapshot_block.get_page(ssm_lc_id) - assert snapshot_page is not None, ( - "Last matched block must have SSM snapshot after truncation" - ) - ssm_level = snapshot_page.cache_level - snapshot_holder = snapshot_page.hold() - self._ssm_blocks[DEFAULT_BEAM_INDEX][ssm_lc_id] = snapshot_holder - self._finalize_cached_tokens_by_level(num_tokens, attention_levels, ssm_level) - if record_manager_stats and ssm_lc_id is not None: - changed = self._pending_stats.record_ssm_snapshot_lookup( - ssm_lc_id, - lookup_tokens=match.num_lookup_tokens, - reused_tokens=num_tokens, - tokens_per_block=tokens_per_block, - ) - if changed: - self.manager.mark_stats_dirty(self.id) - self._num_committed_blocks = BlockOrdinal(len(self._committed_tokens) // tokens_per_block) - for beam_indices in self._base_page_indices: - for indices in beam_indices: - if type(indices) is array.array: - indices.extend([BAD_PAGE_INDEX] * (self.num_blocks - len(indices))) - else: - assert len(indices) >= self.num_blocks - - def _free_scratch_slots(self) -> None: - """Free all scratch slots back to the storage manager.""" - for lc in typed_range(self.manager._life_cycles.size): - for lock in self._scratch_slots[lc]: - lock.unlock() - self._scratch_slots[lc].clear() - - def _clear_blocks(self) -> None: - # drop the last block first - while self._blocks: - self._blocks.pop() - self._free_scratch_slots() - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - if ssm_lc_id is not None: - for beam_block in self._ssm_blocks: - beam_block[ssm_lc_id] = None - - @contextmanager - def _record_event(self) -> Iterator[None]: - assert self._finish_event is None - if self._cuda_stream is None: - # Cache was never resumed — no GPU work was performed, - # so no CUDA event synchronization is needed. Blocks - # only contain _PageHolders (not _SharedPageLocks) and - # their destructors do not read finish_event. - yield - return - self._finish_event = CachedCudaEvent(self.cuda_stream) - try: - yield - finally: - self._finish_event = None - - def _update_base_page_index( - self, beam_idx: BeamIndex, ordinal: BlockOrdinal, lc: LifeCycleId, page_index: PageIndex - ) -> PageIndex: - if ordinal == BAD_BLOCK_ORDINAL: - return PageIndex(BAD_PAGE_INDEX) - indices = self._base_page_indices[beam_idx][lc] - old = PageIndex(indices[ordinal]) - indices[ordinal] = page_index - return old - - def _get_base_page_indices_ref( - self, lc: LifeCycleId, beam_id: BeamIndex = DEFAULT_BEAM_INDEX - ) -> Iterator[int | None]: - assert beam_id < self.beam_width - assert self.is_active - return self.get_aggregated_page_indices(lc, beam_id) - - def _shortcut_set_capacity(self, capacity: int) -> bool: - "Shortcut for cases without side effects. Just for better performance." - tokens_per_block = self.tokens_per_block - if div_up(capacity, tokens_per_block) == div_up(self._capacity, tokens_per_block): - self._capacity = capacity - return True - return False - - def _shortcut_set_history_length(self, history_length: int) -> bool: - "Shortcut for cases without side effects. Just for better performance." - tokens_per_block = self.tokens_per_block - - def no_side_effect(lc: LifeCycle) -> bool: - if type(lc) is SsmLifeCycle: - # history_length change does not impact blocks at all. - return True - assert type(lc) is AttnLifeCycle - window = lc.window_size - return window is None or lc.get_stale_range( - history_length, tokens_per_block - ) == lc.get_stale_range(self.history_length, tokens_per_block) - - if all(no_side_effect(lc) for lc in self.manager._life_cycles): - self._history_length = history_length - return True - return False diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py deleted file mode 100644 index d75316d13fea..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ /dev/null @@ -1,1165 +0,0 @@ -# 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. - -import time -import warnings -from collections import defaultdict -from collections.abc import Callable, Iterable, Sequence -from copy import deepcopy -from dataclasses import dataclass -from typing import TYPE_CHECKING, Iterator, cast - -from .. import rawref -from .._block_radix_tree import Block, BlockRadixTree, ReuseMatch, ReuseScope, RootBlock -from .._common import ( - BAD_PAGE_INDEX, - GPU_LEVEL, - NDEBUG, - PRIORITY_DEFAULT, - BlockOrdinal, - CacheLevel, - CacheTier, - LayerId, - MemAddress, - PageIndexMode, - PageStatus, - Priority, - TokenIdExt, -) -from .._config import DataRole, KVCacheManagerConfig -from .._exceptions import LogicError -from .._life_cycle_registry import LayerGroupId, LifeCycle, LifeCycleId, LifeCycleRegistry -from .._page import Page, _PageHolder -from .._stats import ( - CountsByLevel, - KVCacheIterationStatsDelta, - KVCacheStatsDelta, - ReusedBlocksByLevel, - SsmSnapshotIterationStatsDelta, - add_counts_by_level, -) -from .._storage._config import BufferId, SlotDesc, create_storage_config -from .._storage._core import PoolGroupIndex, PoolIndex, SlotId -from .._storage_manager import StorageManager, StorageStatistics -from .._utils import ( - HalfOpenRange, - HomoTuple, - TypedIndexList, - div_up, - exact_div, - filled_list, - init_cuda_once, - make_typed, - typed_enumerate, - typed_range, - unwrap_rawref, -) -from ._kv_cache import _KVCache -from ._moving_average import MovingAverage - -if TYPE_CHECKING: - from .._event_manager import KVCacheEventManager - - -@dataclass(slots=True, frozen=True) -class PoolDesc: - pool_index: PoolIndex - base_address: MemAddress - slot_bytes: int - - -@dataclass(slots=True, frozen=True) -class PoolGroupDesc: - pool_group_index: PoolGroupIndex - num_slots: int - slot_desc: SlotDesc - pools: TypedIndexList[PoolIndex, PoolDesc] - - -@dataclass(slots=True, frozen=True) -class Range: - start: int - end: int - - def __add__(self, offset: int) -> "Range": - return Range(self.start + offset, self.end + offset) - - def __radd__(self, offset: int) -> "Range": - return self + offset - - -@dataclass(slots=True, frozen=True) -class ExpandedBuffer: - id: BufferId - expansion: int # expansion factor of page due to heterogeneous tokens_per_block - - -@dataclass(slots=True, frozen=True) -class AggregatedPageDesc: - """ - The data you need would be in the following byte ranges: - (base + stride * i + Range(0, size) for i in aggregated_page_indices) - """ - - base: MemAddress - size: int - stride: int - layer_group_id: LayerGroupId - buffers: Sequence[ExpandedBuffer] - - -@dataclass(slots=True, frozen=True) -class ScratchDesc: - """Scratch metadata for one layer group of one sequence. - - Scratch blocks are blocks whose KV data is ephemeral (only needed during one - step's attention). Their pages are stored in shared coalesced slots rather than - per-block slots. - """ - - range: HalfOpenRange[BlockOrdinal] # block ordinal range [beg, end) - slot_ids: Sequence[int] # scratch slot IDs, length = ceil(num_scratch_blocks / scale) - - def __bool__(self) -> bool: - return bool(self.range) - - -@dataclass(slots=True, frozen=True) -class PageIndexConverter: - scale: int - expansion: int - layer_offset: int # sub-page offset within coalesced slot - scratch_pages_per_block: int = 1 - - def __call__( - self, - base_indices: Sequence[int], - index_mode: PageIndexMode | None = None, - scratch: "ScratchDesc | None" = None, - ) -> list[int]: - """ - Convert from base page indices to per-layer page indices expected by operators/kernels. - This is a reference implementation. Users are encouraged to do it with a CUDA kernel. - - When index_mode is PageIndexMode.PER_LAYER, the converted indices include the layer's - position within the coalesced slot. The caller should use the pool group base address. - - When index_mode is PageIndexMode.SHARED, the converted indices do not include any - layer offset — the caller's base pointer (from get_mem_pool_base_address) already - incorporates it. - - Args: - base_indices: Per-block base page indices (slot IDs), from get_base_page_indices(). - index_mode: Page index mode. None defaults to SHARED; must be explicit when - scratch is active (scratch requires PER_LAYER). - scratch: Optional scratch metadata from _KVCache.get_scratch_desc(). - """ - if index_mode is None: - assert not scratch, "index_mode must be provided when scratch is active" - index_mode = PageIndexMode.SHARED - - scale = self.scale - expansion = self.expansion - applied_layer_offset = self.layer_offset if index_mode == PageIndexMode.PER_LAYER else 0 - scratch_pages = self.scratch_pages_per_block - result = list[int]() - - for ordinal, base_index in enumerate(base_indices): - index: int - if scratch and ordinal in scratch.range: - # Scratch block: slot IDs come from ScratchDesc, not base_indices - block_pos = ordinal - scratch.range.beg - total_offset = block_pos * scratch_pages - slot_idx = total_offset // scale - slot_id = scratch.slot_ids[slot_idx] - offset = total_offset % scale - index = slot_id * scale + (offset + applied_layer_offset) % scale - elif base_index == BAD_PAGE_INDEX: - index = BAD_PAGE_INDEX - else: - index = base_index * scale + applied_layer_offset - for i in range(expansion): - result.append(index * expansion + i if index != BAD_PAGE_INDEX else BAD_PAGE_INDEX) - return result - - -@dataclass(slots=True, frozen=True) -class PoolGroupPeakBlockStats: - available: int - unavailable: int - evictable: int - - -class KVCacheManager: - __slots__ = ( - "_init_config", - "_life_cycles", - "_storage", - "_radix_tree", - "_living_kv_caches", - "_avg_reused_length", - "_avg_sqr_capacity", - "_avg_sqr_history_length", - "_target_ratio_list_gpu", - "_target_ratio_list_other", - "_num_created_kv_caches", - "_num_sampled_kv_caches", - "_last_adjustment_time", - "_last_update_num_sampled_kv_caches", - "_event_manager", - "_stats_enabled", - "_committed_stats", - "_iteration_stats_by_life_cycle", - "_ssm_snapshot_iteration_stats_by_life_cycle", - "_iteration_peak_num_blocks_by_cache_level", - "_dirty_stats_kv_cache_ids", - "_stats_excluded_kv_cache_ids", - "_iter_suspended_requests", - "_iter_resumed_requests", - "_iter_disk_prefetch_blocks", - "_iter_cached_tokens_by_level", - "_iter_reused_blocks_by_level", - ) - _init_config: KVCacheManagerConfig - _life_cycles: LifeCycleRegistry - _storage: StorageManager - _radix_tree: BlockRadixTree - _living_kv_caches: set[rawref.ref[_KVCache]] - # Eventually we should let the eviction controller evict associated pages together, i.e. - # when a page eviction makes other pages in the same cache level useless, it should also - # evict those pages. When we have that, we can simply decide capacity ratio based on - # memory pool utilization. For now, we use a simpler approach based on sequence length. - # But this ignores the fact that some pages are shared among multiple sequences. - _avg_reused_length: MovingAverage - # use squared because longer requests also lives for longer but we only update this on - # request closing, use squared average to compensate that. - _avg_sqr_capacity: MovingAverage - _avg_sqr_history_length: MovingAverage - _target_ratio_list_gpu: TypedIndexList[PoolGroupIndex, float] - _target_ratio_list_other: TypedIndexList[PoolGroupIndex, float] - _num_created_kv_caches: int - _num_sampled_kv_caches: int - _last_adjustment_time: float - _last_update_num_sampled_kv_caches: int - _event_manager: "KVCacheEventManager | None" - _stats_enabled: bool - _committed_stats: KVCacheStatsDelta - _iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] - _ssm_snapshot_iteration_stats_by_life_cycle: dict[LifeCycleId, SsmSnapshotIterationStatsDelta] - _iteration_peak_num_blocks_by_cache_level: TypedIndexList[ - CacheLevel, TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats] - ] - _dirty_stats_kv_cache_ids: set[int] - _stats_excluded_kv_cache_ids: set[int] - _iter_suspended_requests: int - _iter_resumed_requests: int - - def __init__( - self, - config: KVCacheManagerConfig, - event_manager: "KVCacheEventManager | None" = None, - cold_page_codec: object | None = None, - ) -> None: - if cold_page_codec is not None: - raise NotImplementedError("Cold-page codecs require the C++ KVCacheManagerV2 backend") - init_cuda_once() - config = deepcopy(config) - self._init_config = config - self._living_kv_caches = set[rawref.ref[_KVCache]]() - self._life_cycles = LifeCycleRegistry(config) - storage_config = create_storage_config(config) - storage = StorageManager( - self._life_cycles, - storage_config, - config.tokens_per_block, - config.swa_scratch_reuse, - typical_batch=config.typical_step, - constraints=config.constraints, - initial_pool_ratio=config.initial_pool_ratio, - event_manager=event_manager, - max_util_for_resume=config.max_util_for_resume, - ) - radix_tree = BlockRadixTree(self._life_cycles, config.tokens_per_block, event_manager) - self._storage = storage - self._radix_tree = radix_tree - decay = 0.9999 - self._avg_reused_length = MovingAverage(decay) - self._avg_sqr_capacity = MovingAverage(decay) - self._avg_sqr_history_length = MovingAverage(decay) - self._target_ratio_list_gpu = self._current_gpu_ratio - self._target_ratio_list_other = self._current_other_ratios - self._num_created_kv_caches = 0 - self._num_sampled_kv_caches = 0 - self._last_adjustment_time = time.monotonic() - self._last_update_num_sampled_kv_caches = 0 - self._event_manager = event_manager - self._stats_enabled = config.enable_stats - self._committed_stats = KVCacheStatsDelta() - self._iteration_stats_by_life_cycle = {} - self._ssm_snapshot_iteration_stats_by_life_cycle = {} - self._reset_iteration_peak_num_blocks() - self._dirty_stats_kv_cache_ids = set() - self._stats_excluded_kv_cache_ids = set() - self._iter_suspended_requests = 0 - self._iter_resumed_requests = 0 - self._iter_disk_prefetch_blocks = 0 - self._iter_cached_tokens_by_level = [] - self._iter_reused_blocks_by_level = {} - - def __del__(self) -> None: - try: - self.shutdown() - except LogicError as e: - warnings.warn(str(e)) - - def _check_no_living_kv_caches(self, api: str) -> None: - """Raise unless every KV cache has been closed. - - `api` names the caller so the message points at the mistake rather than at - whatever breaks later. Entries are dropped by `_KVCache.close()`, so this counts - sequences that are still open, not merely un-collected objects. - """ - if self._living_kv_caches: - raise LogicError( - f"{api} with {len(self._living_kv_caches)} KV cache(s) still open; " - "close them (or drain the engine) first" - ) - - def shutdown(self) -> None: - self._check_no_living_kv_caches("shutdown()") - # A failed constructor may leave either owner unset. Release tree pages - # before destroying storage whenever the corresponding objects exist. - radix_tree = getattr(self, "_radix_tree", None) - if radix_tree is not None: - radix_tree.clear() - storage = getattr(self, "_storage", None) - if storage is not None: - storage.destroy() - - def clear_reusable_blocks(self) -> None: - self._check_no_living_kv_caches("clear_reusable_blocks()") - self._radix_tree.clear() - - def get_mem_pool_base_address( - self, layer_id: LayerId, data_role: DataRole, index_mode: PageIndexMode | None = None - ) -> MemAddress: - """ - Get the base address of the memory pool holding pages for the given layer and data role. - - When index_mode is PageIndexMode.PER_LAYER, returns the pool group base address - (without per-layer offset), since PageIndexConverter includes the layer offset in - the converted indices. Otherwise, returns the per-layer base address (with the - layer offset baked in). - """ - storage = self._storage - attr = storage.get_buffer_attr(layer_id, data_role) - - if index_mode is None: - if self.enable_swa_scratch_reuse: - raise ValueError("index_mode must be provided when SWA scratch reuse is enabled") - index_mode = PageIndexMode.SHARED - - pg_idx = storage.get_pool_group_index(attr.life_cycle_id) - addr = storage.get_mem_pool_base_address(pg_idx, attr.pool_index) - if index_mode == PageIndexMode.SHARED: - addr = MemAddress(addr + attr.offset) - return addr - - # Currently always equals to page size. In the future, that will change when kernels support page stride. - def get_page_stride(self, layer_id: LayerId, data_role: DataRole) -> int: - attr = self._storage.get_buffer_attr(layer_id, data_role) - return exact_div(attr.size, attr.expansion) - - def get_page_index_upper_bound(self, layer_id: LayerId, data_role: DataRole) -> int: - """ - The upper bound of page indices for the given layer and data role. - Note that this is not the same as the max number of pages available for this layer and data role. - Internally, multiple buffers may share one memory pool. The purpose of this API is just in case - users want to wrap the memory pool as a tensor with known shape. - """ - storage = self._storage - lc_id = storage._layer_to_life_cycle_ids[layer_id] - pg_idx = storage.get_pool_group_index(lc_id) - pool_group = storage._levels[GPU_LEVEL].storage._pool_groups[pg_idx] - num_slots = pool_group.num_slots - attr = storage.get_buffer_attr(layer_id, data_role) - pool_idx = attr.pool_index - slot_size = pool_group.slot_size[pool_idx] - return ( - exact_div(slot_size, attr.size) * num_slots - exact_div(attr.offset, attr.size) - ) * attr.expansion - - def get_page_index_scale(self, layer_id: LayerId, data_role: DataRole) -> int: - """ - Deprecated. Use get_page_index_converter instead. - - The multiplier to convert from base page indices to page indices expected by operators/kernels. - - For layers in the same layer group, users are encouraged to share the computed page indices - between buffers of these layers, if the page index scale for these buffers are the same. - """ - storage = self._storage - attr = storage.get_buffer_attr(layer_id, data_role) - return storage._slot_to_page_indices[attr.life_cycle_id][attr.pool_index] - - def get_page_index_converter( - self, layer_id: LayerId, data_role: DataRole - ) -> PageIndexConverter: - """ - Get the converter to convert from base page indices to per-layer page indices - expected by operators/kernels. - - The returned converter is constant and usable by all kv cache instances. - """ - storage = self._storage - attr = storage.get_buffer_attr(layer_id, data_role) - layer_attr = storage.get_layer_attr(layer_id) - scale = storage._slot_to_page_indices[attr.life_cycle_id][attr.pool_index] - layer_offset = exact_div(attr.offset, attr.size) - return PageIndexConverter( - scale, attr.expansion, layer_offset, layer_attr.slot_util[attr.pool_index] - ) - - def create_kv_cache( - self, - reuse_scope: ReuseScope | None = None, - input_tokens: Sequence[TokenIdExt] | None = None, - id: int | None = None, - custom_priority_callback: Callable[[BlockOrdinal, LifeCycle], Priority] = lambda _, - __: PRIORITY_DEFAULT, - expected_prompt_length: int | None = None, - text_only: bool | None = None, - enable_request_stats: bool = False, - ) -> _KVCache: - """ - Args: - reuse_scope: Namespace to match before matching any tokens. - input_tokens: Optional initial tokens used for reuse matching. - id: Optional cache identifier. - custom_priority_callback: Takes a block index and layer sliding-window - size and returns a priority. Reused blocks are updated when the - returned priority is higher than their existing priority. - expected_prompt_length: Optional token count marking the - prefill-to-generation boundary. Once history length reaches it, - subsequent capacity growth is recorded as generation-phase - allocation statistics. Defaults to the length of ``input_tokens`` - and does not affect allocation, reuse, or correctness. - text_only: Optional per-cache override for the manager setting. ``True`` - enables digest-free fast paths and requires all tokens to be text - token IDs; ``False`` permits digest tokens but is invalid when the - manager is configured with ``text_only=True``; ``None`` inherits - the manager setting. - enable_request_stats: Whether to collect request-level allocation and - reuse counters for this cache. Manager-level global and iteration - statistics remain controlled by ``KVCacheManagerConfig.enable_stats``. - - Newly created KV cache is suspended. You need to call resume() with a cuda stream to make it active - & ready in that stream. - Returns None if suspended=False and we don't have enough resource. - This call will attempt to reuse KV cache blocks. - It's user responsibility to remove the last token from prompts if we need to re-compute the token - generated by prefill. - """ - if reuse_scope is None: - reuse_scope = ReuseScope() - assert type(reuse_scope) is ReuseScope - reuse_match = ( - self._match_reuse(reuse_scope, input_tokens) if input_tokens is not None else None - ) - if expected_prompt_length is None and input_tokens is not None: - expected_prompt_length = len(input_tokens) - return _KVCache( - self, - reuse_scope, - reuse_match, - id, - custom_priority_callback, - expected_prompt_length, - text_only, - enable_request_stats, - ) - - def _match_reuse( - self, reuse_scope: ReuseScope, input_tokens: Sequence[TokenIdExt] - ) -> ReuseMatch: - return self._radix_tree.match( - reuse_scope, - input_tokens, - self.enable_partial_match, - self.init_config.reuse_match_backoff, - ) - - def probe_reuse( - self, - reuse_scope: ReuseScope | None = None, - input_tokens: Sequence[TokenIdExt] | None = None, - ) -> int: - """ - Return the currently reusable prefix length without holding pages. - - The returned length is advisory because no page ownership is acquired. - """ - if reuse_scope is None: - reuse_scope = ReuseScope() - assert type(reuse_scope) is ReuseScope - if input_tokens is None: - input_tokens = () - return self._match_reuse(reuse_scope, input_tokens).num_tokens - - def probe_first_new_block_key( - self, - reuse_scope: ReuseScope | None = None, - input_tokens: Sequence[TokenIdExt] | None = None, - ) -> bytes | None: - """Return the first full block's key past the currently reusable prefix. - - Read-only and advisory, like ``probe_reuse``. Reuse the preceding full - block's key from the same fresh match instead of rehashing the prefix. - """ - if reuse_scope is None: - reuse_scope = ReuseScope() - assert type(reuse_scope) is ReuseScope - if input_tokens is None: - return None - match = self._match_reuse(reuse_scope, input_tokens) - block_index = match.num_tokens // self.tokens_per_block - begin = block_index * self.tokens_per_block - end = begin + self.tokens_per_block - if end > len(input_tokens): - return None - # Use the final, pruned match. Its last block can be partial and have a - # different suffix; only a full predecessor has the query's exact key. - previous_key = ( - RootBlock.make_key(reuse_scope) - if block_index == 0 - else match.blocks[block_index - 1].key - ) - return Block.make_key(previous_key, input_tokens[begin:end]) - - def resize(self, cache_level: CacheLevel, quota: int, best_efforts: bool = False) -> bool: - """ - When calling resize, all KV caches must be suspended. - If best_efforts is True, we will try to resize the quota to the largest possible value that is - still <= quota, and returns False only when we cannot resize the quota at all. - If best_efforts is False, we will resize the quota to the exact value of quota, and give up - if not possible. - For now, best_efforts=True is not yet implemented. - """ - if best_efforts: - raise NotImplementedError("Not implemented") - else: - try: - self._adjust_level(cache_level, quota) - return True - except Exception as e: - print(f"Failed to resize cache level {cache_level} to {quota}: {e}") - return False - - def get_quota(self, cache_level: CacheLevel) -> int: - return self._storage._levels[cache_level].storage.total_quota - - def _current_block_stats_by_cache_level( - self, - ) -> TypedIndexList[CacheLevel, TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats]]: - def collect( - cache_level: CacheLevel, - ) -> TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats]: - stats_by_pool_group = self._storage.get_statistics(cache_level) - return make_typed( - lambda pool_group_index: PoolGroupPeakBlockStats( - available=stats_by_pool_group[pool_group_index].available, - unavailable=stats_by_pool_group[pool_group_index].unavailable, - evictable=stats_by_pool_group[pool_group_index].evictable, - ), - self._storage.num_pool_groups, - ) - - return make_typed(collect, self._storage.num_cache_levels) - - def _reset_iteration_peak_num_blocks(self, cache_level: CacheLevel | None = None) -> None: - if cache_level is None: - self._iteration_peak_num_blocks_by_cache_level = ( - self._current_block_stats_by_cache_level() - ) - return - stats_by_pool_group = self._storage.get_statistics(cache_level) - self._iteration_peak_num_blocks_by_cache_level[cache_level] = make_typed( - lambda pool_group_index: PoolGroupPeakBlockStats( - available=stats_by_pool_group[pool_group_index].available, - unavailable=stats_by_pool_group[pool_group_index].unavailable, - evictable=stats_by_pool_group[pool_group_index].evictable, - ), - self._storage.num_pool_groups, - ) - - def _update_iteration_peak_num_blocks(self) -> None: - current = self._current_block_stats_by_cache_level() - for cache_level in typed_range(self._storage.num_cache_levels): - peak = self._iteration_peak_num_blocks_by_cache_level[cache_level] - current_level = current[cache_level] - for pool_group_index in typed_range(self._storage.num_pool_groups): - peak_stats = peak[pool_group_index] - current_stats = current_level[pool_group_index] - peak[pool_group_index] = PoolGroupPeakBlockStats( - available=max(peak_stats.available, current_stats.available), - unavailable=max(peak_stats.unavailable, current_stats.unavailable), - evictable=max(peak_stats.evictable, current_stats.evictable), - ) - - def commit_stats( - self, - stats: KVCacheStatsDelta, - iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] | None = None, - ) -> None: - if not self._stats_enabled: - return - self._update_iteration_peak_num_blocks() - self._committed_stats.add(stats) - if iteration_stats_by_life_cycle is None: - return - for life_cycle, iteration_stats in iteration_stats_by_life_cycle.items(): - if iteration_stats.empty: - continue - destination = self._iteration_stats_by_life_cycle.setdefault( - life_cycle, KVCacheIterationStatsDelta() - ) - destination.add(iteration_stats) - - def get_committed_stats(self) -> KVCacheStatsDelta: - return self._committed_stats.copy() - - def get_and_reset_iteration_stats(self) -> dict[LifeCycleId, KVCacheIterationStatsDelta]: - stats = { - life_cycle: delta.copy() - for life_cycle, delta in self._iteration_stats_by_life_cycle.items() - if not delta.empty - } - self._iteration_stats_by_life_cycle.clear() - return stats - - def _commit_ssm_snapshot_iteration_stats( - self, - iteration_stats_by_life_cycle: dict[LifeCycleId, SsmSnapshotIterationStatsDelta], - ) -> None: - if not self._stats_enabled: - return - for life_cycle, iteration_stats in iteration_stats_by_life_cycle.items(): - if iteration_stats.empty: - continue - destination = self._ssm_snapshot_iteration_stats_by_life_cycle.setdefault( - life_cycle, SsmSnapshotIterationStatsDelta() - ) - destination.add(iteration_stats) - - def get_and_reset_ssm_snapshot_iteration_stats( - self, - ) -> dict[LifeCycleId, SsmSnapshotIterationStatsDelta]: - stats = { - life_cycle: delta.copy() - for life_cycle, delta in self._ssm_snapshot_iteration_stats_by_life_cycle.items() - if not delta.empty - } - self._ssm_snapshot_iteration_stats_by_life_cycle.clear() - return stats - - def _commit_reused_blocks_by_level( - self, by_life_cycle: dict[LifeCycleId, ReusedBlocksByLevel] - ) -> None: - """Commit the per-cache-level split of the reuse block counts. - - Committed alongside the scalar iteration stats so both views cover exactly the same - requests: a cache whose pending stats are discarded contributes to neither. - """ - if not self._stats_enabled: - return - for life_cycle, by_level in by_life_cycle.items(): - if by_level.empty: - continue - self._iter_reused_blocks_by_level.setdefault(life_cycle, ReusedBlocksByLevel()).add( - by_level - ) - - def get_and_reset_iteration_reused_blocks_by_level( - self, - ) -> dict[LifeCycleId, ReusedBlocksByLevel]: - """Return and reset the per-cache-level reuse block counts for this iteration.""" - by_life_cycle = self._iter_reused_blocks_by_level - self._iter_reused_blocks_by_level = {} - return by_life_cycle - - def record_request_suspended(self) -> None: - """Count one ACTIVE->SUSPENDED transition for the current iteration window.""" - if not self._stats_enabled: - return - self._iter_suspended_requests += 1 - - def record_request_resumed(self) -> None: - """Count one preemption recovery for the current iteration window. - - Only a previously-ACTIVE cache that was suspended and then successfully - resumed counts. A freshly-created cache is activated by its first resume() - call, but that is an admission, not a recovery, and is not counted. - """ - if not self._stats_enabled: - return - self._iter_resumed_requests += 1 - - def get_and_reset_iteration_suspend_resume_stats(self) -> tuple[int, int]: - """Return (suspended, resumed) request counts since the last drain and reset them. - - Suspend/resume is a per-request, manager-level event (not per-pool-group), so it - is drained alongside get_and_reset_iteration_stats once per iteration-stats fetch. - - Both counters track the same population, so they are directly comparable: - the running (suspended - resumed) total is the number of requests still - parked in the SUSPENDED state. - """ - suspended = self._iter_suspended_requests - resumed = self._iter_resumed_requests - self._iter_suspended_requests = 0 - self._iter_resumed_requests = 0 - return suspended, resumed - - def record_disk_prefetch_blocks(self, num_blocks: int) -> None: - """Count the blocks a prefetch call actually migrated from disk to host.""" - assert num_blocks >= 0 - if self._stats_enabled: - self._iter_disk_prefetch_blocks += num_blocks - - def get_and_reset_iteration_disk_prefetch_blocks(self) -> int: - """Return and reset disk-to-host prefetch blocks for this iteration.""" - num_blocks = self._iter_disk_prefetch_blocks - self._iter_disk_prefetch_blocks = 0 - return num_blocks - - def _commit_cached_tokens_by_level(self, counts: CountsByLevel) -> None: - """Accumulate a request's initial cached-token attribution, by cache level, into this - iteration. Committed alongside the scalar iteration stats so both views cover exactly the - same requests.""" - assert NDEBUG or all(count >= 0 for count in counts) - if self._stats_enabled: - self._iter_cached_tokens_by_level = add_counts_by_level( - self._iter_cached_tokens_by_level, counts - ) - - def get_and_reset_iteration_cached_tokens_by_level(self) -> CountsByLevel: - """Return the per-cache-level cached-token counts since the last drain and reset them.""" - counts = self._iter_cached_tokens_by_level - self._iter_cached_tokens_by_level = [] - return counts - - def get_storage_statistics( - self, cache_level: CacheLevel = GPU_LEVEL - ) -> list[StorageStatistics]: - """Return independent per-pool values; this backend requires serialized access.""" - return deepcopy(list(self._storage.get_statistics(cache_level))) - - def get_life_cycle_pool_group_indices( - self, cache_level: CacheLevel = GPU_LEVEL - ) -> list[PoolGroupIndex]: - """Return lifecycle-to-pool indices; this backend shares the hot grouping at all levels.""" - return [ - self._storage.get_pool_group_index(life_cycle) - for life_cycle in typed_range(self._storage.num_life_cycles) - ] - - def get_and_reset_iteration_peak_block_stats( - self, cache_level: CacheLevel - ) -> TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats]: - self._update_iteration_peak_num_blocks() - peak = make_typed( - lambda pool_group_index: PoolGroupPeakBlockStats( - available=self._iteration_peak_num_blocks_by_cache_level[cache_level][ - pool_group_index - ].available, - unavailable=self._iteration_peak_num_blocks_by_cache_level[cache_level][ - pool_group_index - ].unavailable, - evictable=self._iteration_peak_num_blocks_by_cache_level[cache_level][ - pool_group_index - ].evictable, - ), - self._storage.num_pool_groups, - ) - self._reset_iteration_peak_num_blocks(cache_level) - return peak - - def get_and_reset_iteration_peak_block_stats_by_level( - self, - ) -> TypedIndexList[CacheLevel, TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats]]: - """Drain every level at once. - - The peaks are already tracked as one per-level record, so a caller that wants all of them - should not take that record apart one level at a time. - """ - self._update_iteration_peak_num_blocks() - peak = self._iteration_peak_num_blocks_by_cache_level - self._reset_iteration_peak_num_blocks() - return peak - - def mark_stats_dirty(self, kv_cache_id: int | None) -> None: - if kv_cache_id is not None: - self._dirty_stats_kv_cache_ids.add(kv_cache_id) - - def clear_stats_dirty(self, kv_cache_id: int | None) -> None: - if kv_cache_id is not None: - self._dirty_stats_kv_cache_ids.discard(kv_cache_id) - - def get_dirty_stats_kv_cache_ids(self) -> set[int]: - return self._dirty_stats_kv_cache_ids.copy() - - def mark_stats_excluded(self, kv_cache_id: int | None) -> None: - if kv_cache_id is not None: - self._stats_excluded_kv_cache_ids.add(kv_cache_id) - self.clear_stats_dirty(kv_cache_id) - - def clear_stats_excluded(self, kv_cache_id: int | None) -> None: - if kv_cache_id is not None: - self._stats_excluded_kv_cache_ids.discard(kv_cache_id) - - def is_stats_excluded(self, kv_cache_id: int | None) -> bool: - return kv_cache_id is not None and kv_cache_id in self._stats_excluded_kv_cache_ids - - # sorted by CacheLevel from warm to cold - @property - def cache_tier_list(self) -> HomoTuple[CacheTier]: - return self._storage.cache_tiers - - @property - def tokens_per_block(self) -> int: - return self._radix_tree.tokens_per_block - - @property - def event_manager(self) -> "KVCacheEventManager | None": - return self._event_manager - - @property - def allow_seq_rebasing(self) -> bool: - """ - If True, when we commit a full block, we will try to find a existing reusable block with the - same tokens and reuse that block instead to save some memory. Intra-batch reuse will be enabled - if this is True. - """ - return True - - @property - def enable_partial_match(self) -> bool: - return self._init_config.enable_partial_reuse - - @property - def enable_swa_scratch_reuse(self) -> bool: - return self._init_config.enable_swa_scratch_reuse - - def supports_index_mode(self, mode: PageIndexMode) -> bool | None: - """Whether managed KV caches support the given page index mode. - - Returns: - True — the mode is supported by every KV cache. - False — the mode is not supported by any KV cache. - None — support is per-instance; check _KVCache.supports_index_mode(). - """ - match mode: - case PageIndexMode.PER_LAYER: - return True - case PageIndexMode.SHARED: - return None if self.enable_swa_scratch_reuse else True - - @property - def num_layers(self) -> int: - return len(self._storage._layer_to_life_cycle_ids) - - @property - def layer_ids(self) -> Iterator[LayerId]: - return iter(self._storage._layer_to_life_cycle_ids.keys()) - - def get_layer_group_id(self, layer_id: LayerId) -> LayerGroupId: - return self._storage._layer_to_life_cycle_ids[layer_id] - - @property - def layer_grouping(self) -> HomoTuple[HomoTuple[LayerId]]: - """ - Layers are divided into multiple groups. - Buffers in the same layer group for the same token block are always allocated/deallocated together. - - NOTE: the iteration order of the layer lists (and of the groups) is NOT part of - the API contract and may differ across backends/runs. Do not rely on it for - buffer/pool memory order -- query ``pool_group_descs`` (PoolGroupDesc.pools[i] - .base_address + coalesced_buffers) for that. - """ - layer_to_life_cycle_ids = self._storage._layer_to_life_cycle_ids - num_life_cycles = self._life_cycles.size - grouping = dict[LifeCycleId, list[LayerId]]({i: [] for i in typed_range(num_life_cycles)}) - for layer_id, life_cycle_id in layer_to_life_cycle_ids.items(): - grouping[life_cycle_id].append(layer_id) - return tuple(tuple(grouping[i]) for i in typed_range(num_life_cycles)) - - @property - def all_buffer_ids(self) -> Iterator[BufferId]: - return iter(self._storage._buffer_attr.keys()) - - def get_aggregated_pages(self, buffers: Iterable[BufferId]) -> Iterator[AggregatedPageDesc]: - """ - Internally, we concatenate buffers into larger buffers. - This API takes a iterable of buffers (unordered), and try to find those that can form - contiguous aggregated buffers. - When we need data transfer, this helps us improve performance. - Args: - buffers: iterable of buffers to aggregate. Order does not matter. - Returns: - A iterator of aggregated buffers. - """ - groups = defaultdict[tuple[LifeCycleId, PoolIndex], list[tuple[Range, ExpandedBuffer]]]( - list[tuple[Range, ExpandedBuffer]] - ) - buffer_attr_map = self._storage._buffer_attr - for buffer in buffers: - attr = buffer_attr_map[buffer] - start = attr.offset - key = (attr.life_cycle_id, attr.pool_index) - groups[key].append( - (Range(start, start + attr.size), ExpandedBuffer(buffer, attr.expansion)) - ) - - storage = self._storage._levels[GPU_LEVEL].storage - lc2pg = self._storage._life_cycle_grouping - for (lc, pool_idx), group in groups.items(): - pg_idx = lc2pg[lc] - group.sort(key=lambda item: item[0].start) - current_start, current_end, current_buffers = ( - group[0][0].start, - group[0][0].end, - [group[0][1]], - ) - stride = storage.slot_size(pg_idx)[pool_idx] - pool_base = int(cast(int, storage.slot_address(pg_idx, pool_idx, SlotId(0)))) - for next_range, next_buffer in group[1:]: - if next_range.start == current_end: - current_end = next_range.end - current_buffers.append(next_buffer) - continue - - base = MemAddress(pool_base + current_start) - yield AggregatedPageDesc( - base, current_end - current_start, stride, lc, tuple(current_buffers) - ) - current_start, current_end, current_buffers = ( - next_range.start, - next_range.end, - [next_buffer], - ) - base = MemAddress(pool_base + current_start) - yield AggregatedPageDesc( - base, current_end - current_start, stride, lc, tuple(current_buffers) - ) - - @property - def pool_group_descs(self) -> TypedIndexList[PoolGroupIndex, PoolGroupDesc]: - storage = self._storage - - def get_pool_group_desc(pg_idx: PoolGroupIndex) -> PoolGroupDesc: - slot_size_list = storage.slot_size(pg_idx) - pools = make_typed( - lambda pool_idx: PoolDesc( - pool_index=pool_idx, - base_address=storage.get_mem_pool_base_address(pg_idx, pool_idx), - slot_bytes=slot_size_list[pool_idx], - ), - storage.num_pools(pg_idx), - ) - return PoolGroupDesc( - pool_group_index=pg_idx, - num_slots=storage.num_slots(pg_idx), - slot_desc=storage._slot_desc_list[pg_idx], - pools=pools, - ) - - return make_typed(get_pool_group_desc, storage.num_pool_groups) - - @property - def _current_gpu_ratio(self) -> TypedIndexList[PoolGroupIndex, float]: - return self._storage.get_ratio_list(GPU_LEVEL) - - @property - def _current_other_ratios(self) -> TypedIndexList[PoolGroupIndex, float]: - storage = self._storage - num_cache_levels = storage.num_cache_levels - if num_cache_levels == 1: - return self._current_gpu_ratio - num_pool_groups = storage.num_pool_groups - other_ratios = [ - storage.get_ratio_list(i) for i in typed_range(CacheLevel(1), num_cache_levels) - ] - other_ratio = filled_list(0.0, num_pool_groups) - for j in typed_range(num_pool_groups): - for i in range(1, num_cache_levels): - other_ratio[j] += other_ratios[i - 1][j] - other_ratio[j] /= num_cache_levels - 1 - return other_ratio - - def _get_target_ratio_list(self, level: CacheLevel) -> TypedIndexList[PoolGroupIndex, float]: - return self._target_ratio_list_gpu if level == GPU_LEVEL else self._target_ratio_list_other - - def _need_adjustment(self, level: CacheLevel) -> bool: - def check_mismatch( - a: TypedIndexList[PoolGroupIndex, float], - b: TypedIndexList[PoolGroupIndex, float], - thres: float, - ) -> bool: - return any(not (1 / thres < x / y < thres) for x, y in zip(a, b, strict=True)) - - if level == GPU_LEVEL: - return check_mismatch(self._target_ratio_list_gpu, self._current_gpu_ratio, 1.25) - else: - return check_mismatch(self._target_ratio_list_other, self._current_other_ratios, 1.25) - - def _adjust_level(self, level: CacheLevel, new_quota: int | None = None) -> None: - new_ratio_list = self._get_target_ratio_list(level) - storage = self._storage - num_cache_levels = storage.num_cache_levels - # held and not evictable as they are already in the last level cache. - persistent_pages: TypedIndexList[PoolGroupIndex, list[Page]] | None = None - if level == num_cache_levels - 1: - persistent_pages = self._gather_persistent_pages() - storage.adjust_cache_level(level, new_quota, new_ratio_list, persistent_pages) - - def _gather_persistent_pages(self) -> TypedIndexList[PoolGroupIndex, list[Page]]: - last_level = self._storage.num_cache_levels - 1 - lc2pg = self._storage._life_cycle_grouping - ret = make_typed(lambda _: list[Page](), self._storage.num_pool_groups) - for r in self._living_kv_caches: - kv_cache = unwrap_rawref(r) - assert kv_cache.status == _KVCache.Status.SUSPENDED - for block in kv_cache._blocks: - for beam in block.pages: - for lc, holder in typed_enumerate(beam): - if holder is None: - continue - assert type(holder) is _PageHolder - page = holder.page - assert page.status == PageStatus.HELD - assert page.scheduled_for_eviction == (page.cache_level != last_level) - if not page.scheduled_for_eviction: - ret[lc2pg[lc]].append(holder.page) - return ret - - @property - def need_adjustment(self) -> bool: - if self._num_sampled_kv_caches < 2000: - return False - if time.monotonic() - self._last_adjustment_time < 120: - return False - return self._need_adjustment(GPU_LEVEL) or self._need_adjustment( - CacheLevel(self._storage.num_cache_levels - 1) - ) - - def adjust(self) -> None: - """ - Adjust the cache level and ratio list. - This function should be called periodically to ensure the cache level and ratio list are - adjusted to the optimal values. All KV caches must be suspended before calling this function. - """ - assert all( - unwrap_rawref(c).status == _KVCache.Status.SUSPENDED for c in self._living_kv_caches - ) - storage = self._storage - for level in typed_range(storage.num_cache_levels): - if self._need_adjustment(level): - self._adjust_level(level) - self._last_adjustment_time = time.monotonic() - - def _try_update_target_ratios(self) -> None: - if self._num_sampled_kv_caches - self._last_update_num_sampled_kv_caches < 100: - return - self._last_update_num_sampled_kv_caches = self._num_sampled_kv_caches - tokens_per_block = self.tokens_per_block - storage = self._storage - - avg_reused_length: int = round(self._avg_reused_length.value) - avg_capacity: int = round(self._avg_sqr_capacity.value**0.5) - avg_history_length: int = round(self._avg_sqr_history_length.value**0.5) - if avg_capacity > 0: - life_cycle_ratio = storage.ratio_from_length( - tokens_per_block, avg_history_length, avg_capacity - ) - pool_group_ratio = storage.pool_group_ratio(life_cycle_ratio) - self._target_ratio_list_gpu = storage.constrain_pool_group_ratio(pool_group_ratio) - if avg_reused_length > 0: - life_cycle_ratio = storage.ratio_from_length( - tokens_per_block, avg_reused_length, avg_reused_length - ) - self._target_ratio_list_other = storage.pool_group_ratio(life_cycle_ratio) - - # @TODO: need updating when dynamic resizing is supported. - def clamp_max_seq_len_for_mem(self, batch_size: int, token_num_upper_bound: int) -> int: - "Get the max possible sequence length limited by the GPU memory pools." - assert batch_size > 0 - tokens_per_block = self.tokens_per_block - life_cycles = self._life_cycles - storage = self._storage - num_pool_groups = storage.num_pool_groups - remaining_slots = cast( - TypedIndexList[PoolGroupIndex, int], - [storage.num_slots(pg) for pg in typed_range(num_pool_groups)], - ) - lc_to_pg_idx = storage._life_cycle_grouping - - def get_num_slots(seq_len: int) -> TypedIndexList[PoolGroupIndex, int]: - ret = filled_list(0, num_pool_groups) - for lc_id, lc in life_cycles.items(): - stale_range = _KVCache._get_stale_range(tokens_per_block, seq_len, lc) - num_stale_blocks = stale_range[1] - stale_range[0] - num_slots = div_up(seq_len, tokens_per_block) - num_stale_blocks - pg_idx = lc_to_pg_idx[lc_id] - ret[pg_idx] += num_slots - return ret - - for pg in typed_range(num_pool_groups): - remaining_slots[pg] -= get_num_slots(1)[pg] * (batch_size - 1) - if remaining_slots[pg] < 0: - return 0 - - def is_enough(num_blocks: int) -> bool: - return all( - cnt <= rem - for cnt, rem in zip( - get_num_slots(num_blocks * tokens_per_block), remaining_slots, strict=True - ) - ) - - if not is_enough(1): - return 0 - lb = 1 - ub = div_up(token_num_upper_bound, tokens_per_block) - if is_enough(ub): - return token_num_upper_bound - while lb < ub - 1: - mid = (lb + ub) // 2 - if is_enough(mid): - lb = mid - else: - ub = mid - return min(lb * tokens_per_block, token_num_upper_bound) - - @property - def init_config(self) -> KVCacheManagerConfig: - return self._init_config - - @property - def commit_min_snapshot(self) -> bool: - return self.init_config.commit_min_snapshot - - @property - def text_only(self) -> bool: - return self.init_config.text_only diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_moving_average.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_moving_average.py deleted file mode 100644 index faee5d7aa0ba..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_moving_average.py +++ /dev/null @@ -1,56 +0,0 @@ -# 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. - - -class MovingAverage: - __slots__ = ("decay", "avg", "weight", "num_updates") - decay: float - avg: float - weight: float - num_updates: int - - def __init__(self, decay: float = 0.9999): - self.decay = decay - self.avg = 0.0 - self.weight = 0.0 - self.num_updates = 0 - - def update(self, value: int | float) -> float: - self.weight = 1.0 + self.decay * self.weight - self.avg += (value - self.avg) / self.weight - self.num_updates += 1 - return self.avg - - @property - def value(self) -> float: - return self.avg - - -class Average: - __slots__ = ("sum", "count") - sum: float - count: int - - def __init__(self): - self.sum = 0.0 - self.count = 0 - - def update(self, value: int | float) -> None: - self.sum += value - self.count += 1 - - @property - def value(self) -> float: - return self.sum / self.count diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py deleted file mode 100644 index c5b4f18fae9f..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py +++ /dev/null @@ -1,317 +0,0 @@ -# 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. - -from dataclasses import dataclass, field - -from .._common import NDEBUG, BlockOrdinal, CacheLevel -from .._life_cycle_registry import LifeCycleId -from .._stats import ( - CountsByLevel, - KVCacheIterationStatsDelta, - KVCacheStatsDelta, - ReusedBlocksByLevel, - SsmSnapshotIterationStatsDelta, - add_counts_by_level, -) - - -@dataclass(slots=True) -class _PendingAllocationSegment: - life_cycle: LifeCycleId - block_begin: BlockOrdinal - block_end: BlockOrdinal - beam_width: int - count_as_missed: bool - count_as_generation: bool - record_manager_stats: bool - record_request_stats: bool - - -@dataclass(slots=True) -class _PendingStatsDelta: - global_stats: KVCacheStatsDelta - request_stats: KVCacheStatsDelta - iteration_stats: KVCacheIterationStatsDelta - life_cycle: LifeCycleId | None = None - - @property - def empty(self) -> bool: - return self.global_stats.empty and self.request_stats.empty and self.iteration_stats.empty - - -@dataclass(slots=True) -class _PendingStats: - request_stats: KVCacheStatsDelta = field(default_factory=KVCacheStatsDelta) - global_stats: KVCacheStatsDelta = field(default_factory=KVCacheStatsDelta) - iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] = field( - default_factory=dict - ) - ssm_snapshot_iteration_stats_by_life_cycle: dict[ - LifeCycleId, SsmSnapshotIterationStatsDelta - ] = field(default_factory=dict) - reused_blocks_by_level_by_life_cycle: dict[LifeCycleId, ReusedBlocksByLevel] = field( - default_factory=dict - ) - # Cached-token attribution for the sequence's reuse match, indexed by cache level. - # - # Unlike the reuse counters this is a manager-global quantity rather than a per-lifecycle one: - # a match spans every lifecycle at once (the final SSM checkpoint summarizes the whole recurrent - # prefix, so its tier applies to every matched token), leaving no single lifecycle to attribute - # it to. It still rides the pending-stats lifecycle so it is committed or discarded together - # with the counters it was derived from -- in particular, a dummy sequence's attribution is - # dropped by the same discard_pending_stats() that drops its reuse counters. - cached_tokens_by_level: CountsByLevel = field(default_factory=list) - allocation_segments: list[_PendingAllocationSegment] = field(default_factory=list) - - @property - def empty(self) -> bool: - return ( - self.request_stats.empty - and self.global_stats.empty - and not self.iteration_stats_by_life_cycle - and not self.ssm_snapshot_iteration_stats_by_life_cycle - and not any(self.cached_tokens_by_level) - ) - - def clear(self) -> None: - self.request_stats.clear() - self.global_stats.clear() - self.iteration_stats_by_life_cycle.clear() - self.ssm_snapshot_iteration_stats_by_life_cycle.clear() - self.reused_blocks_by_level_by_life_cycle.clear() - self.cached_tokens_by_level = [] - self.allocation_segments.clear() - - def record_cached_tokens_by_level(self, counts: CountsByLevel) -> bool: - if not any(counts): - return False - self.cached_tokens_by_level = add_counts_by_level(self.cached_tokens_by_level, counts) - return True - - def limit_cached_tokens_by_level(self, level: CacheLevel, max_tokens: int) -> None: - """Limit staged attribution without restoring counts already dropped or committed.""" - # No attribution is staged when stats are disabled, discarded, or already committed. - if not self.cached_tokens_by_level: - return - assert NDEBUG or level < len(self.cached_tokens_by_level) - if level >= len(self.cached_tokens_by_level): - return - assert NDEBUG or max_tokens >= 0 - self.cached_tokens_by_level[level] = min( - self.cached_tokens_by_level[level], max(0, max_tokens) - ) - - def add(self, delta: _PendingStatsDelta) -> bool: - if delta.empty: - return False - if not delta.global_stats.empty: - self.global_stats.add(delta.global_stats) - if not delta.request_stats.empty: - self.request_stats.add(delta.request_stats) - if not delta.iteration_stats.empty: - assert delta.life_cycle is not None - pending = self.iteration_stats_by_life_cycle.setdefault( - delta.life_cycle, KVCacheIterationStatsDelta() - ) - pending.add(delta.iteration_stats) - return True - - def subtract(self, delta: _PendingStatsDelta) -> bool: - if delta.empty: - return False - if not delta.global_stats.empty: - self.global_stats.subtract(delta.global_stats) - if not delta.request_stats.empty: - self.request_stats.subtract(delta.request_stats) - if not delta.iteration_stats.empty: - assert delta.life_cycle is not None - pending = self.iteration_stats_by_life_cycle.get(delta.life_cycle) - if pending is not None: - pending.subtract(delta.iteration_stats) - if pending.empty: - del self.iteration_stats_by_life_cycle[delta.life_cycle] - return True - - @staticmethod - def _allocation_delta( - segment: _PendingAllocationSegment, - block_begin: BlockOrdinal, - block_end: BlockOrdinal, - ) -> _PendingStatsDelta: - num_blocks = max(0, int(block_end) - int(block_begin)) * segment.beam_width - manager_stats = ( - KVCacheStatsDelta( - alloc_total_blocks=num_blocks, - alloc_new_blocks=num_blocks, - missed_blocks=num_blocks if segment.count_as_missed else 0, - ) - if segment.record_manager_stats - else KVCacheStatsDelta() - ) - request_stats = ( - KVCacheStatsDelta( - alloc_total_blocks=num_blocks, - alloc_new_blocks=num_blocks, - missed_blocks=num_blocks if segment.count_as_missed else 0, - ) - if segment.record_request_stats - else KVCacheStatsDelta() - ) - iteration_stats = ( - KVCacheIterationStatsDelta( - iter_alloc_total_blocks=num_blocks, - iter_alloc_new_blocks=num_blocks, - iter_missed_blocks=num_blocks if segment.count_as_missed else 0, - iter_gen_alloc_blocks=num_blocks if segment.count_as_generation else 0, - ) - if segment.record_manager_stats - else KVCacheIterationStatsDelta() - ) - return _PendingStatsDelta(manager_stats, request_stats, iteration_stats, segment.life_cycle) - - def record_allocation_range( - self, - life_cycle: LifeCycleId, - block_begin: BlockOrdinal, - block_end: BlockOrdinal, - *, - beam_width: int, - count_as_missed: bool, - count_as_generation: bool = False, - record_manager_stats: bool, - record_request_stats: bool, - ) -> bool: - if block_begin >= block_end or not (record_manager_stats or record_request_stats): - return False - segment = _PendingAllocationSegment( - life_cycle=life_cycle, - block_begin=block_begin, - block_end=block_end, - beam_width=beam_width, - count_as_missed=count_as_missed, - count_as_generation=count_as_generation, - record_manager_stats=record_manager_stats, - record_request_stats=record_request_stats, - ) - if not self.add(self._allocation_delta(segment, block_begin, block_end)): - return False - self.allocation_segments.append(segment) - return True - - def record_reuse( - self, - life_cycle: LifeCycleId, - *, - full_reused_blocks: int, - partial_reused_blocks: int, - by_level: ReusedBlocksByLevel | None = None, - record_manager_stats: bool, - record_request_stats: bool, - ) -> bool: - """Record reuse counts for one life cycle. - - ``by_level`` splits the same full/partial counts across the cache levels the reused - pages were resident on. It rides along with the scalar counters so both are committed - or discarded together; reuse is never rolled back (only allocation ranges are), so - add-only is enough. - """ - reused_blocks = full_reused_blocks + partial_reused_blocks - if reused_blocks == 0 or not (record_manager_stats or record_request_stats): - return False - if record_manager_stats and by_level is not None: - self.reused_blocks_by_level_by_life_cycle.setdefault( - life_cycle, ReusedBlocksByLevel() - ).add(by_level) - return self.add( - _PendingStatsDelta( - global_stats=( - KVCacheStatsDelta(reused_blocks=reused_blocks) - if record_manager_stats - else KVCacheStatsDelta() - ), - request_stats=( - KVCacheStatsDelta(reused_blocks=reused_blocks) - if record_request_stats - else KVCacheStatsDelta() - ), - iteration_stats=( - KVCacheIterationStatsDelta( - iter_reused_blocks=reused_blocks, - iter_full_reused_blocks=full_reused_blocks, - iter_partial_reused_blocks=partial_reused_blocks, - ) - if record_manager_stats - else KVCacheIterationStatsDelta() - ), - life_cycle=life_cycle, - ) - ) - - def record_ssm_snapshot_lookup( - self, - life_cycle: LifeCycleId, - *, - lookup_tokens: int, - reused_tokens: int, - tokens_per_block: int, - ) -> bool: - if lookup_tokens == 0: - return False - assert lookup_tokens > 0 - assert 0 <= reused_tokens <= lookup_tokens - assert tokens_per_block > 0 - - is_hit = reused_tokens > 0 - # Alignment describes the reusable snapshot boundary, not whether the - # state itself is complete. Every hit represents one complete SSM - # snapshot; token counters carry the benefit of that single lookup. - delta = SsmSnapshotIterationStatsDelta( - iter_snapshot_lookups=1, - iter_snapshot_hits=int(is_hit), - iter_snapshot_misses=int(not is_hit), - iter_reused_tokens=reused_tokens, - iter_unreused_tokens=lookup_tokens - reused_tokens, - iter_aligned_snapshot_hits=int(is_hit and reused_tokens % tokens_per_block == 0), - iter_unaligned_snapshot_hits=int(is_hit and reused_tokens % tokens_per_block != 0), - ) - pending = self.ssm_snapshot_iteration_stats_by_life_cycle.setdefault( - life_cycle, SsmSnapshotIterationStatsDelta() - ) - pending.add(delta) - return True - - def subtract_allocation_range(self, block_begin: BlockOrdinal, block_end: BlockOrdinal) -> bool: - if block_begin >= block_end or not self.allocation_segments: - return False - changed = False - idx = len(self.allocation_segments) - 1 - while idx >= 0: - segment = self.allocation_segments[idx] - if segment.block_end <= block_begin: - break - removed_begin = max(block_begin, segment.block_begin) - removed_end = min(block_end, segment.block_end) - if removed_begin >= removed_end: - idx -= 1 - continue - changed = True - self.subtract(self._allocation_delta(segment, removed_begin, removed_end)) - if removed_begin <= segment.block_begin: - del self.allocation_segments[idx] - else: - assert removed_end == segment.block_end - segment.block_end = removed_begin - idx -= 1 - return changed diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py deleted file mode 100644 index b138ac2f2632..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py +++ /dev/null @@ -1,228 +0,0 @@ -# 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. - -from typing import Type - -import cuda.bindings.driver as drv - -from ._common import MemAddress -from ._exceptions import CuError -from ._utils import ItemHolderWithSharedPool, PooledFactoryBase, _unwrap, div_up - - -def _is_prop_supported(prop: drv.CUmemAllocationProp) -> bool: - err, handle = drv.cuMemCreate(2 << 20, prop, 0) - err_int = int(err) - if err_int == int(drv.CUresult.CUDA_SUCCESS): - _unwrap(drv.cuMemRelease(handle)) - return True - # Note: OOM is intentionally not caught here — OOM on a 2 MiB probe - # indicates a fundamental resource problem, not an unsupported property. - elif err_int in ( - int(drv.CUresult.CUDA_ERROR_NOT_PERMITTED), - int(drv.CUresult.CUDA_ERROR_NOT_SUPPORTED), - int(drv.CUresult.CUDA_ERROR_INVALID_DEVICE), - int(drv.CUresult.CUDA_ERROR_INVALID_VALUE), - ): - return False - else: - raise CuError(err) - - -# Physical memory -class NativePhysMemAllocator: - __slots__ = ("_device_id", "_size", "_prop", "_outstanding_handles") - - _device_id: int - _size: int - _prop: drv.CUmemAllocationProp - _outstanding_handles: set[int] # allocated but not released - - def __init__(self, size: int) -> None: - self._device_id = int(_unwrap(drv.cuCtxGetDevice())) # pyright: ignore - self._size = size - prop = drv.CUmemAllocationProp() - prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED - prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - prop.location.id = self._device_id - # Prefer shareable handle types so UCX cuda_ipc can export the memory - # for intra-node zero-copy transfers: - # FABRIC (MNNVL) > POSIX_FILE_DESCRIPTOR (pidfd-based, UCX >= 1.22) > NONE, - # and within each handle type prefer gpuDirectRDMACapable over not. - handle_type_preference = ( - drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, - drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, - drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE, - ) - supported = False - for handle_type in handle_type_preference: - for rdma_capable in (1, 0): - prop.requestedHandleTypes = handle_type - prop.allocFlags.gpuDirectRDMACapable = rdma_capable - if _is_prop_supported(prop): - supported = True - break - if supported: - break - if not supported: - raise ValueError("Failed to create physical memory allocation property") - self._prop = prop - self._outstanding_handles = set() - - def allocate(self) -> drv.CUmemGenericAllocationHandle: - handle: drv.CUmemGenericAllocationHandle = _unwrap( - drv.cuMemCreate(self._size, self._prop, 0) - ) - int_handle = int(handle) # pyright: ignore - assert (int_handle not in self._outstanding_handles) and int_handle != 0 - self._outstanding_handles.add(int_handle) - return handle - - def release(self, handle: drv.CUmemGenericAllocationHandle) -> None: - if handle == drv.CUmemGenericAllocationHandle(0): - return - assert int(handle) in self._outstanding_handles - self._outstanding_handles.remove(int(handle)) - try: - _unwrap(drv.cuMemRelease(handle)) - except: - print( - f"Failed to release handle {handle}. num_outstanding = {len(self._outstanding_handles)}" - ) - raise - - @property - def device_id(self) -> int: - return self._device_id - - @property - def size(self) -> int: - return self._size - - -class PhysMem(ItemHolderWithSharedPool[drv.CUmemGenericAllocationHandle]): - __slots__ = () - - -class PooledPhysMemAllocator(PooledFactoryBase[drv.CUmemGenericAllocationHandle, PhysMem]): - _Holder: Type[PhysMem] = PhysMem - __slots__ = ("device_id", "phys_mem_size") - device_id: int - phys_mem_size: int - - def __init__(self, phys_mem_size: int) -> None: - """phys_mem_size is the size of each physical memory chunk.""" - raw_alloc = NativePhysMemAllocator(phys_mem_size) - self.device_id = raw_alloc.device_id - self.phys_mem_size = phys_mem_size - super().__init__(lambda: raw_alloc.allocate(), lambda handle: raw_alloc.release(handle)) - - -# Virtual memory -class VirtMem: - __slots__ = ("_vm_size", "_allocator", "_address", "_pm_stack", "_access_desc") - _vm_size: int - _allocator: PooledPhysMemAllocator - _address: drv.CUdeviceptr - _pm_stack: list[PhysMem] - _access_desc: drv.CUmemAccessDesc - - def __init__( - self, vm_size: int, phys_mem_allocator: PooledPhysMemAllocator, init_num_phys_mem: int = 0 - ): - assert vm_size % phys_mem_allocator.phys_mem_size == 0 - self._allocator = phys_mem_allocator - device_id = phys_mem_allocator.device_id - self._address = _unwrap(drv.cuMemAddressReserve(vm_size, 0, 0, 0)) - self._vm_size = vm_size - self._pm_stack = [] - self._access_desc = drv.CUmemAccessDesc() - self._access_desc.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - self._access_desc.location.id = device_id - self._access_desc.flags = drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE - self.extend(init_num_phys_mem) - - @property - def phys_mem_size(self) -> int: - return self._allocator.phys_mem_size - - def destroy(self) -> None: - if self._vm_size == 0: - return - _unwrap(drv.cuCtxSynchronize()) - while self._pm_stack: - self._pop().close() - _unwrap(drv.cuMemAddressFree(self._address, self._vm_size)) - self._address = drv.CUdeviceptr(0) - self._vm_size = 0 - - def __del__(self) -> None: - self.destroy() - - def extend(self, num_phys_mem: int) -> None: - old_num_phys_mem = self.num_phys_mem - try: - for _ in range(num_phys_mem): - self._push(self._allocator.create()) - except ( - Exception - ): # to make realloc behave like normal realloc, we need to rollback if out of memory - while self.num_phys_mem > old_num_phys_mem: - self._pop().close() - raise - - def shrink(self, num_phys_mem: int) -> None: - _unwrap(drv.cuCtxSynchronize()) - for _ in range(num_phys_mem): - self._pop().close() - - # Different from normal realloc, this function never changes the pointer. - def realloc(self, num_bytes: int) -> None: - required_num_phys_mem = div_up(num_bytes, self.phys_mem_size) - if required_num_phys_mem > self.num_phys_mem: - self.extend(required_num_phys_mem - self.num_phys_mem) - elif required_num_phys_mem < self.num_phys_mem: - self.shrink(self.num_phys_mem - required_num_phys_mem) - - def _push(self, phy_mem: PhysMem) -> None: - phys_mem_size = self.phys_mem_size - assert phys_mem_size * (len(self._pm_stack) + 1) <= self._vm_size - vm_ptr = drv.CUdeviceptr(self.address + phys_mem_size * len(self._pm_stack)) - _unwrap(drv.cuMemMap(vm_ptr, phys_mem_size, 0, phy_mem.handle, 0)) - _unwrap(drv.cuMemSetAccess(vm_ptr, phys_mem_size, (self._access_desc,), 1)) - self._pm_stack.append(phy_mem) - - def _pop(self) -> PhysMem: - assert self._pm_stack - phys_mem_size = self.phys_mem_size - vm_ptr = drv.CUdeviceptr(self.address + phys_mem_size * (len(self._pm_stack) - 1)) - _unwrap(drv.cuMemUnmap(vm_ptr, phys_mem_size)) - return self._pm_stack.pop() - - @property - def mapped_bytes(self) -> int: - return self.phys_mem_size * self.num_phys_mem - - @property - def virtual_bytes(self) -> int: - return self._vm_size - - @property - def num_phys_mem(self) -> int: - return len(self._pm_stack) - - @property - def address(self) -> MemAddress: - return MemAddress(int(self._address)) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py deleted file mode 100644 index ff3e55f7b6ed..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py +++ /dev/null @@ -1,701 +0,0 @@ -# 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. - -from __future__ import annotations - -import sys -import time -from collections import deque -from collections.abc import Iterable, Sequence -from dataclasses import dataclass, field, replace - -# avoid importing the whole tensorrt_llm module, which takes time during debugging. -from importlib.util import find_spec -from pathlib import Path -from threading import Condition -from typing import Any, Callable - -from ._utils import temporary_sys_path - -if "tensorrt_llm" in sys.modules: - from tensorrt_llm.logger import logger - from tensorrt_llm.runtime.kv_cache_hash import ( - KV_CACHE_HASH_ALGO_AUTO, - KV_CACHE_HASH_ALGO_DEFAULT, - KV_CACHE_HASH_ALGO_V1, - KV_CACHE_HASH_ALGO_V2, - KV_CACHE_HASH_ALGO_V2_SHA256_64, - NonTextTokenHashError, - hash_v1_block_key, - truncate_sha256_hash_to_int64, - ) -else: - # fast path for dev, avoids importing the whole tensorrt_llm module - import logging - - logger = logging.getLogger("tensorrt_llm") - - spec = find_spec("kv_cache_manager_v2") - assert spec is not None and spec.origin is not None - with temporary_sys_path(str(Path(spec.origin).parent.parent)): - from kv_cache_hash import ( # noqa - KV_CACHE_HASH_ALGO_AUTO, - KV_CACHE_HASH_ALGO_DEFAULT, - KV_CACHE_HASH_ALGO_V1, - KV_CACHE_HASH_ALGO_V2, - KV_CACHE_HASH_ALGO_V2_SHA256_64, - NonTextTokenHashError, - hash_v1_block_key, - truncate_sha256_hash_to_int64, - ) - -from ._common import GPU_LEVEL, PRIORITY_DEFAULT, CacheLevel, Priority, TokenIdExt - -EventBlockHash = int | str -BlockHashLike = bytes | EventBlockHash -BlockHashesLike = BlockHashLike | Iterable[BlockHashLike] -LayerGroupId = int | None -EventTokenId = int | str -MmKey = tuple[bytes, int] | tuple[bytes, int, str | None] -AttentionDpGatherFn = Callable[[list["KVCacheEvent"]], list[list["KVCacheEvent"]]] - - -@dataclass(slots=True, frozen=True) -class UniqueToken: - token_id: EventTokenId - token_extra_id: int = 0 - - -@dataclass(slots=True, frozen=True) -class KVCacheCreatedData: - num_blocks_per_cache_level: list[int] - - -@dataclass(slots=True, frozen=True) -class KVCacheStoredBlockData: - block_hash: EventBlockHash - tokens: list[UniqueToken] - cache_level: int - priority: int - mm_keys: list[MmKey] = field(default_factory=list) - cache_salt: str | None = None - - -@dataclass(slots=True, frozen=True) -class KVCacheStoredData: - parent_hash: EventBlockHash | None - blocks: list[KVCacheStoredBlockData] - - -@dataclass(slots=True, frozen=True) -class KVCacheRemovedData: - block_hashes: list[EventBlockHash] - - -@dataclass(slots=True, frozen=True) -class KVCacheEventDiff: - old_value: int - new_value: int - - -@dataclass(slots=True, frozen=True) -class KVCacheUpdatedData: - block_hash: EventBlockHash - cache_level: KVCacheEventDiff | None - priority: KVCacheEventDiff | None - - -@dataclass(slots=True, frozen=True) -class KVCacheEvent: - event_id: int - data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData - window_size: int - hash_algo: str | None = None - attention_dp_rank: int | None = None - layer_group_id: int | None = None - - -@dataclass(slots=True) -class _StoredBlockState: - block_hash: EventBlockHash - life_cycle_ids: set[int] - - -class KVCacheEventManager: - """Python event queue matching the C++ KV cache event serializer contract. - - Setting ``mm_token_id_offset`` enables MM keys derived from item-digest tokens - and continuation IDs equal to that offset plus the item-local token index. - It must match the offset passed to ``gen_multimodal_cache_key_tokens``. - """ - - def __init__( - self, - max_kv_event_entries: int, - *, - window_size: int = 0, - attention_dp_rank: int | None = None, - attention_dp_gather: AttentionDpGatherFn | None = None, - hash_algo: str = KV_CACHE_HASH_ALGO_V2, - window_size_by_layer_group: dict[int, int] | None = None, - mm_token_id_offset: int | None = None, - ) -> None: - if mm_token_id_offset is not None and mm_token_id_offset < 0: - raise ValueError("mm_token_id_offset must be nonnegative") - if hash_algo == KV_CACHE_HASH_ALGO_AUTO: - hash_algo = KV_CACHE_HASH_ALGO_DEFAULT - elif hash_algo not in ( - KV_CACHE_HASH_ALGO_V1, - KV_CACHE_HASH_ALGO_V2, - KV_CACHE_HASH_ALGO_V2_SHA256_64, - ): - raise ValueError(f"Unsupported V2 KV cache event hash algorithm: {hash_algo}") - self._max_kv_event_entries = max_kv_event_entries - self._window_size = window_size - self._window_size_by_layer_group = dict(window_size_by_layer_group or {}) - self._attention_dp_rank = attention_dp_rank - self._attention_dp_gather = attention_dp_gather - self._hash_algo = hash_algo - self._mm_token_id_offset = mm_token_id_offset - self._next_event_id = 0 - self._stored_blocks: dict[bytes, _StoredBlockState] = {} - self._latest_stored_events: dict[LayerGroupId, KVCacheEvent] = {} - self._latest_removed_block_hashes: dict[LayerGroupId, list[EventBlockHash]] = {} - self._pending_events: list[KVCacheEvent] = [] - self._events: deque[KVCacheEvent] = deque() - self._condition = Condition() - self._v1_hash_by_block_key: dict[bytes, int] = {} - self._v1_hash_compatible_keys: set[bytes] = set() - self._v1_root_attrs_by_block_key: dict[bytes, tuple[int | None, int | None]] = {} - self._warned_v1_hash_fallback = False - - def add_created_event( - self, - num_blocks_per_cache_level: Sequence[int], - layer_group_ids: Sequence[int] | None = None, - ) -> None: - data = KVCacheCreatedData(list(num_blocks_per_cache_level)) - if layer_group_ids is None: - self._add_event(data) - return - for layer_group_id in layer_group_ids: - self._add_event(data, layer_group_id=int(layer_group_id)) - - def set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: - with self._condition: - self._window_size_by_layer_group = dict(window_sizes) - - def needs_token_digest_context(self) -> bool: - """Whether radix blocks need to retain the latest item digest for events.""" - return self._max_kv_event_entries > 0 and self._mm_token_id_offset is not None - - def add_stored_event( - self, - parent_hash: EventBlockHash | None, - blocks: Sequence[KVCacheStoredBlockData], - layer_group_id: int | None = None, - ) -> None: - if not blocks: - return - self._flush_removed_events(layer_group_id) - self._add_stored_event( - KVCacheStoredData(parent_hash, list(blocks)), - layer_group_id=layer_group_id, - ) - - def add_stored_block_event_from_block(self, block: Any) -> None: - life_cycle_ids = self._life_cycle_ids_from_radix_block(block) - if not life_cycle_ids: - return - parent_hash = self._parent_hash_from_radix_block(block) - self._stored_blocks[block.key] = _StoredBlockState( - block_hash=self._hash_from_radix_block(block), - life_cycle_ids=set(life_cycle_ids), - ) - for life_cycle_id in sorted(life_cycle_ids): - block_data = self._stored_block_from_radix_block(block, life_cycle_ids={life_cycle_id}) - if block_data is not None: - self.add_stored_event(parent_hash, [block_data], life_cycle_id) - - def add_stored_life_cycle_event_from_block(self, block: Any, life_cycle_id: int) -> None: - state = self._stored_blocks.get(block.key) - life_cycle_id = int(life_cycle_id) - if state is not None: - if life_cycle_id in state.life_cycle_ids: - return - block_data = self._stored_block_from_radix_block(block, life_cycle_ids={life_cycle_id}) - if block_data is None: - return - state.life_cycle_ids.add(life_cycle_id) - self.add_stored_event( - self._parent_hash_from_radix_block(block), - [block_data], - layer_group_id=life_cycle_id, - ) - return - self.add_stored_block_event_from_block(block) - - def add_removed_event(self, block_hashes: BlockHashesLike) -> None: - removed_block_hashes_by_layer_group: dict[int, list[EventBlockHash]] = {} - removed_block_hashes_without_layer_group: list[EventBlockHash] = [] - for block_hash in self._iter_block_hashes(block_hashes): - removed_state = self._pop_stored_block_state(block_hash) - if removed_state is None: - continue - normalized_hash, life_cycle_ids = removed_state - if life_cycle_ids: - for life_cycle_id in sorted(life_cycle_ids): - removed_block_hashes_by_layer_group.setdefault(life_cycle_id, []).append( - normalized_hash - ) - else: - removed_block_hashes_without_layer_group.append(normalized_hash) - - if removed_block_hashes_without_layer_group: - self._enqueue_removed_event(removed_block_hashes_without_layer_group) - for layer_group_id, removed_block_hashes in sorted( - removed_block_hashes_by_layer_group.items() - ): - self._enqueue_removed_event(removed_block_hashes, layer_group_id=layer_group_id) - - def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: - removed_state = self._pop_stored_life_cycle_block_state(block_hash, life_cycle_id) - if removed_state is None: - return - normalized_hash, removed_life_cycle_id, _ = removed_state - self._enqueue_removed_event( - [normalized_hash], - layer_group_id=removed_life_cycle_id, - ) - - def add_updated_event( - self, - block_hash: BlockHashLike, - *, - cache_level: KVCacheEventDiff | None = None, - priority: KVCacheEventDiff | None = None, - layer_group_id: int | None = None, - ) -> None: - if cache_level is None and priority is None: - return - normalized_block_hash = self._get_stored_block_hash(block_hash) - if normalized_block_hash is None: - return - self._add_event( - KVCacheUpdatedData( - block_hash=normalized_block_hash, - cache_level=cache_level, - priority=priority, - ), - layer_group_id=layer_group_id, - ) - - def flush_iteration_events(self) -> None: - if self._attention_dp_gather is not None: - with self._condition: - local_events = self._drain_pending_events_unlocked() - local_events = self._trim_events(local_events, self._max_kv_event_entries) - gathered_events = self._attention_dp_gather(local_events) - if self._attention_dp_rank != 0: - return - events = [ - event - for rank_events in gathered_events - for event in self._trim_events(rank_events, self._max_kv_event_entries) - ] - with self._condition: - self._publish_events_unlocked( - events, - max_kv_event_entries=( - self._max_kv_event_entries * max(1, len(gathered_events)) - ), - ) - self._condition.notify_all() - return - - with self._condition: - self._publish_events_unlocked(self._drain_pending_events_unlocked()) - self._condition.notify_all() - - def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: - with self._condition: - if not self._events and timeout_ms is None: - while not self._events: - self._condition.wait() - elif not self._events and timeout_ms > 0: - deadline = time.monotonic() + timeout_ms / 1000 - while not self._events: - remaining = deadline - time.monotonic() - if remaining <= 0: - break - self._condition.wait(timeout=remaining) - - events = list(self._events) - self._events.clear() - return events - - def _add_event( - self, - data: (KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData), - layer_group_id: LayerGroupId = None, - ) -> None: - if self._max_kv_event_entries <= 0: - return - with self._condition: - self._add_event_unlocked(data, layer_group_id) - - def _add_stored_event( - self, - data: KVCacheStoredData, - layer_group_id: LayerGroupId = None, - ) -> None: - if self._max_kv_event_entries <= 0: - return - with self._condition: - has_pending_removed_events = bool(self._latest_removed_block_hashes) - latest_event = self._latest_stored_events.get(layer_group_id) - if ( - not has_pending_removed_events - and latest_event is not None - and isinstance(latest_event.data, KVCacheStoredData) - ): - latest_blocks = latest_event.data.blocks - if latest_blocks and latest_blocks[-1].block_hash == data.parent_hash: - merged_data = replace( - latest_event.data, - blocks=[*latest_blocks, *data.blocks], - ) - merged_event = replace(latest_event, data=merged_data) - self._replace_pending_event_unlocked(latest_event, merged_event) - self._latest_stored_events[layer_group_id] = merged_event - return - - event = self._add_event_unlocked(data, layer_group_id) - self._latest_stored_events[layer_group_id] = event - - def _replace_pending_event_unlocked( - self, - old_event: KVCacheEvent, - new_event: KVCacheEvent, - ) -> None: - for event_idx in range(len(self._pending_events) - 1, -1, -1): - if self._pending_events[event_idx] is old_event: - self._pending_events[event_idx] = new_event - return - raise RuntimeError("Stored event coalescing lost the pending event") - - def _enqueue_removed_event( - self, - block_hashes: Sequence[EventBlockHash], - layer_group_id: LayerGroupId = None, - ) -> None: - if not block_hashes or self._max_kv_event_entries <= 0: - return - with self._condition: - self._latest_removed_block_hashes.setdefault(layer_group_id, []).extend(block_hashes) - self._latest_stored_events.pop(layer_group_id, None) - - def _flush_removed_events(self, layer_group_id: LayerGroupId) -> None: - if self._max_kv_event_entries <= 0: - return - with self._condition: - self._flush_removed_events_unlocked(layer_group_id) - - def _flush_removed_events_unlocked(self, layer_group_id: LayerGroupId) -> None: - block_hashes = self._latest_removed_block_hashes.pop(layer_group_id, None) - if not block_hashes: - return - self._add_event_unlocked( - KVCacheRemovedData(block_hashes), - layer_group_id=layer_group_id, - ) - - def _flush_all_removed_events_unlocked(self) -> None: - layer_group_ids = list(self._latest_removed_block_hashes) - for layer_group_id in layer_group_ids: - self._flush_removed_events_unlocked(layer_group_id) - - def _add_event_unlocked( - self, - data: (KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData), - layer_group_id: LayerGroupId = None, - ) -> KVCacheEvent: - if not isinstance(data, KVCacheRemovedData): - self._flush_all_removed_events_unlocked() - event = KVCacheEvent( - event_id=self._next_event_id, - data=data, - window_size=self._get_window_size(layer_group_id), - hash_algo=self._hash_algo, - attention_dp_rank=self._attention_dp_rank, - layer_group_id=layer_group_id, - ) - self._next_event_id += 1 - self._pending_events.append(event) - if not isinstance(data, KVCacheStoredData): - self._latest_stored_events.pop(layer_group_id, None) - return event - - def _drain_pending_events_unlocked(self) -> list[KVCacheEvent]: - self._flush_all_removed_events_unlocked() - events = self._pending_events - self._pending_events = [] - self._latest_stored_events.clear() - return events - - def _publish_events_unlocked( - self, - events: Sequence[KVCacheEvent], - *, - max_kv_event_entries: int | None = None, - ) -> None: - if not events: - return - if max_kv_event_entries is None: - max_kv_event_entries = self._max_kv_event_entries - self._events.extend(events) - while len(self._events) > max_kv_event_entries: - self._events.popleft() - - @staticmethod - def _trim_events( - events: Sequence[KVCacheEvent], max_kv_event_entries: int - ) -> list[KVCacheEvent]: - if max_kv_event_entries <= 0: - return [] - if len(events) <= max_kv_event_entries: - return list(events) - return list(events[-max_kv_event_entries:]) - - def _get_window_size(self, layer_group_id: LayerGroupId) -> int: - if layer_group_id is None: - return self._window_size - return self._window_size_by_layer_group.get(int(layer_group_id), self._window_size) - - @staticmethod - def _iter_block_hashes(block_hashes: BlockHashesLike) -> Iterable[BlockHashLike]: - if isinstance(block_hashes, (bytes, str, int)): - return (block_hashes,) - return block_hashes - - def _normalize_block_hash(self, block_hash: BlockHashLike) -> EventBlockHash: - if isinstance(block_hash, bytes): - if self._hash_algo == KV_CACHE_HASH_ALGO_V2_SHA256_64: - return truncate_sha256_hash_to_int64(block_hash) - return block_hash.hex() - return block_hash - - def _get_stored_block_hash(self, block_hash: BlockHashLike) -> EventBlockHash | None: - if isinstance(block_hash, bytes): - state = self._stored_blocks.get(block_hash) - return None if state is None else state.block_hash - return block_hash - - def _pop_stored_block_state( - self, block_hash: BlockHashLike - ) -> tuple[EventBlockHash, set[int]] | None: - if isinstance(block_hash, bytes): - state = self._stored_blocks.pop(block_hash, None) - if state is None: - return None - self._drop_hash_cache(block_hash) - return state.block_hash, set(state.life_cycle_ids) - return block_hash, set() - - def _pop_stored_life_cycle_block_state( - self, block_hash: bytes, life_cycle_id: int - ) -> tuple[EventBlockHash, int, bool] | None: - state = self._stored_blocks.get(block_hash) - if state is None or not state.life_cycle_ids: - return None - - life_cycle_id = int(life_cycle_id) - if life_cycle_id not in state.life_cycle_ids: - return None - - state.life_cycle_ids.remove(life_cycle_id) - is_last_life_cycle = not state.life_cycle_ids - if is_last_life_cycle: - self._stored_blocks.pop(block_hash, None) - self._drop_hash_cache(block_hash) - return state.block_hash, life_cycle_id, is_last_life_cycle - - def _drop_hash_cache(self, block_hash: bytes) -> None: - self._v1_hash_by_block_key.pop(block_hash, None) - self._v1_hash_compatible_keys.discard(block_hash) - self._v1_root_attrs_by_block_key.pop(block_hash, None) - - @staticmethod - def _normalize_token(token: TokenIdExt) -> UniqueToken: - if isinstance(token, bytes): - return UniqueToken(token.hex()) - return UniqueToken(int(token)) - - def _mm_keys_from_radix_block(self, block: Any) -> list[MmKey]: - """Decode MM runs from actual block tokens and inherited item context.""" - if not self.needs_token_digest_context(): - return [] - id_offset = self._mm_token_id_offset - assert id_offset is not None - parent = block.prev - digest = parent.last_token_digest if parent.ordinal >= 0 else None - in_mm_run = False - mm_keys: list[MmKey] = [] - for token in block.tokens: - if isinstance(token, bytes): - digest = token - mm_keys.append((digest, 0)) - in_mm_run = True - elif token > id_offset and digest is not None: - if not in_mm_run: - mm_keys.append((digest, int(token) - id_offset)) - in_mm_run = True - else: - # Text may separate runs of the same item, so retain its digest. - in_mm_run = False - return mm_keys - - def _stored_block_from_radix_block( - self, block: Any, life_cycle_ids: set[int] | None = None - ) -> KVCacheStoredBlockData | None: - cache_level: CacheLevel = GPU_LEVEL - priority: Priority = PRIORITY_DEFAULT - found_page = False - for life_cycle_id in range(len(block.storage)): - if life_cycle_ids is not None and life_cycle_id not in life_cycle_ids: - continue - page = block.get_page(life_cycle_id) - if page is None or page.num_tokens_in_block < len(block.tokens): - continue - cache_level = page.cache_level - priority = page.priority - found_page = True - break - - if life_cycle_ids is not None and not found_page: - return None - - return KVCacheStoredBlockData( - block_hash=self._hash_from_radix_block(block), - tokens=[self._normalize_token(token) for token in block.tokens], - cache_level=int(cache_level), - priority=int(priority), - mm_keys=self._mm_keys_from_radix_block(block), - ) - - @staticmethod - def _life_cycle_ids_from_radix_block(block: Any) -> set[int]: - return { - life_cycle_id - for life_cycle_id in range(len(block.storage)) - if (page := block.get_page(life_cycle_id)) is not None - and page.num_tokens_in_block >= len(block.tokens) - } - - def _parent_hash_from_radix_block(self, block: Any) -> EventBlockHash | None: - parent = block.prev - if getattr(parent, "ordinal", -1) == -1: - return None - return self._hash_from_radix_block(parent) - - def _hash_from_radix_block(self, block: Any) -> EventBlockHash: - if self._hash_algo == KV_CACHE_HASH_ALGO_V1: - return self._v1_hash_from_radix_block(block) - return self._normalize_block_hash(getattr(block, "event_key", block.key)) - - def _v1_hash_from_radix_block(self, block: Any) -> int: - key = bytes(block.key) - cached = self._v1_hash_by_block_key.get(key) - if cached is not None: - return cached - - chain: list[Any] = [] - current = block - while self._is_radix_block(current): - current_key = bytes(current.key) - cached = self._v1_hash_by_block_key.get(current_key) - if cached is not None: - parent_hash = cached - parent_is_v1_compatible = current_key in self._v1_hash_compatible_keys - root_attrs = self._v1_root_attrs_by_block_key[current_key] - break - chain.append(current) - current = current.prev - - if not self._is_radix_block(current): - parent_hash = 0 - parent_is_v1_compatible = True - root_attrs = self._root_attrs_from_root_block(current) - - lora_task_id, cache_salt_id = root_attrs - for current in reversed(chain): - current_key = bytes(current.key) - if parent_is_v1_compatible: - try: - parent_hash = self._hash_block_key( - current.tokens, - parent_hash, - lora_task_id, - cache_salt_id, - ) - self._v1_hash_compatible_keys.add(current_key) - except NonTextTokenHashError: - parent_hash = self._fallback_v1_hash(current_key) - parent_is_v1_compatible = False - else: - parent_hash = self._fallback_v1_hash(current_key) - self._v1_hash_by_block_key[current_key] = parent_hash - self._v1_root_attrs_by_block_key[current_key] = root_attrs - - return parent_hash - - def _fallback_v1_hash(self, block_key: bytes) -> int: - if not self._warned_v1_hash_fallback: - logger.warning( - "V2 KV cache event hash algorithm %s only matches v1 for " - "text-token radix blocks. Falling back to truncated V2 block " - "hash for unsupported blocks.", - KV_CACHE_HASH_ALGO_V1, - ) - self._warned_v1_hash_fallback = True - return truncate_sha256_hash_to_int64(block_key) - - @staticmethod - def _is_radix_block(value: Any) -> bool: - return hasattr(value, "tokens") and hasattr(value, "key") - - @staticmethod - def _root_attrs_from_root_block(root: Any) -> tuple[int | None, int | None]: - # Read from the new location first; fall back to the legacy attribute - # names so V1-compat event hashes still resolve for any in-memory - # RootBlock predating the refactor. - scope = getattr(root, "reuse_scope", None) - if scope is not None: - return getattr(scope, "lora_id", None), getattr(scope, "salt", None) - return getattr(root, "lora_task_id", None), getattr(root, "cache_salt_id", None) - - @staticmethod - def _hash_block_key( - tokens: Sequence[TokenIdExt], - parent_hash: int, - lora_task_id: int | None, - cache_salt_id: int | None, - ) -> int: - return hash_v1_block_key( - tokens, - parent_hash=parent_hash, - lora_task_id=lora_task_id, - cache_salt_id=cache_salt_id, - ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/__init__.py deleted file mode 100644 index 0e22e1e7c6a8..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# 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. - -from ._eviction_controller import ( # noqa: E402 - EvictablePage, - EvictionPolicy, - NodeRef, - PerLevelEvictionController, -) - -__all__ = ["EvictionPolicy", "PerLevelEvictionController", "EvictablePage", "NodeRef"] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py deleted file mode 100644 index cf1035c221eb..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py +++ /dev/null @@ -1,242 +0,0 @@ -# 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. - -from typing import Callable, Iterator, Protocol, cast - -from llist import dllist, dllistnode - -from .._common import NDEBUG, CacheLevel, PageStatus, Priority -from .._exceptions import OutOfPagesError -from .._life_cycle_registry import LifeCycleId -from .._storage._core import PoolGroupIndex -from .._utils import ( - TypedIndexList, - assert_critical, - make_typed, - noexcept, - typed_enumerate, - typed_len, - unwrap_optional, -) - - -# @runtime_checkable -class EvictablePage(Protocol): - @property - def cache_level(self) -> CacheLevel: ... - - @property - def priority(self) -> Priority: ... - - @property - def life_cycle(self) -> LifeCycleId: ... - - @property - def status(self) -> PageStatus: ... - - def is_committed(self) -> bool: ... - - node_ref: "NodeRef | None" - - -# @runtime_checkable -class NodeRef(Protocol): - @property - def value(self) -> EvictablePage: ... - - -# @runtime_checkable -class EvictionPolicy(Protocol): - def push(self, page: EvictablePage, evict_first: bool = False) -> NodeRef: ... - - def pop(self) -> EvictablePage: ... - - # Remove a node so we no longer consider it for eviction. Like pop() but allow removing a node - # that is not the first. - def remove(self, node: NodeRef) -> EvictablePage: ... - - def __len__(self) -> int: ... - - def __iter__(self) -> Iterator[EvictablePage]: ... - - -class LRUEvictionPolicy: - __slots__ = ("_queue",) - _queue: dllist - - def __init__(self) -> None: - self._queue = dllist() - - def push(self, page: EvictablePage, evict_first: bool = False) -> dllistnode: - assert page.node_ref is None - return self._queue.appendleft(page) if evict_first else self._queue.append(page) - - def pop(self) -> EvictablePage: - victim = self._queue.first - assert victim is not None - page = victim.value - self.remove(victim) - return page - - def remove(self, node: dllistnode) -> EvictablePage: - # assert isinstance(node, NodeRef) # mypyc does not support runtime_checkable - assert node == node.value.node_ref - return self._queue.remove(node) - - def __len__(self) -> int: - return len(self._queue) - - def __iter__(self) -> Iterator[EvictablePage]: - return iter(self._queue) - - -# helper class to help add support for priority-based eviction -class PrioritizedEvictionPolicy: - __slots__ = ( - "_policy_creator", - "_policies", - ) - _policy_creator: Callable[[Priority], EvictionPolicy] - _policies: dict[Priority, EvictionPolicy] - - def __init__(self, policy_creator: Callable[[Priority], EvictionPolicy]) -> None: - self._policy_creator = policy_creator - self._policies = {} - - def __len__(self) -> int: - return sum(len(policy) for policy in self._policies.values()) - - def get_policy(self, priority: Priority) -> EvictionPolicy: - if priority not in self._policies: - self._policies[priority] = self._policy_creator(priority) - self._policies = dict(sorted(self._policies.items())) - return self._policies[priority] - - def _front_policy(self) -> EvictionPolicy: - return next(iter(self._policies.values())) - - def push(self, page: EvictablePage, evict_first: bool = False) -> NodeRef: - return self.get_policy(page.priority).push(page, evict_first) - - def pop(self) -> EvictablePage: - return self._front_policy().pop() - - def remove(self, node: NodeRef) -> EvictablePage: - page = node.value - policy = self._policies[page.priority] - policy.remove(node) - if not policy: - self._policies.pop(page.priority) - return page - - def __iter__(self) -> Iterator[EvictablePage]: - for p in self._policies.values(): - yield from p - - -class PrioritizedLRUEvictionPolicy(PrioritizedEvictionPolicy): - __slots__ = () - - def __init__(self) -> None: - super().__init__(lambda priority: LRUEvictionPolicy()) - - -class PerLevelEvictionController: # for one cache level - __slots__ = ("_life_cycle_grouping", "_policies", "_cache_level") - _life_cycle_grouping: TypedIndexList[LifeCycleId, PoolGroupIndex] - _policies: TypedIndexList[PoolGroupIndex, EvictionPolicy] - _cache_level: CacheLevel - - def __init__( - self, - life_cycle_grouping: TypedIndexList[LifeCycleId, PoolGroupIndex], - cache_level: CacheLevel, - ): - self._cache_level = cache_level - self._life_cycle_grouping = life_cycle_grouping - num_pool_groups = max(life_cycle_grouping) + 1 - assert num_pool_groups == len(set(life_cycle_grouping)) - self._policies = cast( - TypedIndexList, [PrioritizedLRUEvictionPolicy() for _ in range(num_pool_groups)] - ) - - def __del__(self) -> None: - if not NDEBUG: - assert_critical( - all(len(p) == 0 for p in self._policies), "Eviction controller is not empty" - ) - - def _get_policy(self, life_cycle: LifeCycleId) -> EvictionPolicy: - pg_idx = self._life_cycle_grouping[life_cycle] - return self._policies[pg_idx] - - def schedule_for_eviction(self, page: EvictablePage, evict_first: bool = False): - assert page.node_ref is None and page.cache_level == self._cache_level - page.node_ref = self._get_policy(page.life_cycle).push(page, evict_first) - assert unwrap_optional(page.node_ref).value is page - - # If evicting a node makes some other nodes useless, those nodes will be returned as well. - # One example: for SWA, if the number of blocks just makes up one window size, then evicting any of - # them makes the remaining blocks useless. - # Raise if no enough pages to evict. In this case, pages are returned to the eviction queue. - def evict( - self, min_num_pages: TypedIndexList[PoolGroupIndex, int] - ) -> TypedIndexList[PoolGroupIndex, list[EvictablePage]]: - assert NDEBUG or len(min_num_pages) == self.num_pool_groups - ret = make_typed(lambda _: list[EvictablePage](), self.num_pool_groups) - try: - for pg_idx, count in typed_enumerate(min_num_pages): - if count < 0: - raise ValueError("Eviction count must be non-negative") - policy = self._policies[pg_idx] - if (len(policy) + len(ret[pg_idx])) < count: - raise OutOfPagesError(f"Not enough pages to evict in group {pg_idx}") - while len(ret[pg_idx]) < count: - page = policy.pop() - page.node_ref = None - ret[pg_idx].append(page) - for a, b in zip(ret, self._evict_dependencies(page)): - a.extend(b) - except Exception: - for p in reversed(sum(ret, [])): - self.schedule_for_eviction(p, evict_first=True) - raise - assert all(p.cache_level == self._cache_level for p in sum(ret, [])), ( - "Corrupted eviction controller" - ) - return ret - - def remove(self, node: NodeRef) -> None: - page = node.value - assert page.node_ref == node - self._get_policy(page.life_cycle).remove(node) - page.node_ref = None - - # @TODO: implement this - @noexcept - def _evict_dependencies( - self, page: EvictablePage - ) -> TypedIndexList[PoolGroupIndex, list[EvictablePage]]: - return make_typed(lambda _: list[EvictablePage](), self.num_pool_groups) - - def num_evictable_pages(self, pg_idx: PoolGroupIndex) -> int: - return len(self._policies[pg_idx]) - - @property - def num_pool_groups(self) -> PoolGroupIndex: - return typed_len(self._policies) - - def page_iterator(self, pg_idx: PoolGroupIndex) -> Iterator[EvictablePage]: - return iter(self._policies[pg_idx]) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py deleted file mode 100644 index a425342350ca..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py +++ /dev/null @@ -1,71 +0,0 @@ -# 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. - -import cuda.bindings.driver as drv - - -class OutOfMemoryError(Exception): - pass - - -class HostOOMError(OutOfMemoryError): - pass - - -class DiskOOMError(OutOfMemoryError): - pass - - -class CuOOMError(OutOfMemoryError): - pass - - -class LogicError(Exception): - """ - This exception indicates a bug in the code. - """ - - def __init__(self, message: str) -> None: - super().__init__(message) - - -class CorruptedError(Exception): - """A broken invariant was detected; this process refuses further KV cache work. - - Only the C++ backend has the poison latch that raises this, so the pure-Python backend - never does. It is defined here so that callers can catch it under either backend. - """ - - -class CuError(Exception): - error_code: drv.CUresult - - def __init__(self, error_code: drv.CUresult) -> None: - self.error_code = error_code - err, err_str = drv.cuGetErrorString(error_code) - if err != drv.CUresult.CUDA_SUCCESS: - err_str = "" - super().__init__(f"CUDA driver error: {error_code} ({err_str})") - - def __reduce__(self) -> tuple[type["CuError"], tuple[drv.CUresult]]: - return (self.__class__, (self.error_code,)) - - -class ResourceBusyError(Exception): - pass - - -class OutOfPagesError(Exception): - pass diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index e18d4f4d6bc7..e4d7cb7080e1 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -13,53 +13,61 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Internal KVCacheManagerV2 surface: no stability promise, may change with the C++ side. + +Everything here forwards to the native ``_introspection`` submodule of the bindings. The +indirection exists so callers import one Python path rather than reaching into +``tensorrt_llm.bindings`` directly, and so return values are normalised to plain Python +containers. Most hooks are white-box helpers for tests and accuracy harnesses, but a few +back production internals -- the stats report and the disaggregated bounce buffer -- so +treat this as internal API rather than test-only scaffolding. +""" + from __future__ import annotations import sys from typing import Any -class _TestBlock: - """Own a real radix-tree block and its real committed pages for a test.""" +def _cpp_introspection_module() -> Any | None: + package = sys.modules.get(__package__) + if package is None: + return None + return getattr(package, "_cpp_introspection", None) - __slots__ = ("block", "pages") - def __init__(self, block: Any, pages: list[Any]) -> None: - self.block = block - self.pages = pages +def _cpp() -> Any: + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is None: + raise RuntimeError( + "KVCacheManagerV2 introspection requires the native _introspection module, " + "which is missing from this build of tensorrt_llm.bindings" + ) + return cpp_introspection - def close(self) -> None: - for page in self.pages: - self.block.unlink_page(page.life_cycle, page) - self.pages.clear() - def __del__(self) -> None: - self.close() +#: CUDA virtual-memory primitives, forwarded rather than wrapped: they are classes the +#: caller constructs, so there is nothing to normalise. Used by the native disaggregated +#: bounce buffer to reserve one contiguous fabric region. +_FORWARDED_TYPES = ("PooledPhysMemAllocator", "VirtMem") -def _cpp_introspection_module() -> Any | None: - package = sys.modules.get(__package__) - if package is None: - return None - return getattr(package, "_cpp_introspection", None) +def __getattr__(name: str) -> Any: + # Resolved on first access, not at import, so importing this module never depends on + # the native submodule being present -- the same contract the hooks below follow. + if name in _FORWARDED_TYPES: + return getattr(_cpp(), name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def poison_for_testing(reason: str) -> None: """Set the poison latch directly, to exercise the refusal paths.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is None: - raise RuntimeError("the poison latch requires the C++ backend") - cpp_introspection.poison_for_testing(reason) + _cpp().poison_for_testing(reason) -def create_test_padding_cold_page_codec( - cold_page_bytes_by_layer: dict[int, int], -) -> Any: +def create_test_padding_cold_page_codec(cold_page_bytes_by_layer: dict[int, int]) -> Any: """Create the private native padding codec used by cold-tier end-to-end tests.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is None: - raise RuntimeError("the test padding cold-page codec requires the C++ backend") - return cpp_introspection.create_test_padding_cold_page_codec(cold_page_bytes_by_layer) + return _cpp().create_test_padding_cold_page_codec(cold_page_bytes_by_layer) def make_test_block( @@ -70,129 +78,38 @@ def make_test_block( reuse_scope: Any = None, ) -> Any: """Build a real block with real GPU pages at the requested lifecycle coverage.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return cpp_introspection.make_test_block( - manager, list(tokens), coverage_per_lc, parent, reuse_scope - ) - - from . import rawref - from ._block_radix_tree import ReuseScope, _add_or_get_existing - from ._common import GPU_LEVEL, PRIORITY_DEFAULT - from ._page import CommittedPage - - token_list = list(tokens) - if not token_list: - raise ValueError("make_test_block requires at least one token") - if len(coverage_per_lc) != manager._life_cycles.size: - raise ValueError( - "coverage_per_lc length must match the manager lifecycle count " - f"({manager._life_cycles.size})" - ) - if any(coverage < 0 or coverage > len(token_list) for coverage in coverage_per_lc): - raise ValueError("lifecycle coverage must be between zero and the block token count") - - if parent is None: - parent_block = manager._radix_tree.add_or_get_existing(reuse_scope or ReuseScope()) - else: - parent_block = parent.block - block = _add_or_get_existing(parent_block, token_list) - if block is None: - raise ValueError("make_test_block could not add the requested block") - - counts = [int(coverage > 0) for coverage in coverage_per_lc] - slots = manager._storage.new_gpu_slots(counts) - pages = [] - for life_cycle, coverage in enumerate(coverage_per_lc): - if coverage == 0: - continue - page = CommittedPage( - manager._storage, - block, - life_cycle, - GPU_LEVEL, - slots[life_cycle].pop(), - coverage, - PRIORITY_DEFAULT, - ) - block.storage[life_cycle] = rawref.ref(page) - pages.append(page) - return _TestBlock(block, pages) + return _cpp().make_test_block(manager, list(tokens), coverage_per_lc, parent, reuse_scope) def close_test_block(block: Any) -> None: """Unlink and release a test block pages before its manager shuts down.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - block.close() - return block.close() def set_test_block_page_cache_level(block: Any, life_cycle_id: int, cache_level: int) -> None: """Set synthetic page residency for attribution tests.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - cpp_introspection.set_test_block_page_cache_level(block, life_cycle_id, cache_level) - return - - from ._common import CacheLevel - - page = next((page for page in block.pages if page.life_cycle == life_cycle_id), None) - if page is None: - raise ValueError(f"test block has no page for life cycle {life_cycle_id}") - if page.scheduled_for_eviction: - page.manager.exclude_from_eviction(page) - page.cache_level = CacheLevel(cache_level) + _cpp().set_test_block_page_cache_level(block, life_cycle_id, cache_level) def test_block_key(block: Any) -> bytes: """Return a test block real radix-tree key.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return bytes(cpp_introspection.test_block_key(block)) - return bytes(block.block.key) + return bytes(_cpp().test_block_key(block)) def event_manager_add_stored_block(event_manager: Any, block: Any) -> None: """Derive and enqueue stored events from a real test block.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - cpp_introspection.event_manager_add_stored_block(event_manager, block) - return - event_manager.add_stored_block_event_from_block(block.block) + _cpp().event_manager_add_stored_block(event_manager, block) def event_manager_add_stored_life_cycle(event_manager: Any, block: Any, life_cycle_id: int) -> None: """Derive and enqueue one lifecycle stored event from a real test block.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - cpp_introspection.event_manager_add_stored_life_cycle(event_manager, block, life_cycle_id) - return - event_manager.add_stored_life_cycle_event_from_block(block.block, life_cycle_id) + _cpp().event_manager_add_stored_life_cycle(event_manager, block, life_cycle_id) def active_page_stats(kv_cache: Any) -> tuple[list[int], list[int]]: """Return active pages and unscheduled evictable active pages by cache level.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - counts, unscheduled_evictable = cpp_introspection.active_page_stats(kv_cache) - return list(counts), list(unscheduled_evictable) - - storage = kv_cache.manager._storage - counts = [0] * storage.num_cache_levels - unscheduled_evictable = [0] * storage.num_cache_levels - for ordinal, beam_idx, lc_idx in kv_cache._active_pages(): - block_page = kv_cache._page(ordinal, beam_idx, lc_idx) - if block_page is None: - continue - - page = block_page.page - level = page.cache_level - counts[level] += 1 - if storage.is_evictable(page) and not page.scheduled_for_eviction: - unscheduled_evictable[level] += 1 - return counts, unscheduled_evictable + counts, unscheduled_evictable = _cpp().active_page_stats(kv_cache) + return list(counts), list(unscheduled_evictable) def committed_page_is_linked(kv_cache: Any, ordinal: int, lc_id: int) -> bool | None: @@ -203,82 +120,37 @@ def committed_page_is_linked(kv_cache: Any, ordinal: int, lc_id: int) -> bool | free, and freed-but-mapped memory reads back plausibly enough that only a sanitizer build catches the fault itself. """ - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return cpp_introspection.committed_page_is_linked(kv_cache, ordinal, lc_id) - - from ._common import DEFAULT_BEAM_INDEX - from ._page import CommittedPage - - block_page = kv_cache._page(ordinal, DEFAULT_BEAM_INDEX, lc_id) - if block_page is None: - return None - page = block_page.page - if not isinstance(page, CommittedPage): - return None - return page.block() is not None + return _cpp().committed_page_is_linked(kv_cache, ordinal, lc_id) def all_tree_pages_droppable(manager: Any) -> bool: """Return whether every page reachable from the radix tree is droppable.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return bool(cpp_introspection.all_tree_pages_droppable(manager)) - - from ._block_radix_tree import traverse_post_order - from ._common import PageStatus - from ._utils import unwrap_rawref - - for root_block in manager._radix_tree.next.values(): - for block0 in root_block.next.values(): - for block in traverse_post_order(block0): - for page in block.storage: - if page is not None and unwrap_rawref(page).status != PageStatus.DROPPABLE: - return False - return True + return bool(_cpp().all_tree_pages_droppable(manager)) def is_commit_allowed(kv_cache: Any) -> bool: """Return whether the KV cache still allows token commits.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return bool(cpp_introspection.is_commit_allowed(kv_cache)) - return kv_cache._commit_state == kv_cache.CommitState.ALLOWED + return bool(_cpp().is_commit_allowed(kv_cache)) def current_gpu_ratio(manager: Any) -> list[float]: """Return the current GPU pool-group ratio list.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return list(cpp_introspection.current_gpu_ratio(manager)) - return list(manager._current_gpu_ratio) + return list(_cpp().current_gpu_ratio(manager)) def set_num_sampled_kv_caches(manager: Any, value: int) -> None: """Set the sampled-KV-cache counter that gates auto-tuner rebalancing.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - cpp_introspection.set_num_sampled_kv_caches(manager, value) - return - manager._num_sampled_kv_caches = value + _cpp().set_num_sampled_kv_caches(manager, value) def set_last_adjustment_time(manager: Any, value: float) -> None: - """Set the last pool-rebalance timestamp that gates the auto-tuner cooldown.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - cpp_introspection.set_last_adjustment_time(manager, value) - return - manager._last_adjustment_time = value + """Set the auto-tuner's last-adjustment timestamp.""" + _cpp().set_last_adjustment_time(manager, value) def set_target_ratio_list_gpu(manager: Any, ratios: list[float]) -> None: - """Override the target GPU pool-group ratio list (drives the next rebalance).""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - cpp_introspection.set_target_ratio_list_gpu(manager, list(ratios)) - return - manager._target_ratio_list_gpu = list(ratios) + """Set the auto-tuner's target GPU pool-group ratio list.""" + _cpp().set_target_ratio_list_gpu(manager, list(ratios)) def force_rebalance_precondition(manager: Any, skew: float = 2.0) -> None: @@ -288,8 +160,8 @@ def force_rebalance_precondition(manager: Any, skew: float = 2.0) -> None: ratio so it differs from the current ratio by more than the auto-tuner's adjustment threshold. Requires a model with >=2 pool groups (e.g. a VSWA model) and raises ``ValueError`` otherwise, so a future model change can't - silently turn a dependent test into a no-op. Backend-agnostic white-box - hook intended for accuracy tests, not production code. + silently turn a dependent test into a no-op. White-box hook intended for + accuracy tests, not production code. """ current = current_gpu_ratio(manager) if len(current) < 2: @@ -306,33 +178,17 @@ def force_rebalance_precondition(manager: Any, skew: float = 2.0) -> None: def storage_utilization(manager: Any, cache_level: int = 0) -> list[float]: """Return storage utilization by pool group for a cache level.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return list(cpp_introspection.storage_utilization(manager, cache_level)) - return list(manager._storage.get_utilization(cache_level)) + return list(_cpp().storage_utilization(manager, cache_level)) def grains_for_slots(num_slots: int, slot_size_list: list[int], granularity: int) -> int: """Return the grain count required for a pool group slot count.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return int(cpp_introspection.grains_for_slots(num_slots, slot_size_list, granularity)) - - from ._storage._core import CacheLevelStorage - - return int(CacheLevelStorage._grains_for_slots(num_slots, slot_size_list, granularity)) + return int(_cpp().grains_for_slots(num_slots, slot_size_list, granularity)) def grains_to_slots(pg_grains: int, slot_size_list: list[int], granularity: int) -> tuple[int, int]: """Return (slot count, consumed grains) for a pool group grain budget.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - slots, used = cpp_introspection.grains_to_slots(pg_grains, slot_size_list, granularity) - return int(slots), int(used) - - from ._storage._core import CacheLevelStorage - - slots, used = CacheLevelStorage._grains_to_slots(pg_grains, slot_size_list, granularity) + slots, used = _cpp().grains_to_slots(pg_grains, slot_size_list, granularity) return int(slots), int(used) @@ -344,49 +200,24 @@ def ratio_to_slot_count_list( min_slots: list[int], ) -> list[int]: """Return slot counts by pool group for a quota and ratio list.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return list( - cpp_introspection.ratio_to_slot_count_list( - quota, slot_size_lists, ratio_list, granularity, min_slots - ) - ) - - from ._storage._core import CacheLevelStorage - return list( - CacheLevelStorage.ratio_to_slot_count_list( - quota, slot_size_lists, ratio_list, granularity, min_slots - ) + _cpp().ratio_to_slot_count_list(quota, slot_size_lists, ratio_list, granularity, min_slots) ) def attention_life_cycle_ids(manager: Any) -> list[int]: """Return the lifecycle ids of all attention lifecycles, in order.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return list(cpp_introspection.attention_life_cycle_ids(manager)) - return [lc_id for lc_id, _ in manager._life_cycles.attention_life_cycles()] + return list(_cpp().attention_life_cycle_ids(manager)) def swa_life_cycle_ids(manager: Any) -> list[int]: """Return the lifecycle ids of attention lifecycles that use a sliding window.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return list(cpp_introspection.swa_life_cycle_ids(manager)) - return [ - lc_id - for lc_id, lc in manager._life_cycles.attention_life_cycles() - if lc.window_size is not None - ] + return list(_cpp().swa_life_cycle_ids(manager)) def ssm_life_cycle_id(manager: Any) -> int | None: """Return the SSM lifecycle id, or None if there is no SSM lifecycle.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return cpp_introspection.ssm_life_cycle_id(manager) - return manager._life_cycles.ssm_life_cycle_id + return _cpp().ssm_life_cycle_id(manager) def reuse_match_pages( @@ -403,29 +234,17 @@ def reuse_match_pages( See ``CommittedPage.num_tokens_in_block`` for how attention and SSM life cycles interpret the recorded token count. """ - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - num_tokens, raw_pages = cpp_introspection.reuse_match_pages( - manager, reuse_scope, list(tokens), lc_id, enable_partial - ) - pages: list[tuple[int, int | None] | None] = [] - for entry in raw_pages: - if entry is None: - pages.append(None) - else: - slot_id, num_tokens_in_block = entry - pages.append((slot_id, None if num_tokens_in_block < 0 else num_tokens_in_block)) - return num_tokens, pages - - match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) - py_pages: list[tuple[int, int | None] | None] = [] - for block in match.blocks: - page = block.get_page(lc_id) - if page is None: - py_pages.append(None) + num_tokens, raw_pages = _cpp().reuse_match_pages( + manager, reuse_scope, list(tokens), lc_id, enable_partial + ) + pages: list[tuple[int, int | None] | None] = [] + for entry in raw_pages: + if entry is None: + pages.append(None) else: - py_pages.append((page.slot_id, getattr(page, "num_tokens_in_block", None))) - return match.num_tokens, py_pages + slot_id, num_tokens_in_block = entry + pages.append((slot_id, None if num_tokens_in_block < 0 else num_tokens_in_block)) + return num_tokens, pages def reuse_match_planned_drop_counts( @@ -440,18 +259,9 @@ def reuse_match_planned_drop_counts( Returns ``(num_tokens, counts)`` where ``counts[i]`` is ``None`` when block ``i`` holds no page for lifecycle ``lc_id``, otherwise the matched page's ``planned_drop_count``. """ - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return cpp_introspection.reuse_match_planned_drop_counts( - manager, reuse_scope, list(tokens), lc_id, enable_partial - ) - - match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) - counts: list[int | None] = [] - for block in match.blocks: - page = block.get_page(lc_id) - counts.append(None if page is None else page.planned_drop_count) - return match.num_tokens, counts + return _cpp().reuse_match_planned_drop_counts( + manager, reuse_scope, list(tokens), lc_id, enable_partial + ) def compute_slots_for_batch( @@ -461,15 +271,4 @@ def compute_slots_for_batch( swa_scratch_reuse: Any = None, ) -> list[int]: """Return the minimum per-pool-group slot counts to support ``batch``.""" - cpp_introspection = _cpp_introspection_module() - if cpp_introspection is not None: - return list( - cpp_introspection.compute_slots_for_batch( - manager, batch, tokens_per_block, swa_scratch_reuse - ) - ) - return list( - manager._storage._compute_pool_group_slots_for_batch( - batch, tokens_per_block, swa_scratch_reuse - ) - ) + return list(_cpp().compute_slots_for_batch(manager, batch, tokens_per_block, swa_scratch_reuse)) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py deleted file mode 100644 index d0da6b057678..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py +++ /dev/null @@ -1,167 +0,0 @@ -# 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. - -from typing import Iterator, NamedTuple, NewType, TypeAlias, cast - -from ._common import BlockOrdinal, SlidingWindowSize -from ._config import AttentionLayerConfig, KVCacheManagerConfig, LayerConfig, SsmLayerConfig -from ._utils import HalfOpenRange, TypedIndexList, div_up, intersect, typed_enumerate - - -class AttnLifeCycle(NamedTuple): - window_size: SlidingWindowSize - num_sink_blocks: int # div_up(num_sink_tokens, tokens_per_block) - - @staticmethod - def make( - window_size: SlidingWindowSize, num_sink_tokens: int | None, tokens_per_block: int - ) -> "AttnLifeCycle": - assert tokens_per_block > 0 - assert window_size is None or window_size > 0 - assert num_sink_tokens is None or num_sink_tokens >= 0 - assert num_sink_tokens in (None, 0) or window_size is not None - num_sink_blocks = div_up(num_sink_tokens or 0, tokens_per_block) - return AttnLifeCycle(window_size, num_sink_blocks) - - def get_stale_range( - self, history_length: int, tokens_per_block: int - ) -> HalfOpenRange[BlockOrdinal]: - num_blocks = div_up(history_length, tokens_per_block) - start = BlockOrdinal(min(num_blocks, self.num_sink_blocks)) - if self.window_size is None: - return HalfOpenRange(start, start) - # `+ 1` is intentional: attention always runs for >= 1 in-flight input - # token at position `history_length`, so the live window is - # [history_length + 1 - window_size, history_length]. Do not drop it. - return HalfOpenRange( - start, - BlockOrdinal(max(start, (history_length + 1 - self.window_size) // tokens_per_block)), - ) - - -class SsmLifeCycle(NamedTuple): - def get_stale_range( - self, history_length: int, tokens_per_block: int - ) -> HalfOpenRange[BlockOrdinal]: - return HalfOpenRange(BlockOrdinal(0), BlockOrdinal(history_length // tokens_per_block)) - - -ssm_life_cycle = SsmLifeCycle() - -LifeCycleId = NewType("LifeCycleId", int) - -LifeCycle = AttnLifeCycle | SsmLifeCycle - -# For public exposure -LayerGroupId: TypeAlias = LifeCycleId - - -def make_life_cycle(layer: LayerConfig, tokens_per_block: int) -> LifeCycle: - if isinstance(layer, SsmLayerConfig): - return ssm_life_cycle - else: - assert isinstance(layer, AttentionLayerConfig) - return AttnLifeCycle.make(layer.window_size, layer.num_sink_tokens, tokens_per_block) - - -class LifeCycleRegistry: - __slots__ = ("_life_cycle_list", "_life_cycle_id_dict") - _life_cycle_list: TypedIndexList[LifeCycleId, LifeCycle] - _life_cycle_id_dict: dict[LifeCycle, LifeCycleId] - - def __init__(self, config: KVCacheManagerConfig) -> None: - self._life_cycle_list = cast(TypedIndexList[LifeCycleId, LifeCycle], []) - self._life_cycle_id_dict = dict[LifeCycle, LifeCycleId]() - for layer in config.layers: - details = make_life_cycle(layer, config.tokens_per_block) - if details not in self._life_cycle_id_dict: - assert len(self._life_cycle_id_dict) == len(self._life_cycle_list), ( - "corrupted life cycle registry" - ) - self._life_cycle_list.append(details) - self._life_cycle_id_dict[details] = LifeCycleId(len(self._life_cycle_list) - 1) - - def get_life_cycle(self, id: LifeCycleId) -> LifeCycle: - return self._life_cycle_list[id] - - def get_id(self, life_cycle_details: LifeCycle) -> LifeCycleId: - return self._life_cycle_id_dict[life_cycle_details] - - @property - def size(self) -> LifeCycleId: - assert len(self._life_cycle_list) == len(self._life_cycle_id_dict), ( - "corrupted life cycle registry" - ) - return LifeCycleId(len(self._life_cycle_list)) - - def __iter__(self) -> Iterator[LifeCycle]: - return iter(self._life_cycle_list) - - def __getitem__(self, idx: LifeCycleId) -> LifeCycle: - return self._life_cycle_list[idx] - - def items(self) -> Iterator[tuple[LifeCycleId, LifeCycle]]: - return typed_enumerate(self.get()) - - def get(self) -> TypedIndexList[LifeCycleId, LifeCycle]: - return self._life_cycle_list - - def __contains__(self, lc: LifeCycle) -> bool: - return lc in self._life_cycle_id_dict - - @property - def ssm_life_cycle_id(self) -> LifeCycleId | None: - return self._life_cycle_id_dict.get(ssm_life_cycle) - - @property - def has_ssm(self) -> bool: - return ssm_life_cycle in self._life_cycle_id_dict - - def attention_life_cycles(self) -> Iterator[tuple[LifeCycleId, AttnLifeCycle]]: - for lc_id, lc in self.items(): - if isinstance(lc, AttnLifeCycle): - yield lc_id, lc - - -def compute_scratch_range( - life_cycle: LifeCycle, - history_length: int, - capacity: int, - tokens_per_block: int, - max_rewind_len: int, -) -> HalfOpenRange[BlockOrdinal]: - """ - Range of blocks that should use scratch (shared) slots during SWA prefill. - - Scratch = stale_at_capacity ∩ input_blocks, where: - - stale_at_capacity: blocks out-of-window when all non-rewindable capacity tokens - become history. - - input_blocks: [div_up(history_length, tpb), div_up(capacity, tpb)) — new blocks - for the current chunk. Blocks before this range already contain real KV data - from previous chunks and must not be overwritten. - - max_rewind_len protects the speculative tail from scratch reuse. Those tokens may - survive after rejected draft tokens are rewound, so their KV data must remain in - normal per-block pages. - """ - if not isinstance(life_cycle, AttnLifeCycle) or life_cycle.window_size is None: - return HalfOpenRange(BlockOrdinal(0), BlockOrdinal(0)) - non_rewindable_capacity = max(0, capacity - max_rewind_len) - cap_stale = life_cycle.get_stale_range(non_rewindable_capacity, tokens_per_block) - input_range = HalfOpenRange( - BlockOrdinal(div_up(history_length, tokens_per_block)), - BlockOrdinal(div_up(capacity, tokens_per_block)), - ) - return intersect(cap_stale, input_range) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py deleted file mode 100644 index 0ace68f53278..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py +++ /dev/null @@ -1,561 +0,0 @@ -# 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. - -from collections.abc import Callable, Sequence -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, NamedTuple, cast - -from . import rawref -from ._block_radix_tree import Block -from ._common import ( - BAD_BLOCK_ORDINAL, - BAD_PAGE_INDEX, - DEFAULT_BEAM_INDEX, - GPU_LEVEL, - NDEBUG, - BeamIndex, - BlockOrdinal, - CacheLevel, - PageIndex, - PageStatus, - Priority, -) - -if TYPE_CHECKING: - from ._core._kv_cache import _KVCache - from ._storage_manager import StorageManager - - -from ._eviction_controller import NodeRef -from ._exceptions import LogicError -from ._life_cycle_registry import LifeCycleId -from ._storage._core import Slot -from ._utils import ( - CachedCudaEvent, - assert_critical, - filled_list, - get_uniform_attribute, - merge_events, - partition, - stream_wait_events, - unwrap_rawref, -) - -ReferenceType = rawref.ReferenceType - - -# We will have a huge amount of pages for large storage capacity. -# So we prefer inheritance over composition to save some memory. -@dataclass(slots=True) -class Page(Slot): - _manager: ReferenceType["StorageManager"] - life_cycle: LifeCycleId - cache_level: CacheLevel - _priority: Priority - # _holder is either None or a valid rawref. - _holder: ReferenceType["_PageHolder"] | None - node_ref: NodeRef | None - - def __del__(self) -> None: - if not NDEBUG: - assert_critical(self.status == PageStatus.DROPPABLE and not self.scheduled_for_eviction) - if self.has_valid_slot: - self.manager.release_slot(self.life_cycle, self.cache_level, self) - - @property - def manager(self) -> "StorageManager": - return unwrap_rawref(self._manager) - - @property - def priority(self) -> Priority: - return self._priority - - # prevent dropping - def hold(self) -> "_PageHolder": - if self._holder is not None: - return unwrap_rawref(self._holder) - holder = _PageHolder(self) - self._holder = rawref.ref(holder) - controller = self.manager - if self.scheduled_for_eviction and not controller.is_evictable(self): - controller.exclude_from_eviction(self) - assert not self.scheduled_for_eviction - return holder - - # Prevent eviction. You need to migrate the page to GPU later. - def lock( - self, - kv_cache: "_KVCache", - beam_index: BeamIndex, - ordinal: BlockOrdinal, - life_cycle: LifeCycleId, - skip_wait: bool = False, - ) -> "_SharedPageLock": - "If skip wait, you are responsible for making the page ready in kv_cache.cuda_stream." - return self.hold().lock(kv_cache, beam_index, ordinal, life_cycle, skip_wait) - - @property - def status(self) -> PageStatus: - if self._holder is None: - return PageStatus.DROPPABLE - lock_ref = unwrap_rawref(self._holder)._lock - if lock_ref is None: - return PageStatus.HELD - assert unwrap_rawref(lock_ref) is not None - return PageStatus.LOCKED - - @property - def scheduled_for_eviction(self) -> bool: - return self.node_ref is not None - - def is_committed(self) -> bool: - raise LogicError("Unexpected call to this implementation.") - - -@dataclass(slots=True) -class UncommittedPage(Page): - # @TODO: consider move this to _PageHolder - kv_cache: rawref.ref["_KVCache"] - ordinal: BlockOrdinal - beam_index: BeamIndex - - def is_committed(self) -> bool: - return False - - def __init__( - self, - kv_cache: "_KVCache", - ordinal: BlockOrdinal, - life_cycle: LifeCycleId, - cache_level: CacheLevel, - slot: Slot, - beam_index: BeamIndex = DEFAULT_BEAM_INDEX, - ): - self.kv_cache = rawref.ref(kv_cache) - self.ordinal = ordinal - self.beam_index = beam_index - manager = kv_cache.manager - priority = kv_cache._get_priority(ordinal, manager._life_cycles.get_life_cycle(life_cycle)) - Page.__init__( - self, - None, - CachedCudaEvent.NULL, - rawref.ref(manager._storage), - life_cycle, - cache_level, - priority, - None, - None, - ) - self.set_slot(slot) - - def convert_to_committed( - self, block: Block, ready_event: CachedCudaEvent, num_tokens_in_block: int - ) -> "CommittedPage": - """ - Moves the slot to a new committed page and add the new page to the block. - The uncommitted page becomes invalid. - - `num_tokens_in_block` records the page's token count. See - `CommittedPage.num_tokens_in_block` for its attention and SSM interpretations. - """ - assert not self.scheduled_for_eviction - # Check before building: replace_page() below drops the superseded page, so a - # failure in between must not lose a usable snapshot. - assert block.can_replace_page(self.life_cycle, num_tokens_in_block) - # If you hit this assertion failure, it's likely because you are using debugpy, which delayed GC - # for _KVCache._take_uncommitted_page(). Disable breakpoints on exceptions to avoid this issue. - assert self.status == PageStatus.DROPPABLE, "Release holder/lock first" - self.ready_event = ready_event - committed_page = CommittedPage( - self.manager, - block, - self.life_cycle, - self.cache_level, - self, - num_tokens_in_block, - self.priority, - ) - assert not self.has_valid_slot and self.ready_event is CachedCudaEvent.NULL - assert committed_page.has_valid_slot - block.replace_page(self.life_cycle, committed_page) - return committed_page - - def __del__(self) -> None: - def check_page(p: "BlockPage") -> bool: - return p is None or isinstance(p.page, CommittedPage) - - if not NDEBUG: - assert_critical( - self.life_cycle == self.manager.life_cycles.ssm_life_cycle_id - # this may not be true for C++: container size change may happen before or after element destructor. - or len(unwrap_rawref(self.kv_cache)._blocks) <= self.ordinal - or check_page( - unwrap_rawref(self.kv_cache) - ._blocks[self.ordinal] - .pages[self.beam_index][self.life_cycle] - ) - ) - Page.__del__(self) - - -@dataclass(slots=True) -class CommittedPage(Page): - """A committed page is immutable — all access after commit is read-only. - - We intentionally do not add a separate read_event to track read completion. - The inherited Slot.ready_event serves double duty: after commit or migration it - represents write completion; after _UniqPageLock is destroyed it is set to the - merged finish events of all prior readers. This means a new reader may - unnecessarily wait for a prior reader (read-after-read on immutable data), but - this is functionally correct, only occurs when the lock is fully released between - reuses, and saves one event field per committed page — a worthwhile tradeoff given - the potentially huge number of committed pages in the system. - """ - - block: rawref.ref["Block"] - # Token count recorded for this page. It is usually len(block.tokens), but a snapshot - # taken at an earlier token boundary may live in a block that spans more tokens -- see - # Block.__init__ and _KVCache._snapshot_partial_block_to_tree. - # - # Attention and SSM life cycles interpret it differently: - # * for attention pages, it is the number of leading tokens with valid per-token KV, - # so the page is reusable for any prefix up to that count (compare with `>=`); - # * for an SSM page, it is the exact recurrent-state checkpoint, so reuse must be - # truncated to exactly that boundary. - num_tokens_in_block: int - planned_drop_count: int - __rawref__: rawref.ref["CommittedPage"] - - def is_committed(self) -> bool: - return True - - def __init__( - self, - storage: "StorageManager", - block: Block, - life_cycle: LifeCycleId, - cache_level: CacheLevel, - slot: Slot, - num_tokens_in_block: int, - priority: Priority, - ): - assert 0 < num_tokens_in_block <= len(block.tokens) - self.block = rawref.ref(block) - self.num_tokens_in_block = num_tokens_in_block - self.planned_drop_count = 0 - self.__rawref__ = rawref.NULL - Page.__init__( - self, - None, - CachedCudaEvent.NULL, - rawref.ref(storage), - life_cycle, - cache_level, - priority, - None, - None, - ) - self.set_slot(slot) - - def __del__(self) -> None: - block = self.block() - # block may be None when rebase happens, i.e. another block with the same key is committed, - # replacing it, but the page is still used by a _KVCache. - if block is not None and block.unlink_page(self.life_cycle, self) is not None: - Block.clear_stale_blocks_after_page_unlink( - block, - self.life_cycle, - self.manager._life_cycles.get_life_cycle(self.life_cycle), - ) - Page.__del__(self) - self.__rawref__.invalidate() - - -@dataclass(slots=True) -class _PageHolder: - "Prevents pages from being dropped." - - page: Page - _lock: rawref.ref["_UniqPageLock"] | None = None - __rawref__: rawref.ref["_PageHolder"] = field(default_factory=lambda: rawref.NULL) - - def __init__(self, page: Page) -> None: - self.page = page - self._lock = None - self.__rawref__ = rawref.NULL - - def __del__(self) -> None: - if not NDEBUG: - assert_critical(self._lock is None) - page = self.page - page._holder = None - # If a held page was in last level cache, it was not scheduled for eviction. - if page.is_committed(): - page = cast(CommittedPage, page) - if not page.scheduled_for_eviction: - page.manager.schedule_for_eviction(page) - block = page.block() - # A page that no longer sits in its block's slot (orphaned block, or replaced - # by a page with a larger recorded token count) is unreachable for reuse, so - # keeping it in the eviction LRU would just pin a slot until memory pressure hits. - if block is None or block.is_orphan or not block.holds_page(page): - page.manager.exclude_from_eviction(page) - elif page.scheduled_for_eviction: - page = cast(UncommittedPage, self.page) - page.manager.exclude_from_eviction(self.page) - self.__rawref__.invalidate() - - # Prevent eviction. You need to migrate the page to GPU later. - def lock( - self, - kv_cache: "_KVCache", - beam_index: BeamIndex, - ordinal: BlockOrdinal, - life_cycle: LifeCycleId, - skip_wait: bool = False, - ) -> "_SharedPageLock": - if self._lock is None: - lock = _UniqPageLock(self) - self._lock = rawref.ref(lock) - else: - lock = unwrap_rawref(self._lock) - if self.page.scheduled_for_eviction: - manager = self.page.manager - manager.exclude_from_eviction(self.page) - assert not self.page.scheduled_for_eviction - return lock.share(kv_cache, beam_index, ordinal, life_cycle, skip_wait) - - -@dataclass(slots=True) -class _UniqPageLock: - "Locks pages to prevent eviction." - - holder: _PageHolder | None - finish_events: list[CachedCudaEvent] - __rawref__: rawref.ref["_UniqPageLock"] = field(default_factory=lambda: rawref.NULL) - - def __init__(self, holder: _PageHolder) -> None: - if holder.page.cache_level != CacheLevel(0): - raise ValueError("Lock can be applied only on GPU memory pages.") - self.holder = holder - self.finish_events = [] - self.__rawref__ = rawref.NULL - - def share( - self, - kv_cache: "_KVCache", - beam_index: BeamIndex, - ordinal: BlockOrdinal, - life_cycle: LifeCycleId, - skip_wait: bool, - ) -> "_SharedPageLock": - ret = _SharedPageLock(self, kv_cache, beam_index, ordinal, life_cycle, skip_wait) - return ret - - @property - def page(self) -> Page: - assert self.holder is not None - return self.holder.page - - def __del__(self) -> None: - page = self.page - if not NDEBUG: - assert_critical(page.cache_level == CacheLevel(0) and not page.scheduled_for_eviction) - # Set ready_event to the merged finish events of all readers. For committed (read-only) - # pages, this means the next reader will wait for prior reads to complete, which is - # unnecessary but correct. See the CommittedPage docstring for rationale. - page.ready_event = merge_events(self.finish_events) - assert self.holder is not None - self.holder._lock = None - if False: - if page.manager.is_evictable(page): - page.manager.schedule_for_eviction(page) - else: - # Optimized code path: - # delete holder first, so if nobody holds the page elsewhere, it becomes droppable immediately, - # before we hand it over to eviction controller. - self.holder = None - # if it's not droppable, then it means self.holder=None had no impact. We need to schedule it - # for eviction as usual. - if page.status != PageStatus.DROPPABLE and page.manager.is_evictable(page): - page.manager.schedule_for_eviction(page) - self.__rawref__.invalidate() - - def notify_finish(self, event: CachedCudaEvent): - self.finish_events.append(event) - # Avoid unbounded growth for system prompt pages shared by all requests - if len(self.finish_events) > 32: - self.finish_events = [merge_events(self.finish_events)] - - -class LockOwner(NamedTuple): - kv_cache: rawref.ref["_KVCache"] - beam_index: BeamIndex - ordinal: BlockOrdinal - life_cycle: LifeCycleId - - -@dataclass(slots=True, init=False) -class _SharedPageLock: - _uniq_lock: _UniqPageLock | None - _user: LockOwner - - @property - def page(self) -> Page: - assert self._uniq_lock is not None - return self._uniq_lock.page - - @property - def holder(self) -> _PageHolder: - assert self._uniq_lock is not None - assert self._uniq_lock.holder is not None - return self._uniq_lock.holder - - def __hash__(self) -> int: - return hash(id(self)) - - def __eq__(self, other: object) -> bool: - return self is other - - def __init__( - self, - uniq_lock: _UniqPageLock, - kv_cache: "_KVCache", - beam_index: BeamIndex, - ordinal: BlockOrdinal, - life_cycle: LifeCycleId, - skip_wait: bool, - ) -> None: - self._uniq_lock = uniq_lock - if not skip_wait: - self.page.ready_event.wait_in_stream(kv_cache.cuda_stream) - self._user = LockOwner(rawref.ref(kv_cache), beam_index, ordinal, life_cycle) - old_base_index = kv_cache._update_base_page_index( - beam_index, ordinal, life_cycle, PageIndex(self.page.slot_id) - ) - assert old_base_index == BAD_PAGE_INDEX - - def __del__(self) -> None: - if self._uniq_lock is not None: - self.unlock() - - def unlock(self) -> Page: - assert self._uniq_lock is not None - page = self.page - self._uniq_lock.notify_finish(unwrap_rawref(self._user.kv_cache).finish_event) - beam_index = self._user.beam_index - ordinal = self._user.ordinal - life_cycle = self._user.life_cycle - kv_cache = unwrap_rawref(self._user.kv_cache) - new_index = BAD_PAGE_INDEX - old_base_index = kv_cache._update_base_page_index( - beam_index, ordinal, life_cycle, new_index - ) - assert NDEBUG or old_base_index == ( - self._get_base_page_index() if ordinal != BAD_BLOCK_ORDINAL else BAD_PAGE_INDEX - ) - self._uniq_lock = None - return page - - def _get_base_page_index(self) -> PageIndex: - return PageIndex(self.page.slot_id) - - -BlockPage = _SharedPageLock | _PageHolder | None - - -class BatchedLockTarget(NamedTuple): - page: Page - beam_index: BeamIndex - ordinal: BlockOrdinal - life_cycle: LifeCycleId - - -def batched_lock_to_gpu( - kv_cache: "_KVCache", - tasks: Sequence[BatchedLockTarget], - migration_recorder: Callable[[Sequence[Page], Sequence[Slot], CacheLevel, CacheLevel], None] - | None = None, - drop_recorder: Callable[[Sequence[Page], CacheLevel], None] | None = None, -) -> list["_SharedPageLock"]: - "Lock pages after migrating all pages to GPU. If migration fails, no locking happens." - storage = kv_cache.manager._storage - assert not tasks or storage is get_uniform_attribute(tasks, lambda p: p.page.manager) - requirements = filled_list(0, storage.num_pool_groups) - scheduled_for_eviction = [t.page.scheduled_for_eviction for t in tasks] - lc2pg = storage._life_cycle_grouping - for t, e in zip(tasks, scheduled_for_eviction): - if e: - storage.exclude_from_eviction(t.page) - if t.page.cache_level == GPU_LEVEL: - continue - requirements[lc2pg[t.life_cycle]] += 1 - - try: - storage.prepare_free_slots(GPU_LEVEL, requirements, migration_recorder, drop_recorder) - partitioned = partition(tasks, lambda p: (p.page.cache_level, lc2pg[p.life_cycle])) - for (lvl, pg_idx), part in partitioned.items(): - if lvl == GPU_LEVEL: - continue - storage._batched_migrate( - pg_idx, - GPU_LEVEL, - lvl, - [p.page for p in part], - update_src=True, - migration_recorder=migration_recorder, - ) - except Exception: - for t, e in zip(tasks, scheduled_for_eviction): - if e: - storage.schedule_for_eviction(t.page) - raise - stream_wait_events(kv_cache.cuda_stream, (p.page.ready_event for p in tasks)) - return [ - page.lock(kv_cache, beam_index, ordinal, life_cycle, skip_wait=True) - for page, beam_index, ordinal, life_cycle in tasks - ] - - -@dataclass(slots=True) -class ScratchSlotLock: - slot: Slot - owner: rawref.ref["_KVCache"] - life_cycle: LifeCycleId - - def __init__( - self, slot: Slot, kv_cache: "_KVCache", life_cycle: LifeCycleId, skip_wait: bool = False - ): - if not skip_wait: - slot.ready_event.wait_in_stream(kv_cache.cuda_stream) - self.slot = slot.move_to_new_slot() - self.owner = rawref.ref(kv_cache) - self.life_cycle = life_cycle - - def detach_slot(self) -> Slot: - assert self.slot.has_valid_slot - return self.slot.move_to_new_slot() - - def unlock(self): - assert self.slot.has_valid_slot - kv_cache = unwrap_rawref(self.owner) - self.slot.ready_event = kv_cache.finish_event - kv_cache.manager._storage.release_slot(self.life_cycle, GPU_LEVEL, self.slot) - assert not self.slot.has_valid_slot - - def __del__(self): - if self.slot.has_valid_slot: - self.unlock() diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py deleted file mode 100644 index 0bb975ab790d..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py +++ /dev/null @@ -1,142 +0,0 @@ -# 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. - -import dataclasses -from dataclasses import dataclass, fields - -from ._common import CacheLevel -from ._utils import TypedIndexList - - -class _StatsDeltaMixin: - __slots__ = () - - def add(self, other) -> None: - for field in fields(self): - name = field.name - setattr(self, name, getattr(self, name) + getattr(other, name)) - - def subtract(self, other) -> None: - for field in fields(self): - name = field.name - setattr(self, name, getattr(self, name) - getattr(other, name)) - - def clear(self) -> None: - for field in fields(self): - setattr(self, field.name, 0) - - def copy(self): - return type(self)(**{field.name: getattr(self, field.name) for field in fields(self)}) - - @property - def empty(self) -> bool: - return all(getattr(self, field.name) == 0 for field in fields(self)) - - -@dataclass(slots=True) -class KVCacheStatsDelta(_StatsDeltaMixin): - alloc_total_blocks: int = 0 - alloc_new_blocks: int = 0 - reused_blocks: int = 0 - missed_blocks: int = 0 - - -@dataclass(slots=True) -class KVCacheIterationStatsDelta(_StatsDeltaMixin): - iter_alloc_total_blocks: int = 0 - iter_alloc_new_blocks: int = 0 - iter_reused_blocks: int = 0 - iter_full_reused_blocks: int = 0 - iter_partial_reused_blocks: int = 0 - iter_missed_blocks: int = 0 - iter_gen_alloc_blocks: int = 0 - iter_onboard_blocks: int = 0 - iter_onboard_bytes: int = 0 - iter_offload_blocks: int = 0 - iter_offload_bytes: int = 0 - iter_intra_device_copy_blocks: int = 0 - iter_intra_device_copy_bytes: int = 0 - # Host-tier pages released by LRU without ever being onboarded back to GPU - # in the lifetime since they were offloaded. Counted at the drop site in - # _storage_manager._prepare_free_slots when is_last_level(lvl). - iter_host_dropped_blocks: int = 0 - iter_host_dropped_bytes: int = 0 - - @property - def iter_cache_hit_rate(self) -> float: - total = self.iter_reused_blocks + self.iter_missed_blocks - if self.iter_reused_blocks == 0 or total == 0: - return 0.0 - return self.iter_reused_blocks / total - - -CountsByLevel = TypedIndexList[CacheLevel, int] -"""Counters indexed by CacheLevel: entry ``i`` belongs to the i-th configured cache tier. - -The length follows the configured tier list instead of a hard-coded gpu/host/disk split, which -is what lets a deployment with a hot and a cold GPU level report them as two distinct entries. -""" - - -def add_counts_by_level(dst: CountsByLevel, src: CountsByLevel) -> CountsByLevel: - """Element-wise accumulate, widening ``dst`` when ``src`` covers more levels.""" - if len(dst) < len(src): - dst = list(dst) + [0] * (len(src) - len(dst)) - for level, value in enumerate(src): - dst[level] += value - return dst - - -@dataclass(slots=True) -class ReusedBlocksByLevel: - """Reuse block counts split by the cache level the reused pages were resident on. - - Kept outside ``KVCacheIterationStatsDelta`` on purpose: the level count is a runtime - quantity, while that dataclass is a fixed-field record whose field-wise add/subtract - helpers assume scalar members. - """ - - full: CountsByLevel = dataclasses.field(default_factory=list) - partial: CountsByLevel = dataclasses.field(default_factory=list) - - def add(self, other: "ReusedBlocksByLevel") -> None: - self.full = add_counts_by_level(self.full, other.full) - self.partial = add_counts_by_level(self.partial, other.partial) - - @property - def empty(self) -> bool: - return not any(self.full) and not any(self.partial) - - -@dataclass(slots=True) -class SsmSnapshotIterationStatsDelta(_StatsDeltaMixin): - iter_snapshot_lookups: int = 0 - iter_snapshot_hits: int = 0 - iter_snapshot_misses: int = 0 - iter_reused_tokens: int = 0 - iter_unreused_tokens: int = 0 - iter_aligned_snapshot_hits: int = 0 - iter_unaligned_snapshot_hits: int = 0 - - @property - def iter_snapshot_hit_rate(self) -> float: - if self.iter_snapshot_hits == 0 or self.iter_snapshot_lookups == 0: - return 0.0 - return self.iter_snapshot_hits / self.iter_snapshot_lookups - - -_KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple( - field.name for field in fields(KVCacheIterationStatsDelta) -) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/__init__.py deleted file mode 100644 index 8540398cde1a..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# 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. - -from ._config import BufferId -from ._core import CacheLevelStorage - -# These are re-exported for external use -__all__ = ["CacheLevelStorage", "BufferId"] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.py deleted file mode 100644 index 6a62470e582a..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.py +++ /dev/null @@ -1,242 +0,0 @@ -# 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. - -from collections import defaultdict -from dataclasses import dataclass -from fractions import Fraction -from typing import NamedTuple, cast - -from .._common import LayerId -from .._config import CacheTierConfig, DataRole, KVCacheManagerConfig -from .._life_cycle_registry import LayerGroupId, LifeCycleId, LifeCycleRegistry, make_life_cycle -from .._storage._core import PoolGroupIndex, PoolIndex -from .._utils import ( - HomoTuple, - TypedIndexList, - exact_div, - filled_list, - get_uniform_attribute, - is_sorted, - typed_enumerate, - typed_map, - value_or, -) - - -class BufferId(NamedTuple): - layer_id: LayerId - role: DataRole - - -@dataclass(slots=True, frozen=True) -class CoalescedBuffer: - """Each coalesced buffer has multiple buffers with the same size and life cycle.""" - - single_buffer_size: int # identical for all buffers in the same coalesced buffer - buffer_ids: HomoTuple[BufferId] - - @property - def size(self) -> int: - return self.single_buffer_size * len(self.buffer_ids) - - @property - def num_buffers(self) -> int: - return len(self.buffer_ids) - - -@dataclass(slots=True, frozen=True) -class SlotDescVariant: - """ - A group of coalesced buffers for the same life cycle. Each coalesced buffer goes to one different memory pool. - These pools forms a pool group. Allocation / deallocation of buffers in these pools are mirrored. - """ - - life_cycle_id: LifeCycleId - coalesced_buffers: TypedIndexList[PoolIndex, CoalescedBuffer] - - def __post_init__(self) -> None: - assert is_sorted(self.coalesced_buffers, key=lambda s: s.size, reverse=True) - - @property - def layer_group_id(self) -> LayerGroupId: - return self.life_cycle_id - - @property - def slot_size_list(self) -> TypedIndexList[PoolIndex, int]: - return typed_map(self.coalesced_buffers, lambda cb: cb.size) - - -@dataclass(slots=True, frozen=True) -class SlotDesc: - """ - A slot in a memory pool group can have different composition, if they have the same slot_size_list. - """ - - variants: HomoTuple[SlotDescVariant] - - @property - def slot_size_list(self) -> TypedIndexList[PoolIndex, int]: - return get_uniform_attribute(self.variants, lambda s: s.slot_size_list) - - -@dataclass(slots=True, frozen=True) -class BufferAttr: - life_cycle_id: LifeCycleId - pool_index: PoolIndex - offset: int - size: int # expanded size of the buffer - expansion: int # expansion factor of page due to heterogeneous tokens_per_block - - -@dataclass(slots=True) -class LayerAttr: - life_cycle_id: LifeCycleId - # Number of sub-pages within a single coalesced slot that belong to this layer. - # When multiple buffers from the same layer are coalesced into the same pool (e.g., K and V buffers - # of the same size), this equals the number of such buffers. For scratch allocation, each layer - # needs exactly this many sub-pages per block within the slot. - slot_util: TypedIndexList[PoolIndex, int] - - # Fraction of slot_util to total number of buffers in the slot, max over all pools in the slot - slot_util_frac_max: Fraction - - -@dataclass(slots=True, frozen=True) -class StorageConfig: - cache_tiers: HomoTuple[CacheTierConfig] - slot_desc_list: TypedIndexList[PoolGroupIndex, SlotDesc] - expansion: dict[BufferId, int] # expansion factor of page due to heterogeneous tokens_per_block - - @property - def num_life_cycles(self) -> LifeCycleId: - return LifeCycleId(sum(len(pg.variants) for pg in self.slot_desc_list)) - - def life_cycle_grouping(self) -> TypedIndexList[LifeCycleId, PoolGroupIndex]: - ret = filled_list(PoolGroupIndex(-1), self.num_life_cycles) - for pg_idx, pg in enumerate(self.slot_desc_list): - pg_idx = PoolGroupIndex(pg_idx) - for s in pg.variants: - ret[s.life_cycle_id] = pg_idx - return ret - - def buffer_attributes(self) -> dict[BufferId, BufferAttr]: - ret = dict[BufferId, BufferAttr]() - for pg in self.slot_desc_list: - for slot in pg.variants: - life_cycle_id = slot.life_cycle_id - for pool, cb in enumerate(slot.coalesced_buffers): - offset = 0 - for b in cb.buffer_ids: - ret[b] = BufferAttr( - life_cycle_id, - PoolIndex(pool), - offset, - cb.single_buffer_size, - self.expansion.get(b, 1), - ) - offset += cb.single_buffer_size - return ret - - def slot_to_page_indices(self) -> TypedIndexList[LifeCycleId, TypedIndexList[PoolIndex, int]]: - ret = [[]] * self.num_life_cycles - for pg in self.slot_desc_list: - for slot in pg.variants: - life_cycle = slot.life_cycle_id - ret[life_cycle] = [cb.num_buffers for cb in slot.coalesced_buffers] - return cast(TypedIndexList[LifeCycleId, TypedIndexList[PoolIndex, int]], ret) - - def layer_attributes(self) -> dict[LayerId, LayerAttr]: - ret = dict[LayerId, LayerAttr]() - for pg in self.slot_desc_list: - for slot in pg.variants: - life_cycle_id = slot.life_cycle_id - num_pools = len(slot.coalesced_buffers) - for pool, cb in typed_enumerate(slot.coalesced_buffers): - slot_util = defaultdict[LayerId, int](int) - for b in cb.buffer_ids: - slot_util[b.layer_id] += 1 - buffers_per_slots = cb.num_buffers - for layer_id, count in slot_util.items(): - if layer_id not in ret: - ret[layer_id] = LayerAttr( - life_cycle_id, filled_list(0, num_pools), Fraction(0, 1) - ) - attr = ret[layer_id] - assert attr.life_cycle_id == life_cycle_id - attr.slot_util[pool] = count - attr.slot_util_frac_max = max( - attr.slot_util_frac_max, Fraction(count, buffers_per_slots) - ) - return ret - - def layer_to_life_cycle_ids(self) -> dict[LayerId, LifeCycleId]: - map = dict[LayerId, LifeCycleId]() - for (layer_id, _), attr in self.buffer_attributes().items(): - lc_id = map.setdefault(layer_id, attr.life_cycle_id) - assert lc_id == attr.life_cycle_id - return map - - def __post_init__(self) -> None: - groups = [tuple(s.life_cycle_id for s in pg.variants) for pg in self.slot_desc_list] - all_life_cycle_ids = sum((g for g in groups), ()) - assert len(all_life_cycle_ids) == len(set(all_life_cycle_ids)) - - -def create_storage_config(config: KVCacheManagerConfig) -> StorageConfig: - # group buffers first by life cycle, then by single buffer size. - buffer_groups = defaultdict[LifeCycleId, defaultdict[int, list[BufferId]]]( - lambda: defaultdict[int, list[BufferId]](list[BufferId]) - ) - life_cycle_registry = LifeCycleRegistry(config) - tokens_per_block = config.tokens_per_block - expansion_map = dict[BufferId, int]() - for layer in config.layers: - life_cycle = make_life_cycle(layer, tokens_per_block) - life_cycle_id = life_cycle_registry.get_id(life_cycle) - size_to_buffers = buffer_groups[life_cycle_id] - for buffer in layer.buffers: - tokens_per_block_override = value_or(buffer.tokens_per_block_override, tokens_per_block) - expansion = exact_div(tokens_per_block, tokens_per_block_override) - size = buffer.size * expansion - buf_id = BufferId(layer.layer_id, buffer.role) - expansion_map[buf_id] = expansion - size_to_buffers[size].append(buf_id) - # Create one slot group for each life cycle. - # It's possible that buffers with different sizes form coalesced buffers with the same coalesced size. - # @TODO: add test for this case. - slot_groups: list[SlotDescVariant] = [] - for life_cycle_id, size_to_buffers in buffer_groups.items(): - slots = [ - CoalescedBuffer(size, tuple(buffer_ids)) for size, buffer_ids in size_to_buffers.items() - ] - slots.sort(key=lambda p: p.size, reverse=True) - slot_groups.append( - SlotDescVariant(life_cycle_id, cast(TypedIndexList[PoolIndex, CoalescedBuffer], slots)) - ) - # Merge slot groups with the same slot_size_list - pool_groups_by_slot_size_list = defaultdict[HomoTuple[int], list[SlotDescVariant]]( - list[SlotDescVariant] - ) - for slot_group in slot_groups: - pool_groups_by_slot_size_list[tuple(slot_group.slot_size_list)].append(slot_group) - slot_desc_list = cast( - TypedIndexList[PoolGroupIndex, SlotDesc], - [SlotDesc(tuple(slot_groups)) for slot_groups in pool_groups_by_slot_size_list.values()], - ) - return StorageConfig( - cache_tiers=tuple(config.cache_tiers), - slot_desc_list=slot_desc_list, - expansion=expansion_map, - ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py deleted file mode 100644 index cf3c9c42f2a0..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py +++ /dev/null @@ -1,1000 +0,0 @@ -# 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. - -import abc -import errno -import os -import sys -import tempfile -import warnings -from collections import deque -from collections.abc import Sequence -from dataclasses import dataclass -from typing import ClassVar, NewType, final - -if sys.version_info[:2] >= (3, 12): - from typing import override -else: - from typing_extensions import override - -from .._common import ( - BAD_FILE_DESCRIPTOR, - NDEBUG, - Address, - CacheTier, - DiskAddress, - FileDescriptor, - MemAddress, -) -from .._cuda_virt_mem import PooledPhysMemAllocator, VirtMem -from .._exceptions import LogicError, OutOfPagesError -from .._utils import ( - CachedCudaEvent, - DynamicBitset, - HomoTuple, - HostMem, - TypedIndexList, - assert_critical, - div_up, - filled_list, - make_typed, - query_total_gpu_memory, - resize_file, - round_down, - round_up, - typed_enumerate, - typed_len, - typed_map, - typed_range, -) - -PoolGroupIndex = NewType("PoolGroupIndex", int) -PoolIndex = NewType("PoolIndex", int) -SlotId = NewType("SlotId", int) - -# A temporary work-around while migrating to new page index API. -# To be removed later. -PoolIndex0 = PoolIndex(0) - - -class SlotPoolBase(abc.ABC): - _slot_size: int - - @property - def slot_size(self) -> int: - return self._slot_size - - @property - @abc.abstractmethod - def num_slots(self) -> int: ... - - @property - def num_bytes(self) -> int: - return self.slot_size * self.num_slots - - def __init__(self, slot_size: int) -> None: - self._slot_size = slot_size - - @abc.abstractmethod - def destroy(self) -> None: - pass - - @abc.abstractmethod - def resize(self, new_num_slots: int) -> None: - pass - - @abc.abstractmethod - def slot_address(self, slot: SlotId) -> Address: - pass - - def __del__(self) -> None: - self.destroy() - - -@final -class GpuSlotPool(SlotPoolBase): - __slots__ = ("_vm",) - _vm: VirtMem - - def __init__( - self, - slot_size: int, - vm_size: int, - shared_phys_mem_pool: PooledPhysMemAllocator, - num_slots: int, - ): - super().__init__(slot_size) - assert vm_size % shared_phys_mem_pool.phys_mem_size == 0 - self._vm = VirtMem(vm_size, shared_phys_mem_pool) - self.resize(num_slots) - - @override - def destroy(self) -> None: - self._vm.destroy() - - @override - def resize(self, new_num_slots: int) -> None: - new_num_phys_mem = self._compute_num_phys_mem( - self.slot_size, new_num_slots, self._vm.phys_mem_size - ) - self._vm.realloc(self._vm.phys_mem_size * new_num_phys_mem) - - def extend_by_one_phys_mem(self) -> int: - self._vm.extend(1) - return self.num_slots - - @override - def slot_address(self, slot: SlotId) -> MemAddress: - return MemAddress(int(self._vm.address) + self.slot_size * int(slot)) - - @property - @override - def num_slots(self) -> int: - return self._compute_num_slots( - self.slot_size, self._vm.num_phys_mem, self._vm.phys_mem_size - ) - - @staticmethod - def _compute_num_phys_mem(slot_size: int, num_slots: int, phys_mem_size: int) -> int: - return div_up(num_slots * slot_size, phys_mem_size) - - @staticmethod - def _compute_num_slots(slot_size: int, num_phys_mem: int, phys_mem_size: int) -> int: - return num_phys_mem * phys_mem_size // slot_size - - -class HostSlotPool(SlotPoolBase): - __slots__ = ("_host_mem",) - _host_mem: HostMem - - def __init__(self, slot_size: int, num_slots: int) -> None: - super().__init__(slot_size) - self._host_mem = HostMem(self.aligned_size(num_slots)) - - @override - def destroy(self) -> None: - self._host_mem.destroy() - - @override - def resize(self, new_num_slots: int) -> None: - self._host_mem.resize(self.aligned_size(new_num_slots)) - - @override - def slot_address(self, slot: SlotId) -> MemAddress: - return MemAddress(self._host_mem._address + self.slot_size * int(slot)) - - @property - @override - def num_slots(self) -> int: - return self._host_mem.size // self.slot_size - - def aligned_size(self, num_slots: int) -> int: - return round_up(num_slots * self.slot_size, HostMem.ALIGNMENT) - - -class DiskSlotPool(SlotPoolBase): - __slots__ = ("_filename", "_fd") - # Currently only used to get the parent folder where we create temporary files. - # You won't find file with this name. - filename: str - _fd: FileDescriptor - - def __init__(self, filename: str, slot_size: int, num_slots: int) -> None: - super().__init__(slot_size) - self.filename = filename - folder = os.path.dirname(filename) - assert os.path.isdir(folder), f"Folder {folder} does not exist" - try: - fd = os.open(folder, os.O_TMPFILE | os.O_RDWR | os.O_EXCL, 0o664) - except OSError as e: - if e.errno != errno.EOPNOTSUPP: - raise - # Fallback for filesystems/architectures not supporting O_TMPFILE - fd, path = tempfile.mkstemp(dir=folder) - try: - os.unlink(path) - except OSError: - os.close(fd) - raise - self._fd = FileDescriptor(fd) - self.resize(num_slots) - - @override - def destroy(self) -> None: - if self.fd == BAD_FILE_DESCRIPTOR: - return - os.close(self.fd) - self._fd = BAD_FILE_DESCRIPTOR - - @property - def fd(self) -> FileDescriptor: - return self._fd - - @property - def file_size(self) -> int: - return os.lseek(self.fd, 0, os.SEEK_END) - - @override - def resize(self, new_num_slots: int) -> None: - file_size = new_num_slots * self.slot_size - resize_file(self.fd, file_size) - - @override - def slot_address(self, slot: SlotId) -> DiskAddress: - assert slot < self.num_slots - return DiskAddress(self.fd, slot * self.slot_size) - - @property - @override - def num_slots(self) -> int: - return self.file_size // self.slot_size - - -@dataclass(slots=True) -class Slot: - # ready_event indicates whether the slot is ready for use. - # For newly allocated BlockData, it indicates finish of the last usage by the previous owners of the - # slot (who returned the slot to the pool). - # After data migration, it indicates finish of data migration. - # When passed to release(), it indicates finish of usage by the current owners of the slot. - _slot_id: SlotId | None - ready_event: CachedCudaEvent - - @property - def slot_id(self) -> SlotId: - assert self._slot_id is not None - return self._slot_id - - def query_ready(self) -> bool: - ev = self.ready_event - if ev is CachedCudaEvent.NULL: - return True - ret = ev.query_complete() - if ret: - self.ready_event = CachedCudaEvent.NULL - return ret - - @property - def has_valid_slot(self) -> bool: - return self._slot_id is not None - - def move_to_new_slot(self) -> "Slot": - ret = Slot(None, CachedCudaEvent.NULL) - ret.set_slot(self) - return ret - - def set_slot(self, slot: "Slot") -> None: - if self.has_valid_slot: - raise LogicError("Slot is already set.") - self._slot_id = slot.slot_id - self.ready_event = slot.ready_event - slot._slot_id = None - slot.ready_event = CachedCudaEvent.NULL - - def __del__(self) -> None: - if self.has_valid_slot: - warnings.warn("[KVCacheManager] slot is not freed before deletion") - - -class SlotAllocator: - __slots__ = ( - "_capacity", - "_num_active_slots", - "_recycled_slots", - "_num_ready_recycled_slots", - "_occupied_mask", - "_target_capacity", - "_overflow_slots", - ) - _capacity: int - _num_active_slots: int # active slots are either in use or recycled. - _recycled_slots: deque[ - Slot - ] # only store recycled slots to avoid excessive memory usage on program start - _num_ready_recycled_slots: int # number of recycled slots that are ready to be used immediately - # (no need for sync or wait in stream), i.e. their ready events are triggered. - _occupied_mask: DynamicBitset - - # for scheduled shrinking resize - _target_capacity: ( - int # _target_capacity <= _capacity. Inequal if a shrinking resize is in progress. - ) - _overflow_slots: list[ - Slot - ] # slots that will be out-of-range after a in-progress resize. scheduled for removal. - - def __init__(self, capacity: int) -> None: - self._capacity = capacity - self._num_active_slots = 0 - self._recycled_slots = deque[Slot]() - self._num_ready_recycled_slots = 0 - self._occupied_mask = DynamicBitset(capacity) - self._target_capacity = capacity - self._overflow_slots = [] - - def __del__(self) -> None: - assert_critical( - self._num_ready_recycled_slots == len(self._recycled_slots), - "did you call synchronize()?", - ) - assert_critical( - self._target_capacity == self._capacity and not self._overflow_slots, - "resize is in progress", - ) - assert_critical(self._occupied_mask.num_set_bits == 0, "some slots are still in use") - assert_critical( - len(self._recycled_slots) == self._num_active_slots, "some slots are not free" - ) - - @property - def num_free_slots(self) -> int: - return len(self._recycled_slots) + max(self._target_capacity - self._num_active_slots, 0) - - @property - def num_occupied_slots(self) -> int: - return self._occupied_mask.num_set_bits - - def allocate(self) -> Slot: - if self.num_free_slots == 0: - raise OutOfPagesError("No free slots") - self._scrub_events() - # prefererence: ready recycled slots > new slots > recycled slots that are not ready - if self._num_ready_recycled_slots > 0: - assert self._recycled_slots - slot = self._recycled_slots.popleft() - assert slot.has_valid_slot - self._num_ready_recycled_slots -= 1 - assert slot.ready_event is CachedCudaEvent.NULL - elif self._num_active_slots < min(self.num_slots, self._target_capacity): - slot = Slot(SlotId(self._num_active_slots), CachedCudaEvent.NULL) - self._num_active_slots += 1 - else: - slot = self._recycled_slots.popleft() - assert slot.has_valid_slot - self._occupied_mask.set(slot.slot_id) - return slot - - # The reason why we don't use allocate() multiple times is that if what user need is all or none, - # and when we don't have enough free slots, we will free these newly allocated slots by appending - # them to the back of the recycled slot queue, which may impact perf. - def allocate_multiple(self, num_slots: int) -> list[Slot]: - if num_slots < 0: - raise LogicError("SlotAllocator.allocate_multiple: slot count must be non-negative") - if self.num_free_slots < num_slots: - raise OutOfPagesError("Not enough free slots") - return [self.allocate() for _ in range(num_slots)] - - def release(self, slot: Slot) -> None: - assert slot.has_valid_slot - slot = slot.move_to_new_slot() - if slot.slot_id >= self._capacity or not self._occupied_mask.get(slot.slot_id): - raise LogicError(f"Slot {slot.slot_id} is not occupied") - assert type(slot) is Slot and slot.has_valid_slot - if slot.slot_id < self._target_capacity: - self._recycled_slots.append(slot) - else: - self._overflow_slots.append(slot) - self._occupied_mask.clear(slot.slot_id) - self._scrub_events() - assert NDEBUG or self._check() - - @property - def num_slots(self) -> int: - return self._capacity - - def expand(self, new_num_slots: int) -> None: - assert NDEBUG or self._check() - assert self._target_capacity == self._capacity - assert new_num_slots > self._capacity - self._occupied_mask.resize(new_num_slots) - self._capacity = new_num_slots - self._target_capacity = self._capacity - assert NDEBUG or self._check() - - def prepare_for_shrink(self, new_num_slots: int) -> None: - assert NDEBUG or self._check() - assert self._target_capacity == self._capacity - assert new_num_slots < self._capacity - new_recycled_slots = deque[Slot]() - new_num_ready_recycled_slots = 0 - old_num_ready_recycled_slots = self._num_ready_recycled_slots - for i, slot in enumerate(self._recycled_slots): - if slot.slot_id < new_num_slots: - new_recycled_slots.append(slot) - if i < old_num_ready_recycled_slots: - new_num_ready_recycled_slots += 1 - else: - self._overflow_slots.append(slot) - self._recycled_slots = new_recycled_slots - self._num_ready_recycled_slots = new_num_ready_recycled_slots - self._target_capacity = new_num_slots - assert NDEBUG or self._check() - - @property - def shrink_in_progress(self) -> bool: - "Indicates if a scheduled shrink is in progress." - assert self._target_capacity <= self._capacity - return self._target_capacity < self._capacity - - def finish_shrink(self) -> bool: - assert NDEBUG or self._check() - # Overflow-range IDs that were ever issued are exactly - # max(0, _num_active_slots - _target_capacity); the underused case - # (_num_active_slots <= _target_capacity) collapses to zero. - expected_overflow = max(0, self._num_active_slots - self._target_capacity) - if self.shrink_in_progress and len(self._overflow_slots) == expected_overflow: - assert len(set(s.slot_id for s in self._overflow_slots)) == len(self._overflow_slots), ( - "Some slots are still in use." - ) - for ev in set(s.ready_event for s in self._overflow_slots): - ev.synchronize() - for slot in self._overflow_slots: - slot.ready_event = CachedCudaEvent.NULL - slot._slot_id = None - self._overflow_slots.clear() - self._capacity = self._target_capacity - self._num_active_slots = min(self._num_active_slots, self._capacity) - self._scrub_events() - assert NDEBUG or self._check() - return True - raise RuntimeError("shrink can't be finished") - - def get_slots_blocking_shrink(self) -> HomoTuple[SlotId]: - return tuple( - SlotId(id) - for id in range(self._target_capacity, self._capacity) - if self._occupied_mask.get(id) - ) - - def _scrub_events(self) -> None: - self._num_ready_recycled_slots = self._scrub_events_impl( - self._recycled_slots, self._num_ready_recycled_slots - ) - - def _check(self) -> bool: - return ( - self._num_active_slots <= self._capacity - and self._target_capacity <= self._capacity - and (self.shrink_in_progress or len(self._overflow_slots) == 0) - and all( - self._target_capacity <= slot.slot_id < self._capacity - for slot in self._overflow_slots - ) - and len(self._recycled_slots) + len(self._overflow_slots) + self.num_occupied_slots - == self._num_active_slots - ) - - @staticmethod - def _scrub_events_impl(slots: Sequence[Slot], num_ready: int) -> int: - assert num_ready <= len(slots) - for i in range(num_ready, len(slots)): - slot = slots[i] - if slot.ready_event.query_complete(): - slot.ready_event = CachedCudaEvent.NULL - num_ready += 1 - else: - break - return num_ready - - def _synchronize(self) -> None: - "synchronize the events of all unused slots" - while self._num_ready_recycled_slots != len(self._recycled_slots): - self._scrub_events() - - -class PoolGroupBase: - __slots__ = ("_slot_allocator", "_pools", "_destroyed") - - _slot_allocator: SlotAllocator - _pools: TypedIndexList[PoolIndex, SlotPoolBase] - _destroyed: bool - - def __init__(self, num_slots: int) -> None: - self._slot_allocator = SlotAllocator(num_slots) - self._destroyed = False - - def __del__(self) -> None: - self.destroy() - - def destroy(self) -> None: - if self._destroyed: - return - allocator = self._slot_allocator - if allocator.num_slots != 0: - allocator._synchronize() - allocator.prepare_for_shrink(0) - allocator.finish_shrink() - for pool in self._pools: - pool.destroy() - self._destroyed = True - - @property - def num_pools(self) -> PoolIndex: - return PoolIndex(len(self._pools)) - - @property - def num_slots(self) -> int: - num_slots = self._slot_allocator.num_slots - assert num_slots <= self._get_num_slots_from_pools() - return num_slots - - @property - def num_free_slots(self) -> int: - return self._slot_allocator.num_free_slots - - @property - def num_bytes(self) -> int: - return sum(pool.num_bytes for pool in self._pools) - - def resize_pools(self, new_num_slots: int | None) -> None: - """ - Resize the pools, but not the slot allocator. If new_num_slots is None, make pool sizes match - the slot allocator. - If exception is raised, size of pools may be imbalanced. Call resize_pools() again with None or - self._get_num_slots_from_pools() to fix it. - """ - if new_num_slots is None: - new_num_slots = self._slot_allocator.num_slots - for pool in self._pools: - pool.resize(new_num_slots) - assert NDEBUG or self._check(True) - - def allocate(self) -> Slot: - return self._slot_allocator.allocate() - - def allocate_multiple(self, num_slots: int) -> list[Slot]: - return self._slot_allocator.allocate_multiple(num_slots) - - def release(self, slot: Slot) -> None: - self._slot_allocator.release(slot) - - def slot_address(self, slot_id: SlotId) -> HomoTuple[Address]: - return tuple(pool.slot_address(slot_id) for pool in self._pools) - - @property - def slot_size(self) -> TypedIndexList[PoolIndex, int]: - return typed_map(self._pools, lambda pg: pg.slot_size) - - def _check(self, allow_mismatch: bool = False) -> bool: - pool_num_slots = self._get_num_slots_from_pools() - return ( - self._slot_allocator.num_slots <= pool_num_slots - if allow_mismatch - else self._slot_allocator.num_slots == pool_num_slots - ) - - def _get_num_slots_from_pools(self) -> int: - return min(p.num_slots for p in self._pools) - - @staticmethod - def _compute_num_phys_mem( - slot_size_list: Sequence[int], num_slots: int, phys_mem_size: int - ) -> HomoTuple[int]: - return tuple( - GpuSlotPool._compute_num_phys_mem(slot_size, num_slots, phys_mem_size) - for slot_size in slot_size_list - ) - - -class GpuPoolGroup(PoolGroupBase): - __slots__ = () - - def __init__( - self, - num_slots: int, - slot_size_list: TypedIndexList[PoolIndex, int], - shared_phys_mem_pool: PooledPhysMemAllocator, - ): - super().__init__(num_slots) - total_gpu_memory = query_total_gpu_memory() - max_slot_size = max(slot_size_list) - phys_mem_size = shared_phys_mem_pool.phys_mem_size - self._pools = typed_map( - slot_size_list, - lambda slot_size: GpuSlotPool( - slot_size, - max( - round_down(int(total_gpu_memory * slot_size / max_slot_size), phys_mem_size), - round_up(num_slots * slot_size, phys_mem_size), - ), - shared_phys_mem_pool, - num_slots, - ), - ) - - -class HostPoolGroup(PoolGroupBase): - __slots__ = () - - def __init__(self, num_slots: int, slot_size_list: TypedIndexList[PoolIndex, int]): - super().__init__(num_slots) - self._pools = typed_map( - slot_size_list, lambda slot_size: HostSlotPool(slot_size, num_slots) - ) - - -class DiskPoolGroup(PoolGroupBase): - __slots__ = () - - def __init__( - self, num_slots: int, slot_size_list: TypedIndexList[PoolIndex, int], filename_template: str - ): - super().__init__(num_slots) - num_pools = typed_len(slot_size_list) - self._pools = make_typed( - lambda pool_idx: DiskSlotPool( - filename_template.format(pool_idx), slot_size_list[pool_idx], num_slots - ), - num_pools, - ) - - -class CacheLevelStorage: - TIER: ClassVar[CacheTier] - __slots__ = "_pool_groups" - # _total_quota: int # fixme: remove _total_quota and _ratio_list and compute from _pool_groups - # _ratio_list: TypedIndexList[PoolGroupIndex, float] - _pool_groups: TypedIndexList[PoolGroupIndex, PoolGroupBase] - - def __init__(self) -> None: - if not hasattr(self.__class__, "TIER"): - raise ValueError(f"{self.__class__.__name__} must define 'TIER' as a class variable") - - def __del__(self) -> None: - self.destroy() - - @property - def cache_tier(self) -> CacheTier: - return self.TIER - - def destroy(self) -> None: - for pg in self._pool_groups: - pg.destroy() - - def allocate(self, pool_group_index: PoolGroupIndex) -> Slot: - return self._pool_groups[pool_group_index].allocate() - - def allocate_multiple(self, pool_group_index: PoolGroupIndex, num_slots: int) -> list[Slot]: - return self._pool_groups[pool_group_index].allocate_multiple(num_slots) - - def release(self, pool_group_index: PoolGroupIndex, slot: Slot) -> None: - self._pool_groups[pool_group_index].release(slot) - - @property - def total_quota(self) -> int: - granularity = self.pool_size_granularity - quota = 0 - for pg in self._pool_groups: - for p in pg._pools: - quota += round_up(p.num_bytes, granularity) - return quota - - @property - def ratio_list(self) -> TypedIndexList[PoolGroupIndex, float]: - num_pool_groups = self.num_pool_groups - ret = filled_list(0.0, num_pool_groups) - total = 0 - for i, pg in typed_enumerate(self._pool_groups): - size = pg.num_bytes - total += size - ret[i] = size - assert total > 0 - for i in typed_range(num_pool_groups): - ret[i] /= total - return ret - - def num_slots(self, pool_group_index: PoolGroupIndex) -> int: - return self._pool_groups[pool_group_index].num_slots - - def get_num_free_slots(self, pool_group_index: PoolGroupIndex) -> int: - return self._pool_groups[pool_group_index].num_free_slots - - @property - def slot_count_list(self) -> TypedIndexList[PoolGroupIndex, int]: - """ - The number of slots in each pool group. - """ - return typed_map(self._pool_groups, lambda pg: pg.num_slots) - - def slot_size(self, pool_group_index: PoolGroupIndex) -> TypedIndexList[PoolIndex, int]: - """ - The slot sizes of each pool in the pool group. - """ - return self._pool_groups[pool_group_index].slot_size - - @property - def slot_size_lists(self) -> TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]]: - """ - A tuple of tuples, each containing the slot sizes for a pool group. - """ - return typed_map(self._pool_groups, lambda pg: typed_map(pg._pools, lambda p: p.slot_size)) - - @property - def num_pool_groups(self) -> PoolGroupIndex: - return typed_len(self._pool_groups) - - def slot_address( - self, pool_group_index: PoolGroupIndex, pool_index: PoolIndex, slot_id: SlotId - ) -> Address: - return self._pool(pool_group_index, pool_index).slot_address(slot_id) - - def post_resize(self) -> None: - pass - - def _pool(self, pool_group_index: PoolGroupIndex, pool_index: PoolIndex) -> SlotPoolBase: - return self._pool_groups[pool_group_index]._pools[pool_index] - - # Calculate how many slots will there be in each pool group with the given total_quota and - # ratio_list. Use ratio_to_slot_count_list for initialization. - def compute_slot_count_list( - self, - ratio_list: TypedIndexList[PoolGroupIndex, float], - min_slots: TypedIndexList[PoolGroupIndex, int], - total_quota: int | None = None, - ) -> TypedIndexList[PoolGroupIndex, int]: - if total_quota is None: - total_quota = self.total_quota - assert len(ratio_list) == len(self._pool_groups), ( - f"Wrong ratio_list length. Expected {len(self._pool_groups)}, got {len(ratio_list)}" - ) - return self.ratio_to_slot_count_list( - total_quota, self.slot_size_lists, ratio_list, self.pool_size_granularity, min_slots - ) - - @staticmethod - def _grains_to_slots( - pg_grains: int, - slot_size_list: TypedIndexList[PoolIndex, int], - granularity: int, - ) -> tuple[int, int]: - """Compute the maximum slots that fit in a pool group grain budget. - - Returns (num_slots, grains_consumed). - """ - if granularity <= 0 or any(slot_size <= 0 for slot_size in slot_size_list): - raise ValueError("Cache slot sizes and granularity must be positive") - num_pools = typed_len(slot_size_list) - min_pool_grains = typed_map(slot_size_list, lambda s: div_up(s, granularity)) - if pg_grains < sum(min_pool_grains): - return (0, 0) - num_slots = 1 << 63 - remaining_pg_grains = pg_grains - pool_idx_lst = sorted(typed_range(num_pools), key=lambda i: slot_size_list[i]) - for j, pool in enumerate(pool_idx_lst): - slot_size = slot_size_list[pool] - pool_grains = max( - min_pool_grains[pool], - round( - remaining_pg_grains - * (slot_size / sum(slot_size_list[k] for k in pool_idx_lst[j:])) - ), - ) - num_slots = min(num_slots, pool_grains * granularity // slot_size) - remaining_pg_grains -= pool_grains - assert remaining_pg_grains == 0 - assert num_slots > 0 - _s2g = CacheLevelStorage._grains_for_slots - lo = num_slots - step = 1 - hi = lo + step - while _s2g(hi, slot_size_list, granularity) <= pg_grains: - lo = hi - step *= 2 - hi = lo + step - while lo + 1 < hi: - mid = (lo + hi) // 2 - if _s2g(mid, slot_size_list, granularity) <= pg_grains: - lo = mid - else: - hi = mid - used = _s2g(lo, slot_size_list, granularity) - assert used <= pg_grains - assert _s2g(lo + 1, slot_size_list, granularity) > pg_grains - return lo, used - - @staticmethod - def _grains_for_slots( - num_slots: int, - slot_size_list: TypedIndexList[PoolIndex, int], - granularity: int, - ) -> int: - """Compute the minimum grains needed for num_slots in a pool group.""" - return sum(div_up(num_slots * s, granularity) for s in slot_size_list) - - @staticmethod - def ratio_to_slot_count_list( - total_quota: int, - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - ratio_list: TypedIndexList[PoolGroupIndex, float], - pool_size_granularity: int, - min_slots: TypedIndexList[PoolGroupIndex, int], - ) -> TypedIndexList[PoolGroupIndex, int]: - num_pool_groups = typed_len(ratio_list) - assert all(x > 0 for x in ratio_list) - assert num_pool_groups == typed_len(slot_size_lists) - assert total_quota % pool_size_granularity == 0 - total_grains = total_quota // pool_size_granularity - assert total_grains >= sum(len(sizes) for sizes in slot_size_lists) - g = pool_size_granularity - _g2s = CacheLevelStorage._grains_to_slots - _s2g = CacheLevelStorage._grains_for_slots - - slot_cnt_list = filled_list(0, num_pool_groups) - remaining_grains = total_grains - active_pgs = list(typed_range(num_pool_groups)) - - # Iteratively peel off constrained PGs until all active PGs are - # unconstrained: - # 1. Distribute remaining quota among active PGs by ratio. - # 2. Any PG with slots <= min_slots is constrained — pin it to - # min_slots and subtract its grains from the budget. - # 3. Repeat with the remaining PGs and re-normalized ratios. - # Each iteration removes at least one PG, so this terminates. - while active_pgs: - # Distribute remaining_grains among active PGs by ratio. - active_ratio = [ratio_list[pg] for pg in active_pgs] - slots_for_active = filled_list(0, len(active_pgs)) - grains_for_active = filled_list(0, len(active_pgs)) - budget = remaining_grains - idx_lst = sorted(range(len(active_pgs)), key=lambda i: active_ratio[i]) - for i, idx in enumerate(idx_lst): - pct = active_ratio[idx] / sum(active_ratio[j] for j in idx_lst[i:]) - slots, used = _g2s(round(budget * pct), slot_size_lists[active_pgs[idx]], g) - slots_for_active[idx] = slots - grains_for_active[idx] = used - budget -= used - assert budget >= 0 - - # Identify constrained PGs (slots <= min_slots). - constrained = [] - unconstrained = [] - for idx in range(len(active_pgs)): - pg = active_pgs[idx] - if slots_for_active[idx] <= min_slots[pg]: - constrained.append(idx) - else: - unconstrained.append(idx) - - if not constrained: - # All active PGs are unconstrained — accept their allocations. - for idx in range(len(active_pgs)): - slot_cnt_list[active_pgs[idx]] = slots_for_active[idx] - break - - # Pin constrained PGs to min_slots and subtract from budget. - for idx in constrained: - pg = active_pgs[idx] - min_grains = _s2g(min_slots[pg], slot_size_lists[pg], g) - slots, used = _g2s(min_grains, slot_size_lists[pg], g) - slot_cnt_list[pg] = slots - remaining_grains -= used - - if not unconstrained: - # All PGs are constrained — nothing left to redistribute. - break - - if remaining_grains <= 0: - raise ValueError("Insufficient quota to satisfy min_slots constraints") - - # Continue with unconstrained PGs only. - active_pgs = [active_pgs[idx] for idx in unconstrained] - - # _g2s may under-count slots due to imperfect grain distribution - # across pools. Try bumping each PG's slot count while it still fits - # within the same grain budget. - for pg in typed_range(num_pool_groups): - grains_now = _s2g(slot_cnt_list[pg], slot_size_lists[pg], g) - while _s2g(slot_cnt_list[pg] + 1, slot_size_lists[pg], g) <= grains_now: - slot_cnt_list[pg] += 1 - - return slot_cnt_list - - @property - def pool_size_granularity(self) -> int: - return 2 << 20 - - -class GpuCacheLevelStorage(CacheLevelStorage): - TIER: ClassVar[CacheTier] = CacheTier.GPU_MEM - __slots__ = ("shared_phys_mem_pool",) - shared_phys_mem_pool: PooledPhysMemAllocator - - def __init__( - self, - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - slot_count_list: TypedIndexList[PoolGroupIndex, int], - phys_mem_size: int, - ): - num_pool_groups = typed_len(slot_size_lists) - assert num_pool_groups == typed_len(slot_count_list), ( - "slot_size_lists and slot_count_list must have the same length" - ) - super().__init__() - self.shared_phys_mem_pool = PooledPhysMemAllocator(phys_mem_size) - self._pool_groups = make_typed( - lambda pg_idx: GpuPoolGroup( - slot_count_list[pg_idx], slot_size_lists[pg_idx], self.shared_phys_mem_pool - ), - num_pool_groups, - ) - - @override - def post_resize(self) -> None: - super().post_resize() - self.shared_phys_mem_pool.clear() # clear cached unused phys mem - - @property - def pool_size_granularity(self) -> int: - return self.shared_phys_mem_pool.phys_mem_size - - @override - def destroy(self) -> None: - super().destroy() - self.shared_phys_mem_pool.clear() - - -class HostCacheLevelStorage(CacheLevelStorage): - TIER: ClassVar[CacheTier] = CacheTier.HOST_MEM - POOL_SIZE_GRANULARITY: ClassVar[int] = HostMem.ALIGNMENT - __slots__ = () - - def __init__( - self, - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - slot_count_list: TypedIndexList[PoolGroupIndex, int], - ): - num_pool_groups = typed_len(slot_size_lists) - assert num_pool_groups == typed_len(slot_count_list), ( - "slot_size_lists and slot_count_list must have the same length" - ) - super().__init__() - self._pool_groups = make_typed( - lambda pg_idx: HostPoolGroup(slot_count_list[pg_idx], slot_size_lists[pg_idx]), - num_pool_groups, - ) - - @property - def pool_size_granularity(self) -> int: - return self.POOL_SIZE_GRANULARITY - - -class DiskCacheLevelStorage(CacheLevelStorage): - __slots__ = () - TIER: ClassVar[CacheTier] = CacheTier.DISK - POOL_SIZE_GRANULARITY: ClassVar[int] = 2 << 20 - - def __init__( - self, - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - slot_count_list: TypedIndexList[PoolGroupIndex, int], - filename_template: str, - ): - num_pool_groups = typed_len(slot_size_lists) - assert num_pool_groups == typed_len(slot_count_list), ( - "slot_size_lists and slot_count_list must have the same length" - ) - super().__init__() - self._pool_groups = make_typed( - lambda pg_idx: DiskPoolGroup( - slot_count_list[pg_idx], - slot_size_lists[pg_idx], - filename_template.format(pg_idx, "{}"), - ), - num_pool_groups, - ) - - @property - def pool_size_granularity(self) -> int: - return self.POOL_SIZE_GRANULARITY diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py deleted file mode 100644 index 8032502677c7..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ /dev/null @@ -1,1175 +0,0 @@ -# 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. - -import math -import os -import warnings -from collections import deque -from dataclasses import dataclass -from fractions import Fraction -from typing import TYPE_CHECKING, Callable, Iterator, Sequence, cast - -from . import rawref -from ._common import ( - GPU_LEVEL, - NDEBUG, - Address, - BlockOrdinal, - CacheLevel, - CacheTier, - LayerId, - MemAddress, - PageStatus, -) -from ._config import ( - BatchDesc, - CacheTierConfig, - DataRole, - DiskCacheTierConfig, - KVCacheDesc, - SwaScratchReuseConfig, -) -from ._copy_engine import CopyTask, batched_copy -from ._event_manager import KVCacheEventDiff -from ._eviction_controller import EvictablePage, PerLevelEvictionController -from ._exceptions import LogicError, OutOfPagesError -from ._life_cycle_registry import ( - AttnLifeCycle, - LifeCycleId, - LifeCycleRegistry, - compute_scratch_range, -) -from ._page import CommittedPage, Page -from ._storage import CacheLevelStorage -from ._storage._config import BufferAttr, BufferId, LayerAttr, SlotDesc, StorageConfig -from ._storage._core import ( - DiskCacheLevelStorage, - GpuCacheLevelStorage, - HostCacheLevelStorage, - PoolGroupBase, - PoolGroupIndex, - PoolIndex, - Slot, - SlotId, -) -from ._utils import ( - Array2D, - CachedCudaEvent, - HalfOpenRange, - HomoTuple, - TemporaryCudaStream, - TypedIndexList, - div_up, - filled_array2d, - filled_list, - get_uniform_attribute, - intersect, - make_typed, - partition, - remove_if, - round_up, - typed_enumerate, - typed_len, - typed_map, - typed_range, - unwrap_optional, -) - -if TYPE_CHECKING: - from ._event_manager import KVCacheEventManager - - -class CacheLevelManager: - __slots__ = ("cache_level", "storage", "controller") - cache_level: CacheLevel - storage: CacheLevelStorage - controller: PerLevelEvictionController - - @property - def cache_tier(self) -> CacheTier: - return self.storage.cache_tier - - def __init__( - self, - life_cycle_grouping: TypedIndexList[LifeCycleId, PoolGroupIndex], - cache_level: CacheLevel, - config: CacheTierConfig, - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - slot_count_list: TypedIndexList[PoolGroupIndex, int], - ): - self.cache_level = cache_level - self.storage = self._create_cache_level_storage(config, slot_size_lists, slot_count_list) - self.controller = PerLevelEvictionController(life_cycle_grouping, cache_level) - - @property - def num_pool_groups(self) -> PoolGroupIndex: - assert self.storage.num_pool_groups == self.controller.num_pool_groups - return self.storage.num_pool_groups - - @staticmethod - def cache_tier_granularity(tier: CacheTier, quota: int) -> int: - """Compute pool size granularity for a given cache tier and quota.""" - match tier: - case CacheTier.GPU_MEM: - page_size = 2 << 20 - return page_size << min(4, max(0, int(math.log(quota / (page_size * 512), 2)))) - case CacheTier.HOST_MEM: - return HostCacheLevelStorage.POOL_SIZE_GRANULARITY - case CacheTier.DISK: - return DiskCacheLevelStorage.POOL_SIZE_GRANULARITY - case _: - raise ValueError(f"Invalid cache tier: {tier}") - - @staticmethod - def _create_cache_level_storage( - config: CacheTierConfig, - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - slot_count_list: TypedIndexList[PoolGroupIndex, int], - ) -> CacheLevelStorage: - match config.tier: - case CacheTier.GPU_MEM: - granularity = CacheLevelManager.cache_tier_granularity( - CacheTier.GPU_MEM, config.quota - ) - return GpuCacheLevelStorage(slot_size_lists, slot_count_list, granularity) - case CacheTier.HOST_MEM: - return HostCacheLevelStorage(slot_size_lists, slot_count_list) - case CacheTier.DISK: - assert isinstance(config, DiskCacheTierConfig) - assert os.path.isdir(config.path), ( - f"Disk path {config.path} does not exist or is not a directory" - ) - filename_template = os.path.join(config.path, "g{}p{}.bin") - return DiskCacheLevelStorage(slot_size_lists, slot_count_list, filename_template) - case _: - raise ValueError(f"Invalid cache tier: {config.tier}") - - -@dataclass(slots=True, frozen=True) -class StorageStatistics: - "All in number of slots, for one pool group" - - slot_size: TypedIndexList[PoolIndex, int] - total: int - free: int - evictable: int - - @property - def slot_sizes(self) -> list[int]: - return list(self.slot_size) - - @property - def available(self) -> int: - return self.free + self.evictable - - @property - def unavailable(self) -> int: - return self.total - self.available - - -MigrationRecorder = Callable[[Sequence[Page], Sequence[Slot], CacheLevel, CacheLevel], None] -# Invoked when pages at the last cache level are released to free their slots -# without being migrated to any further tier (i.e. dropped from the cache hierarchy). -DropRecorder = Callable[[Sequence[Page], CacheLevel], None] - - -class StorageManager: - __slots__ = ( - "_life_cycles", - "_layer_to_life_cycle_ids", - "_slot_to_page_indices", - "_buffer_attr", - "_layer_attributes", - "_slot_util_frac_max", - "_life_cycle_grouping", - "_slot_desc_list", - "_levels", - "_min_slots", - "_event_manager", - "__rawref__", - ) - _life_cycles: LifeCycleRegistry - _layer_to_life_cycle_ids: dict[LayerId, LifeCycleId] - _slot_to_page_indices: TypedIndexList[LifeCycleId, TypedIndexList[PoolIndex, int]] - _slot_util_frac_max: TypedIndexList[LifeCycleId, Fraction] - _buffer_attr: dict[BufferId, BufferAttr] - _layer_attributes: dict[LayerId, LayerAttr] - _life_cycle_grouping: TypedIndexList[LifeCycleId, PoolGroupIndex] - _slot_desc_list: TypedIndexList[PoolGroupIndex, SlotDesc] - _levels: TypedIndexList[CacheLevel, CacheLevelManager] - _min_slots: TypedIndexList[PoolGroupIndex, int] - _event_manager: "KVCacheEventManager | None" - __rawref__: rawref.ref["StorageManager"] - - def __init__( - self, - life_cycles: LifeCycleRegistry, - config: StorageConfig, - tokens_per_block: int, - swa_scratch_reuse: SwaScratchReuseConfig | None, - typical_batch: BatchDesc | None = None, - constraints: list[BatchDesc] | None = None, - initial_pool_ratio: list[float] | None = None, - event_manager: "KVCacheEventManager | None" = None, - max_util_for_resume: float = 1.0, - ) -> None: - self.__rawref__ = rawref.NULL - self._event_manager = event_manager - assert config.cache_tiers[GPU_LEVEL].tier == CacheTier.GPU_MEM, ( - "The first cache tier must be GPU memory" - ) - self._life_cycles = life_cycles - self._layer_to_life_cycle_ids = config.layer_to_life_cycle_ids() - self._slot_to_page_indices = config.slot_to_page_indices() - self._layer_attributes = config.layer_attributes() - self._slot_util_frac_max = filled_list(Fraction(0, 1), life_cycles.size) - for attr in self._layer_attributes.values(): - if attr.slot_util_frac_max > self._slot_util_frac_max[attr.life_cycle_id]: - self._slot_util_frac_max[attr.life_cycle_id] = attr.slot_util_frac_max - self._buffer_attr = config.buffer_attributes() - self._life_cycle_grouping = config.life_cycle_grouping() - self._slot_desc_list = config.slot_desc_list - assert all(pg < self.num_pool_groups for pg in self._life_cycle_grouping) - assert self.num_pool_groups == PoolGroupIndex(len(set(self._life_cycle_grouping))) - slot_size_lists = typed_map(self._slot_desc_list, lambda pg: pg.slot_size_list) - - gpu_quota = config.cache_tiers[GPU_LEVEL].quota - gpu_granularity = CacheLevelManager.cache_tier_granularity(CacheTier.GPU_MEM, gpu_quota) - - # Constraints are hot-level feasibility floors. Other levels need only one - # structural slot per pool group. - self._min_slots = self._compute_pool_group_min_slots_from_constraints( - constraints or [], tokens_per_block, swa_scratch_reuse, max_util_for_resume - ) - - # Derive one lifecycle ratio, then project it onto each level's pool grouping. - life_cycle_ratio: TypedIndexList[LifeCycleId, float] - if initial_pool_ratio is not None: - if len(initial_pool_ratio) != self.num_life_cycles: - raise ValueError( - f"initial_pool_ratio length must match number of layer groups " - f"({self.num_life_cycles}), got {len(initial_pool_ratio)}" - ) - if any(r <= 0 for r in initial_pool_ratio): - raise ValueError("initial_pool_ratio values must be positive") - if not math.isclose(sum(initial_pool_ratio), 1.0, rel_tol=0.0, abs_tol=1e-6): - raise ValueError("initial_pool_ratio values must sum to 1.0") - life_cycle_ratio = cast(TypedIndexList[LifeCycleId, float], list(initial_pool_ratio)) - elif typical_batch is not None: - life_cycle_ratio = self.ratio_from_batch( - typical_batch, tokens_per_block, swa_scratch_reuse, gpu_granularity - ) - elif constraints: - life_cycle_slots = self._compute_slots_from_constraints( - constraints, tokens_per_block, swa_scratch_reuse, max_util_for_resume - ) - life_cycle_bytes = self._slots_to_bytes(life_cycle_slots, gpu_granularity) - total = sum(life_cycle_bytes) - life_cycle_ratio = typed_map(life_cycle_bytes, lambda x: x / total) - else: - life_cycle_ratio = self.ratio_from_batch( - BatchDesc([KVCacheDesc(capacity=2049, history_length=2048)]), - tokens_per_block, - swa_scratch_reuse, - gpu_granularity, - ) - init_ratio = self.pool_group_ratio(life_cycle_ratio) - - num_levels = CacheLevel(len(config.cache_tiers)) - self._levels = cast( - TypedIndexList, - [ - CacheLevelManager( - self._life_cycle_grouping, - i, - config.cache_tiers[i], - slot_size_lists, - self._compute_slot_count_for_level( - i, config.cache_tiers[i], slot_size_lists, init_ratio - ), - ) - for i in typed_range(num_levels) - ], - ) - assert self.num_pool_groups == get_uniform_attribute( - self._levels, lambda level: level.storage.num_pool_groups - ) - - def __del__(self) -> None: - self.destroy() - - def destroy(self) -> None: - if self.__rawref__.is_valid: - self.__rawref__.invalidate() - for lvl in self._levels: - lvl.storage.destroy() - - def get_pool_group_index(self, life_cycle: LifeCycleId) -> PoolGroupIndex: - return self._life_cycle_grouping[life_cycle] - - def new_gpu_slots( - self, - num_slots: TypedIndexList[LifeCycleId, int], - migration_recorder: MigrationRecorder | None = None, - drop_recorder: DropRecorder | None = None, - ) -> TypedIndexList[LifeCycleId, list[Slot]]: - return self.new_slots(GPU_LEVEL, num_slots, migration_recorder, drop_recorder) - - def new_slots( - self, - level: CacheLevel, - num_slots: TypedIndexList[LifeCycleId, int], - migration_recorder: MigrationRecorder | None = None, - drop_recorder: DropRecorder | None = None, - ) -> TypedIndexList[LifeCycleId, list[Slot]]: - lc2pg = self._life_cycle_grouping - pg_num_slots = filled_list(0, self.num_pool_groups) - for lc in typed_range(self.num_life_cycles): - if num_slots[lc] < 0: - raise LogicError("StorageManager.new_slots: slot count must be non-negative") - pg_num_slots[lc2pg[lc]] += num_slots[lc] - storage = self._levels[level].storage - if any( - pg_num_slots[pg] > storage.get_num_free_slots(pg) - for pg in typed_range(self.num_pool_groups) - ): - self.prepare_free_slots(level, pg_num_slots, migration_recorder, drop_recorder) - assert all( - pg_num_slots[pg] <= storage.get_num_free_slots(pg) - for pg in typed_range(self.num_pool_groups) - ) - ret = filled_list(list[Slot](), self.num_life_cycles) - try: - for life_cycle in typed_range(self.num_life_cycles): - pg_idx = lc2pg[life_cycle] - ret[life_cycle] = storage.allocate_multiple(pg_idx, num_slots[life_cycle]) - except Exception: - warnings.warn("Exception not expected here. Please report a bug.") - for lc, slots in typed_enumerate(ret): - pg_idx = lc2pg[lc] - for s in slots: - storage.release(pg_idx, s) - raise - return ret - - def new_slots_for_pool_group( - self, - level: CacheLevel, - pg_idx: PoolGroupIndex, - num_slots: int, - migration_recorder: MigrationRecorder | None = None, - drop_recorder: DropRecorder | None = None, - ) -> list[Slot]: - if num_slots < 0: - raise LogicError( - "StorageManager.new_slots_for_pool_group: slot count must be non-negative" - ) - storage = self._levels[level].storage - if num_slots > storage.get_num_free_slots(pg_idx): - num_slots_list = filled_list(0, self.num_pool_groups) - num_slots_list[pg_idx] = num_slots - self.prepare_free_slots(level, num_slots_list, migration_recorder, drop_recorder) - assert num_slots <= storage.get_num_free_slots(pg_idx) - try: - return storage.allocate_multiple(pg_idx, num_slots) - except Exception: - warnings.warn("Exception not expected here. Please report a bug.") - raise - - @property - def life_cycles(self) -> LifeCycleRegistry: - return self._life_cycles - - @property - def num_life_cycles(self) -> LifeCycleId: - return typed_len(self._life_cycle_grouping) - - @property - def num_pool_groups(self) -> PoolGroupIndex: - return typed_len(self._slot_desc_list) - - @property - def num_cache_levels(self) -> CacheLevel: - return CacheLevel(len(self._levels)) - - def is_last_level(self, level: CacheLevel) -> bool: - return level == self.num_cache_levels - 1 - - @property - def cache_tiers(self) -> HomoTuple[CacheTier]: - return tuple(cache_level.cache_tier for cache_level in self._levels) - - def get_ratio_list(self, level: CacheLevel) -> TypedIndexList[PoolGroupIndex, float]: - return self._levels[level].storage.ratio_list - - def is_evictable(self, page: EvictablePage, level: CacheLevel | None = None) -> bool: - """ - Check if a page is evictable. If level is specified, check if the page will be evictable after - migrating to the given level. - """ - status = page.status - level = page.cache_level if level is None else level - # droppable pages that are not committed should be dropped immediately. - # held pages in last level cache can't be evicted. - return (status == PageStatus.DROPPABLE and page.is_committed()) or ( - status == PageStatus.HELD and level < self.num_cache_levels - 1 - ) - - def prepare_free_slots( - self, - level: CacheLevel, - requirements: TypedIndexList[PoolGroupIndex, int], - migration_recorder: MigrationRecorder | None = None, - drop_recorder: DropRecorder | None = None, - ) -> None: - goals = filled_array2d(self.num_cache_levels, self.num_pool_groups, 0) - for pg in typed_range(self.num_pool_groups): - goals[level, pg] = requirements[pg] - fallen_pages = make_typed(lambda _: list[Page](), self.num_pool_groups) - self._prepare_free_slots(goals, level, fallen_pages, migration_recorder, drop_recorder) - - def force_evict( - self, - level: CacheLevel, - min_num_pages: TypedIndexList[PoolGroupIndex, int], - drop_recorder: DropRecorder | None = None, - ) -> None: - # If we break inside this function with debugpy, pages in `evicted` won't be - # released even after the function returns. This is a debugpy artifact. - evicted = self._levels[level].controller.evict(min_num_pages) - if int(level) == self.num_cache_levels - 1: - assert all(p.status == PageStatus.DROPPABLE for pages in evicted for p in pages), ( - "Corrupted eviction controller" - ) - if drop_recorder is not None: - for pg_idx in typed_range(self.num_pool_groups): - if evicted[pg_idx]: - drop_recorder(evicted[pg_idx], level) - return - next_lvl = CacheLevel(level + 1) - goals = filled_array2d(self.num_cache_levels, self.num_pool_groups, 0) - self._prepare_free_slots( - goals, - next_lvl, - cast(TypedIndexList[PoolGroupIndex, list[Page]], evicted), - drop_recorder=drop_recorder, - ) - - def _prepare_free_slots( - self, - goals: Array2D[CacheLevel, PoolGroupIndex, int], - lvl_id: CacheLevel, - fallen_pages: TypedIndexList[PoolGroupIndex, list[Page]], - migration_recorder: MigrationRecorder | None = None, - drop_recorder: DropRecorder | None = None, - ) -> None: - assert NDEBUG or goals.rows == self.num_cache_levels and goals.cols == self.num_pool_groups - assert NDEBUG or all( - all(p.cache_level < lvl_id for p in pages) for pages in fallen_pages - ), "Fallen pages must come from upper cache levels" - storage = self._levels[lvl_id].storage - ctrl = self._levels[lvl_id].controller - num_to_evict = filled_list(0, self.num_pool_groups) - held_pages = make_typed(lambda _: list[Page](), self.num_pool_groups) - for pg_idx in typed_range(self.num_pool_groups): - goal = goals[lvl_id, pg_idx] - fallen = len(fallen_pages[pg_idx]) - old_free_cnt = storage.get_num_free_slots(pg_idx) - evictable_cnt = ctrl.num_evictable_pages(pg_idx) - num_to_evict[pg_idx] = max(0, min(goal + fallen - old_free_cnt, evictable_cnt)) - fallen_held_cnt = 0 # fallen held pages we must accept in the current level. - if self.is_last_level(lvl_id): - held_pages[pg_idx] = remove_if( - fallen_pages[pg_idx], lambda p: p.status == PageStatus.HELD - ) - fallen_held_cnt = len(held_pages[pg_idx]) - if fallen_held_cnt > old_free_cnt + evictable_cnt: - # Do we need to revert the eviction we did before? Maybe not. - raise OutOfPagesError( - "Too many held pages are being evicted to the last-level cache for group {pg_idx}" - ) - if old_free_cnt + evictable_cnt - fallen_held_cnt < goal: - raise OutOfPagesError( - "Impossible to meet the goal ({goal} free slots) for group {pg_idx}" - ) - evicted = ctrl.evict(num_to_evict) - accepted_pages = make_typed(lambda _: list[Page](), self.num_pool_groups) - is_last_level = self.is_last_level(lvl_id) - if is_last_level: - for pg_idx in typed_range(self.num_pool_groups): - old_free_cnt = storage.get_num_free_slots(pg_idx) - num_evicted = len(evicted[pg_idx]) - assert NDEBUG or all(p.status == PageStatus.DROPPABLE for p in evicted[pg_idx]) - if not NDEBUG: - dbg_rawrefs = [rawref.ref(p) for p in evicted[pg_idx]] - # Record the drop event before releasing — these pages are leaving the - # cache hierarchy entirely without being onboarded back to GPU. - if drop_recorder is not None and num_evicted > 0: - drop_recorder(evicted[pg_idx], lvl_id) - evicted[pg_idx].clear() - if not NDEBUG: - assert all(p() is None for p in dbg_rawrefs) # pyright: ignore - new_free_cnt = storage.get_num_free_slots(pg_idx) - # GC of some pages may trigger removal of radix tree blocks and some other pages. - assert new_free_cnt >= num_evicted + old_free_cnt - assert len(held_pages[pg_idx]) <= new_free_cnt - fallen_pages[pg_idx].extend(held_pages[pg_idx]) - held_pages[pg_idx].clear() - goal = goals[lvl_id, pg_idx] - num_accepted = min(new_free_cnt - goal, len(fallen_pages[pg_idx])) - assert num_accepted >= 0 - accepted_pages[pg_idx] = ( - fallen_pages[pg_idx][-num_accepted:] if num_accepted > 0 else [] - ) - fallen_pages[pg_idx].clear() - else: - assert all(len(g) == 0 for g in held_pages) - for pg_idx in typed_range(self.num_pool_groups): - old_free_cnt = storage.get_num_free_slots(pg_idx) - e = evicted[pg_idx] - num_evicted = len(e) - fallen_pages[pg_idx][:0] = cast(list[Page], e) - e.clear() - num_accepted = min( - old_free_cnt + num_evicted - goals[lvl_id, pg_idx], len(fallen_pages[pg_idx]) - ) - assert num_accepted >= 0 - if num_accepted > 0: - accepted_pages[pg_idx] = fallen_pages[pg_idx][-num_accepted:] - del fallen_pages[pg_idx][-num_accepted:] - self._prepare_free_slots( - goals, - CacheLevel(lvl_id + 1), - fallen_pages, - migration_recorder, - drop_recorder, - ) - assert all(len(f) == 0 for f in fallen_pages) - # migrate pages - for pg_idx in typed_range(self.num_pool_groups): - partitioned = partition( - accepted_pages[pg_idx], - lambda p: (p.cache_level, self.get_pool_group_index(p.life_cycle)), - ) - accepted_pages[pg_idx].clear() - for (src_lvl, pg_idx), pages in partitioned.items(): - dst_lvl = lvl_id - self._batched_migrate( - pg_idx, - dst_lvl, - src_lvl, - pages, - update_src=True, - migration_recorder=migration_recorder, - ) - for p in pages: - if is_last_level and p.status == PageStatus.HELD: - continue - self._levels[dst_lvl].controller.schedule_for_eviction(p) - return - - def _batched_migrate( - self, - pool_group_index: PoolGroupIndex, - dst_level: CacheLevel, - src_level: CacheLevel, - src_pages: Sequence[Page], - update_src: bool, - migration_recorder: MigrationRecorder | None = None, - defrag: bool = False, # we are doing defragmentation - ) -> Sequence[Slot] | None: - "Free slots must be prepared before calling this function." - assert defrag or dst_level != src_level, ( - "dst_level and src_level must be different unless performing defragmentation" - ) - num_slots = len(src_pages) - num_pools = self.num_pools(pool_group_index) - src_pool_group = self._pool_group(src_level, pool_group_index) - dst_pool_group = self._pool_group(dst_level, pool_group_index) - if dst_pool_group.num_free_slots < num_slots: - raise OutOfPagesError("Not enough free slots") - dst_slots = dst_pool_group.allocate_multiple(num_slots) - try: - assert len(dst_slots) == num_slots - prior_events: set[CachedCudaEvent] = set() - tasks_per_pool: TypedIndexList[PoolIndex, list[CopyTask]] = make_typed( - lambda _: list[CopyTask](), num_pools - ) - for src, dst in zip(src_pages, dst_slots): - assert defrag or src.node_ref is None - prior_events.update((dst.ready_event, src.ready_event)) - dst_addresses = dst_pool_group.slot_address(dst.slot_id) - src_addresses = src_pool_group.slot_address(src.slot_id) - for pool_idx in typed_range(num_pools): - tasks_per_pool[pool_idx].append( - CopyTask(dst_addresses[pool_idx], src_addresses[pool_idx]) - ) - dst_tier = self._levels[dst_level].cache_tier - src_tier = self._levels[src_level].cache_tier - with TemporaryCudaStream(prior_events) as stream: - slot_sizes = self.slot_size(pool_group_index) - for pool_idx, tasks in typed_enumerate(tasks_per_pool): - batched_copy(dst_tier, src_tier, slot_sizes[pool_idx], tasks, stream.get()) - finish_event = stream.take_finish_event() - emit_cache_level_updates = ( - update_src - and not defrag - and src_level != dst_level - and self._event_manager is not None - ) - emitted_update_keys: set[tuple[bytes, LifeCycleId]] = set() - if migration_recorder is not None and not defrag: - migration_recorder(src_pages, dst_slots, src_level, dst_level) - for src, dst in zip(src_pages, dst_slots): - dst.ready_event = finish_event - src.ready_event = ( - finish_event # compulsory for the next owner getting this slot from the pool. - ) - if update_src: - scheduled_for_eviction = src.scheduled_for_eviction - if scheduled_for_eviction: - self.exclude_from_eviction(src) - src_pool_group.release(src) - src.set_slot(dst) - src.cache_level = dst_level - if emit_cache_level_updates: - self._emit_cache_level_updated_event( - src, src_level, dst_level, emitted_update_keys - ) - if scheduled_for_eviction: - self.schedule_for_eviction(src) - return None if update_src else dst_slots - except Exception: - for s in dst_slots: - dst_pool_group.release(s) - raise - - def _emit_cache_level_updated_event( - self, - page: Page, - old_level: CacheLevel, - new_level: CacheLevel, - emitted_keys: set[tuple[bytes, LifeCycleId]], - ) -> None: - if self._event_manager is None or not isinstance(page, CommittedPage): - return - - block = page.block() - if block is None or block.is_orphan: - return - - event_key = (block.key, page.life_cycle) - if event_key in emitted_keys: - return - - emitted_keys.add(event_key) - self._event_manager.add_updated_event( - block.key, - cache_level=KVCacheEventDiff( - old_value=int(old_level), - new_value=int(new_level), - ), - layer_group_id=int(page.life_cycle), - ) - - def _pool_group( - self, cache_level: CacheLevel, pool_group_index: PoolGroupIndex - ) -> PoolGroupBase: - return self._levels[cache_level].storage._pool_groups[pool_group_index] - - def num_pools(self, pool_group_index: PoolGroupIndex) -> PoolIndex: - return get_uniform_attribute( - self._levels, lambda level: level.storage._pool_groups[pool_group_index].num_pools - ) - - def slot_size(self, pool_group_index: PoolGroupIndex) -> TypedIndexList[PoolIndex, int]: - return self._slot_desc_list[pool_group_index].slot_size_list - - def num_slots( - self, pool_group_index: PoolGroupIndex, cache_level: CacheLevel = GPU_LEVEL - ) -> int: - return self._levels[cache_level].storage.num_slots(pool_group_index) - - def release_slot(self, life_cycle: LifeCycleId, cache_level: CacheLevel, slot: Slot) -> None: - pg_idx = self.get_pool_group_index(life_cycle) - self._levels[cache_level].storage.release(pg_idx, slot) - - def schedule_for_eviction(self, page: EvictablePage) -> None: - if self.is_evictable(page): - self._levels[page.cache_level].controller.schedule_for_eviction(page) - - def exclude_from_eviction(self, page: EvictablePage) -> None: - assert page.node_ref is not None - self._levels[page.cache_level].controller.remove(page.node_ref) - - def get_mem_pool_base_address(self, pg_idx: PoolGroupIndex, pool_idx: PoolIndex) -> MemAddress: - storage = self._levels[GPU_LEVEL].storage - return MemAddress(cast(int, storage.slot_address(pg_idx, pool_idx, SlotId(0)))) - - def get_buffer_attr(self, layer_id: LayerId, data_role: DataRole) -> BufferAttr: - return self._buffer_attr[BufferId(layer_id, data_role)] - - def get_layer_attr(self, layer_id: LayerId) -> LayerAttr: - return self._layer_attributes[layer_id] - - def slot_address( - self, level: CacheLevel, pg_idx: PoolGroupIndex, slot_id: SlotId, pool_idx: PoolIndex - ) -> Address: - return self._levels[level].storage.slot_address(pg_idx, pool_idx, slot_id) - - def get_statistics( - self, level: CacheLevel = GPU_LEVEL - ) -> TypedIndexList[PoolGroupIndex, StorageStatistics]: - ret = make_typed( - lambda pg_idx: StorageStatistics(filled_list(0, self.num_pools(pg_idx)), 0, 0, 0), - self.num_pool_groups, - ) - for pg_idx in typed_range(self.num_pool_groups): - pg = self._pool_group(level, pg_idx) - evictable_cnt = self._levels[level].controller.num_evictable_pages(pg_idx) - ret[pg_idx] = StorageStatistics( - pg.slot_size, pg.num_slots, pg.num_free_slots, evictable_cnt - ) - return ret - - def get_utilization( - self, level: CacheLevel = GPU_LEVEL - ) -> TypedIndexList[PoolGroupIndex, float]: - ret = filled_list(0.0, self.num_pool_groups) - stats = self.get_statistics(level) - for pg_idx in typed_range(self.num_pool_groups): - ret[pg_idx] = stats[pg_idx].unavailable / stats[pg_idx].total - return ret - - def get_overall_utilization(self, level: CacheLevel = GPU_LEVEL) -> float: - stats = self.get_statistics(level) - return sum(sum(s.slot_size) * s.unavailable for s in stats) / sum( - sum(s.slot_size) * s.total for s in stats - ) - - def shrink_pool_group( - self, - level: CacheLevel, - pg_idx: PoolGroupIndex, - new_num_slots: int, - persistent_pages: list[Page], - ) -> None: - """Move pages to eliminate overflow slots then shrink the pool group""" - lc2pg = self._life_cycle_grouping - assert len(persistent_pages) <= new_num_slots and all( - p.cache_level == level and lc2pg[p.life_cycle] == pg_idx for p in persistent_pages - ), "Not enough slots" - pool_group = self._levels[level].storage._pool_groups[pg_idx] - assert new_num_slots < pool_group.num_slots, "Not required for expansion of pools" - allocator = pool_group._slot_allocator - # Fast path: when no slot id has ever been issued in the to-be-removed - # range [new_num_slots, _capacity), there is nothing to migrate. - # _num_active_slots is a monotone high-water mark of issued ids. - if allocator._num_active_slots <= new_num_slots: - allocator.prepare_for_shrink(new_num_slots) - allocator.finish_shrink() - pool_group.resize_pools(new_num_slots) - return - ctrl = self._levels[level].controller - # pages with overflow slots and their indices in the eviction queue. - overflow_slots = deque[tuple[int, Page]]() - for i, p in enumerate(cast(Iterator[Page], ctrl.page_iterator(pg_idx))): - if p.slot_id >= new_num_slots: - overflow_slots.append((i, p)) - overflow_persistent_pages = [p for p in persistent_pages if p.slot_id >= new_num_slots] - num_overflow_persistent = len(overflow_persistent_pages) - if num_overflow_persistent > new_num_slots: - raise OutOfPagesError("Not enough slots to hold all persistent pages") - # prevent allocating slots with id >= new_num_slots - allocator.prepare_for_shrink(new_num_slots) - min_num_evicted = 0 - # Need this because evicted overflow pages won't become free, because only free - # non-overflow slots can be used for defragmentation. - num_evicted_overflow_slots = 0 - while overflow_slots and len(overflow_slots) + num_overflow_persistent > min( - new_num_slots, - overflow_slots[0][0] + allocator.num_free_slots - num_evicted_overflow_slots, - ): - min_num_evicted = overflow_slots.popleft()[0] + 1 - num_evicted_overflow_slots += 1 - self.force_evict( - level, make_typed(lambda i: min_num_evicted if i == pg_idx else 0, self.num_pool_groups) - ) - # These are the pages that will remain in the cache level and require defragmentation. - overflow_pages = [s[1] for s in overflow_slots] + overflow_persistent_pages - requirements = filled_list(0, self.num_pool_groups) - requirements[pg_idx] = len(overflow_pages) - self.prepare_free_slots(level, requirements) - assert NDEBUG or all(p.cache_level == level for p in overflow_pages), ( - "Some pages are not overflowed" - ) - self._batched_migrate(pg_idx, level, level, overflow_pages, update_src=True, defrag=True) - assert ( - len(allocator._overflow_slots) - == allocator._num_active_slots - allocator._target_capacity - ) - allocator.finish_shrink() - pool_group.resize_pools(new_num_slots) - - def expand_pool_group( - self, level: CacheLevel, pg_idx: PoolGroupIndex, new_num_slots: int - ) -> None: - pool_group = self._levels[level].storage._pool_groups[pg_idx] - assert new_num_slots > pool_group.num_slots - pool_group.resize_pools(new_num_slots) - pool_group._slot_allocator.expand(new_num_slots) - - def adjust_cache_level( - self, - level: CacheLevel, - new_quota: int | None, - new_ratio_list: TypedIndexList[PoolGroupIndex, float], - persistent_pages: TypedIndexList[PoolGroupIndex, list[Page]] | None = None, - ) -> None: - """Adapt the cache level by adjusting the ratio list. Persistent pages are those held and not evictable.""" - num_cache_levels = self.num_cache_levels - lvl_storage = self._levels[level].storage - old_num_slots = lvl_storage.slot_count_list - new_quota = ( - lvl_storage.total_quota - if new_quota is None - else round_up(new_quota, lvl_storage.pool_size_granularity) - ) - min_slots = self._min_slots_for_level(level) - min_quota = self._min_quota_for_level( - lvl_storage.slot_size_lists, lvl_storage.pool_size_granularity, min_slots - ) - if new_quota < min_quota: - raise ValueError( - f"Quota {new_quota} is insufficient for min_slots constraints " - f"(requires at least {min_quota})" - ) - new_num_slots = lvl_storage.compute_slot_count_list(new_ratio_list, min_slots, new_quota) - if level != num_cache_levels - 1: - assert persistent_pages is None, ( - "Persistent pages should be None for non-last level cache" - ) - # shrink first - for pg_idx in typed_range(self.num_pool_groups): - if new_num_slots[pg_idx] >= old_num_slots[pg_idx]: - continue - pages = persistent_pages[pg_idx] if persistent_pages is not None else [] - self.shrink_pool_group(level, pg_idx, new_num_slots[pg_idx], pages) - # then expand - for pg_idx in typed_range(self.num_pool_groups): - if new_num_slots[pg_idx] <= old_num_slots[pg_idx]: - continue - self.expand_pool_group(level, pg_idx, new_num_slots[pg_idx]) - lvl_storage.post_resize() - - def ratio_from_length( - self, tokens_per_block: int, history_length: int, capacity: int - ) -> TypedIndexList[LifeCycleId, float]: - if capacity < history_length: - warnings.warn("Bad sampling for capacity and history_length") - capacity = history_length - num_blocks = div_up(capacity, tokens_per_block) - num_bytes = filled_list(0, self.num_life_cycles) - ssm_lc_idx = self._life_cycles.ssm_life_cycle_id - for life_cycle, lc in typed_enumerate(self._life_cycles.get()): - pool_group = self.get_pool_group_index(life_cycle) - if life_cycle == ssm_lc_idx: - num_required_blocks = 1 - else: - stale = lc.get_stale_range(history_length, tokens_per_block) - num_required_blocks = max(num_blocks - len(stale), 1) - num_bytes[life_cycle] = num_required_blocks * sum(self.slot_size(pool_group)) - total = sum(num_bytes) - assert total > 0 - return typed_map(num_bytes, lambda x: x / total) - - def pool_group_ratio( - self, life_cycle_ratio: TypedIndexList[LifeCycleId, float] - ) -> TypedIndexList[PoolGroupIndex, float]: - assert len(life_cycle_ratio) == self.num_life_cycles - pool_group_ratio = filled_list(0.0, self.num_pool_groups) - for life_cycle in typed_range(self.num_life_cycles): - pool_group_ratio[self.get_pool_group_index(life_cycle)] += life_cycle_ratio[life_cycle] - total = sum(pool_group_ratio) - assert total > 0 - return typed_map(pool_group_ratio, lambda x: x / total) - - def ratio_from_batch( - self, - batch: BatchDesc, - tokens_per_block: int, - swa_scratch_reuse: SwaScratchReuseConfig | None, - granularity: int, - ) -> TypedIndexList[LifeCycleId, float]: - num_slots = self._compute_slots_for_batch(batch, tokens_per_block, swa_scratch_reuse) - num_bytes = self._slots_to_bytes(num_slots, granularity) - total = sum(num_bytes) - assert total > 0 - return typed_map(num_bytes, lambda x: x / total) - - def _compute_slots_from_constraints( - self, - constraints: list[BatchDesc], - tokens_per_block: int, - swa_scratch_reuse: SwaScratchReuseConfig | None, - max_util_for_resume: float, - ) -> TypedIndexList[LifeCycleId, int]: - if not 0 < max_util_for_resume <= 1: - raise ValueError(f"max_util_for_resume must be in (0, 1], got {max_util_for_resume}") - max_slots = filled_list(0, self.num_life_cycles) - - def swa_floor_blocks(lc: AttnLifeCycle) -> int: - window = unwrap_optional(lc.window_size) - return lc.num_sink_blocks + (window + tokens_per_block - 2) // tokens_per_block + 1 - - floor_num_blocks = 1 - for _, lc in self.life_cycles.attention_life_cycles(): - if lc.window_size is not None: - floor_num_blocks = max(floor_num_blocks, swa_floor_blocks(lc)) - for life_cycle, lc in self.life_cycles.items(): - if not isinstance(lc, AttnLifeCycle): - max_slots[life_cycle] = 1 - elif lc.window_size is not None: - max_slots[life_cycle] = swa_floor_blocks(lc) - else: - max_slots[life_cycle] = floor_num_blocks - - for batch in constraints: - slots = self._compute_slots_for_batch(batch, tokens_per_block, swa_scratch_reuse) - for life_cycle in typed_range(self.num_life_cycles): - scaled_slots = math.ceil(slots[life_cycle] / max_util_for_resume) - max_slots[life_cycle] = max(max_slots[life_cycle], scaled_slots) - return max_slots - - def _compute_pool_group_min_slots_from_constraints( - self, - constraints: list[BatchDesc], - tokens_per_block: int, - swa_scratch_reuse: SwaScratchReuseConfig | None, - max_util_for_resume: float, - ) -> TypedIndexList[PoolGroupIndex, int]: - """Compute the minimum slots per pool group across all constraints (element-wise max). - - All returned elements are positive. Constraint-derived floors include - headroom for the utilization gate checked by ``_KVCache.resume``. - """ - if not 0 < max_util_for_resume <= 1: - raise ValueError(f"max_util_for_resume must be in (0, 1], got {max_util_for_resume}") - max_slots = filled_list(0, self.num_pool_groups) - life_cycle_floors = self._compute_slots_from_constraints( - [], tokens_per_block, swa_scratch_reuse, max_util_for_resume - ) - for life_cycle in typed_range(self.num_life_cycles): - max_slots[self.get_pool_group_index(life_cycle)] += life_cycle_floors[life_cycle] - - for batch in constraints: - slots = self._compute_pool_group_slots_for_batch( - batch, tokens_per_block, swa_scratch_reuse - ) - for pg_idx in typed_range(self.num_pool_groups): - scaled_slots = math.ceil(slots[pg_idx] / max_util_for_resume) - max_slots[pg_idx] = max(max_slots[pg_idx], scaled_slots) - return max_slots - - def _compute_slots_for_batch( - self, - batch: BatchDesc, - tokens_per_block: int, - swa_scratch_reuse: SwaScratchReuseConfig | None, - ) -> TypedIndexList[LifeCycleId, int]: - """Compute the minimum number of slots per lifecycle to support a BatchDesc.""" - num_slots = filled_list(0, self.num_life_cycles) - ssm_lc_idx = self._life_cycles.ssm_life_cycle_id - sys_blocks = batch.system_prompt_length // tokens_per_block - for lc_idx, lc in typed_enumerate(self._life_cycles.get()): - if lc_idx == ssm_lc_idx: - # SSM: always 1 dedicated block per request, never shared. - num_slots[lc_idx] += len(batch.kv_caches) - continue - # Shared sys blocks (counted once): union of non-stale sys blocks across all requests. - # A sys block needs memory if it's non-stale for ANY request. - # = sys_blocks - (blocks stale for ALL requests within [0, sys_blocks)) - sys_range = HalfOpenRange(BlockOrdinal(0), BlockOrdinal(sys_blocks)) - # Intersection of per-request stale ranges, clamped to sys_range. - stale_intersection = sys_range - for kv in batch.kv_caches: - stale = lc.get_stale_range(kv.history_length, tokens_per_block) - stale_intersection = intersect(stale_intersection, stale) - num_slots[lc_idx] += sys_blocks - len(stale_intersection) - # Per-request unique blocks (excluding shared sys blocks already counted above). - for kv in batch.kv_caches: - total_blocks = div_up(kv.capacity, tokens_per_block) - stale = lc.get_stale_range(kv.history_length, tokens_per_block) - non_stale = total_blocks - len(stale) - # Non-stale sys blocks for this request. - non_stale_sys = sys_blocks - len(intersect(stale, sys_range)) - unique_non_stale = max(0, non_stale - non_stale_sys) - if swa_scratch_reuse is not None: - scratch = compute_scratch_range( - lc, - kv.history_length, - kv.capacity, - tokens_per_block, - swa_scratch_reuse.max_rewind_len, - ) - # Scratch blocks are always input blocks, so they never - # overlap with shared sys blocks (which are history). - num_scratch = len(scratch) - frac_max = self._slot_util_frac_max[lc_idx] - num_slots[lc_idx] += (unique_non_stale - num_scratch) + math.ceil( - num_scratch * frac_max - ) - else: - num_slots[lc_idx] += unique_non_stale - return num_slots - - def _compute_pool_group_slots_for_batch( - self, - batch: BatchDesc, - tokens_per_block: int, - swa_scratch_reuse: SwaScratchReuseConfig | None, - ) -> TypedIndexList[PoolGroupIndex, int]: - """Compute the minimum number of slots per hot pool group.""" - life_cycle_slots = self._compute_slots_for_batch(batch, tokens_per_block, swa_scratch_reuse) - num_slots = filled_list(0, self.num_pool_groups) - for life_cycle in typed_range(self.num_life_cycles): - num_slots[self.get_pool_group_index(life_cycle)] += life_cycle_slots[life_cycle] - return num_slots - - def _slots_to_bytes( - self, num_slots: TypedIndexList[LifeCycleId, int], granularity: int - ) -> TypedIndexList[LifeCycleId, int]: - num_bytes = filled_list(0, self.num_life_cycles) - for life_cycle in typed_range(self.num_life_cycles): - pool_group = self.get_pool_group_index(life_cycle) - for pool_size in self.slot_size(pool_group): - num_bytes[life_cycle] += round_up(num_slots[life_cycle] * pool_size, granularity) - return num_bytes - - def _pool_group_slots_to_bytes( - self, num_slots: TypedIndexList[PoolGroupIndex, int], granularity: int - ) -> TypedIndexList[PoolGroupIndex, int]: - """Convert slot counts to bytes, rounding up each pool to granularity.""" - num_bytes = filled_list(0, self.num_pool_groups) - for pg_idx in typed_range(self.num_pool_groups): - for pool_size in self.slot_size(pg_idx): - num_bytes[pg_idx] += round_up(num_slots[pg_idx] * pool_size, granularity) - return num_bytes - - def _min_slots_for_level(self, level: CacheLevel) -> TypedIndexList[PoolGroupIndex, int]: - if level == GPU_LEVEL: - return self._min_slots - return filled_list(1, self.num_pool_groups) - - @staticmethod - def _min_quota_for_level( - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - granularity: int, - min_slots: TypedIndexList[PoolGroupIndex, int], - ) -> int: - """Minimum quota in bytes required by the supplied slot floors.""" - return sum( - round_up(ms * s, granularity) - for ms, sizes in zip(min_slots, slot_size_lists) - for s in sizes - ) - - def _compute_slot_count_for_level( - self, - level: CacheLevel, - tier_config: CacheTierConfig, - slot_size_lists: TypedIndexList[PoolGroupIndex, TypedIndexList[PoolIndex, int]], - ratio: TypedIndexList[PoolGroupIndex, float], - ) -> TypedIndexList[PoolGroupIndex, int]: - """Compute slot counts for a cache level from its tier config and ratio. - - Applies hot constraint floors or a one-slot structural floor on colder levels. - """ - granularity = CacheLevelManager.cache_tier_granularity(tier_config.tier, tier_config.quota) - min_slots = self._min_slots_for_level(level) - quota = max( - self._min_quota_for_level(slot_size_lists, granularity, min_slots), - round_up(tier_config.quota, granularity), - ) - return CacheLevelStorage.ratio_to_slot_count_list( - quota, slot_size_lists, ratio, granularity, min_slots - ) - - def constrain_pool_group_ratio( - self, - ratio: TypedIndexList[PoolGroupIndex, float], - ) -> TypedIndexList[PoolGroupIndex, float]: - """Apply the stored min_slots constraint to a ratio list for GPU level. - - Converts ratio to slot counts (with min_slots floor), - then converts back to a bytes-based ratio. - """ - gpu_storage = self._levels[GPU_LEVEL].storage - granularity = gpu_storage.pool_size_granularity - slot_count_list = gpu_storage.compute_slot_count_list(ratio, self._min_slots) - num_bytes = self._pool_group_slots_to_bytes(slot_count_list, granularity) - total = sum(num_bytes) - assert total > 0 - return typed_map(num_bytes, lambda x: x / total) - - def prefetch( - self, - dst_lvl: CacheLevel, - pages: TypedIndexList[PoolGroupIndex, TypedIndexList[CacheLevel, list[Page]]], - ) -> int: - """Dispatch page migration to the destination cache level. - - Args: - dst_lvl: Destination cache level for pages currently in lower tiers. - pages: Pages grouped by pool group and current cache level. - - Returns: - How many pages it moved off the disk tier, counted per migrated batch rather than per - page. A raise reports nothing, which in practice means slot preparation failed before - anything moved. - - Raises: - OutOfPagesError: If there are not enough pages available for the prefetch hint. - """ - num_slots = filled_list(0, self.num_pool_groups) - disk_blocks_migrated = 0 - scheduled = list[Page]() - try: - for pg_idx, pg_pages in typed_enumerate(pages): - for lvl, lvl_pages in typed_enumerate(pg_pages): - assert lvl >= dst_lvl or not lvl_pages - for p in lvl_pages: - if p.scheduled_for_eviction: - self.exclude_from_eviction(p) - scheduled.append(p) - elif self.is_evictable(p, dst_lvl): - scheduled.append(p) - assert lvl >= dst_lvl - if lvl == dst_lvl: - continue - num_slots[pg_idx] += 1 - self.prepare_free_slots(dst_lvl, num_slots) - for pg_idx, pg_tasks in typed_enumerate(pages): - for lvl in typed_range(CacheLevel(dst_lvl + 1), self.num_cache_levels): - lvl_tasks = pg_tasks[lvl] - self._batched_migrate(pg_idx, dst_lvl, lvl, lvl_tasks, True) - # Per batch, after it landed: the tasks are already grouped by source level, so - # this costs nothing per page and never credits a batch that did not run. - if self.cache_tiers[lvl] == CacheTier.DISK: - disk_blocks_migrated += len(lvl_tasks) - finally: - for p in scheduled: - self.schedule_for_eviction(p) - return disk_blocks_migrated diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py deleted file mode 100644 index 2f970ca347b7..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py +++ /dev/null @@ -1,1126 +0,0 @@ -# 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. - -import array -import concurrent.futures -import ctypes -import errno -import functools -import itertools -import operator -import os -import platform -import sys -import traceback -import warnings -import weakref -from abc import ABC, abstractmethod -from collections import defaultdict, deque -from collections.abc import Set -from contextlib import contextmanager -from ctypes.util import find_library -from itertools import pairwise -from typing import ( - Any, - Callable, - ClassVar, - Final, - Generic, - Iterable, - Iterator, - MutableSequence, - Protocol, - Reversible, - Sequence, - Type, - TypeVar, - cast, -) - -import cuda.bindings.driver as drv -import cuda.bindings.runtime as cudart - -from . import rawref -from ._common import NDEBUG, CudaStream -from ._exceptions import CuError, CuOOMError, DiskOOMError, HostOOMError - -T = TypeVar("T") -U = TypeVar("U") -Index = TypeVar("Index", bound=int, contravariant=True) -IndexO = TypeVar("IndexO", bound=int, covariant=True) -Row = TypeVar("Row", bound=int) -Col = TypeVar("Col", bound=int) - - -def _unwrap( - ret: drv.CUresult - | tuple[ - drv.CUresult, - T, - ] - | tuple[drv.CUresult, T, U], -): - if isinstance(ret, drv.CUresult): - if int(ret) != int(drv.CUresult.CUDA_SUCCESS): # pyright: ignore - if int(ret) == int(drv.CUresult.CUDA_ERROR_OUT_OF_MEMORY): # pyright: ignore - raise CuOOMError() - raise CuError(ret) - else: - _unwrap(ret[0]) - return ret[1] if len(ret) == 2 else ret[1:] - - -def div_up(x: int, y: int) -> int: - return (x + y - 1) // y - - -def round_up(x: int, y: int) -> int: - return div_up(x, y) * y - - -def round_down(x: int, y: int) -> int: - return x // y * y - - -def in_range(x: int, lower: int, upper: int) -> bool: - return lower <= x < upper - - -def exact_div(x: int, y: int) -> int: - assert x % y == 0 - return x // y - - -Idx = TypeVar("Idx", bound=int) - - -class HalfOpenRange(tuple[Idx, Idx], Generic[Idx]): - """A half-open range [beg, end). Falsy when empty (beg >= end). - Generic over index type. Supports unpacking into (beg, end).""" - - __slots__ = () - - def __new__(cls, beg: Idx, end: Idx) -> "HalfOpenRange[Idx]": - return tuple.__new__(cls, (beg, end)) - - @property - def beg(self) -> Idx: - return self[0] - - @property - def end(self) -> Idx: - return self[1] - - def __eq__(self, other: object) -> bool: - if not isinstance(other, HalfOpenRange): - return NotImplemented - return (not self and not other) or tuple.__eq__(self, other) - - def __hash__(self) -> int: - return hash((0, 0)) if not self else tuple.__hash__(self) - - def __bool__(self) -> bool: - return self[0] < self[1] - - def __len__(self) -> int: - return max(0, self[1] - self[0]) - - def __contains__(self, item: Any) -> bool: - return self[0] <= item < self[1] - - -def intersect(a: HalfOpenRange[Idx], b: HalfOpenRange[Idx]) -> HalfOpenRange[Idx]: - """Returns the intersection of two half-open ranges [beg, end). - The result may be empty (beg >= end), which is safe to chain into further intersections.""" - return HalfOpenRange(max(a[0], b[0]), min(a[1], b[1])) - - -def value_or(opt: T | None, default: T) -> T: - return default if opt is None else opt - - -def unwrap_optional(value: T | None) -> T: - if value is not None: - return value - raise ValueError("Expected non-None value") - - -def unwrap_weakref(ref: weakref.ref[T]) -> T: - obj = ref() - if obj is not None: - return obj - raise ValueError("Dereferencing a dangling weakref") - - -def unwrap_rawref(ref: rawref.ref[T]) -> T: - obj = ref() - if obj is not None: - return obj - raise ValueError("Dereferencing a dangling rawref") - - -def map_optional(value: T | None, func: Callable[[T], U]) -> U | None: - return func(value) if value is not None else None - - -def remove_if(original: MutableSequence[T], predicate: Callable[[T], bool]) -> list[T]: - "Remove items from original that satisfy the predicate and return the removed items." - removed = [] - for idx, item in enumerate(original): - if predicate(item): - removed.append(item) - else: - original[idx - len(removed)] = item - del original[len(original) - len(removed) :] - return removed - - -def chunked(iterable: Iterable[T], size: int) -> Iterator[list[T]]: - iterator = iter(iterable) - while True: - chunk = list(itertools.islice(iterator, size)) - if not chunk: - break - yield chunk - - -def partition(original: Iterable[T], classifier: Callable[[T], U]) -> defaultdict[U, list[T]]: - ret = defaultdict(list) - for item in original: - ret[classifier(item)].append(item) - return ret - - -def get_uniform_attribute(iterable: Iterable[T], attribute_func: Callable[[T], U]) -> U: - ret = attribute_func(next(iter(iterable))) - assert NDEBUG or all(attribute_func(item) == ret for item in iterable) - return ret - - -def assert_critical(condition: bool, message: str | None = None) -> None: - "Similar to assert, but instead of raising an exception, it terminates the process, even if inside __del__()." - if not condition: - warnings.warn(value_or(message, "Critical assertion failed")) - traceback.print_stack() - os._exit(1) - - -def noexcept(func: Callable[..., T]) -> Callable[..., T]: - @functools.wraps(func) - def wrapper(*args, **kwargs) -> T: - try: - return func(*args, **kwargs) - except Exception as e: - raise AssertionError(f"Function {func.__name__} raised an exception: {e}") from e - - return wrapper - - -def not_implemented(func: Callable[..., T]) -> Callable[..., T]: - @functools.wraps(func) - def wrapper(*args, **kwargs) -> T: - raise NotImplementedError(f"The function '{func.__name__}' is not implemented yet.") - - return wrapper - - -def expect_type(ExpectedType: Type[T], value: Any) -> T: - "Similar to typing.cast, but does runtime checking with assert." - assert isinstance(value, ExpectedType), f"Expected {ExpectedType}, got {type(value)}" - return value - - -def is_sorted( - iterable: Iterable[T], key: Callable[[T], Any] = lambda x: x, reverse: bool = False -) -> bool: - comp = operator.ge if reverse else operator.le - return all(comp(key(a), key(b)) for a, b in pairwise(iterable)) - - -HomoTuple = tuple[T, ...] - - -class TypedIndexList(Protocol[Index, T]): - """ - A protocol representing a list-like container with a strongly typed integer index. - Useful for enforcing index types like NewType wrappers around int. - """ - - def __getitem__(self, index: Index) -> T: ... - - def __setitem__(self, index: Index, value: T) -> None: ... - - def __delitem__(self, index: Index | slice) -> None: ... - - def __iter__(self) -> Iterator[T]: ... - - def __len__(self) -> int: ... - - def __reversed__(self) -> Iterator[T]: ... - - def clear(self) -> None: ... - - def pop(self) -> T: ... - - def append(self, value: T) -> None: ... - - -# @TODO: use this where applicable. -def to_typed(index_type: Callable[[Any], Index], lst: list[T]) -> TypedIndexList[Index, T]: - """ - Casts a standard list to a TypedIndexList with a strongly typed integer index. - - Parameters: - index_type: A type alias for the NewType index, e.g. type(BlockOrdinal(0)) or a concrete class derived from int. - lst: The list to cast - - Returns: - A TypedIndexList[Index, T] with the specified index type - """ - return cast(TypedIndexList[Index, T], lst) - - -def typed_range(*args: Index) -> Reversible[Index]: - return cast(Reversible[Index], range(*args)) - - -def filled_list(value: T, count: Index) -> TypedIndexList[Index, T]: - "Note that all elements will be the same value. Do not use mutable values." - return cast(TypedIndexList[Index, T], [value] * int(count)) - - -def make_typed(generator: Callable[[Index], T], count: Index) -> TypedIndexList[Index, T]: - return cast(TypedIndexList[Index, T], [generator(Index) for Index in typed_range(count)]) - - -def typed_len(iterable: TypedIndexList[IndexO, T]) -> IndexO: - return cast(IndexO, len(iterable)) - - -def typed_enumerate(iterable: TypedIndexList[Index, T]) -> Iterator[tuple[Index, T]]: - return cast(Iterator[tuple[Index, T]], enumerate(iterable)) - - -def typed_map( - iterable: TypedIndexList[Index, T], func: Callable[[T], U] -) -> TypedIndexList[Index, U]: - return cast(TypedIndexList[Index, U], [func(item) for item in iterable]) - - -class Array2D(Generic[Row, Col, T]): - __slots__ = ("_data", "_cols") - _data: list[T] - _cols: int - - def __init__(self, rows: Row, cols: Col, init_val: Iterable[T]) -> None: - self._data = list(init_val) - self._cols = cols - - def __getitem__(self, index: tuple[Row, Col]) -> T: - return self._data[index[0] * self._cols + index[1]] - - def __setitem__(self, index: tuple[Row, Col], value: T) -> None: - self._data[index[0] * self._cols + index[1]] = value - - @property - def rows(self) -> int: - assert len(self._data) % self._cols == 0 - return len(self._data) // self._cols - - def row(self, row: Row) -> TypedIndexList[Col, T]: - return cast(TypedIndexList[Col, T], self._data[row * self._cols : (row + 1) * self._cols]) - - def col(self, col: Col) -> TypedIndexList[Row, T]: - return cast(TypedIndexList[Row, T], self._data[col :: self._cols]) - - @property - def cols(self) -> int: - return self._cols - - def __len__(self) -> int: - return len(self._data) - - def __iter__(self) -> Iterator[T]: - return iter(self._data) - - def __reversed__(self) -> Iterator[T]: - return reversed(self._data) - - -def filled_array2d(rows: Row, cols: Col, val: T) -> Array2D[Row, Col, T]: - return Array2D(rows, cols, [val] * rows * cols) - - -def find(seq: Sequence[T], predicate: Callable[[T], bool], default: U) -> T | U: - return next((item for item in seq if predicate(item)), default) - - -def find_index(seq: Iterable[T], predicate: Callable[[T], bool]) -> int: - i = 0 - for i, item in enumerate(seq): - if predicate(item): - return i - return i + 1 - - -# mmap constants (Linux x86_64) -MAP_PRIVATE: Final[int] = 0x02 -MAP_ANONYMOUS: Final[int] = 0x20 -PROT_READ: Final[int] = 0x1 -PROT_WRITE: Final[int] = 0x2 -MREMAP_MAYMOVE: Final[int] = 1 -MAP_FAILED: Final[int] = -1 - -_libc = ctypes.CDLL(find_library("c"), use_errno=True) -_libc.mmap.restype = ctypes.c_void_p -_libc.mmap.argtypes = [ - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_int, - ctypes.c_int, - ctypes.c_int, - ctypes.c_longlong, -] -_libc.munmap.restype = ctypes.c_int -_libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t] -_libc.mremap.restype = ctypes.c_void_p -_libc.mremap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_size_t, ctypes.c_int] -_libc.madvise.restype = ctypes.c_int -_libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] -_libc.posix_fallocate.restype = ctypes.c_int -_libc.posix_fallocate.argtypes = [ctypes.c_int, ctypes.c_longlong, ctypes.c_longlong] - -MADV_HUGEPAGE: Final[int] = 14 -MADV_NOHUGEPAGE: Final[int] = 15 -MADV_POPULATE_WRITE: Final[int] = 23 - -# TLLM_KV_CACHE_MANAGER_V2_THP=0 backs host pools with regular 4KB pages -# (MADV_NOHUGEPAGE). On nodes with fragmented physical memory and THP -# defrag=madvise, every 2MB THP fault stalls in direct compaction that -# rarely succeeds, slowing pool population from GB/s to GB/min. -USE_THP: Final[bool] = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_THP", "1") == "1" -# TLLM_KV_CACHE_MANAGER_V2_PREFAULT_THREADS=0 disables prefaulting; pages are -# then faulted in lazily, single-threaded, inside cuMemHostRegister. -PREFAULT_THREADS: Final[int] = int( - os.environ.get( - "TLLM_KV_CACHE_MANAGER_V2_PREFAULT_THREADS", str(min(64, (os.cpu_count() or 32) // 2)) - ) -) - - -def _madvise(ptr: int, size: int, advice: int) -> None: - if os.name == "nt": - return - ret = _libc.madvise(ctypes.c_void_p(ptr), ctypes.c_size_t(size), ctypes.c_int(advice)) - if ret != 0: - error_code = ctypes.get_errno() - # Advisory, so just warn - warnings.warn( - f"madvise failed with errno {error_code}: {errno.errorcode.get(error_code, 'Unknown error')}" - ) - - -def _mmap(size: int) -> int: - """ - Allocates size bytes using mmap (anonymous, private). - Returns the address as an integer. - Raises HostOOMError on failure. - """ - prot = PROT_READ | PROT_WRITE - flags = MAP_PRIVATE | MAP_ANONYMOUS - - ptr = _libc.mmap( - None, - ctypes.c_size_t(size), - ctypes.c_int(prot), - ctypes.c_int(flags), - ctypes.c_int(-1), - ctypes.c_longlong(0), - ) - - ptr_int = int(ptr) if ptr is not None else 0 - if ptr_int == -1 or ptr_int == 0xFFFFFFFFFFFFFFFF or ptr_int == 0: - error_code = ctypes.get_errno() - raise HostOOMError(f"mmap failed with errno {error_code}") - - return ptr_int - - -def _munmap(ptr: int, size: int) -> None: - ret = _libc.munmap(ctypes.c_void_p(ptr), ctypes.c_size_t(size)) - if ret != 0: - error_code = ctypes.get_errno() - warnings.warn(f"munmap failed with errno {error_code}") - - -def _mremap(ptr: int, old_size: int, new_size: int) -> int: - """ - Remaps memory using mremap. - Returns the new address as an integer. - Raises HostOOMError on failure. - """ - ptr_new = _libc.mremap( - ctypes.c_void_p(ptr), - ctypes.c_size_t(old_size), - ctypes.c_size_t(new_size), - ctypes.c_int(MREMAP_MAYMOVE), - ) - ptr_int = int(ptr_new) if ptr_new is not None else 0 - if ptr_int == -1 or ptr_int == 0xFFFFFFFFFFFFFFFF or ptr_int == 0: - error_code = ctypes.get_errno() - raise HostOOMError(f"mremap failed with errno {error_code}") - return ptr_int - - -def _posix_fallocate(fd: int, offset: int, length: int) -> None: - ret = _libc.posix_fallocate( - ctypes.c_int(fd), ctypes.c_longlong(offset), ctypes.c_longlong(length) - ) - if ret != 0: - raise DiskOOMError(ret, "posix_fallocate failed") - - -class HostMem: - ALIGNMENT: ClassVar[int] = 4096 # 4KB - """ - Host memory, reallocable for low-cost resizing and registered to CUDA as page-locked memory. - Uses MADV_HUGEPAGE to opportunistically use Transparent Huge Pages (THP) where possible. - Resizing will keep the original memory content, like `realloc` in C. - """ - __slots__ = ("_address", "_size", "_num_registered_chunks") - _address: int - _size: int - # If True and _size > 2GB, use multiple chunks to register pinned memory due to a Linux kernel - # 6.11/6.12/6.13 bug preventing pinning more than 2GB of host memory in one operation. - _CHUNKED_REGISTRATION: ClassVar[bool] = platform.system() == "Linux" and platform.release()[ - :4 - ] in ["6.11", "6.12", "6.13"] - _CHUNK_SIZE: ClassVar[int] = 2 << 30 - _num_registered_chunks: int - - @property - def address(self) -> int: - return self._address - - @property - def size(self) -> int: - return self._size - - def __init__(self, size: int) -> None: - self._num_registered_chunks = 0 - if size == 0: - self._address = 0 - self._size = 0 - return - - # Allocate with standard mmap (4KB alignment) - self._address = _mmap(size) - assert self._address % self.ALIGNMENT == 0 - self._size = size - - # Opportunistically advise huge pages for the whole range. - # The kernel will use huge pages for aligned 2MB chunks within this range. - _madvise(self._address, self._size, MADV_HUGEPAGE if USE_THP else MADV_NOHUGEPAGE) - - if PREFAULT_THREADS > 0: - self._parallel_prefault(PREFAULT_THREADS) - self._register_to_cuda() - - def resize(self, new_size: int) -> None: - self._unregister_from_cuda() - try: - self._address = _mremap(self._address, self._size, new_size) - assert self._address % self.ALIGNMENT == 0 - self._size = new_size - - # Re-advise the configured page mode for the new range. - _madvise(self._address, self._size, MADV_HUGEPAGE if USE_THP else MADV_NOHUGEPAGE) - finally: - self._register_to_cuda() - - def destroy(self) -> None: - if self._address == 0: - return - self._unregister_from_cuda() - _munmap(self._address, self._size) - self._address = 0 - self._size = 0 - - def __del__(self) -> None: - self.destroy() - - def _parallel_prefault(self, nthreads: int) -> None: - """Fault in all pages with parallel threads before cuMemHostRegister, - so registration only pins pages (never allocates them). - - Lazy faulting inside cuMemHostRegister is single-threaded and, under - memory pressure or THP compaction stalls, can take minutes for - multi-hundred-GiB pools. MADV_POPULATE_WRITE populates in bulk; small - chunks keep mmap_lock hold times short (one giant madvise per thread - serializes every other thread behind it) and let threads - load-balance. - """ - chunk = 512 << 20 - - def populate(off: int) -> None: - ln = min(chunk, self._size - off) - if ln <= 0: - return - ret = _libc.madvise( - ctypes.c_void_p(self._address + off), - ctypes.c_size_t(ln), - ctypes.c_int(MADV_POPULATE_WRITE), - ) - if ret != 0: - error_code = ctypes.get_errno() - if error_code in (errno.EINVAL, getattr(errno, "ENOSYS", -1)): - # MADV_POPULATE_WRITE requires Linux >= 5.14; on older - # kernels fall back to touching every page. - ctypes.memset(self._address + off, 0, ln) - return - error_name = errno.errorcode.get(error_code, "Unknown error") - if error_code == errno.ENOMEM: - # Surface real allocation failures instead of masking them - # with a memset that would trigger a system OOM kill. - raise HostOOMError( - f"madvise(MADV_POPULATE_WRITE) failed with errno {error_code}: {error_name}" - ) - raise OSError( - error_code, - f"madvise(MADV_POPULATE_WRITE) failed: {error_name}", - ) - - with concurrent.futures.ThreadPoolExecutor(max_workers=nthreads) as executor: - list(executor.map(populate, range(0, self._size, chunk))) - - def _register_to_cuda(self) -> None: - assert self._num_registered_chunks == 0 - for addr, size in self._iterate_chunks(): - _unwrap( - drv.cuMemHostRegister( - addr, size, drv.CU_MEMHOSTREGISTER_PORTABLE | drv.CU_MEMHOSTREGISTER_DEVICEMAP - ) - ) - self._num_registered_chunks += 1 - - def _unregister_from_cuda(self) -> None: - for addr, _ in self._iterate_chunks(): - if self._num_registered_chunks == 0: - break - _unwrap(drv.cuMemHostUnregister(addr)) - self._num_registered_chunks -= 1 - assert self._num_registered_chunks == 0 - - def _iterate_chunks(self) -> Iterator[tuple[int, int]]: - start = self._address - end = start + self._size - chunk_size = self._CHUNK_SIZE if self._CHUNKED_REGISTRATION else self._size - for addr in range(start, end, chunk_size): - yield addr, min(end - addr, chunk_size) - - -def resize_file(fd: int, new_size: int) -> None: - old_size = os.lseek(fd, 0, os.SEEK_END) - if new_size > old_size: - _posix_fallocate(fd, old_size, new_size - old_size) - elif new_size < old_size: - os.truncate(fd, new_size) - - -class DynamicBitset: - """ - A memory efficient bitset that can be resized. - """ - - __slots__ = ("_bits", "_num_set_bits") - _bits: array.array - _num_set_bits: int - - TYPE_CODE: ClassVar[str] = "Q" - ALL_SET_MASK: ClassVar[int] = (1 << 64) - 1 - - def __init__(self, capacity: int) -> None: - self._bits = array.array(self.TYPE_CODE, [0] * (div_up(capacity, 64))) - self._num_set_bits = 0 - - def set(self, index: int) -> None: - if not self.get(index): - self._bits[index // 64] |= 1 << (index % 64) - self._num_set_bits += 1 - - def get(self, index: int) -> bool: - return self._bits[index // 64] & (1 << (index % 64)) != 0 - - def clear(self, index: int) -> None: - if self.get(index): - self._bits[index // 64] &= ~(1 << (index % 64)) - self._num_set_bits -= 1 - - @property - def num_set_bits(self) -> int: - return self._num_set_bits - - def resize(self, new_capacity: int) -> None: - old_elems = len(self._bits) - new_elems = div_up(new_capacity, 64) - - # When the capacity shrinks, every set bit at or above new_capacity is - # dropped. Account for those bits so num_set_bits stays accurate, and - # mask the retained partial word so any_set() cannot observe stale bits. - # This covers both fewer-words and same-word-count (new_elems == - # old_elems with a smaller new_capacity) shrinks. - if new_elems <= old_elems: - for w in range(new_elems, old_elems): - self._num_set_bits -= bin(self._bits[w]).count("1") - if new_elems >= 1 and new_capacity % 64 != 0: - keep_mask = self.ALL_SET_MASK >> (64 - (new_capacity % 64)) - word = self._bits[new_elems - 1] - dropped = word & ~keep_mask & self.ALL_SET_MASK - self._num_set_bits -= bin(dropped).count("1") - self._bits[new_elems - 1] = word & keep_mask - - if new_elems > old_elems: - self._bits.extend(array.array(self.TYPE_CODE, [0] * (new_elems - old_elems))) - elif new_elems < old_elems: - self._bits = self._bits[:new_elems] - - # check if any bit in the range [start, end) is set - def any_set(self, start: int, end: int) -> bool: - if start >= end: - return False - start_word_mask = self.ALL_SET_MASK << (start % 64) - end_word_mask = self.ALL_SET_MASK >> (64 - (end % 64)) - if start // 64 == end // 64: - if (start_word_mask & end_word_mask & self._bits[start // 64]) != 0: - return True - else: - if (start_word_mask & self._bits[start // 64]) != 0 or ( - end % 64 != 0 and end_word_mask & self._bits[end // 64] - ) != 0: - return True - return any(self._bits[i] != 0 for i in range(start // 64 + 1, end // 64)) - - -@functools.cache -def init_cuda_once() -> None: - (err,) = cudart.cudaFree(0) - assert int(err) == int(cudart.cudaError_t.cudaSuccess) - - -class SimplePool(Generic[T]): - __slots__ = ( - "_create_func", - "_destroy_func", - "_init_size", - "_max_size", - "_outstanding_count", - "_items", - ) - _create_func: Callable[[], T] - _destroy_func: Callable[[T], None] - _init_size: int - _max_size: int | None - _items: deque[T] | None - _outstanding_count: ( - int # number of items currently we gave out but not returned, i.e. get() but not put() - ) - - def __init__( - self, - create_func: Callable[[], T], - destroy_func: Callable[[T], None], - init_size: int = 0, - max_size: int | None = None, - ): - self._create_func = create_func - self._destroy_func = destroy_func - self._init_size = init_size - self._max_size = max_size - self._items = None - self._outstanding_count = 0 - - def clear(self) -> None: - while self.items: - self._destroy_func(self.items.popleft()) - - def __del__(self) -> None: - self.clear() - - @property - def items(self) -> deque[T]: - if self._items is None: - self._items = deque[T]( - (self._create_func() for _ in range(self._init_size)), maxlen=self._max_size - ) - return self._items - - def get(self) -> T: - ret = self.items.popleft() if self.items else self._create_func() - self._outstanding_count += 1 - return ret - - def put(self, item: T) -> None: - self._outstanding_count -= 1 - if self._max_size is not None and len(self.items) >= self._max_size: - self._destroy_func(item) - else: - self.items.append(item) - - @property - def outstanding_count(self) -> int: - "number of items currently we get() but not put()" - return self._outstanding_count - - @property - def cached_count(self) -> int: - "number of items currently in the pool" - return len(self.items) - - @property - def total_count(self) -> int: - "total number of items created, including both outstanding and cached" - return self.outstanding_count + self.cached_count - - -class ItemHolderBase(Generic[T], ABC): - __slots__ = ("_item",) - _item: T | None - - def __init__(self) -> None: - self._item = self.pool.get() - - def close(self) -> None: - # Manually inlined for better performance. - item = self._item - if item is not None: - self.pool.put(item) - self._item = None - - def __del__(self) -> None: - self.close() - - def is_closed(self) -> bool: - return self._item is None - - def get(self) -> T: - # Manually inlined for better performance. - item = self._item - assert item is not None - return item - - @property - def handle(self) -> T: - # Manually inlined for better performance. - item = self._item - assert item is not None - return item - - @property - @abstractmethod - def pool(self) -> SimplePool[T]: ... - - -class CachedCudaEvent(ItemHolderBase[drv.CUevent]): - """ - A cached CUDA event without support for timing. Recorded to a stream when created. - """ - - __slots__ = () - _pool: ClassVar[SimplePool[drv.CUevent] | None] = None - NULL: ClassVar["_NullCudaEvent"] - - def __init__(self, stream: CudaStream) -> None: - super().__init__() - self._record(stream) - - def query_complete(self) -> bool: - """ - Query the event. If complete, also close the event. Closed events are always considered complete. - """ - # Manually inlined for better performance. - ev = self._item - if ev is None: - return True - (err,) = drv.cuEventQuery(ev) - if int(err) == int(drv.CUresult.CUDA_SUCCESS): - self.close() - return True - elif int(err) == int(drv.CUresult.CUDA_ERROR_NOT_READY): - return False - else: - raise CuError(err) - - def synchronize(self) -> None: - # Manually inlined for better performance. - ev = self._item - if ev is None: - return - _unwrap(drv.cuEventSynchronize(ev)) - self.close() - - def wait_in_stream(self, stream: CudaStream) -> None: - # Manually inlined for better performance. - ev = self._item - if ev is None: - return - _unwrap(drv.cuStreamWaitEvent(stream, ev, 0)) - - def _record(self, stream: CudaStream) -> None: - """ - Prefer new event instead of recording an existing event. - """ - # Manually inlined for better performance. - ev = self._item - assert ev is not None - _unwrap(drv.cuEventRecord(ev, stream)) - - @property - def pool(self) -> SimplePool[drv.CUevent]: - if CachedCudaEvent._pool is None: - CachedCudaEvent._pool = SimplePool[drv.CUevent]( - lambda: _unwrap(drv.cuEventCreate(drv.CUevent_flags.CU_EVENT_DISABLE_TIMING)), - lambda ev: _unwrap(drv.cuEventDestroy(ev)), # pyright: ignore - init_size=1024, - ) - return CachedCudaEvent._pool - - -class _NullCudaEvent(CachedCudaEvent): - """ - A null CUDA event that is closed (and always complete). - """ - - __slots__ = () - - def __init__(self) -> None: - # do not call super().__init__(). We don't need an event here. - self._item = None - - -CachedCudaEvent.NULL = _NullCudaEvent() - - -# @TODO: consider do this in a single batch call to C++. -def stream_wait_events(stream: CudaStream, events: Iterable[CachedCudaEvent]) -> None: - "Batched wait for multiple events with deduplication first." - if not isinstance(events, Set): - events = set(events) - for ev in events: - ev.wait_in_stream(stream) - - -class CachedCudaStream(ItemHolderBase[CudaStream]): - """ - A cached non-blocking CUDA stream. - """ - - __slots__ = () - _pool: ClassVar[SimplePool[CudaStream] | None] = None - - def __init__(self) -> None: - super().__init__() - - def wait_event(self, event: drv.CUevent) -> None: - _unwrap(drv.cuStreamWaitEvent(self.get(), event, drv.CU_STREAM_WAIT_VALUE_COMPLETED)) - - def wait_events(self, events: Sequence[CachedCudaEvent] | set[CachedCudaEvent]) -> None: - """ - Wait for events with deduplication first. - """ - stream_wait_events(self.get(), events) - - def record_event(self) -> CachedCudaEvent: - return CachedCudaEvent(self.get()) - - def __cuda_stream__(self) -> tuple[int, int]: - return 0, int(self.get()) - - def synchronize(self) -> None: - _unwrap(drv.cuStreamSynchronize(self.handle)) - - @property - def pool(self) -> SimplePool[CudaStream]: - if CachedCudaStream._pool is None: - CachedCudaStream._pool = SimplePool[CudaStream]( - lambda: CudaStream( - int(_unwrap(drv.cuStreamCreate(drv.CUstream_flags.CU_STREAM_NON_BLOCKING))) # pyright: ignore - ), - lambda stream: _unwrap(drv.cuStreamDestroy(stream)), # pyright: ignore - init_size=128, - ) - return CachedCudaStream._pool - - -class TemporaryCudaStream(CachedCudaStream): - """ - A cached non-blocking CUDA stream. Mainly used as temporary worker streams. - Requires a list of prior events to wait for dependencies. A finish event is recorded when exiting - normally. Call take_finish_event() to consume the finish event, otherwise you get a warning. - """ - - __slots__ = "_finish_event" - _finish_event: CachedCudaEvent | None - - def __init__(self, prior_events: Sequence[CachedCudaEvent] | set[CachedCudaEvent]): - super().__init__() - self.wait_events(prior_events) - self._finish_event = None - - def __del__(self) -> None: - if self._finish_event is not None: - warnings.warn("[KVCacheManager] finish event recorded but not taken") - super().__del__() - - def take_finish_event(self) -> CachedCudaEvent: - ret = unwrap_optional(self._finish_event) - self._finish_event = None - return ret - - def __enter__(self) -> "TemporaryCudaStream": - return self - - def __exit__(self, exc_type, exc_value, traceback) -> None: - if not exc_type: - self._finish_event = self.record_event() - - -def merge_events(events: Sequence[CachedCudaEvent] | set[CachedCudaEvent]) -> CachedCudaEvent: - if len(events) == 0: - return CachedCudaEvent.NULL - if len(events) == 1: - ev = next(iter(events)) - return ev if not ev.is_closed() else CachedCudaEvent.NULL - with TemporaryCudaStream(events) as stream: - pass - return stream.take_finish_event() - - -class MultiStreamExecutor: - __slots__ = ("_prior_event", "_streams", "_finish_event") - _prior_event: CachedCudaEvent - _streams: list[TemporaryCudaStream] - _finish_event: CachedCudaEvent | None - - def __init__(self, prior_events: Sequence[CachedCudaEvent] | set[CachedCudaEvent]): - self._prior_event = merge_events(prior_events) - self._streams = [] - self._finish_event = None - - def __enter__(self) -> "MultiStreamExecutor": - return self - - def __exit__(self, exc_type, exc_value, traceback) -> None: - events = [s.take_finish_event() for s in self._streams] - self._streams.clear() - self._finish_event = merge_events(events) - - def __del__(self) -> None: - assert_critical(self._finish_event is None, "finish event not taken") - - def new_stream(self) -> TemporaryCudaStream: - stream = TemporaryCudaStream((self._prior_event,)) - self._streams.append(stream) - return stream - - def take_finish_event(self) -> CachedCudaEvent: - ret = unwrap_optional(self._finish_event) - self._finish_event = None - return ret - - -class SharedPoolProvider(Generic[T]): - _pool: SimplePool[T] - - def __init__(self, pool: SimplePool[T]): - self._pool = pool - - def pool(self) -> SimplePool[T]: - return self._pool - - -class ItemHolderWithSharedPool(ItemHolderBase[T]): - __slots__ = ("_pool",) - _pool: SimplePool[T] - - def __init__(self, pool: SimplePool[T]) -> None: - self._pool = pool - super().__init__() - - def __del__(self) -> None: - self.close() - - @property - def pool(self) -> SimplePool[T]: - return self._pool - - -HolderT = TypeVar("HolderT", bound=ItemHolderWithSharedPool) - - -# For subclassing if holder needs to be customized -class PooledFactoryBase(Generic[T, HolderT]): - _Holder: Type[HolderT] # subclasses must initialize this static attribute - __slots__ = ("_pool",) - _pool: SimplePool[T] - - def __init__( - self, - create_func: Callable[[], T], - destroy_func: Callable[[T], None], - init_size: int = 0, - max_cache_size: int | None = None, - ): - self._pool = SimplePool[T](create_func, destroy_func, init_size, max_cache_size) - - def create(self) -> HolderT: - return self._Holder(self._pool) - - def clear(self) -> None: - self._pool.clear() - - -def query_total_gpu_memory() -> int: - _, total = _unwrap(drv.cuMemGetInfo()) # pyright: ignore - return total - - -def query_free_gpu_memory() -> int: - free, _ = _unwrap(drv.cuMemGetInfo()) # pyright: ignore - return free - - -class CudaStreamWrapper: - "Just a wrapper to make it compatible with IsStreamT protocol. Does not own the stream." - - __slots__ = ("_stream",) - _stream: CudaStream - - def __init__(self, stream: CudaStream) -> None: - self._stream = stream - - def __cuda_stream__(self) -> tuple[int, int]: - return 0, int(self._stream) - - -@contextmanager -def temporary_sys_path(path: str) -> Iterator[None]: - already_in_path = path in sys.path - if not already_in_path: - sys.path.insert(0, path) - try: - yield - finally: - if not already_in_path: - sys.path.remove(path) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/mypy_mypyc.ini b/tensorrt_llm/runtime/kv_cache_manager_v2/mypy_mypyc.ini deleted file mode 100644 index 1b798b09ee0a..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/mypy_mypyc.ini +++ /dev/null @@ -1,51 +0,0 @@ -; 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. - -[mypy] -# Only check files explicitly listed - don't follow any imports -follow_imports = skip -follow_imports_for_stubs = False - -# Ignore all missing imports -ignore_missing_imports = True - -# Allow untyped code in dependencies -allow_untyped_calls = True -allow_untyped_defs = True -check_untyped_defs = False - -# Disable various warnings to reduce noise -warn_return_any = False -warn_unused_ignores = False -warn_unreachable = False -no_implicit_optional = False - -# Don't check .pyi files outside our target -exclude = (?x)( - ^(?!tensorrt_llm/runtime/kv_cache_manager_v2/) -) - -# Ignore errors in any imported modules -[mypy-tensorrt_llm.executor.*] -ignore_errors = True -follow_imports = skip - -[mypy-tensorrt_llm.bindings.*] -ignore_errors = True -follow_imports = skip - -[mypy-torch.*] -ignore_errors = True -follow_imports = skip diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/README.md b/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/README.md deleted file mode 100644 index 6128f1a7b6ec..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/README.md +++ /dev/null @@ -1,140 +0,0 @@ - - -# rawref - Mutable Reference C Extension - -A C extension that provides a mutable reference class similar to `weakref.ref` for holding weak-like references to Python objects. - -## Features - -- **`ref[T]`**: A generic reference class (like `weakref.ref`) that stores an object's ID -- **Singleton pattern**: `ref(obj)` returns the same reference if `obj.__rawref__` is valid -- **Dereferencing**: Call `r()` to get the object, or `None` if invalid -- **Invalidation**: Call `r.invalidate()` to mark the reference as invalid -- **NULL constant**: Use `NULL` to initialize `__rawref__` attributes -- **Type-safe**: Comes with `.pyi` stub file for proper type checking -- **API compatible with weakref**: Use `ref` for both object creation and type hints - -## Building - -From the `rawref` directory: - -```bash -python setup.py build_ext --inplace -``` - -Or install it: - -```bash -pip install -e . -``` - -## Usage - -```python -from rawref import ref, NULL - -class MyClass: - # Class attribute: default value for __rawref__ - # Each instance will get its own __rawref__ instance attribute when ref() is called - __rawref__ = NULL - - def __init__(self, value): - self.value = value - - def __del__(self): - # self.__rawref__ is an instance attribute (set by ref()) - # Invalidate the canonical reference when object is destroyed - if self.__rawref__.is_valid: - self.__rawref__.invalidate() - -# Create an object and a reference to it (just like weakref.ref) -obj = MyClass(42) -r1 = ref(obj) - -# The reference is automatically stored as an instance attribute obj.__rawref__ -print(obj.__rawref__ is r1) # True - -# Singleton pattern: creating another ref returns the same one -r2 = ref(obj) -print(r1 is r2) # True - -# Dereference to get the object back -print(r1()) # -print(r1().value) # 42 - -# Check validity -print(r1.is_valid) # True - -# After invalidation -r1.invalidate() -print(r1()) # None -print(r1.is_valid) # False - -# Creating a new ref after invalidation creates a new reference -r3 = ref(obj) -print(r1 is r3) # False -print(r3.is_valid) # True -``` - -## Type Hints - -Like `weakref.ref`, you can use `ref` for both object creation and type hints: - -```python -from rawref import ref, NULL - -class MyClass: - __rawref__ = NULL - -# Create and type a reference -r: ref[MyClass] = ref(MyClass()) - -# Alternative: use ReferenceType directly -from rawref import ReferenceType -r: ReferenceType[MyClass] = ReferenceType(MyClass()) -``` - -## Warning - -This implementation uses raw object IDs (memory addresses) and attempts to dereference them. This is inherently unsafe and should be used with caution. The reference does not keep the object alive (unlike a strong reference), so care must be taken to ensure the object is not garbage collected while references exist. - -## API - -### Classes and Constants -- `ReferenceType`: The main reference class -- `ref`: Alias for `ReferenceType` (like `weakref.ref`) -- `NULL`: An invalid reference constant for initialization - -### Creation -- `ref(obj)`: Create a reference to `obj`, or return existing valid reference from `obj.__rawref__` - -### Properties -- `r.is_valid`: Check if the reference is still valid (read-only) - -### Methods -- `r()`: Dereference to get the object, or `None` if invalid -- `r.invalidate()`: Mark the reference as invalid - -## Singleton Pattern - -The `ref()` function implements a singleton pattern: -1. When `ref(obj)` is called, it checks if `obj.__rawref__` (instance attribute) exists and is valid -2. If yes, it returns the existing reference -3. If no, it creates a new reference and sets `obj.__rawref__` as an instance attribute - -**Note**: The class attribute `__rawref__ = NULL` is just a default value. When `ref(obj)` is called, it creates an **instance attribute** `obj.__rawref__` that shadows the class attribute. Each instance gets its own `__rawref__` instance attribute, ensuring each object has at most one canonical reference at a time. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.py deleted file mode 100644 index 672e53603a5e..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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. - -"""rawref - Mutable reference with singleton pattern. - -This module provides a C extension for creating mutable references to Python -objects, similar to weakref.ref but with manual invalidation control and a -singleton pattern via __rawref__. - -The main purpose is to work around the issue that mypyc does not support -weakref. - -Main exports: -- ReferenceType: The reference class -- ref: Alias for ReferenceType (recommended, like weakref.ref) -- NULL: Invalid reference constant for initialization -""" - -from ._rawref import NULL, ReferenceType, ref - -__all__ = ["ReferenceType", "ref", "NULL"] - -__version__ = "2.0.0" diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.pyi deleted file mode 100644 index 54b3e9549d09..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.pyi +++ /dev/null @@ -1,82 +0,0 @@ -# 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. - -from typing import Generic, Optional, TypeVar - -T = TypeVar("T") - -class ReferenceType(Generic[T]): - """A mutable reference holder that stores an object ID. - - This class holds a reference to an object via its ID and allows - dereferencing it. The reference can be invalidated. - - Like weakref.ref, but stores raw object IDs instead of proper weak references. - - Implements a singleton pattern: calling ref(obj) multiple times returns the - same reference if obj.__rawref__ exists and is valid. - """ - - @property - def is_valid(self) -> bool: - """Check if the reference is still valid (read-only).""" - ... - - def __init__(self, obj: T) -> None: - """Initialize a ReferenceType with an object. - - If obj.__rawref__ exists and is valid, returns that instead. - Otherwise creates a new reference and sets obj.__rawref__ to it. - - Args: - obj: The object to reference. - """ - ... - - def __call__(self) -> Optional[T]: - """Dereference the object. - - Returns: - The referenced object, or None if the reference is invalid. - """ - ... - - def __hash__(self) -> int: - """Return the hash of the referenced object (its ID). - - Raises: - RuntimeError: If the reference is invalid. - """ - ... - - def invalidate(self) -> None: - """Invalidate the reference. - - After calling this method, __call__() will return None. - This should be called from T.__del__ to invalidate the reference. - """ - ... - -# Alias 'ref' to 'ReferenceType' (like weakref.ref is an alias to weakref.ReferenceType) -ref = ReferenceType - -# NULL is an invalid reference constant that can be used to initialize __rawref__ -NULL: ReferenceType - -# For type hints, you can use either: -# r: ref[MyClass] = ref(obj) -# or: -# r: ReferenceType[MyClass] = ReferenceType(obj) -# Both are equivalent. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/rawrefmodule.c b/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/rawrefmodule.c deleted file mode 100644 index a283d697f55a..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/rawrefmodule.c +++ /dev/null @@ -1,250 +0,0 @@ -/* - * 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. - */ - -#define PY_SSIZE_T_CLEAN -#include -#include - -/* ReferenceType object structure */ -typedef struct -{ - PyObject_HEAD Py_ssize_t object_id; /* ID of the referenced object. 0 if invalid. */ -} ReferenceTypeObject; - -/* Forward declarations */ -static PyTypeObject ReferenceTypeType; - -/* Cached attribute name for faster lookups */ -static PyObject* rawref_attr_name = NULL; - -/* ReferenceType.__new__ - implements singleton pattern via __rawref__ */ -static PyObject* ReferenceType_new(PyTypeObject* type, PyObject* args, PyObject* kwds) -{ - PyObject* obj = NULL; - static char* kwlist[] = {"obj", NULL}; - - /* Parse arguments to get the object */ - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &obj)) - { - return NULL; - } - - /* Try to get existing __rawref__ using cached attribute name (faster) */ - PyObject* existing_ref = PyObject_GetAttr(obj, rawref_attr_name); - if (existing_ref != NULL) - { - /* Check if it's a ReferenceType instance and is valid */ - if (PyObject_TypeCheck(existing_ref, &ReferenceTypeType)) - { - ReferenceTypeObject* ref_obj = (ReferenceTypeObject*) existing_ref; - if (ref_obj->object_id != 0) - { - /* Return existing valid reference */ - return existing_ref; - } - } - Py_DECREF(existing_ref); - } - else - { - /* Clear the AttributeError if __rawref__ doesn't exist */ - PyErr_Clear(); - } - - /* Create new reference */ - ReferenceTypeObject* self; - self = (ReferenceTypeObject*) type->tp_alloc(type, 0); - if (self != NULL) - { - self->object_id = (Py_ssize_t) obj; - - /* Set obj.__rawref__ to this new reference using cached attr name */ - if (PyObject_SetAttr(obj, rawref_attr_name, (PyObject*) self) < 0) - { - /* If we can't set __rawref__, just clear the error and continue */ - PyErr_Clear(); - } - } - return (PyObject*) self; -} - -/* ReferenceType.__init__ */ -static int ReferenceType_init(ReferenceTypeObject* self, PyObject* args, PyObject* kwds) -{ - /* __new__ already did all the work, including setting object_id */ - /* Skip argument parsing since __new__ already validated them */ - /* This saves ~5-10% overhead on object creation */ - return 0; -} - -/* ReferenceType.__hash__ */ -static Py_hash_t ReferenceType_hash(ReferenceTypeObject* self) -{ - if (self->object_id == 0) - { - PyErr_SetString(PyExc_RuntimeError, "Reference is invalid"); - return -1; - } - Py_hash_t const h = (Py_hash_t) self->object_id; - /* Follow CPython hash implementation to remap -1 to -2, as -1 is reserved to signal - an error. */ - return h == -1 ? -2 : h; -} - -/* ReferenceType.__call__() - dereference the object */ -static PyObject* ReferenceType_call(ReferenceTypeObject* self, PyObject* args, PyObject* kwds) -{ - PyObject* obj; - - if (self->object_id == 0) - { - Py_RETURN_NONE; - } - - /* Use _PyObject_FromStackRefSteal or ctypes approach */ - /* We need to find the object by its id */ - /* This is the tricky part - we need to convert id back to object */ - - /* Use ctypes.cast to convert id to PyObject* */ - obj = (PyObject*) self->object_id; - - /* Check if the object is still alive by verifying ref count > 0 */ - /* This is somewhat unsafe but matches the intended behavior */ - if (Py_REFCNT(obj) > 0) - { - Py_INCREF(obj); - return obj; - } - - /* Object no longer valid */ - self->object_id = 0; - Py_RETURN_NONE; -} - -/* ReferenceType.invalidate() */ -static PyObject* ReferenceType_invalidate(ReferenceTypeObject* self, PyObject* Py_UNUSED(ignored)) -{ - self->object_id = 0; - Py_RETURN_NONE; -} - -/* ReferenceType.is_valid property getter */ -static PyObject* ReferenceType_is_valid(ReferenceTypeObject* self, void* closure) -{ - return PyBool_FromLong(self->object_id != 0); -} - -/* ReferenceType.__class_getitem__() - support for generic type subscripting */ -static PyObject* ReferenceType_class_getitem(PyObject* cls, PyObject* item) -{ - /* Just return the class itself, ignore the type parameter */ - /* This allows rawref.ref[T] to work at runtime like weakref.ref[T] */ - Py_INCREF(cls); - return cls; -} - -/* Method definitions */ -static PyMethodDef ReferenceType_methods[] = { - {"invalidate", (PyCFunction) ReferenceType_invalidate, METH_NOARGS, - "Invalidate the reference, making it return None on dereference."}, - {"__class_getitem__", (PyCFunction) ReferenceType_class_getitem, METH_O | METH_CLASS, - "Support for generic type subscripting (e.g., ref[T])."}, - {NULL} /* Sentinel */ -}; - -/* Property definitions */ -static PyGetSetDef ReferenceType_getsetters[] = { - {"is_valid", (getter) ReferenceType_is_valid, NULL, "Check if the reference is still valid (read-only).", NULL}, - {NULL} /* Sentinel */ -}; - -/* Type definition */ -static PyTypeObject ReferenceTypeType = { - PyVarObject_HEAD_INIT(NULL, 0).tp_name = "_rawref.ReferenceType", - .tp_doc = "A mutable reference holder that stores an object ID.", - .tp_basicsize = sizeof(ReferenceTypeObject), - .tp_itemsize = 0, - .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, - .tp_new = ReferenceType_new, - .tp_init = (initproc) ReferenceType_init, - .tp_call = (ternaryfunc) ReferenceType_call, - .tp_hash = (hashfunc) ReferenceType_hash, - .tp_methods = ReferenceType_methods, - .tp_getset = ReferenceType_getsetters, -}; - -/* Module definition */ -static PyModuleDef rawrefmodule = { - PyModuleDef_HEAD_INIT, - .m_name = "_rawref", - .m_doc = "C extension providing mutable reference class ReferenceType (internal module).", - .m_size = -1, -}; - -/* Module initialization */ -PyMODINIT_FUNC PyInit__rawref(void) -{ - PyObject* m; - ReferenceTypeObject* null_ref; - - if (PyType_Ready(&ReferenceTypeType) < 0) - return NULL; - - m = PyModule_Create(&rawrefmodule); - if (m == NULL) - return NULL; - - /* Cache the __rawref__ attribute name for faster lookups */ - rawref_attr_name = PyUnicode_InternFromString("__rawref__"); - if (rawref_attr_name == NULL) - { - Py_DECREF(m); - return NULL; - } - - Py_INCREF(&ReferenceTypeType); - if (PyModule_AddObject(m, "ReferenceType", (PyObject*) &ReferenceTypeType) < 0) - { - Py_DECREF(&ReferenceTypeType); - Py_DECREF(m); - return NULL; - } - - /* Add 'ref' as an alias for 'ReferenceType' (like weakref.ref) */ - Py_INCREF(&ReferenceTypeType); - if (PyModule_AddObject(m, "ref", (PyObject*) &ReferenceTypeType) < 0) - { - Py_DECREF(&ReferenceTypeType); - Py_DECREF(m); - return NULL; - } - - /* Create NULL constant - an invalid reference */ - null_ref = (ReferenceTypeObject*) ReferenceTypeType.tp_alloc(&ReferenceTypeType, 0); - if (null_ref != NULL) - { - null_ref->object_id = 0; - if (PyModule_AddObject(m, "NULL", (PyObject*) null_ref) < 0) - { - Py_DECREF(null_ref); - Py_DECREF(m); - return NULL; - } - } - - return m; -} diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/setup.py b/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/setup.py deleted file mode 100644 index 48a6634d3781..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/setup.py +++ /dev/null @@ -1,28 +0,0 @@ -# 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. - -from setuptools import Extension, setup - -rawref_module = Extension( - '_rawref', - sources=['rawrefmodule.c'], -) - -setup( - name='rawref', - version='1.0', - description='C extension providing mutable reference class Ref[T]', - ext_modules=[rawref_module], -) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/test_rawref.py b/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/test_rawref.py deleted file mode 100644 index fab925869fac..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/rawref/test_rawref.py +++ /dev/null @@ -1,269 +0,0 @@ -# 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. - -"""Test script for the rawref module.""" - -import os -import sys - -# Add parent directory to path to import the rawref package -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -try: - from rawref import NULL, ReferenceType, ref -except ImportError as e: - print(f"Error importing rawref: {e}") - print("Make sure to build the extension first with: python setup.py build_ext --inplace") - sys.exit(1) - - -class TestObject: - """Test class with __del__ that invalidates references.""" - - def __init__(self, value): - self.value = value - self.refs = [] # Store references to invalidate - - def __del__(self): - print(f"TestObject({self.value}).__del__ called, invalidating {len(self.refs)} references") - for ref in self.refs: - ref.invalidate() - - -def test_basic_reference(): - """Test basic reference creation and dereferencing.""" - print("\n=== Test 1: Basic Reference ===") - - obj = TestObject(42) - r = ref(obj) - obj.refs.append(r) - - print(f"Created object with value: {obj.value}") - print(f"Reference is_valid: {r.is_valid}") - - # Dereference - dereferenced = r() - print(f"Dereferenced object: {dereferenced}") - if dereferenced: - print(f"Dereferenced value: {dereferenced.value}") - - assert r.is_valid, "Reference should be valid" - assert dereferenced is obj, "Dereferenced object should be the same as original" - - -def test_invalidation(): - """Test manual invalidation.""" - print("\n=== Test 2: Manual Invalidation ===") - - obj = TestObject(123) - r = ref(obj) - - print(f"Before invalidation - is_valid: {r.is_valid}") - print(f"Dereferenced: {r()}") - - r.invalidate() - - print(f"After invalidation - is_valid: {r.is_valid}") - print(f"Dereferenced: {r()}") - - assert not r.is_valid, "Reference should be invalid after invalidate()" - assert r() is None, "Dereferencing invalid reference should return None" - - -def test_del_invalidation(): - """Test invalidation via __del__.""" - print("\n=== Test 3: Invalidation via __del__ ===") - - r = None - - # Create object in a scope - def create_and_ref(): - obj = TestObject(999) - nonlocal r - r = ref(obj) - obj.refs.append(r) - - print(f"Inside scope - is_valid: {r.is_valid}") - print(f"Inside scope - dereferenced: {r()}") - - create_and_ref() - - # Object should be deleted and reference invalidated - print(f"After scope - is_valid: {r.is_valid}") - print(f"After scope - dereferenced: {r()}") - - assert not r.is_valid, "Reference should be invalid after object deletion" - assert r() is None, "Dereferencing should return None after object deletion" - - -def test_multiple_references(): - """Test that singleton pattern returns same reference.""" - print("\n=== Test 4: Singleton Pattern - Same Reference ===") - - obj = TestObject(555) - r1 = ref(obj) - r2 = ref(obj) - obj.refs.extend([r1, r2]) - - print(f"r1 is r2: {r1 is r2}") - - assert r1 is r2, "With singleton pattern, should return the same reference" - - # Invalidate r1 (which is the same as r2) - r1.invalidate() - - print("After invalidating r1:") - print(f" r1 is_valid: {r1.is_valid}, dereferenced: {r1()}") - print(f" r2 is_valid: {r2.is_valid}, dereferenced: {r2()}") - - # Since they're the same object, both are invalid - assert not r1.is_valid, "r1 should be invalid" - assert not r2.is_valid, "r2 should also be invalid (same object)" - - -def test_alias_equivalence(): - """Test that ref and ReferenceType are the same.""" - print("\n=== Test 5: ref and ReferenceType are equivalent ===") - - obj = TestObject(777) - r1 = ref(obj) - r2 = ReferenceType(obj) - - print(f"ref is ReferenceType: {ref is ReferenceType}") - print(f"type(r1): {type(r1)}") - print(f"type(r2): {type(r2)}") - print(f"r1 is r2: {r1 is r2}") - - assert ref is ReferenceType, "ref should be an alias for ReferenceType" - assert r1 is r2, "Should return the same reference object (singleton pattern)" - - -def test_singleton_pattern(): - """Test that ref(obj) returns the same reference if __rawref__ is valid.""" - print("\n=== Test 6: Singleton Pattern ===") - - obj = TestObject(888) - - # First call creates a new reference - r1 = ref(obj) - print(f"First ref: {r1}, valid: {r1.is_valid}") - - # Second call returns the same reference - r2 = ref(obj) - print(f"Second ref: {r2}, valid: {r2.is_valid}") - print(f"r1 is r2: {r1 is r2}") - - assert r1 is r2, "Should return the same reference object" - - # After invalidation, a new call creates a new reference - r1.invalidate() - print(f"After invalidation: r1.valid={r1.is_valid}, r2.valid={r2.is_valid}") - - r3 = ref(obj) - print(f"Third ref (after invalidation): {r3}, valid: {r3.is_valid}") - print(f"r1 is r3: {r1 is r3}") - - assert r1 is not r3, "Should create a new reference after invalidation" - assert r3.is_valid, "New reference should be valid" - - -def test_null_constant(): - """Test the NULL constant.""" - print("\n=== Test 7: NULL Constant ===") - - print(f"NULL: {NULL}") - print(f"NULL.is_valid: {NULL.is_valid}") - print(f"NULL(): {NULL()}") - print(f"type(NULL): {type(NULL)}") - - assert not NULL.is_valid, "NULL should be invalid" - assert NULL() is None, "NULL() should return None" - - # Test using NULL to initialize __rawref__ - class MyClass: - __rawref__ = NULL - - obj = MyClass() - r = ref(obj) - print("After creating ref for obj with __rawref__=NULL:") - print(f" obj.__rawref__ is r: {obj.__rawref__ is r}") - print(f" r.is_valid: {r.is_valid}") - - assert obj.__rawref__ is r, "Should update __rawref__ to new reference" - assert r.is_valid, "New reference should be valid" - - -def test_hidden_object_id(): - """Test that object_id is hidden.""" - print("\n=== Test 8: Hidden object_id ===") - - obj = TestObject(123) - r = ref(obj) - - # Try to access object_id - should raise AttributeError - try: - _ = r.object_id - assert False, "Should not be able to access object_id" - except AttributeError: - print("✓ object_id is hidden (AttributeError raised)") - - print() - - -def test_is_valid_readonly(): - """Test that is_valid is read-only.""" - print("\n=== Test 9: is_valid is read-only ===") - - obj = TestObject(456) - r = ref(obj) - - print(f"r.is_valid: {r.is_valid}") - - # Try to set is_valid - should raise AttributeError - try: - r.is_valid = False - assert False, "Should not be able to set is_valid" - except AttributeError: - print("✓ is_valid is read-only (AttributeError raised)") - - print() - - -if __name__ == "__main__": - print("Testing rawref module...") - - try: - test_basic_reference() - test_invalidation() - test_del_invalidation() - test_multiple_references() - test_alias_equivalence() - test_singleton_pattern() - test_null_constant() - test_hidden_object_id() - test_is_valid_readonly() - - print("\n" + "=" * 50) - print("All tests passed!") - print("=" * 50) - except AssertionError as e: - print(f"\n❌ Test failed: {e}") - sys.exit(1) - except Exception as e: - print(f"\n❌ Unexpected error: {e}") - import traceback - - traceback.print_exc() - sys.exit(1) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py b/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py deleted file mode 100644 index 4e40456a4dbe..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py +++ /dev/null @@ -1,141 +0,0 @@ -# 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. - -"""Setup script for compiling kv_cache_manager_v2 with mypyc. - -Usage (from project root): - python tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py build_ext --inplace - -Or use the build script: - ./tensorrt_llm/runtime/kv_cache_manager_v2/build_mypyc.sh -""" - -import os -import sys - -from mypyc.build import mypycify -from setuptools import setup - -# Set environment variables BEFORE importing mypyc -os.environ["MYPY_FORCE_COLOR"] = "0" - -# Write a strict mypy config that won't check external files -mypy_config_path = os.path.abspath("mypy_mypyc_build.ini") -with open(mypy_config_path, "w") as f: - f.write("""[mypy] -# Critical: Don't follow any imports outside the specified files -follow_imports = skip -follow_imports_for_stubs = False - -# Ignore missing imports completely -ignore_missing_imports = True - -# Allow all untyped code -allow_untyped_calls = True -allow_untyped_defs = True -allow_incomplete_defs = True -allow_untyped_globals = True -check_untyped_defs = False - -# Disable all warnings that might cause errors -disallow_untyped_calls = False -disallow_untyped_defs = False -disallow_incomplete_defs = False -warn_return_any = False -warn_unused_ignores = False - -# Disable type validation errors (for external types like drv.CUstream) -disable_error_code = valid-type -""") - -# Point mypy to this config by adding to sys.argv before mypyc runs -sys.argv.extend(["--config-file", mypy_config_path]) - -# List all Python modules in kv_cache_manager_v2 to compile -# -# EXCLUDED FILES: -# - _exceptions.py: inherits from builtin Exception classes (mypyc limitation) -# -modules = [ - # Main module files - "kv_cache_manager_v2/__init__.py", - "kv_cache_manager_v2/_block_radix_tree.py", - "kv_cache_manager_v2/_common.py", - "kv_cache_manager_v2/_config.py", - "kv_cache_manager_v2/_copy_engine.py", - "kv_cache_manager_v2/_cuda_virt_mem.py", - "kv_cache_manager_v2/_event_manager.py", - "kv_cache_manager_v2/_exceptions.py", - "kv_cache_manager_v2/_introspection.py", - "kv_cache_manager_v2/_life_cycle_registry.py", - "kv_cache_manager_v2/_page.py", - "kv_cache_manager_v2/_storage_manager.py", - "kv_cache_manager_v2/_utils.py", - # _core submodule - "kv_cache_manager_v2/_core/__init__.py", - "kv_cache_manager_v2/_core/_kv_cache_manager.py", - "kv_cache_manager_v2/_core/_kv_cache.py", - "kv_cache_manager_v2/_core/_pending_stats.py", - # _eviction_controller submodule - "kv_cache_manager_v2/_eviction_controller/__init__.py", - "kv_cache_manager_v2/_eviction_controller/_eviction_controller.py", - # _storage submodule - "kv_cache_manager_v2/_storage/__init__.py", - "kv_cache_manager_v2/_storage/_config.py", - "kv_cache_manager_v2/_storage/_core.py", -] - -print(f"Compiling {len(modules)} modules with mypyc...") -print("Excluded: None") -print("") - -try: - ext_modules = mypycify( - modules, - opt_level="3", # Maximum optimization - multi_file=True, # Allow cross-module references (needed for inheritance) - verbose=True, # Show what's being compiled - separate=False, # Compile into single .so (required for cross-module inheritance) - strip_asserts=False, # Keep assertions for debugging - ) - -except Exception as e: - print(f"Error during mypyc compilation: {e}") - sys.exit(1) -finally: - # Cleanup temp config - if os.path.exists(mypy_config_path): - try: - os.remove(mypy_config_path) - except OSError: - pass - - # Remove --config-file arguments from sys.argv before calling setup() - # This prevents setuptools from seeing arguments it doesn't understand - while "--config-file" in sys.argv: - idx = sys.argv.index("--config-file") - sys.argv.pop(idx) # Remove '--config-file' - if idx < len(sys.argv): # Remove the path that follows it - sys.argv.pop(idx) - -setup( - name="kv_cache_manager_v2_compiled", - ext_modules=ext_modules, - packages=["kv_cache_manager_v2.rawref"], - package_data={ - "kv_cache_manager_v2": ["*.pyi", "**/*.pyi"], - }, - python_requires=">=3.8", -) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 5cc532b9e8a8..6deb648730e1 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -34,8 +34,7 @@ KV_CACHE_HASH_ALGO_DEFAULT, KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2, KV_CACHE_HASH_ALGO_V2_SHA256_64, KV_CACHE_HASH_ALGOS, BlockHash, BlockHashMixin, OpenAIRequest, block_key_hasher, get_cache_salt_id, - get_request_num_tokens, hash_v1_block_key, truncate_sha256_hash_to_int64, - v2_sha256_block_hasher) + get_request_num_tokens, hash_v1_block_key, truncate_sha256_hash_to_int64) _MSGPACK_HEADERS = {"Content-Type": "application/msgpack"} COORDINATOR_SELECT_MAX_ATTEMPTS = 2 diff --git a/tensorrt_llm/serve/router_utils.py b/tensorrt_llm/serve/router_utils.py index 7e416f5e7782..45b3ca6829ab 100644 --- a/tensorrt_llm/serve/router_utils.py +++ b/tensorrt_llm/serve/router_utils.py @@ -25,9 +25,7 @@ from tensorrt_llm.bindings.internal.batch_manager import BlockKeyHasher as _NativeBlockKeyHasher from tensorrt_llm.logger import logger from tensorrt_llm.runtime import kv_cache_hash -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import Block as V2Block -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ReuseScope -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import RootBlock as V2RootBlock +from tensorrt_llm.runtime.kv_cache_manager_v2 import ReuseScope, sequence_to_blockchain_keys from tensorrt_llm.serve.chat_tokenization import ( resolve_model_type_from_config, tokenize_chat_request_for_serving, @@ -65,7 +63,6 @@ "BlockHash", "get_request_num_tokens", "block_key_hasher", - "v2_sha256_block_hasher", "BlockHashMixin", "PrefixBlockSet", ] @@ -170,17 +167,6 @@ def block_key_hasher( return hash_v1_block_key(token_ids, parent_hash=parent, cache_salt_id=cache_salt_id) -def v2_sha256_block_hasher( - token_ids: list[int], parent_hash: Optional[str] = None, cache_salt_id: Optional[int] = None -) -> str: - parent_key = ( - V2RootBlock.make_key(ReuseScope(salt=cache_salt_id)) - if parent_hash is None - else bytes.fromhex(parent_hash) - ) - return V2Block.make_key(parent_key, token_ids).hex() - - class BlockHashMixin: """Shared tokenization and block-hash computation. @@ -327,19 +313,24 @@ def _compute_block_hashes( ) -> list[list[BlockHash]]: if hash_algo == KV_CACHE_HASH_ALGO_V1: block_hasher = block_key_hasher - elif hash_algo == KV_CACHE_HASH_ALGO_V2: - block_hasher = v2_sha256_block_hasher - elif hash_algo == KV_CACHE_HASH_ALGO_V2_SHA256_64: + elif hash_algo in (KV_CACHE_HASH_ALGO_V2, KV_CACHE_HASH_ALGO_V2_SHA256_64): + # V2 keys chain from a reuse-scope root, which is exactly what + # sequence_to_blockchain_keys yields; its first pair is the root itself, so + # skip it. In KvCacheManager the last token is not part of any block key. reuse_scope = ReuseScope(salt=cache_salt_id) + to_hash = truncate_sha256_hash_to_int64 + if hash_algo == KV_CACHE_HASH_ALGO_V2: + + def to_hash(key: bytes) -> BlockHash: + return key.hex() + block_hashes: list[list[BlockHash]] = [] for token_list in token_lists: - hash_list = [] - parent_key = V2RootBlock.make_key(reuse_scope) - for t in range(0, len(token_list) - 1, self._tokens_per_block): - t_end = min(t + self._tokens_per_block, len(token_list) - 1) - parent_key = V2Block.make_key(parent_key, token_list[t:t_end]) - hash_list.append(truncate_sha256_hash_to_int64(parent_key)) - block_hashes.append(hash_list) + keys = sequence_to_blockchain_keys( + self._tokens_per_block, reuse_scope, token_list[:-1] + ) + next(keys, None) # the root key labels no block + block_hashes.append([to_hash(key) for _, key in keys]) return block_hashes else: raise ValueError(f"Unsupported KV cache hash algorithm: {hash_algo}") diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index 0ab917aff209..216c318a953d 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -37,14 +37,12 @@ environment: trtllm_repo: '' build_wheel: false work_dir: - # TLLM_KV_CACHE_MANAGER_V2_BACKEND / TRTLLM_ENABLE_PDL / - # NCCL_GRAPH_MIXING_SUPPORT act on the model runtime, so + # TRTLLM_ENABLE_PDL / NCCL_GRAPH_MIXING_SUPPORT act on the model runtime, so # they belong on the ctx and gen workers only; the disagg proxy server does # not read them. TRTLLM_SERVER_DISABLE_GC / TRTLLM_WORKER_DISABLE_GC are set # on both sides. worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 - TRTLLM_ENABLE_PDL=1 NCCL_GRAPH_MIXING_SUPPORT=0 TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp - ENROOT_ALLOW_DEV=yes + TRTLLM_ENABLE_PDL=1 NCCL_GRAPH_MIXING_SUPPORT=0 ENROOT_ALLOW_DEV=yes server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index fda64c2f673d..a7a7ded3e688 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -52,12 +52,12 @@ ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, BatchDesc, KVCacheDesc, PageIndexMode, _introspection, ) -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX _RequestCache = Dict[ Tuple[int, DeepseekV4AttentionType], # (layer index, attention type) diff --git a/tests/unittest/_torch/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py index 0aba2e052faf..2df28efc7de3 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_attention.py +++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py @@ -27,7 +27,7 @@ from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.functional import AttentionMaskType from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX class TestingFlashInferAttentionMetadata(FlashInferAttentionMetadata): diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py index e60a759b662b..7e1e82f71047 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py @@ -11,9 +11,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( - gen_multimodal_cache_key_tokens, -) +from tensorrt_llm.runtime.kv_cache_manager_v2 import gen_multimodal_cache_key_tokens pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py similarity index 99% rename from tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py rename to tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 55c917fd6f9e..9fa6593a9091 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -67,7 +67,6 @@ LayerId, SsmLayerConfig, ) -from tensorrt_llm.runtime.kv_cache_manager_v2._utils import init_cuda_once TOKENS_PER_BLOCK = 4 MAX_SEQ_LEN = 16 @@ -392,7 +391,7 @@ def test_zero_size_layers_are_removed_before_pool_allocation( ) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") - init_cuda_once() + torch.cuda.init() manager = KVCacheManagerV2( KvCacheConfig( max_gpu_total_bytes=16 << 20, @@ -1445,7 +1444,7 @@ def test_external_draft_estimated_quota_supports_allocation_and_resume( ) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") - init_cuda_once() + torch.cuda.init() spec = config_cls(max_draft_len=4, speculative_model="draft") model_config = SimpleNamespace( quant_config=None, @@ -1519,7 +1518,7 @@ def max_num_turns() -> int: def manager(max_num_turns: int) -> KVCacheManagerV2: if not torch.cuda.is_available(): pytest.skip("requires CUDA") - init_cuda_once() + torch.cuda.init() manager = KVCacheManagerV2( KvCacheConfig( enable_block_reuse=True, @@ -1793,7 +1792,9 @@ def test_per_conversation_policy_ignores_overlapping_request( def test_live_storage_stats_use_the_manager_api() -> None: - init_cuda_once() + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() core = KVCacheManager( KVCacheManagerConfig( tokens_per_block=TOKENS_PER_BLOCK, @@ -1934,7 +1935,9 @@ def test_disagg_gen_init_drops_only_the_partial_block() -> None: @pytest.mark.parametrize("enable_stats", [True, False]) def test_disagg_partial_attribution_survives_admission_retry(enable_stats: bool) -> None: """A failed resume preserves the cache, so retrying must not exclude its tail again.""" - init_cuda_once() + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() stream = torch.cuda.Stream() core = KVCacheManager( KVCacheManagerConfig( diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 3b6bf453b8f1..319458b69691 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -50,7 +50,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantAlgo -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX if TYPE_CHECKING: from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -1685,9 +1685,6 @@ def test_e4b_like_config(self): # ---- VSWA (Variable Sliding Window Attention) page index tests ---- @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_vswa_per_pool_page_indices(self): """VSWA: FlashInfer metadata builds separate page indices per pool. @@ -1775,9 +1772,6 @@ def test_vswa_per_pool_page_indices(self): kv_cache_manager.shutdown() @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_vswa_page_index_bounds(self): """VSWA: page indices must be within each layer's pool buffer bounds. @@ -1837,9 +1831,6 @@ def test_vswa_page_index_bounds(self): kv_cache_manager.shutdown() @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_vswa_swap_restores_correct_pool(self): """VSWA: swapping indices between pools and back produces correct data. @@ -2128,9 +2119,6 @@ def test_vswa_pool_cache_not_aliased(self): kv_cache_manager.shutdown() @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_vswa_evicted_page_indices_are_sanitized(self) -> None: """FlashInfer metadata replaces evicted SWA page markers.""" from tensorrt_llm._torch.attention.backends.utils import get_attention_backend @@ -2878,9 +2866,6 @@ def _expected_decode_block_table( @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_shared_kv_draft_view(self) -> None: """The draft view advances lengths without modifying target KV.""" kv_cache_manager, layers, metadata, queries, _, _ = self._make_trtllm_gen_decode_case( @@ -2920,9 +2905,6 @@ def test_shared_kv_draft_view(self) -> None: @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_trtllm_gen_block_table_transitions(self) -> None: """Shrinking the active rectangle clears stale rows and columns.""" initial_page_counts = [8, 5, 3, 2] @@ -2958,9 +2940,6 @@ def test_cuda_graph_trtllm_gen_block_table_transitions(self) -> None: @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> None: """Crossing 64 pages grows host staging without moving the graph buffer.""" initial_page_counts = [63, 2] @@ -3009,9 +2988,6 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_trtllm_gen_request_turnover_matches_eager(self) -> None: """A captured graph remains correct when long requests are replaced by short ones.""" initial_page_counts = [8, 4] @@ -3080,9 +3056,6 @@ def test_cuda_graph_trtllm_gen_request_turnover_matches_eager(self) -> None: ) @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_decode_hybrid_headdim(self): """CUDA graph decode with hybrid head_dim (VSWA). @@ -3261,9 +3234,6 @@ def test_cuda_graph_decode_hybrid_headdim(self): kv_cache_manager.shutdown() @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_multi_step_decode(self): """CUDA graph multi-step decode with hybrid head_dim. @@ -3435,9 +3405,6 @@ def test_cuda_graph_multi_step_decode(self): @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_decode_high_gqa(self) -> None: """CUDA graph decode with GQA=8 and real head_dim (E2B-like). @@ -3602,9 +3569,6 @@ def test_cuda_graph_decode_high_gqa(self) -> None: kv_cache_manager.shutdown() @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def _run_cuda_graph_real_headdim( self, config_dict: dict, @@ -3804,17 +3768,11 @@ def _run_cuda_graph_real_headdim( kv_cache_manager.shutdown() @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_decode_real_headdim(self): """E2B-like: GQA=8, hd=256/512, non-K=V.""" self._run_cuda_graph_real_headdim(deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG), "E2B") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) @unittest.skipUnless( torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0), "FA2 split-K schedule refresh is Hopper-specific", @@ -3833,26 +3791,17 @@ def test_cuda_graph_split_kv_schedule_refresh(self) -> None: ) @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_decode_31b_like(self): """31B-like: mixed GQA (2 sliding, 8 full K=V), hd=256/512.""" self._run_cuda_graph_real_headdim(deepcopy(GEMMA4_31B_REAL_DIMS_CONFIG), "31B") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_decode_26b_like(self): """26B-like: GQA=2, K=V, hd=256/512.""" self._run_cuda_graph_real_headdim(deepcopy(GEMMA4_26B_REAL_DIMS_CONFIG), "26B") @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() - @unittest.mock.patch( - "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None - ) def test_cuda_graph_multi_step_trtllm_gen(self) -> None: """Multi-step CG decode with trtllm-gen (hd=256/512). diff --git a/tests/unittest/executor/test_stats_serializer.py b/tests/unittest/executor/test_stats_serializer.py index c6e4d9a127f5..5cb680a3de7a 100644 --- a/tests/unittest/executor/test_stats_serializer.py +++ b/tests/unittest/executor/test_stats_serializer.py @@ -121,36 +121,6 @@ def _make_mock_kv_iter_stats( return {window_size: s} -class _FakeStorageStatistics(SimpleNamespace): - @property - def unavailable(self): - return self.total - self.available - - -class _FakePeakStorage: - num_pool_groups = 2 - num_cache_levels = 2 - - def __init__(self): - self._levels = [] - self.primary_stats = [ - _FakeStorageStatistics(total=10, available=8, evictable=1), - _FakeStorageStatistics(total=10, available=9, evictable=0), - ] - self.secondary_stats = [ - _FakeStorageStatistics(total=5, available=4, evictable=1), - _FakeStorageStatistics(total=5, available=5, evictable=0), - ] - - def get_statistics(self, level): - if int(level) == 0: - return self.primary_stats - return self.secondary_stats - - def destroy(self): - pass - - class TestStatsSerializer: def test_serializer_without_kv_iter_stats(self): """Legacy 2-tuple and 3-tuple with None should produce same output.""" @@ -532,100 +502,3 @@ def test_serializer_emits_v2_suspend_resume_counters(self) -> None: pool_group = d["kvCacheIterationStatsByPoolGroup"]["0"] assert "iterSuspendedRequests" not in pool_group assert "iterResumedRequests" not in pool_group - - def test_v2_peak_block_stats_reset_tracks_interval_peak(self): - """Peak block stats should cover the interval since the previous reset.""" - from tensorrt_llm.runtime.kv_cache_manager_v2._common import GPU_LEVEL, CacheLevel - from tensorrt_llm.runtime.kv_cache_manager_v2._core._kv_cache_manager import KVCacheManager - - storage = _FakePeakStorage() - manager = object.__new__(KVCacheManager) - manager._living_kv_caches = set() - manager._storage = storage - manager._radix_tree = SimpleNamespace(clear=lambda: []) - manager._reset_iteration_peak_num_blocks() - - # Some gauges rise above the reset baseline, then fall before drain. - storage.primary_stats[0].available = 5 # primary used = 5 - storage.primary_stats[0].evictable = 3 - storage.primary_stats[1].available = 6 # primary used = 4 - storage.primary_stats[1].evictable = 4 - storage.secondary_stats[0].available = 2 # secondary used = 3 - storage.secondary_stats[0].evictable = 2 - manager._update_iteration_peak_num_blocks() - storage.primary_stats[0].available = 7 # primary used = 3 - storage.primary_stats[0].evictable = 1 - storage.secondary_stats[0].available = 4 # secondary used = 1 - storage.secondary_stats[0].evictable = 1 - - primary_peak = manager.get_and_reset_iteration_peak_block_stats(GPU_LEVEL) - secondary_peak = manager.get_and_reset_iteration_peak_block_stats(CacheLevel(1)) - assert [stats.available for stats in primary_peak] == [8, 9] - assert [stats.unavailable for stats in primary_peak] == [5, 4] - assert [stats.evictable for stats in primary_peak] == [3, 4] - assert [stats.available for stats in secondary_peak] == [4, 5] - assert [stats.unavailable for stats in secondary_peak] == [3, 0] - assert [stats.evictable for stats in secondary_peak] == [2, 0] - - # The next interval starts from current usage, not zero. - primary_peak = manager.get_and_reset_iteration_peak_block_stats(GPU_LEVEL) - secondary_peak = manager.get_and_reset_iteration_peak_block_stats(CacheLevel(1)) - assert [stats.available for stats in primary_peak] == [7, 6] - assert [stats.unavailable for stats in primary_peak] == [3, 4] - assert [stats.evictable for stats in primary_peak] == [1, 4] - assert [stats.available for stats in secondary_peak] == [4, 5] - assert [stats.unavailable for stats in secondary_peak] == [1, 0] - assert [stats.evictable for stats in secondary_peak] == [1, 0] - - def test_v2_peak_block_stats_by_level_matches_per_level_drain(self): - """One by-level drain reports and resets exactly what the per-level drains would.""" - from tensorrt_llm.runtime.kv_cache_manager_v2._common import CacheLevel - from tensorrt_llm.runtime.kv_cache_manager_v2._core._kv_cache_manager import KVCacheManager - - def make_manager(): - storage = _FakePeakStorage() - manager = object.__new__(KVCacheManager) - manager._living_kv_caches = set() - manager._storage = storage - manager._radix_tree = SimpleNamespace(clear=lambda: []) - manager._reset_iteration_peak_num_blocks() - # Same gauge movement as the per-level test above: rise, then fall before drain. - storage.primary_stats[0].available = 5 - storage.primary_stats[0].evictable = 3 - storage.primary_stats[1].available = 6 - storage.primary_stats[1].evictable = 4 - storage.secondary_stats[0].available = 2 - storage.secondary_stats[0].evictable = 2 - manager._update_iteration_peak_num_blocks() - storage.primary_stats[0].available = 7 - storage.primary_stats[0].evictable = 1 - storage.secondary_stats[0].available = 4 - storage.secondary_stats[0].evictable = 1 - return manager - - def as_tuples(stats_by_pool_group): - return [ - (stats.available, stats.unavailable, stats.evictable) - for stats in stats_by_pool_group - ] - - per_level = make_manager() - expected = [ - as_tuples(per_level.get_and_reset_iteration_peak_block_stats(CacheLevel(level))) - for level in range(_FakePeakStorage.num_cache_levels) - ] - expected_after_reset = [ - as_tuples(per_level.get_and_reset_iteration_peak_block_stats(CacheLevel(level))) - for level in range(_FakePeakStorage.num_cache_levels) - ] - - by_level = make_manager() - assert [ - as_tuples(stats) - for stats in by_level.get_and_reset_iteration_peak_block_stats_by_level() - ] == expected - # Draining every level at once must also reset every level. - assert [ - as_tuples(stats) - for stats in by_level.get_and_reset_iteration_peak_block_stats_by_level() - ] == expected_after_reset diff --git a/tests/unittest/kv_cache_manager_v2_tests/cuda_test_utils.py b/tests/unittest/kv_cache_manager_v2_tests/cuda_test_utils.py new file mode 100644 index 000000000000..2c8286a9ba94 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/cuda_test_utils.py @@ -0,0 +1,472 @@ +# 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. + +"""CUDA driver and arithmetic helpers used by the KVCacheManagerV2 test-suite. + +These live in the test directory rather than in the ``kv_cache_manager_v2`` package: they +are test-side scaffolding (raw-driver stream/event pooling, small index helpers, a +``sys.path`` context manager for sibling-module imports) with no role in the shipped +Python surface, which only re-exports the C++ implementation. +""" + +import functools +import os +import sys +import warnings +from abc import ABC, abstractmethod +from collections import deque +from collections.abc import Set +from contextlib import contextmanager +from typing import ( + Any, + Callable, + ClassVar, + Final, + Generic, + Iterable, + Iterator, + MutableSequence, + NewType, + Reversible, + Sequence, + TypeVar, + cast, +) + +import cuda.bindings.driver as drv +import cuda.bindings.runtime as cudart + +T = TypeVar("T") +U = TypeVar("U") + +NDEBUG: Final[bool] = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] != "1" + +CudaStream = NewType("CudaStream", int) + + +class OutOfMemoryError(Exception): + pass + + +class CuOOMError(OutOfMemoryError): + pass + + +class CuError(Exception): + error_code: drv.CUresult + + def __init__(self, error_code: drv.CUresult) -> None: + self.error_code = error_code + err, err_str = drv.cuGetErrorString(error_code) + if err != drv.CUresult.CUDA_SUCCESS: + err_str = "" + super().__init__(f"CUDA driver error: {error_code} ({err_str})") + + def __reduce__(self) -> tuple[type["CuError"], tuple[drv.CUresult]]: + return (self.__class__, (self.error_code,)) + + +def _unwrap( + ret: drv.CUresult + | tuple[ + drv.CUresult, + T, + ] + | tuple[drv.CUresult, T, U], +): + if isinstance(ret, drv.CUresult): + if int(ret) != int(drv.CUresult.CUDA_SUCCESS): # pyright: ignore + if int(ret) == int(drv.CUresult.CUDA_ERROR_OUT_OF_MEMORY): # pyright: ignore + raise CuOOMError() + raise CuError(ret) + else: + _unwrap(ret[0]) + return ret[1] if len(ret) == 2 else ret[1:] + + +def div_up(x: int, y: int) -> int: + return (x + y - 1) // y + + +def round_up(x: int, y: int) -> int: + return div_up(x, y) * y + + +def exact_div(x: int, y: int) -> int: + assert x % y == 0 + return x // y + + +Idx = TypeVar("Idx", bound=int) +Index = TypeVar("Index", bound=int, contravariant=True) + + +class HalfOpenRange(tuple[Idx, Idx], Generic[Idx]): + """A half-open range [beg, end), falsy when empty (beg >= end). + + Generic over index type. Supports unpacking into (beg, end). + """ + + __slots__ = () + + def __new__(cls, beg: Idx, end: Idx) -> "HalfOpenRange[Idx]": + return tuple.__new__(cls, (beg, end)) + + @property + def beg(self) -> Idx: + return self[0] + + @property + def end(self) -> Idx: + return self[1] + + def __eq__(self, other: object) -> bool: + if not isinstance(other, HalfOpenRange): + return NotImplemented + return (not self and not other) or tuple.__eq__(self, other) + + def __hash__(self) -> int: + return hash((0, 0)) if not self else tuple.__hash__(self) + + def __bool__(self) -> bool: + return self[0] < self[1] + + def __len__(self) -> int: + return max(0, self[1] - self[0]) + + def __contains__(self, item: Any) -> bool: + return self[0] <= item < self[1] + + +def intersect(a: HalfOpenRange[Idx], b: HalfOpenRange[Idx]) -> HalfOpenRange[Idx]: + """Return the intersection of two half-open ranges [beg, end). + + The result may be empty (beg >= end), which is safe to chain into further intersections. + """ + return HalfOpenRange(max(a[0], b[0]), min(a[1], b[1])) + + +def value_or(opt: T | None, default: T) -> T: + return default if opt is None else opt + + +def unwrap_optional(value: T | None) -> T: + if value is not None: + return value + raise ValueError("Expected non-None value") + + +def remove_if(original: MutableSequence[T], predicate: Callable[[T], bool]) -> list[T]: + """Remove items from original that satisfy the predicate and return the removed items.""" + removed = [] + for idx, item in enumerate(original): + if predicate(item): + removed.append(item) + else: + original[idx - len(removed)] = item + del original[len(original) - len(removed) :] + return removed + + +def get_uniform_attribute(iterable: Iterable[T], attribute_func: Callable[[T], U]) -> U: + ret = attribute_func(next(iter(iterable))) + assert NDEBUG or all(attribute_func(item) == ret for item in iterable) + return ret + + +def typed_range(*args: Index) -> Reversible[Index]: + return cast(Reversible[Index], range(*args)) + + +@functools.cache +def init_cuda_once() -> None: + (err,) = cudart.cudaFree(0) + assert int(err) == int(cudart.cudaError_t.cudaSuccess) + + +class SimplePool(Generic[T]): + __slots__ = ( + "_create_func", + "_destroy_func", + "_init_size", + "_max_size", + "_outstanding_count", + "_items", + ) + _create_func: Callable[[], T] + _destroy_func: Callable[[T], None] + _init_size: int + _max_size: int | None + _items: deque[T] | None + _outstanding_count: ( + int # number of items currently we gave out but not returned, i.e. get() but not put() + ) + + def __init__( + self, + create_func: Callable[[], T], + destroy_func: Callable[[T], None], + init_size: int = 0, + max_size: int | None = None, + ): + self._create_func = create_func + self._destroy_func = destroy_func + self._init_size = init_size + self._max_size = max_size + self._items = None + self._outstanding_count = 0 + + def clear(self) -> None: + while self.items: + self._destroy_func(self.items.popleft()) + + def __del__(self) -> None: + self.clear() + + @property + def items(self) -> deque[T]: + if self._items is None: + self._items = deque[T]( + (self._create_func() for _ in range(self._init_size)), maxlen=self._max_size + ) + return self._items + + def get(self) -> T: + ret = self.items.popleft() if self.items else self._create_func() + self._outstanding_count += 1 + return ret + + def put(self, item: T) -> None: + self._outstanding_count -= 1 + if self._max_size is not None and len(self.items) >= self._max_size: + self._destroy_func(item) + else: + self.items.append(item) + + @property + def outstanding_count(self) -> int: + """Number of items acquired with get() and not yet returned with put().""" + return self._outstanding_count + + @property + def cached_count(self) -> int: + """Number of items currently in the pool.""" + return len(self.items) + + @property + def total_count(self) -> int: + """Total number of items created, both outstanding and cached.""" + return self.outstanding_count + self.cached_count + + +class ItemHolderBase(Generic[T], ABC): + __slots__ = ("_item",) + _item: T | None + + def __init__(self) -> None: + self._item = self.pool.get() + + def close(self) -> None: + # Manually inlined for better performance. + item = self._item + if item is not None: + self.pool.put(item) + self._item = None + + def __del__(self) -> None: + self.close() + + def is_closed(self) -> bool: + return self._item is None + + def get(self) -> T: + # Manually inlined for better performance. + item = self._item + assert item is not None + return item + + @property + def handle(self) -> T: + # Manually inlined for better performance. + item = self._item + assert item is not None + return item + + @property + @abstractmethod + def pool(self) -> SimplePool[T]: ... + + +class CachedCudaEvent(ItemHolderBase[drv.CUevent]): + """A cached CUDA event without support for timing. Recorded to a stream when created.""" + + __slots__ = () + _pool: ClassVar[SimplePool[drv.CUevent] | None] = None + NULL: ClassVar["_NullCudaEvent"] + + def __init__(self, stream: CudaStream) -> None: + super().__init__() + self._record(stream) + + def query_complete(self) -> bool: + """Query the event. If complete, also close the event. Closed events are always considered complete.""" + # Manually inlined for better performance. + ev = self._item + if ev is None: + return True + (err,) = drv.cuEventQuery(ev) + if int(err) == int(drv.CUresult.CUDA_SUCCESS): + self.close() + return True + elif int(err) == int(drv.CUresult.CUDA_ERROR_NOT_READY): + return False + else: + raise CuError(err) + + def synchronize(self) -> None: + # Manually inlined for better performance. + ev = self._item + if ev is None: + return + _unwrap(drv.cuEventSynchronize(ev)) + self.close() + + def wait_in_stream(self, stream: CudaStream) -> None: + # Manually inlined for better performance. + ev = self._item + if ev is None: + return + _unwrap(drv.cuStreamWaitEvent(stream, ev, 0)) + + def _record(self, stream: CudaStream) -> None: + """Prefer new event instead of recording an existing event.""" + # Manually inlined for better performance. + ev = self._item + assert ev is not None + _unwrap(drv.cuEventRecord(ev, stream)) + + @property + def pool(self) -> SimplePool[drv.CUevent]: + if CachedCudaEvent._pool is None: + CachedCudaEvent._pool = SimplePool[drv.CUevent]( + lambda: _unwrap(drv.cuEventCreate(drv.CUevent_flags.CU_EVENT_DISABLE_TIMING)), + lambda ev: _unwrap(drv.cuEventDestroy(ev)), # pyright: ignore + init_size=1024, + ) + return CachedCudaEvent._pool + + +class _NullCudaEvent(CachedCudaEvent): + """A null CUDA event that is closed (and always complete).""" + + __slots__ = () + + def __init__(self) -> None: + # do not call super().__init__(). We don't need an event here. + self._item = None + + +CachedCudaEvent.NULL = _NullCudaEvent() + + +def stream_wait_events(stream: CudaStream, events: Iterable[CachedCudaEvent]) -> None: + """Batched wait for multiple events with deduplication first.""" + if not isinstance(events, Set): + events = set(events) + for ev in events: + ev.wait_in_stream(stream) + + +class CachedCudaStream(ItemHolderBase[CudaStream]): + """A cached non-blocking CUDA stream.""" + + __slots__ = () + _pool: ClassVar[SimplePool[CudaStream] | None] = None + + def __init__(self) -> None: + super().__init__() + + def wait_event(self, event: drv.CUevent) -> None: + _unwrap(drv.cuStreamWaitEvent(self.get(), event, drv.CU_STREAM_WAIT_VALUE_COMPLETED)) + + def wait_events(self, events: Sequence[CachedCudaEvent] | set[CachedCudaEvent]) -> None: + """Wait for events with deduplication first.""" + stream_wait_events(self.get(), events) + + def record_event(self) -> CachedCudaEvent: + return CachedCudaEvent(self.get()) + + def __cuda_stream__(self) -> tuple[int, int]: + return 0, int(self.get()) + + def synchronize(self) -> None: + _unwrap(drv.cuStreamSynchronize(self.handle)) + + @property + def pool(self) -> SimplePool[CudaStream]: + if CachedCudaStream._pool is None: + CachedCudaStream._pool = SimplePool[CudaStream]( + lambda: CudaStream( + int(_unwrap(drv.cuStreamCreate(drv.CUstream_flags.CU_STREAM_NON_BLOCKING))) # pyright: ignore + ), + lambda stream: _unwrap(drv.cuStreamDestroy(stream)), # pyright: ignore + init_size=128, + ) + return CachedCudaStream._pool + + +class TemporaryCudaStream(CachedCudaStream): + """A cached non-blocking CUDA stream, used as a temporary worker stream. + + Requires a list of prior events to wait for dependencies. A finish event is recorded when exiting + normally. Call take_finish_event() to consume the finish event, otherwise you get a warning. + """ + + __slots__ = "_finish_event" + _finish_event: CachedCudaEvent | None + + def __init__(self, prior_events: Sequence[CachedCudaEvent] | set[CachedCudaEvent]): + super().__init__() + self.wait_events(prior_events) + self._finish_event = None + + def __del__(self) -> None: + if self._finish_event is not None: + warnings.warn("[KVCacheManager] finish event recorded but not taken") + super().__del__() + + def take_finish_event(self) -> CachedCudaEvent: + ret = unwrap_optional(self._finish_event) + self._finish_event = None + return ret + + def __enter__(self) -> "TemporaryCudaStream": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + if not exc_type: + self._finish_event = self.record_event() + + +@contextmanager +def temporary_sys_path(path: str) -> Iterator[None]: + already_in_path = path in sys.path + if not already_in_path: + sys.path.insert(0, path) + try: + yield + finally: + if not already_in_path: + sys.path.remove(path) diff --git a/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py b/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py index 70a4e48d6c85..08c6e1aa04eb 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py +++ b/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py @@ -13,6 +13,8 @@ # limitations under the License. import itertools +import os +import sys from collections.abc import Sequence from functools import cached_property from importlib.util import find_spec @@ -20,46 +22,45 @@ if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: from kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, + NDEBUG, AttentionLayerConfig, BeamIndex, CudaStream, DataRole, KVCacheManagerConfig, LayerId, + MemAddress, + PageIndexMode, SsmLayerConfig, TokenIdExt, _KVCache, ) - from kv_cache_manager_v2._common import BAD_PAGE_INDEX, NDEBUG, MemAddress, PageIndexMode - from kv_cache_manager_v2._utils import ( - HalfOpenRange, - div_up, - exact_div, - get_uniform_attribute, - intersect, - temporary_sys_path, - typed_range, - value_or, - ) else: from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, + NDEBUG, AttentionLayerConfig, BeamIndex, CudaStream, DataRole, KVCacheManagerConfig, LayerId, + MemAddress, + PageIndexMode, SsmLayerConfig, TokenIdExt, _KVCache, ) - from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( - BAD_PAGE_INDEX, - NDEBUG, - MemAddress, - PageIndexMode, - ) - from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +# cuda_test_utils supplies temporary_sys_path, so its own path entry is added and +# removed by hand here; every later sibling import goes through that helper. +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import ( # noqa: E402 HalfOpenRange, div_up, exact_div, @@ -69,10 +70,11 @@ typed_range, value_or, ) +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) -import os - -with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): +with temporary_sys_path(_TEST_DIR): from kernels import check_values, fill_values diff --git a/tests/unittest/kv_cache_manager_v2_tests/kernels.py b/tests/unittest/kv_cache_manager_v2_tests/kernels.py index d5955026a52f..0f7e3a9bb01e 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/kernels.py +++ b/tests/unittest/kv_cache_manager_v2_tests/kernels.py @@ -14,6 +14,8 @@ import contextlib import ctypes +import os +import sys from collections.abc import Sequence from functools import lru_cache from importlib.util import find_spec @@ -28,16 +30,21 @@ from cuda.core.experimental._module import ObjectCode if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: - from kv_cache_manager_v2._common import CudaStream, LayerId, MemAddress, TokenIdExt - from kv_cache_manager_v2._utils import _unwrap, div_up, exact_div + from kv_cache_manager_v2 import CudaStream, LayerId, MemAddress, TokenIdExt else: - from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( - CudaStream, - LayerId, - MemAddress, - TokenIdExt, - ) - from tensorrt_llm.runtime.kv_cache_manager_v2._utils import _unwrap, div_up, exact_div + from tensorrt_llm.runtime.kv_cache_manager_v2 import CudaStream, LayerId, MemAddress, TokenIdExt + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +# cuda_test_utils supplies temporary_sys_path, so its own path entry is added and +# removed by hand here; every later sibling import goes through that helper. +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import _unwrap, div_up, exact_div # noqa: E402 +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) _SLEEP_TIME_NS: int = 0 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py b/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py index 489896279247..fa33e88bdb37 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py @@ -16,6 +16,7 @@ import gc import itertools import os +import sys import unittest from importlib.util import find_spec from typing import TYPE_CHECKING, cast @@ -27,17 +28,10 @@ DataRole, KVCacheManager, LayerId, + MemAddress, TokenId, TokenIdExt, ) - from kv_cache_manager_v2._common import MemAddress - from kv_cache_manager_v2._utils import ( - TemporaryCudaStream, - exact_div, - init_cuda_once, - round_up, - temporary_sys_path, - ) else: from tensorrt_llm.runtime.kv_cache_manager_v2 import ( DEFAULT_BEAM_INDEX, @@ -45,19 +39,30 @@ DataRole, KVCacheManager, LayerId, + MemAddress, TokenId, TokenIdExt, ) - from tensorrt_llm.runtime.kv_cache_manager_v2._common import MemAddress - from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +# cuda_test_utils supplies temporary_sys_path, so its own path entry is added and +# removed by hand here; every later sibling import goes through that helper. +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import ( # noqa: E402 TemporaryCudaStream, exact_div, init_cuda_once, round_up, temporary_sys_path, ) +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) -with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): +with temporary_sys_path(_TEST_DIR): from kernels import check_values, fill_values from test_kv_cache_manager_v2 import create_config diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_first_new_block_probe.py b/tests/unittest/kv_cache_manager_v2_tests/test_first_new_block_probe.py index bee8d0179e2a..0352c30199e3 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_first_new_block_probe.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_first_new_block_probe.py @@ -30,6 +30,7 @@ import gc import os +import sys import unittest from collections.abc import Sequence from dataclasses import replace @@ -46,24 +47,33 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import sequence_to_blockchain_keys if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: - from kv_cache_manager_v2 import CudaStream, KVCacheManager, TokenId, _introspection - from kv_cache_manager_v2._block_radix_tree import ReuseScope - from kv_cache_manager_v2._utils import TemporaryCudaStream, init_cuda_once, temporary_sys_path + from kv_cache_manager_v2 import CudaStream, KVCacheManager, ReuseScope, TokenId, _introspection else: from tensorrt_llm.runtime.kv_cache_manager_v2 import ( CudaStream, KVCacheManager, + ReuseScope, TokenId, _introspection, ) - from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ReuseScope - from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +# cuda_test_utils supplies temporary_sys_path, so its own path entry is added and +# removed by hand here; every later sibling import goes through that helper. +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import ( # noqa: E402 TemporaryCudaStream, init_cuda_once, temporary_sys_path, ) +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) -with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): +with temporary_sys_path(_TEST_DIR): from test_kv_cache_manager_v2 import create_config TOKENS_PER_BLOCK = 4 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py index 30239eb3021e..1f1e47650589 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py @@ -28,7 +28,6 @@ import faulthandler import itertools -import os import threading import time @@ -45,17 +44,7 @@ KVCacheManagerConfig, ) -KV_CACHE_MANAGER_V2_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() - -pytestmark = [ - pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"), - # The pure-Python backend is not thread-safe: it relies on the GIL and has no API lock, so - # these tests would race against it rather than exercise the property they assert. - pytest.mark.skipif( - KV_CACHE_MANAGER_V2_BACKEND != "cpp", - reason="GIL-free API locking exists only in the C++ backend", - ), -] +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") # Generous relative to the work each test does (milliseconds); tight enough that a deadlock fails # the run in reasonable time rather than hanging it. diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index 16854c737b9e..f427394b56f2 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -16,6 +16,7 @@ import gc import os import pickle +import sys import threading import time from importlib.util import find_spec @@ -29,20 +30,7 @@ KV_CACHE_HASH_ALGO_V2_SHA256_64, truncate_sha256_hash_to_int64, ) -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheCreatedData as NativeKVCacheCreatedData -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheEvent as NativeKVCacheEvent -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheEventDiff as NativeKVCacheEventDiff from tensorrt_llm.runtime.kv_cache_manager_v2 import ( - KVCacheEventManager as NativeKVCacheEventManager, -) -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheRemovedData as NativeKVCacheRemovedData -from tensorrt_llm.runtime.kv_cache_manager_v2 import ( - KVCacheStoredBlockData as NativeKVCacheStoredBlockData, -) -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheStoredData as NativeKVCacheStoredData -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheUpdatedData as NativeKVCacheUpdatedData -from tensorrt_llm.runtime.kv_cache_manager_v2 import UniqueToken as NativeUniqueToken -from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import ( KVCacheCreatedData, KVCacheEvent, KVCacheEventDiff, @@ -55,48 +43,45 @@ ) if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: - from kv_cache_manager_v2 import CacheLevel, CudaStream, KVCacheManager, TokenId, _introspection - from kv_cache_manager_v2._block_radix_tree import ( - Block, - BlockRadixTree, + from kv_cache_manager_v2 import ( + CacheLevel, + CudaStream, + KVCacheManager, ReuseScope, - RootBlock, - detach_next, + TokenId, + _introspection, + sequence_to_blockchain_keys, ) - from kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry - from kv_cache_manager_v2._utils import CachedCudaStream, init_cuda_once, temporary_sys_path else: from tensorrt_llm.runtime.kv_cache_manager_v2 import ( CacheLevel, CudaStream, KVCacheManager, + ReuseScope, TokenId, _introspection, - ) - from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( - Block, - BlockRadixTree, - ReuseScope, - RootBlock, - detach_next, - ) - from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry - from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( - CachedCudaStream, - init_cuda_once, - temporary_sys_path, + sequence_to_blockchain_keys, ) +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +# cuda_test_utils supplies temporary_sys_path, so its own path entry is added and +# removed by hand here; every later sibling import goes through that helper. +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import CachedCudaStream, init_cuda_once, temporary_sys_path # noqa: E402 +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) + try: import torch except ImportError: torch = None -_USING_CPP_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() != "python" - - -with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): +with temporary_sys_path(_TEST_DIR): from test_kv_cache_manager_v2 import create_config @@ -164,6 +149,13 @@ def _block_key(block): return _introspection.test_block_key(block) +def _blockchain_keys(tokens_per_block, tokens, reuse_scope=None): + """Block keys for a token sequence, without the reuse-scope root the chain starts from.""" + keys = sequence_to_blockchain_keys(tokens_per_block, reuse_scope or ReuseScope(), tokens) + next(keys, None) + return [key for _, key in keys] + + def _add_stored_block(event_manager, block): _introspection.event_manager_add_stored_block(event_manager, block) @@ -181,13 +173,6 @@ def _flush_serialized_events(event_manager): return KVCacheEventSerializer.serialize(event_manager.get_latest_events(0)) -def _flush_contract_events(event_manager): - events = _flush_serialized_events(event_manager) - for event in events: - event["hash_algo"] = "" - return events - - def _stored_events(events): return [event for event in events if event["data"]["type"] == "stored"] @@ -219,9 +204,8 @@ def _commit_and_close(manager, stream, tokens, *, input_tokens=None, reuse_scope gc.collect() -@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="requires the native C++ event manager") -def test_native_event_manager_queue_and_stored_coalescing(): - event_manager = NativeKVCacheEventManager( +def test_event_manager_queue_and_stored_coalescing(): + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, hash_algo="v2_sha256", @@ -229,9 +213,9 @@ def test_native_event_manager_queue_and_stored_coalescing(): event_manager.add_stored_event( None, [ - NativeKVCacheStoredBlockData( + KVCacheStoredBlockData( "block0", - [NativeUniqueToken(1), NativeUniqueToken(2)], + [UniqueToken(1), UniqueToken(2)], cache_level=0, priority=35, ) @@ -240,9 +224,9 @@ def test_native_event_manager_queue_and_stored_coalescing(): event_manager.add_stored_event( "block0", [ - NativeKVCacheStoredBlockData( + KVCacheStoredBlockData( "block1", - [NativeUniqueToken(3), NativeUniqueToken(4)], + [UniqueToken(3), UniqueToken(4)], cache_level=0, priority=35, ) @@ -261,27 +245,14 @@ def test_native_event_manager_queue_and_stored_coalescing(): ] -def test_native_event_manager_v1_hash_matches_legacy_cpp_hasher(): - _tb = pytest.importorskip("tensorrt_llm.bindings") - block_key = _tb.internal.batch_manager.BlockKey - block_key_hasher = _tb.internal.batch_manager.BlockKeyHasher - parent_hash = block_key_hasher.hash(block_key([1, 2, 3, 4])) - - assert NativeKVCacheEventManager._hash_block_key([1, 2, 3, 4], 0, None, None) == parent_hash - assert NativeKVCacheEventManager._hash_block_key( - [5, 6], parent_hash, None, None - ) == block_key_hasher.hash(block_key([5, 6]), parent_hash) - - -@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="requires the native C++ event manager") -def test_native_event_manager_attention_dp_gather_callback(): +def test_event_manager_attention_dp_gather_callback(): gathered_events = [] def gather(local_events): gathered_events.extend(local_events) return [pickle.loads(pickle.dumps(local_events))] - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=2, window_size=128, attention_dp_rank=0, @@ -296,172 +267,98 @@ def gather(local_events): assert events[0]["data"]["num_blocks_per_cache_level"] == [4] -@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") -def test_native_event_data_value_semantics_match_python_reference(): - native_token = NativeUniqueToken("token", 7) - native_block = NativeKVCacheStoredBlockData( - "block", - [native_token], - cache_level=1, - priority=35, - mm_keys=[(b"short-mm-key", 3), (b"another-key", 5, "uuid")], - cache_salt="salt", - ) - native_diff = NativeKVCacheEventDiff(0, 1) - native_objects = [ - native_token, - NativeKVCacheCreatedData([2, 3]), - native_block, - NativeKVCacheStoredData(None, [native_block]), - NativeKVCacheRemovedData(["block"]), - native_diff, - NativeKVCacheUpdatedData("block", native_diff, None), - ] - native_event = NativeKVCacheEvent( - 4, - native_objects[3], - 128, - "same-hash-label", - 1, - 2, - ) - native_objects.append(native_event) - - python_token = UniqueToken("token", 7) - python_block = KVCacheStoredBlockData( +def test_event_data_pickle_round_trip_preserves_value_semantics(): + token = UniqueToken("token", 7) + block = KVCacheStoredBlockData( "block", - [python_token], + [token], cache_level=1, priority=35, mm_keys=[(b"short-mm-key", 3), (b"another-key", 5, "uuid")], cache_salt="salt", ) - python_diff = KVCacheEventDiff(0, 1) - python_objects = [ - python_token, + diff = KVCacheEventDiff(0, 1) + stored = KVCacheStoredData(None, [block]) + values = [ + token, KVCacheCreatedData([2, 3]), - python_block, - KVCacheStoredData(None, [python_block]), + block, + stored, KVCacheRemovedData(["block"]), - python_diff, - KVCacheUpdatedData("block", python_diff, None), + diff, + KVCacheUpdatedData("block", diff, None), + KVCacheEvent(4, stored, 128, "same-hash-label", 1, 2), ] - python_objects.append( - KVCacheEvent( - 4, - python_objects[3], - 128, - "same-hash-label", - 1, - 2, - ) - ) - for value, python_value in zip(native_objects, python_objects): + for value in values: assert pickle.loads(pickle.dumps(value)) == value - assert repr(value) == repr(python_value) - - python_event = KVCacheEventManager(max_kv_event_entries=1) - native_manager = NativeKVCacheEventManager(max_kv_event_entries=1) - python_event.add_stored_event(None, [python_block], layer_group_id=2) - native_manager.add_stored_event(None, [native_block], layer_group_id=2) - assert _flush_contract_events(native_manager) == _flush_contract_events(python_event) -@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") -def test_native_event_manager_public_methods_match_python_reference(): - python_manager = KVCacheEventManager( +def test_event_manager_public_methods_drain_queue(): + manager = KVCacheEventManager( max_kv_event_entries=16, window_size=128, window_size_by_layer_group=None, ) - native_manager = NativeKVCacheEventManager( - max_kv_event_entries=16, - window_size=128, - window_size_by_layer_group=None, + manager.set_layer_group_window_sizes({0: 64, 1: 96}) + + manager.add_created_event([2, 3], layer_group_ids=[0, 1]) + manager.add_stored_event( + None, + [KVCacheStoredBlockData("block-0", [UniqueToken(1)], cache_level=0, priority=35)], + layer_group_id=0, ) - python_manager.set_layer_group_window_sizes({0: 64, 1: 96}) - native_manager.set_layer_group_window_sizes({0: 64, 1: 96}) - - for manager, block_type, token_type, diff_type in ( - (python_manager, KVCacheStoredBlockData, UniqueToken, KVCacheEventDiff), - ( - native_manager, - NativeKVCacheStoredBlockData, - NativeUniqueToken, - NativeKVCacheEventDiff, - ), - ): - manager.add_created_event([2, 3], layer_group_ids=[0, 1]) - manager.add_stored_event( - None, - [block_type("block-0", [token_type(1)], cache_level=0, priority=35)], - layer_group_id=0, - ) - manager.add_stored_event( - "block-0", - [block_type("block-1", [token_type(2)], cache_level=0, priority=35)], - layer_group_id=0, - ) - manager.add_removed_event(block_hash for block_hash in ("removed-0", "removed-1")) - manager.add_updated_event( - "updated", - cache_level=diff_type(0, 1), - priority=diff_type(35, 50), - layer_group_id=1, - ) + manager.add_stored_event( + "block-0", + [KVCacheStoredBlockData("block-1", [UniqueToken(2)], cache_level=0, priority=35)], + layer_group_id=0, + ) + manager.add_removed_event(block_hash for block_hash in ("removed-0", "removed-1")) + manager.add_updated_event( + "updated", + cache_level=KVCacheEventDiff(0, 1), + priority=KVCacheEventDiff(35, 50), + layer_group_id=1, + ) + + assert _flush_serialized_events(manager) + assert manager.get_latest_events(0) == [] + assert manager.get_latest_events(-1) == [] + + +def test_event_manager_unknown_raw_keys_are_noops(): + manager = KVCacheEventManager(max_kv_event_entries=4) - assert _flush_contract_events(native_manager) == _flush_contract_events(python_manager) - assert native_manager.get_latest_events(0) == python_manager.get_latest_events(0) == [] - assert native_manager.get_latest_events(-1) == python_manager.get_latest_events(-1) == [] - - -@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") -def test_native_event_manager_unknown_raw_keys_match_python_noop_behavior(): - python_manager = KVCacheEventManager(max_kv_event_entries=4) - native_manager = NativeKVCacheEventManager(max_kv_event_entries=4) - - for manager, diff_type in ( - (python_manager, KVCacheEventDiff), - (native_manager, NativeKVCacheEventDiff), - ): - manager.add_removed_event([b"short", b"still-not-a-radix-key"]) - manager.add_removed_life_cycle_event(b"short", 0) - manager.add_updated_event(b"short") - manager.add_updated_event(b"short", cache_level=diff_type(0, 1)) - - assert _flush_contract_events(native_manager) == _flush_contract_events(python_manager) == [] - - -@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") -def test_native_event_manager_disabled_queue_and_blocking_read_match_python_reference(): - for manager_type in (KVCacheEventManager, NativeKVCacheEventManager): - disabled = manager_type(max_kv_event_entries=0) - disabled.add_created_event([1]) - disabled.flush_iteration_events() - assert disabled.get_latest_events(0) == [] - - manager = manager_type(max_kv_event_entries=1) - result = [] - reader = threading.Thread(target=lambda: result.extend(manager.get_latest_events())) - reader.start() - time.sleep(0.05) - assert reader.is_alive() - manager.add_created_event([1]) - manager.flush_iteration_events() - reader.join(timeout=1) - assert not reader.is_alive() - assert len(result) == 1 - - -@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") -def test_native_event_manager_constructor_errors_match_python_reference(): - for manager_type in (KVCacheEventManager, NativeKVCacheEventManager): - with pytest.raises(ValueError, match="Unsupported V2 KV cache event hash algorithm"): - manager_type(max_kv_event_entries=1, hash_algo="unsupported") + manager.add_removed_event([b"short", b"still-not-a-radix-key"]) + manager.add_removed_life_cycle_event(b"short", 0) + manager.add_updated_event(b"short") + manager.add_updated_event(b"short", cache_level=KVCacheEventDiff(0, 1)) + assert _flush_serialized_events(manager) == [] + + +def test_event_manager_disabled_queue_and_blocking_read(): + disabled = KVCacheEventManager(max_kv_event_entries=0) + disabled.add_created_event([1]) + disabled.flush_iteration_events() + assert disabled.get_latest_events(0) == [] + + manager = KVCacheEventManager(max_kv_event_entries=1) + result = [] + reader = threading.Thread(target=lambda: result.extend(manager.get_latest_events())) + reader.start() + time.sleep(0.05) + assert reader.is_alive() + manager.add_created_event([1]) + manager.flush_iteration_events() + reader.join(timeout=1) + assert not reader.is_alive() + assert len(result) == 1 + + +def test_updated_data_requires_a_cache_level_diff(): with pytest.raises(TypeError): - NativeKVCacheUpdatedData("block") + KVCacheUpdatedData("block") def test_v2_kv_cache_event_manager_serialization(): @@ -600,7 +497,7 @@ def test_v2_kv_cache_event_manager_accepts_removed_iterables(): def test_v2_kv_cache_event_manager_coalesces_contiguous_stored_events( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager, num_life_cycles=2, tokens_per_block=2) block0 = make_block([1, 2], [2, 2]) block1 = make_block([3, 4], [2, 2], parent=block0) @@ -632,7 +529,7 @@ def test_v2_kv_cache_event_manager_serializes_layer_group_id(): def test_v2_kv_cache_event_manager_sha256_64_compatibility_mode(real_block_factory): - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V2_SHA256_64, @@ -668,7 +565,7 @@ def test_v2_kv_cache_event_manager_sha256_64_compatibility_mode(real_block_facto def test_v2_kv_cache_event_manager_v1_hash_algo_matches_v1_block_key_hash( real_block_factory, ): - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -691,16 +588,20 @@ def test_v2_kv_cache_event_manager_v1_hash_algo_matches_v1_block_key_hash( def test_v2_root_key_distinguishes_lora_from_cache_salt_id(): - assert RootBlock.make_key(ReuseScope()) != RootBlock.make_key(ReuseScope(lora_id=123)) - assert RootBlock.make_key(ReuseScope()) != RootBlock.make_key(ReuseScope(salt=123)) - assert RootBlock.make_key(ReuseScope(lora_id=123)) != RootBlock.make_key(ReuseScope(salt=123)) - assert RootBlock.make_key(ReuseScope(lora_id=123, salt=456)) != RootBlock.make_key( + def root_key(scope): + # The first pair a blockchain yields is the reuse-scope root itself. + return next(iter(sequence_to_blockchain_keys(1, scope, [])))[1] + + assert root_key(ReuseScope()) != root_key(ReuseScope(lora_id=123)) + assert root_key(ReuseScope()) != root_key(ReuseScope(salt=123)) + assert root_key(ReuseScope(lora_id=123)) != root_key(ReuseScope(salt=123)) + assert root_key(ReuseScope(lora_id=123, salt=456)) != root_key( ReuseScope(lora_id=456, salt=123) ) def test_v2_kv_cache_event_manager_v1_hash_algo_mixes_cache_salt_id(real_block_factory): - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -722,7 +623,7 @@ def test_v2_kv_cache_event_manager_v1_hash_algo_mixes_cache_salt_id(real_block_f def test_v2_kv_cache_event_manager_v1_hash_reads_root_reuse_scope(real_block_factory): def hashes_for(reuse_scope): - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -742,7 +643,7 @@ def hashes_for(reuse_scope): def test_v2_kv_cache_event_manager_v1_hash_recomputes_removed_parent( real_block_factory, ): - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -762,12 +663,12 @@ def test_v2_kv_cache_event_manager_v1_hash_recomputes_removed_parent( def test_v2_kv_cache_event_manager_v1_hash_algo_matches_cpp_hasher(): _tb = pytest.importorskip("tensorrt_llm.bindings") - block_key = _tb.internal.batch_manager.BlockKey + block_key_type = _tb.internal.batch_manager.BlockKey block_key_hasher = _tb.internal.batch_manager.BlockKeyHasher - parent_hash = block_key_hasher.hash(block_key([1, 2, 3, 4])) - child_hash = block_key_hasher.hash(block_key([5, 6]), parent_hash) - lora_hash = block_key_hasher.hash(block_key([1, 2, 3, 4], 123)) + parent_hash = block_key_hasher.hash(block_key_type([1, 2, 3, 4])) + child_hash = block_key_hasher.hash(block_key_type([5, 6]), parent_hash) + lora_hash = block_key_hasher.hash(block_key_type([1, 2, 3, 4], 123)) assert KVCacheEventManager._hash_block_key([1, 2, 3, 4], 0, None, None) == parent_hash assert KVCacheEventManager._hash_block_key([5, 6], parent_hash, None, None) == child_hash @@ -776,9 +677,9 @@ def test_v2_kv_cache_event_manager_v1_hash_algo_matches_cpp_hasher(): def test_v2_kv_cache_event_manager_v1_hash_events_match_cpp_hasher(real_block_factory): _tb = pytest.importorskip("tensorrt_llm.bindings") - block_key = _tb.internal.batch_manager.BlockKey + block_key_type = _tb.internal.batch_manager.BlockKey block_key_hasher = _tb.internal.batch_manager.BlockKeyHasher - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -791,8 +692,8 @@ def test_v2_kv_cache_event_manager_v1_hash_events_match_cpp_hasher(real_block_fa _add_stored_block(event_manager, block1) events = _flush_serialized_events(event_manager) - parent_hash = block_key_hasher.hash(block_key([1, 2, 3, 4])) - child_hash = block_key_hasher.hash(block_key([5, 6]), parent_hash) + parent_hash = block_key_hasher.hash(block_key_type([1, 2, 3, 4])) + child_hash = block_key_hasher.hash(block_key_type([5, 6]), parent_hash) assert _stored_block_hashes(events) == [parent_hash, child_hash] @@ -856,7 +757,7 @@ def test_v1_and_v2_managers_emit_same_v1_hash_stored_events(): manager_v1.flush_iteration_events() v1_events = KVCacheEventSerializer.serialize(manager_v1.get_latest_events(10)) - event_manager_v2 = NativeKVCacheEventManager( + event_manager_v2 = KVCacheEventManager( max_kv_event_entries=event_buffer_max_size, window_size=max_seq_len, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -1026,7 +927,7 @@ def test_v2_kv_cache_event_manager_serializes_updated_event(): def test_v2_kv_cache_event_manager_uses_stored_registry_for_removed_event( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager) block = make_block([1, 2], [2]) block_hash = _block_key(block).hex() @@ -1048,7 +949,7 @@ def test_v2_kv_cache_event_manager_uses_stored_registry_for_removed_event( def test_v2_kv_cache_event_manager_derives_mm_keys_across_blocks( real_block_factory, ancestor_coverage ): - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, mm_token_id_offset=1000 ) make_block = real_block_factory(event_manager) @@ -1088,7 +989,7 @@ def test_v2_kv_cache_event_manager_derives_mm_keys_across_blocks( def test_v2_kv_cache_event_manager_preserves_mm_keys_after_life_cycle_removal(real_block_factory): - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=8, window_size=128, mm_token_id_offset=1000 ) make_block = real_block_factory(event_manager, num_life_cycles=2) @@ -1125,7 +1026,7 @@ def test_v2_kv_cache_event_manager_preserves_mm_keys_after_life_cycle_removal(re def test_v2_kv_cache_event_manager_requires_mm_token_encoding(real_block_factory): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager) mm_hash = bytes(range(32)) block = make_block([mm_hash, 1001], [2]) @@ -1137,44 +1038,10 @@ def test_v2_kv_cache_event_manager_requires_mm_token_encoding(real_block_factory assert [token["token_id"] for token in stored_block["tokens"]] == [mm_hash.hex(), 1001] -def test_python_v2_mm_digest_context_survives_detached_ancestor(): - event_manager = KVCacheEventManager(max_kv_event_entries=8, mm_token_id_offset=1000) - life_cycles = LifeCycleRegistry( - create_config( - tokens_per_block=4, - gpu_quota=16 << 20, - host_quota=0, - disk_quota=0, - num_layers=0, - window_size=None, - sink_tokens=0, - ) - ) - tree = BlockRadixTree(life_cycles, tokens_per_block=4, event_manager=event_manager) - root = tree.add_or_get_existing(ReuseScope()) - digest = bytes(range(32)) - first = Block([0, digest, 1001, 1002], root) - gap = Block([1, 2, 3, 4], first) - continued = Block([1003, 5, 1004, 1005], gap) - parent_ref = gap._prev - try: - assert detach_next(first, gap.key) is gap - assert gap.last_token_digest == digest - assert event_manager._mm_keys_from_radix_block(continued) == [(digest, 3), (digest, 4)] - - gap._prev = parent_ref - first.next[gap.key] = gap - assert event_manager._mm_keys_from_radix_block(continued) == [(digest, 3), (digest, 4)] - finally: - gap._prev = parent_ref - first.next[gap.key] = gap - tree.clear() - - def test_v2_kv_cache_event_manager_omits_partial_life_cycle_coverage( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager, num_life_cycles=2) block = make_block([1, 2], [2, 1]) @@ -1191,7 +1058,7 @@ def test_v2_kv_cache_event_manager_omits_partial_life_cycle_coverage( def test_v2_kv_cache_event_manager_emits_partial_life_cycle_removed_events( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager, num_life_cycles=2) block = make_block([1, 2], [2, 2]) block_hash = _block_key(block).hex() @@ -1222,7 +1089,7 @@ def test_v2_kv_cache_event_manager_emits_partial_life_cycle_removed_events( def test_v2_kv_cache_event_manager_whole_block_removal_clears_life_cycle_state( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager, num_life_cycles=2) block = make_block([1, 2], [2, 2]) block_hash = _block_key(block).hex() @@ -1247,19 +1114,19 @@ def test_v2_kv_cache_event_manager_whole_block_removal_clears_life_cycle_state( def test_v2_kv_cache_event_manager_readds_life_cycle_emits_stored_event( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager, num_life_cycles=2) block = make_block([1, 2], [2, 2]) - block_key = _block_key(block) - block_hash = block_key.hex() + key = _block_key(block) + block_hash = key.hex() _add_stored_block(event_manager, block) event_types = [event["data"]["type"] for event in _flush_serialized_events(event_manager)] assert event_types == ["stored", "stored"] - event_manager.add_removed_life_cycle_event(block_key, 0) + event_manager.add_removed_life_cycle_event(key, 0) _add_stored_life_cycle(event_manager, block, 0) - event_manager.add_removed_life_cycle_event(block_key, 1) + event_manager.add_removed_life_cycle_event(key, 1) events = _flush_serialized_events(event_manager) assert [event["data"]["type"] for event in events] == [ @@ -1272,7 +1139,7 @@ def test_v2_kv_cache_event_manager_readds_life_cycle_emits_stored_event( assert events[1]["data"]["blocks"][0]["block_hash"] == block_hash assert events[2]["data"]["block_hashes"] == [block_hash] - event_manager.add_removed_life_cycle_event(block_key, 0) + event_manager.add_removed_life_cycle_event(key, 0) events = _flush_serialized_events(event_manager) assert [event["data"]["type"] for event in events] == ["removed"] assert events[0]["layer_group_id"] == 0 @@ -1282,7 +1149,7 @@ def test_v2_kv_cache_event_manager_readds_life_cycle_emits_stored_event( def test_v2_kv_cache_event_manager_serializes_only_fully_covered_life_cycle_tokens( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager, num_life_cycles=2) block = make_block([1, 2], [2, 1]) @@ -1299,13 +1166,13 @@ def test_v2_kv_cache_event_manager_serializes_only_fully_covered_life_cycle_toke def test_v2_kv_cache_event_manager_reemits_stored_after_all_life_cycles_were_removed( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager) block = make_block([1, 2], [2]) - block_key = _block_key(block) + key = _block_key(block) _add_stored_block(event_manager, block) - event_manager.add_removed_life_cycle_event(block_key, 0) + event_manager.add_removed_life_cycle_event(key, 0) event_types = [event["data"]["type"] for event in _flush_serialized_events(event_manager)] assert event_types == ["stored", "removed"] @@ -1314,13 +1181,13 @@ def test_v2_kv_cache_event_manager_reemits_stored_after_all_life_cycles_were_rem assert [event["data"]["type"] for event in events] == ["stored"] assert events[0]["layer_group_id"] == 0 - assert events[0]["data"]["blocks"][0]["block_hash"] == block_key.hex() + assert events[0]["data"]["blocks"][0]["block_hash"] == key.hex() def test_v2_kv_cache_event_manager_flushes_removed_before_updated_event( real_block_factory, ): - event_manager = NativeKVCacheEventManager(max_kv_event_entries=8, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) make_block = real_block_factory(event_manager) removed_block = make_block([1, 2], [2]) updated_block = make_block([3, 4], [2]) @@ -1334,7 +1201,7 @@ def test_v2_kv_cache_event_manager_flushes_removed_before_updated_event( event_manager.add_removed_event(removed_key) event_manager.add_updated_event( updated_key, - cache_level=NativeKVCacheEventDiff(old_value=0, new_value=1), + cache_level=KVCacheEventDiff(old_value=0, new_value=1), layer_group_id=0, ) events = _flush_serialized_events(event_manager) @@ -1350,7 +1217,7 @@ def test_v2_stored_events_match_block_hash_chain(): gc.collect() gc.disable() - event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) manager = None try: tokens_per_block = 4 @@ -1365,13 +1232,8 @@ def test_v2_stored_events_match_block_hash_chain(): stored_events = _stored_events(events) assert len(stored_events) == 1 - # Both the Python and C++ backends hash blocks with SHA-256 over the same - # little-endian token encoding, so the block-key chain matches the - # pure-Python Block.make_key implementation on either backend. - root_key = RootBlock.make_key(ReuseScope()) - block0_key = Block.make_key(root_key, tokens[:tokens_per_block]) - block1_key = Block.make_key(block0_key, tokens[tokens_per_block:]) - expected_hashes = [block0_key.hex(), block1_key.hex()] + chain = _blockchain_keys(tokens_per_block, tokens) + expected_hashes = [key.hex() for key in chain] assert _stored_block_hashes(stored_events) == expected_hashes assert stored_events[0]["hash_algo"] == "v2_sha256" @@ -1386,14 +1248,13 @@ def test_v2_stored_events_match_block_hash_chain(): token["token_id"] for token in stored_events[0]["data"]["blocks"][1]["tokens"] ] == list(range(tokens_per_block, 2 * tokens_per_block)) - if _USING_CPP_BACKEND: - event_manager.add_updated_event( - bytes.fromhex(expected_hashes[0]), - cache_level=NativeKVCacheEventDiff(old_value=0, new_value=1), - layer_group_id=0, - ) - updated_events = _flush_serialized_events(event_manager) - assert updated_events[0]["data"]["block_hash"] == expected_hashes[0] + event_manager.add_updated_event( + bytes.fromhex(expected_hashes[0]), + cache_level=KVCacheEventDiff(old_value=0, new_value=1), + layer_group_id=0, + ) + updated_events = _flush_serialized_events(event_manager) + assert updated_events[0]["data"]["block_hash"] == expected_hashes[0] finally: gc.enable() if manager is not None: @@ -1406,7 +1267,7 @@ def test_v2_v1_hash_events_include_cache_salt_from_kv_cache(): gc.collect() gc.disable() - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=16, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -1442,7 +1303,7 @@ def test_v2_reused_prefix_does_not_emit_duplicate_stored_events(): gc.collect() gc.disable() - event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) manager = None try: tokens_per_block = 4 @@ -1466,14 +1327,12 @@ def test_v2_reused_prefix_does_not_emit_duplicate_stored_events(): reuse_events = _flush_serialized_events(event_manager) reused_hashes = _stored_block_hashes(reuse_events) - # SHA-256 block-key chain is identical across both backends. - root_key = RootBlock.make_key(ReuseScope()) - block0_key = Block.make_key(root_key, prefix_tokens[:tokens_per_block]) - block1_key = Block.make_key(block0_key, prefix_tokens[tokens_per_block:]) - block2_key = Block.make_key(block1_key, new_tokens) + block0_key, block1_key, block2_key = _blockchain_keys( + tokens_per_block, [*prefix_tokens, *new_tokens] + ) prefix_keys = [block0_key, block1_key] - assert prefix_hashes == [block_key.hex() for block_key in prefix_keys] + assert prefix_hashes == [key.hex() for key in prefix_keys] assert reused_hashes == [block2_key.hex()] assert not (set(prefix_hashes) & set(reused_hashes)) finally: @@ -1487,7 +1346,7 @@ def test_v2_mm_partial_commit_reuses_actual_longer_block(): init_cuda_once() gc.collect() gc.disable() - event_manager = NativeKVCacheEventManager( + event_manager = KVCacheEventManager( max_kv_event_entries=16, window_size=128, mm_token_id_offset=1000 ) manager = None @@ -1536,7 +1395,7 @@ def test_v2_removed_events_match_stored_hashes(): gc.collect() gc.disable() - event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=128) + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) manager = None try: tokens_per_block = 4 @@ -1593,7 +1452,7 @@ def test_v2_removed_event_emitted_when_last_level_page_is_dropped(): tokens_per_block = 8 window_size = 8 num_blocks = 6 - event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=window_size) + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=window_size) manager = None try: manager = _create_test_manager( @@ -1664,7 +1523,7 @@ def test_v2_kv_cache_event_manager_emits_updated_on_level_migration(): tokens_per_block = 8 window_size = 8 num_blocks = 8 - event_manager = NativeKVCacheEventManager(max_kv_event_entries=64, window_size=window_size) + event_manager = KVCacheEventManager(max_kv_event_entries=64, window_size=window_size) manager = None try: manager = _create_test_manager( diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 4b6d504ee659..c5dff10bcc15 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -20,6 +20,7 @@ import math import os import random +import sys import time import unittest from contextlib import contextmanager @@ -27,13 +28,19 @@ from importlib.util import find_spec from random import randbytes from statistics import median -from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, Sequence, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, Sequence, cast import pytest if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: + from bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + MemToMemTask, + copy_device_to_device, + ) from kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, + GPU_LEVEL, AttentionLayerConfig, BatchDesc, BufferConfig, @@ -50,8 +57,12 @@ KVCacheManagerConfig, LayerGroupId, LayerId, + MemAddress, + OutOfPagesError, + PageIndexMode, PlannedDropHandle, ReuseScope, + SlidingWindowSize, SsmLayerConfig, SwaScratchReuseConfig, TokenId, @@ -64,35 +75,15 @@ sequence_to_blockchain_keys, take_poison, ) - from kv_cache_manager_v2._block_radix_tree import Hasher - from kv_cache_manager_v2._common import ( - BAD_PAGE_INDEX, - GPU_LEVEL, - CacheTier, - MemAddress, - PageIndexMode, - SlidingWindowSize, - ) - from kv_cache_manager_v2._copy_engine import CopyTask, batched_copy - from kv_cache_manager_v2._exceptions import LogicError, OutOfPagesError - from kv_cache_manager_v2._storage._core import CacheLevelStorage, PoolGroupBase, SlotAllocator - from kv_cache_manager_v2._storage_manager import StorageManager - from kv_cache_manager_v2._utils import ( - CachedCudaStream, - HalfOpenRange, - TemporaryCudaStream, - div_up, - exact_div, - init_cuda_once, - intersect, - remove_if, - round_up, - temporary_sys_path, - typed_range, - ) else: + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + MemToMemTask, + copy_device_to_device, + ) from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, + GPU_LEVEL, AttentionLayerConfig, BatchDesc, BufferConfig, @@ -109,8 +100,12 @@ KVCacheManagerConfig, LayerGroupId, LayerId, + MemAddress, + OutOfPagesError, + PageIndexMode, PlannedDropHandle, ReuseScope, + SlidingWindowSize, SsmLayerConfig, SwaScratchReuseConfig, TokenId, @@ -118,26 +113,24 @@ _introspection, _KVCache, gen_multimodal_cache_key_tokens, + num_live_managers, + poison_reason, sequence_to_blockchain_keys, + take_poison, ) - from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import Hasher - from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( - BAD_PAGE_INDEX, - GPU_LEVEL, - CacheTier, - MemAddress, - PageIndexMode, - SlidingWindowSize, - ) - from tensorrt_llm.runtime.kv_cache_manager_v2._copy_engine import CopyTask, batched_copy - from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import LogicError, OutOfPagesError - from tensorrt_llm.runtime.kv_cache_manager_v2._storage._core import ( - CacheLevelStorage, - PoolGroupBase, - SlotAllocator, - ) - from tensorrt_llm.runtime.kv_cache_manager_v2._storage_manager import StorageManager - from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + +from copy import deepcopy + +from parameterized import parameterized + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +# cuda_test_utils supplies temporary_sys_path, so its own path entry is added and +# removed by hand here; every later sibling import goes through that helper. +_ADDED_TEST_DIR = _TEST_DIR not in sys.path +if _ADDED_TEST_DIR: + sys.path.insert(0, _TEST_DIR) +try: + from cuda_test_utils import ( # noqa: E402 CachedCudaStream, HalfOpenRange, TemporaryCudaStream, @@ -150,53 +143,24 @@ temporary_sys_path, typed_range, ) +finally: + if _ADDED_TEST_DIR: + sys.path.remove(_TEST_DIR) -from copy import deepcopy - -from parameterized import parameterized - -with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): +with temporary_sys_path(_TEST_DIR): from fake_engine import FakeEngine, Role, Step from kernels import HostGate, enable_kernel_delay -KV_CACHE_MANAGER_V2_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() - -# Gate for white-box tests that reach into the pure-Python implementation's objects -# (e.g. mutating a CommittedPage field). Prefer `_introspection`, which works on both -# backends; use this only when the behaviour under test cannot be reached through it. -requires_python_backend = unittest.skipIf( - KV_CACHE_MANAGER_V2_BACKEND == "cpp", - "white-box test over pure-Python KVCacheManagerV2 internals", -) - -requires_cpp_backend = unittest.skipUnless( - KV_CACHE_MANAGER_V2_BACKEND == "cpp", - "cold-page codec end-to-end test requires the C++ backend", -) - - def get_cached_cuda_event_type(): - backend = KV_CACHE_MANAGER_V2_BACKEND - if backend == "cpp": - try: - from bindings.internal.batch_manager.kv_cache_manager_v2 import _introspection - - return _introspection.CachedCudaEvent - except ImportError: - from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2 import ( - _introspection, - ) - - return _introspection.CachedCudaEvent - - if find_spec("kv_cache_manager_v2") is not None: - from kv_cache_manager_v2._utils import CachedCudaEvent + try: + from bindings.internal.batch_manager.kv_cache_manager_v2 import _introspection - return CachedCudaEvent - from tensorrt_llm.runtime.kv_cache_manager_v2._utils import CachedCudaEvent + return _introspection.CachedCudaEvent + except ImportError: + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2 import _introspection - return CachedCudaEvent + return _introspection.CachedCudaEvent seed = int.from_bytes(os.urandom(8), "little") @@ -248,28 +212,6 @@ def wrapper(self, *args, **kwargs): return wrapper -class TestTypedSlotIds(unittest.TestCase): - def test_num_slots_accessors_return_int(self) -> None: - self.assertIs(get_type_hints(SlotAllocator.num_slots.fget)["return"], int) - self.assertIs(get_type_hints(SlotAllocator.num_free_slots.fget)["return"], int) - self.assertIs(get_type_hints(SlotAllocator.num_occupied_slots.fget)["return"], int) - self.assertIs(get_type_hints(PoolGroupBase.num_slots.fget)["return"], int) - self.assertIs(get_type_hints(PoolGroupBase.num_free_slots.fget)["return"], int) - self.assertIs(get_type_hints(CacheLevelStorage.num_slots)["return"], int) - self.assertIs(get_type_hints(CacheLevelStorage.get_num_free_slots)["return"], int) - self.assertIs(get_type_hints(StorageManager.num_slots)["return"], int) - - self.assertIs(get_type_hints(SlotAllocator.allocate_multiple)["num_slots"], int) - self.assertIs(get_type_hints(PoolGroupBase.allocate_multiple)["num_slots"], int) - self.assertIs(get_type_hints(CacheLevelStorage.allocate_multiple)["num_slots"], int) - self.assertIs(get_type_hints(StorageManager.new_slots_for_pool_group)["num_slots"], int) - - allocator = SlotAllocator(3) - self.assertEqual(allocator.num_slots, 3) - self.assertEqual(allocator.num_free_slots, 3) - self.assertEqual(allocator.num_occupied_slots, 0) - - class TestCacheLevelStorage(unittest.TestCase): def test_grains_to_slots_rejects_zero_divisors(self) -> None: invalid_inputs = [ @@ -498,8 +440,6 @@ def run_request( tic = time.perf_counter() # prefill num_reused = kv_cache.num_committed_tokens - # workaround a mypyc bug: exception in property setter is not propagated - # kv_cache.capacity = round_up(len(prompt), interval) if not kv_cache.resize(round_up(len(prompt), interval)): raise OutOfPagesError("Not enough pages in GPU memory") capacity = kv_cache.capacity @@ -516,8 +456,6 @@ def run_request( if required_capacity > capacity: if not delay_commit: kv_cache.commit(history[kv_cache.history_length :]) - # workaround a mypyc bug: exception in property setter is not propagated - # kv_cache.capacity = round_up(required_capacity, interval) if not kv_cache.resize(round_up(required_capacity, interval)): raise OutOfPagesError("Not enough pages in GPU memory") capacity = kv_cache.capacity @@ -602,7 +540,6 @@ def _run_cold_page_codec_round_trip( if hasattr(self, "manager"): self.manager.clear_reusable_blocks() - @requires_cpp_backend def test_cold_codec_merges_lifecycles_from_different_hot_pool_groups(self) -> None: """Padding merges full attention with one of two differently-sized SWA LCs.""" unit = 1 << 20 @@ -664,7 +601,6 @@ def test_cold_codec_merges_lifecycles_from_different_hot_pool_groups(self) -> No }, ) - @requires_cpp_backend def test_cold_codec_splits_lifecycles_from_one_hot_pool_group(self) -> None: """Padding one SWA lifecycle splits a shared hot pool group in cold storage.""" unit = 1 << 20 @@ -1079,43 +1015,6 @@ def plan_drop(tokens: list[TokenIdExt]) -> PlannedDropHandle: with self.assertRaisesRegex(RuntimeError, "already been dropped"): long_handle.drop() - @requires_python_backend - def test_planned_drop_handle_rejects_partial_coverage(self) -> None: - # plan_committed_block_drop() rejects via _prune_match, which clamps the match to - # the page's recorded token count, so the endpoint no longer matches exactly. - # Forcing that state needs a direct write to the page, hence the backend gate. - window_size = 8 - tokens_per_block = 8 - self.prepare(16 << 20, 0, 0, 2, window_size, 0, tokens_per_block=tokens_per_block) - tokens = [self.next_token() for _ in range(3 * tokens_per_block)] - - with TemporaryCudaStream([]) as stream_holder: - stream = cast(CudaStream, stream_holder.handle) - kv_cache = self.manager.create_kv_cache(None, tokens) - self.assertTrue(kv_cache.resume(stream)) - self.assertTrue(kv_cache.resize(len(tokens))) - kv_cache.commit(tokens) - kv_cache.stop_committing() - - swa_lc_id = next( - lc_id - for lc_id, lc in self.manager._life_cycles.attention_life_cycles() - if lc.window_size is not None - ) - tree_block = kv_cache._blocks[2].tree_block - assert tree_block is not None - page = tree_block.get_page(swa_lc_id) - assert page is not None - self.assertEqual(page.num_tokens_in_block, len(tree_block.tokens)) - - page.num_tokens_in_block -= 1 - try: - self.assertIsNone(kv_cache.plan_committed_block_drop()) - finally: - page.num_tokens_in_block += 1 - kv_cache.close() - stream_holder.take_finish_event().synchronize() - def test_int32_ndarray_ingest_matches_list(self) -> None: """Zero-copy int32-ndarray ingest must hash identically to the list path. @@ -1128,11 +1027,6 @@ def test_int32_ndarray_ingest_matches_list(self) -> None: bit, blocks committed via the ndarray path would not be found by a list probe (and vice versa), so the equalities below would fail. """ - # The int32-ndarray ingest fast path lives in the C++ binding; the pure-Python - # backend consumes plain lists (the dispatcher hands it get_tokens, not a view). - if os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() == "python": - self.skipTest("int32-ndarray ingest is a C++-backend fast path") - import numpy as np tokens_per_block = 8 @@ -1400,15 +1294,12 @@ class TestLivingKvCacheGuard(TestKVCacheManagerV2): `clear_reusable_blocks()` detaches the whole radix tree and `shutdown()` frees the storage the pages live in. A request that is still open keeps committing into the - detached subtree, which silently discards work on the Python backend and segfaults on - the C++ one, so both entry points must reject the call instead. + detached subtree, which segfaults, so both entry points must reject the call instead. """ - # The two backends raise different types: the Python backend raises its own - # LogicError, while the C++ backend's TLLM_CHECK_WITH_INFO throws a TllmException - # (a std::runtime_error), which nanobind surfaces as RuntimeError. Accept either so - # this test is meaningful under both. - GuardError = (LogicError, RuntimeError) + # TLLM_CHECK_WITH_INFO throws a TllmException (a std::runtime_error), which nanobind + # surfaces as RuntimeError. + GuardError = RuntimeError def _seed_reusable_prompt(self) -> list[TokenIdExt]: """Commit and close a sequence so the radix tree actually holds reusable blocks. @@ -1904,7 +1795,7 @@ def get_rank_and_slice(max_par_size: int, par_size: int, idx: int) -> tuple[int, assert src_tp_slice.num_slices == 1 and dst_tp_slice.num_slices == 1 num_bytes = exact_div(src_page.size, src_tp_slice.num_slices) for i, j in zip(dst_indices, src_indices, strict=True): - task = CopyTask( + task = MemToMemTask( MemAddress(dst_page.base + dst_page.stride * i), MemAddress(src_page.base + src_page.stride * j), ) @@ -1927,12 +1818,12 @@ def get_rank_and_slice(max_par_size: int, par_size: int, idx: int) -> tuple[int, + num_bytes * src_tp_slice.slice_rank ) for b in range(num_buffers): - task = CopyTask( + task = MemToMemTask( MemAddress(dst_base + dst_buf_size * b), MemAddress(src_base + src_buf_size * b), ) tasks.append(task) - batched_copy(CacheTier.GPU_MEM, CacheTier.GPU_MEM, num_bytes, tasks, stream) + copy_device_to_device(tasks, num_bytes, stream) @parameterized.expand([(1, 1, 1, 1), (1, 2, 1, 1), (1, 1, 1, 2), (2, 1, 1, 1), (1, 1, 2, 1)]) def test_disaggregated_serving( @@ -4949,11 +4840,8 @@ def test_page_coverage_only_grows(self) -> None: class TestPoolRebalance(TestKVCacheManagerV2): """Drive the auto-tuner's pool rebalance end to end. - TestSlotAllocatorShrink below pokes the Python SlotAllocator directly, so it - never reaches the backend selected by TLLM_KV_CACHE_MANAGER_V2_BACKEND. This - class goes through the manager instead, covering - need_adjustment -> adjust() -> adjust_cache_level -> shrink/expand_pool_group - on whichever backend is active (C++ by default). + Everything goes through the manager, covering + need_adjustment -> adjust() -> adjust_cache_level -> shrink/expand_pool_group. """ _TOKENS_PER_BLOCK = 32 @@ -5143,43 +5031,6 @@ def test_kv_survives_adjust(self) -> None: self._run_sequence(prompt=prompt, expect_reuse=True) -class TestSlotAllocatorShrink(unittest.TestCase): - def test_shrink_underused_pool(self) -> None: - # Regression for NVBug 6225866: shrinking a pool whose new size is - # still above the slot-ID high-water mark used to assert because - # _num_active_slots - _target_capacity went negative. - allocator = SlotAllocator(capacity=184064) - slots = [allocator.allocate() for _ in range(2048)] - for s in slots: - allocator.release(s) - self.assertEqual(allocator._num_active_slots, 2048) - - allocator.prepare_for_shrink(122624) - self.assertEqual(len(allocator._overflow_slots), 0) - self.assertTrue(allocator.finish_shrink()) - self.assertEqual(allocator._capacity, 122624) - self.assertEqual(allocator._num_active_slots, 2048) - self.assertFalse(allocator.shrink_in_progress) - - def test_shrink_touched_pool(self) -> None: - # Sanity-check that the non-trivial migration path still works: - # all ids are issued, half released, shrink to half. - allocator = SlotAllocator(capacity=16) - slots = [allocator.allocate() for _ in range(16)] - for s in slots[8:]: - allocator.release(s) - self.assertEqual(allocator._num_active_slots, 16) - - allocator.prepare_for_shrink(8) - self.assertEqual(len(allocator._overflow_slots), 8) - self.assertTrue(allocator.finish_shrink()) - self.assertEqual(allocator._capacity, 8) - self.assertEqual(allocator._num_active_slots, 8) - - for s in slots[:8]: - allocator.release(s) - - class TestCachedTokensByTier(TestKVCacheManagerV2): @contextmanager def _tiered_prefix(self): @@ -5310,33 +5161,74 @@ def test_drop_partial_block_keeps_block_aligned_attribution(self) -> None: @pytest.mark.cpu_only class TestBlockKeyHashing(unittest.TestCase): - """Verify Hasher.update produces bit-identical digests to the per-token reference (no GPU needed).""" + """Verify blockchain keys are bit-identical to the per-token reference (no GPU needed).""" + + @staticmethod + def _chain(block: "list[int | bytes]") -> bytes: + """The single block key for one whole-sequence block, and its parent root key.""" + keys = list(sequence_to_blockchain_keys(max(len(block), 1), ReuseScope(), block)) + return keys[-1][1] @staticmethod - def _ref_update(seed: bytes, block: "list[int | bytes]") -> bytes: + def _root() -> bytes: + return next(iter(sequence_to_blockchain_keys(1, ReuseScope(), [])))[1] + + @staticmethod + def _ref_update(parent_key: bytes, block: "list[int | bytes]") -> bytes: h = hashlib.sha256() - h.update(seed) + h.update(parent_key) for item in block: # Normal token ids are packed as 4 little-endian bytes (31-bit range), - # matching the C++ backend's 4-byte TokenIdExt layout. + # matching the 4-byte TokenIdExt layout. h.update(item.to_bytes(4, "little") if type(item) is int else item) return h.digest() - def test_update_int_block_matches_reference(self) -> None: + def test_block_key_int_block_matches_reference(self) -> None: rng = random.Random(123) - seed = b"\xaa\xbb\xcc" - for n in (0, 1, 7, 32, 33, 257): + parent_key = self._root() + for n in (1, 7, 32, 33, 257): block = [rng.randint(0, (1 << 31) - 1) for _ in range(n)] self.assertEqual( - Hasher(seed).update(block).digest, - self._ref_update(seed, block), + self._chain(block), + self._ref_update(parent_key, block), f"int block of length {n}", ) - def test_update_mixed_multimodal_block(self) -> None: + def test_block_keys_chain_across_blocks(self) -> None: + """Each block key hashes its parent's key, so block N depends on blocks 0..N-1.""" + rng = random.Random(4242) + tokens_per_block = 32 + # Three whole blocks plus a short tail: the tail is chunked like any other block, + # so it gets a key of its own. + tokens = [rng.randint(0, (1 << 31) - 1) for _ in range(3 * tokens_per_block + 5)] + + keys = [ + key for _, key in sequence_to_blockchain_keys(tokens_per_block, ReuseScope(), tokens) + ] + self.assertEqual(keys[0], self._root()) + self.assertEqual(len(keys), 5) + + expected = self._root() + for ordinal in range(4): + block = tokens[ordinal * tokens_per_block : (ordinal + 1) * tokens_per_block] + expected = self._ref_update(expected, block) + self.assertEqual(keys[ordinal + 1], expected, f"block {ordinal}") + + # Chaining runs forward only: rewriting the tail must not disturb the keys of the + # blocks before it, which is what makes a matched prefix reusable. + tail_rewritten = tokens[:-1] + [tokens[-1] ^ 1] + rewritten_keys = [ + key + for _, key in sequence_to_blockchain_keys( + tokens_per_block, ReuseScope(), tail_rewritten + ) + ] + self.assertEqual(rewritten_keys[:-1], keys[:-1]) + self.assertNotEqual(rewritten_keys[-1], keys[-1]) + + def test_block_key_mixed_multimodal_block(self) -> None: block = [randbytes(32), 5, 6, randbytes(32)] + list(range(20)) - seed = b"\x01" - self.assertEqual(Hasher(seed).update(block).digest, self._ref_update(seed, block)) + self.assertEqual(self._chain(block), self._ref_update(self._root(), block)) def test_multimodal_digest_requires_sha256_length(self) -> None: for digest_size in (31, 33): diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py index e27596f5ff23..79ecd69176af 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py @@ -15,9 +15,8 @@ """Pure unit tests for KV cache reuse scopes.""" import unittest -from collections.abc import Iterator from importlib.util import find_spec -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import pytest @@ -25,42 +24,24 @@ if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: - from kv_cache_manager_v2 import TokenId - from kv_cache_manager_v2._block_radix_tree import ( - Block, - BlockRadixTree, - ReuseScope, - sequence_to_blockchain_keys, - ) - from kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry + from kv_cache_manager_v2 import ReuseScope, TokenId, sequence_to_blockchain_keys else: - from tensorrt_llm.runtime.kv_cache_manager_v2 import TokenId - from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( - Block, - BlockRadixTree, + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( ReuseScope, + TokenId, sequence_to_blockchain_keys, ) - from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry -class _EmptyLifeCycles: - size = 0 - - @property - def ssm_life_cycle_id(self) -> None: - return None - - def attention_life_cycles(self) -> Iterator[tuple[object, object]]: - return iter(()) +def _root_key(scope): + """The reuse-scope digest a blockchain starts from, which is its first key.""" + return next(iter(sequence_to_blockchain_keys(1, scope, [])))[1] class TestReuseScope(unittest.TestCase): def test_reuse_scope_seeds_distinct_keys(self) -> None: # Distinct reuse scopes -- including the None-vs-0 cases for each field -- - # must seed distinct radix-tree keys. The first blockchain key is the - # root (the reuse-scope digest), so hashing the same tokens under each - # scope isolates the scope's contribution to the key. + # must seed distinct radix-tree keys. scopes = [ ReuseScope(), ReuseScope(lora_id=0), @@ -68,15 +49,11 @@ def test_reuse_scope_seeds_distinct_keys(self) -> None: ReuseScope(lora_id=0, salt=0), ReuseScope(lora_id=7, salt=11), ] - tokens = [TokenId(1), TokenId(2)] - def root_key(scope: "ReuseScope") -> bytes: - return next(iter(sequence_to_blockchain_keys(2, scope, tokens)))[1] - - roots = [root_key(scope) for scope in scopes] + roots = [_root_key(scope) for scope in scopes] self.assertEqual(len(set(roots)), len(scopes)) # Deterministic across repeated derivation. - self.assertEqual(roots, [root_key(scope) for scope in scopes]) + self.assertEqual(roots, [_root_key(scope) for scope in scopes]) def test_blockchain_keys_are_seeded_by_reuse_scope(self) -> None: tokens = [TokenId(1), TokenId(2), TokenId(3), TokenId(4)] @@ -89,26 +66,35 @@ def test_blockchain_keys_are_seeded_by_reuse_scope(self) -> None: self.assertEqual(keys, same_scope_keys) self.assertNotEqual([key for _, key in keys], [key for _, key in different_scope_keys]) - def test_radix_tree_match_is_scoped(self) -> None: - tree = BlockRadixTree(cast(LifeCycleRegistry, _EmptyLifeCycles()), tokens_per_block=2) + def test_blockchain_keys_chain_from_the_scope_root(self) -> None: + # The first blockchain key is the reuse-scope digest, and each subsequent key + # is that chain extended by one block of tokens. + tokens = [TokenId(1), TokenId(2), TokenId(3), TokenId(4)] scope = ReuseScope(lora_id=7, salt=11) - other_scope = ReuseScope(lora_id=7, salt=12) - tokens = [TokenId(1), TokenId(2)] - root = tree.add_or_get_existing(scope) - block = Block(tokens, root) + keys = list(sequence_to_blockchain_keys(2, scope, tokens)) + self.assertEqual([chunk for chunk, _ in keys], [[], [1, 2], [3, 4]]) + self.assertEqual(keys[0][1], _root_key(scope)) - match = tree.match(scope, tokens) - self.assertEqual(match.blocks, [block]) - self.assertEqual(match.num_tokens, len(tokens)) + # Each key covers a prefix of the sequence, so extending the sequence leaves the + # earlier keys untouched. + prefix = list(sequence_to_blockchain_keys(2, scope, tokens[:2])) + self.assertEqual([key for _, key in prefix], [key for _, key in keys[:2]]) - match = tree.match(other_scope, tokens) - self.assertEqual(match.blocks, []) - self.assertEqual(match.num_tokens, 0) + def test_block_keys_are_scoped_through_their_root(self) -> None: + tokens = [TokenId(1), TokenId(2)] - match = tree.match(other_scope, tokens[:1], enable_partial_match=True) - self.assertEqual(match.blocks, []) - self.assertEqual(match.num_tokens, 0) + def block_keys(scope): + return [key for _, key in sequence_to_blockchain_keys(2, scope, tokens)][1:] + + self.assertNotEqual( + block_keys(ReuseScope(lora_id=7, salt=11)), + block_keys(ReuseScope(lora_id=7, salt=12)), + ) + self.assertEqual( + block_keys(ReuseScope(lora_id=7, salt=11)), + block_keys(ReuseScope(7, 11)), + ) if __name__ == "__main__": diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py index f89f50326af2..299062806eef 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os - import pytest import torch @@ -87,9 +85,6 @@ def test_stats_delta_arithmetic() -> None: def test_cpp_stats_types_are_native() -> None: - if os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() != "cpp": - pytest.skip("C++ backend only") - from tensorrt_llm.bindings.internal.batch_manager import kv_cache_manager_v2 as cpp assert KVCacheStatsDelta is cpp.KVCacheStatsDelta @@ -108,9 +103,6 @@ def test_cpp_stats_types_are_native() -> None: def test_native_cold_page_codec_is_consumed_after_failure() -> None: - if os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() != "cpp": - pytest.skip("C++ backend only") - from tensorrt_llm.bindings.internal.batch_manager import kv_cache_manager_v2 as cpp from tensorrt_llm.runtime.kv_cache_manager_v2 import create_default_kv_cache_cold_page_codec @@ -127,9 +119,6 @@ def test_native_cold_page_codec_is_consumed_after_failure() -> None: def test_native_cold_page_codec_rejects_wrong_type() -> None: - if os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() != "cpp": - pytest.skip("C++ backend only") - with pytest.raises(TypeError, match="IKvCacheColdPageCodec instance or None"): KVCacheManager(_make_config(), cold_page_codec=5) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py deleted file mode 100644 index 7c26296a80ec..000000000000 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py +++ /dev/null @@ -1,167 +0,0 @@ -# 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. - -"""TRTLLM-15217: SSM/recurrent life cycles must appear in V2 iteration stats. - -The page-movement recorders used to drop every non-attention life cycle, which -made KDA (Kimi K3) recurrent-state offload / onboard / drop invisible in -iteration statistics. These tests drive the recorders directly with a -duck-typed stand-in so they run without a GPU or an allocated cache. -""" - -from types import SimpleNamespace - -import pytest - -from tensorrt_llm.runtime.kv_cache_manager_v2._common import GPU_LEVEL, CacheLevel -from tensorrt_llm.runtime.kv_cache_manager_v2._core._kv_cache import _KVCache -from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import ( - AttnLifeCycle, - LifeCycleId, - SsmLifeCycle, -) -from tensorrt_llm.runtime.kv_cache_manager_v2._stats import KVCacheIterationStatsDelta - -ATTN_LC = LifeCycleId(0) -SSM_LC = LifeCycleId(1) -PAGE_BYTES = 16 -HOST_LEVEL = CacheLevel(GPU_LEVEL + 1) - - -def _make_recorder(*, manager_stats_enabled: bool = True, request_stats_enabled: bool = False): - """Duck-typed _KVCache exposing only what the stats recorders touch. - - The recording methods are bound off the real class, so the life-cycle - filtering under test is the production implementation. - """ - committed = [] - life_cycles = {ATTN_LC: AttnLifeCycle(None, 0), SSM_LC: SsmLifeCycle()} - manager = SimpleNamespace( - _life_cycles=SimpleNamespace(get_life_cycle=life_cycles.__getitem__), - _storage=SimpleNamespace( - get_pool_group_index=lambda life_cycle: life_cycle, - slot_size=lambda _pool_group: [PAGE_BYTES], - ), - commit_stats=lambda stats, by_life_cycle: committed.append((stats, by_life_cycle)), - ) - recorder = SimpleNamespace(manager=manager) - recorder._should_record_manager_stats = lambda: manager_stats_enabled - recorder._should_record_request_stats = lambda: request_stats_enabled - for name in ( - "_is_attention_life_cycle", - "_record_direct_iteration_stats", - "_record_migrated_slots", - "_record_dropped_pages", - ): - setattr(recorder, name, getattr(_KVCache, name).__get__(recorder)) - return recorder, committed - - -def test_request_only_stats_do_not_record_page_movement() -> None: - """Request-only accounting must not enable manager iteration statistics.""" - recorder, committed = _make_recorder(manager_stats_enabled=False, request_stats_enabled=True) - page = SimpleNamespace(life_cycle=ATTN_LC) - - recorder._record_migrated_slots([page], [object()], GPU_LEVEL, HOST_LEVEL) - recorder._record_dropped_pages([page], HOST_LEVEL) - recorder._record_direct_iteration_stats( - ATTN_LC, - KVCacheIterationStatsDelta(iter_intra_device_copy_blocks=1), - ) - - assert committed == [] - - -@pytest.mark.parametrize("life_cycle", [ATTN_LC, SSM_LC]) -def test_offload_is_recorded_for_every_life_cycle(life_cycle: LifeCycleId) -> None: - recorder, committed = _make_recorder() - page = SimpleNamespace(life_cycle=life_cycle) - - recorder._record_migrated_slots([page], [object()], GPU_LEVEL, HOST_LEVEL) - - assert len(committed) == 1 - _, by_life_cycle = committed[0] - assert set(by_life_cycle) == {life_cycle} - assert by_life_cycle[life_cycle].iter_offload_blocks == 1 - assert by_life_cycle[life_cycle].iter_offload_bytes == PAGE_BYTES - - -@pytest.mark.parametrize("life_cycle", [ATTN_LC, SSM_LC]) -def test_host_drop_is_recorded_for_every_life_cycle(life_cycle: LifeCycleId) -> None: - recorder, committed = _make_recorder() - page = SimpleNamespace(life_cycle=life_cycle) - - recorder._record_dropped_pages([page], HOST_LEVEL) - - assert len(committed) == 1 - _, by_life_cycle = committed[0] - assert set(by_life_cycle) == {life_cycle} - assert by_life_cycle[life_cycle].iter_host_dropped_blocks == 1 - assert by_life_cycle[life_cycle].iter_host_dropped_bytes == PAGE_BYTES - - -@pytest.mark.parametrize("life_cycle", [ATTN_LC, SSM_LC]) -def test_direct_iteration_stats_are_recorded_for_every_life_cycle( - life_cycle: LifeCycleId, -) -> None: - """SSM deferred copies must reach iteration stats. - - The resume() deferred copy reports iter_intra_device_copy_* through this - recorder for SSM life cycles too, matching the C++ backend. - """ - recorder, committed = _make_recorder() - - recorder._record_direct_iteration_stats( - life_cycle, - KVCacheIterationStatsDelta( - iter_intra_device_copy_blocks=1, - iter_intra_device_copy_bytes=PAGE_BYTES, - ), - ) - - assert len(committed) == 1 - _, by_life_cycle = committed[0] - assert set(by_life_cycle) == {life_cycle} - assert by_life_cycle[life_cycle].iter_intra_device_copy_blocks == 1 - assert by_life_cycle[life_cycle].iter_intra_device_copy_bytes == PAGE_BYTES - - -def test_onboard_counts_globally_only_for_attention() -> None: - """Onboard is per-life-cycle; global cache-hit counters are attention-only. - - alloc_total_blocks / alloc_new_blocks feed the global cache-hit rate, which - is defined over attention blocks only. - """ - recorder, committed = _make_recorder() - - recorder._record_migrated_slots( - [SimpleNamespace(life_cycle=SSM_LC)], [object()], HOST_LEVEL, GPU_LEVEL - ) - recorder._record_migrated_slots( - [SimpleNamespace(life_cycle=ATTN_LC)], [object()], HOST_LEVEL, GPU_LEVEL - ) - - assert len(committed) == 2 - ssm_stats, ssm_by_life_cycle = committed[0] - attn_stats, attn_by_life_cycle = committed[1] - - for life_cycle, by_life_cycle in ((SSM_LC, ssm_by_life_cycle), (ATTN_LC, attn_by_life_cycle)): - assert by_life_cycle[life_cycle].iter_onboard_blocks == 1 - assert by_life_cycle[life_cycle].iter_onboard_bytes == PAGE_BYTES - - assert ssm_stats.alloc_total_blocks == 0 - assert ssm_stats.alloc_new_blocks == 0 - assert attn_stats.alloc_total_blocks == 1 - assert attn_stats.alloc_new_blocks == 1 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py b/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py index 8e7a0a36e6a2..b3f48f789795 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py @@ -72,10 +72,6 @@ CHURN_REQUESTS = 100 -@unittest.skipUnless( - kv_test.KV_CACHE_MANAGER_V2_BACKEND == "cpp", - "the Python backend carries its own copy of this logic and is not fixed yet", -) class TestNvBug6625710(unittest.TestCase): """Evicting an unheld SSM snapshot must not detach a still-referenced block.""" diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py index c9c25731e3b6..6866f780864c 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py @@ -29,16 +29,14 @@ validate_streaming_support, ) from tensorrt_llm.llmapi.llm_args import KVEventsConfig -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( - Block, - BlockRadixTree, - ReuseScope, -) -from tensorrt_llm.runtime.kv_cache_manager_v2._config import ( - GpuCacheTierConfig, - KVCacheManagerConfig, + +# Streaming KV cache events have no implementation: StreamingKVCacheEventManager +# construction and validate_streaming_support() both raise. The supported route for KV +# cache events is the buffered one, via kv_cache_config.event_buffer_max_size. +streaming_unsupported = pytest.mark.skip( + reason="streaming KV cache events have no implementation; " + "use the buffered path via kv_cache_config.event_buffer_max_size" ) -from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry _ZMQ_SETUP_ATTEMPTS = 4 _RECEIVE_TIMEOUT_MS = 2_000 @@ -93,52 +91,7 @@ def _run_on_fresh_port(scenario: Callable[[int], None]) -> None: pytest.fail(f"ZeroMQ setup failed after {_ZMQ_SETUP_ATTEMPTS} attempts") -def test_streaming_sink_supports_real_radix_blocks(monkeypatch: pytest.MonkeyPatch) -> None: - """Real radix construction must accept the streaming sink's capability hooks.""" - manager = StreamingKVCacheEventManager( - KVEventsConfig(enable_kv_cache_events=True, publisher="null"), - data_parallel_rank=0, - block_size=4, - max_window_size=128, - ) - life_cycles = LifeCycleRegistry( - KVCacheManagerConfig( - tokens_per_block=4, - cache_tiers=[GpuCacheTierConfig(quota=4096)], - layers=[], - ) - ) - tree = BlockRadixTree(life_cycles, tokens_per_block=4, event_manager=manager) - published: list[KVEventBatch] = [] - monkeypatch.setattr( - manager._publisher, "publish", lambda batch: published.append(batch) or True - ) - try: - root = tree.add_or_get_existing(ReuseScope()) - first = Block([1, 2, 3, 4], root) - second = Block([5, 6, 7, 8], first) - - # Exercise wire event production after the page-coverage gate without - # allocating GPU pages; the radix blocks and sink are real objects. - manager._add_full_block(first) - manager._add_full_block(second) - manager.flush_iteration_events() - - assert len(published) == 1 - decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(published[0])) - assert len(decoded[1]) == 1 - stored = decoded[1][0] - assert stored["type"] == "BlockStored" - assert stored["token_ids"] == [1, 2, 3, 4, 5, 6, 7, 8] - assert stored["parent_block_hash"] is None - assert stored["block_size"] == 4 - assert len(stored["block_hashes"]) == 2 - assert manager.stored_blocks == 2 - finally: - tree.clear() - manager.shutdown() - - +@streaming_unsupported def test_streaming_fast_path_publishes_only_full_max_window_blocks() -> None: """Protect radix hash reuse, filtering, wire format, and shutdown.""" topic = "kv-events" @@ -261,6 +214,7 @@ def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: _run_on_fresh_port(scenario) +@streaming_unsupported def test_streaming_removals_are_never_dropped_by_the_entry_cap() -> None: """Removals must survive the per-iteration cap or the consumer desyncs.""" manager = StreamingKVCacheEventManager( @@ -339,6 +293,7 @@ def test_dropped_batches_leave_a_sequence_gap() -> None: publisher.shutdown() +@streaming_unsupported def test_construction_binds_nothing_until_start() -> None: """A constructed-but-unstarted publisher must hold no socket and no thread.""" port = _unused_tcp_port() @@ -363,6 +318,7 @@ def test_construction_binds_nothing_until_start() -> None: manager.shutdown() +@streaming_unsupported def test_shutdown_without_start_is_safe() -> None: """Tearing down a manager that never started must not raise.""" manager = StreamingKVCacheEventManager( @@ -378,19 +334,15 @@ def test_shutdown_without_start_is_safe() -> None: def test_validate_streaming_support_rejects_unsupported_setups() -> None: config = KVEventsConfig(enable_kv_cache_events=True, endpoint="tcp://*:5557") - supported = dict(pp_size=1, cp_size=1, ranks_per_host=1, data_parallel_size=1, backend="python") - - # The supported baseline must not raise, or the negative cases prove nothing. - validate_streaming_support(config, **supported) + supported = dict(pp_size=1, cp_size=1, ranks_per_host=1, data_parallel_size=1) with pytest.raises(ValueError, match="pipeline parallelism"): validate_streaming_support(config, **{**supported, "pp_size": 2}) with pytest.raises(ValueError, match="context parallelism"): validate_streaming_support(config, **{**supported, "cp_size": 2}) - # The default backend is "cpp", whose nanobind KVCacheManager cannot accept a - # duck-typed Python event sink; the error must name the env var that fixes it. - with pytest.raises(ValueError, match="TLLM_KV_CACHE_MANAGER_V2_BACKEND=python"): - validate_streaming_support(config, **{**supported, "backend": "cpp"}) + # Streaming has no implementation, so even an otherwise supported setup is rejected. + with pytest.raises(ValueError, match="event_buffer_max_size"): + validate_streaming_support(config, **supported) @pytest.mark.parametrize( @@ -427,6 +379,7 @@ def test_validate_endpoint_ranges(endpoint, replay_endpoint, ranks_per_host, ove validate_endpoint_ranges(config, ranks_per_host, ranks_per_host) +@streaming_unsupported def test_partial_target_page_coverage_is_suppressed_until_fully_covered() -> None: """A page adopted from a shorter sibling must not be published as a full block.""" manager = StreamingKVCacheEventManager( @@ -471,6 +424,7 @@ def test_partial_target_page_coverage_is_suppressed_until_fully_covered() -> Non manager.shutdown() +@streaming_unsupported def test_life_cycle_hooks_ignore_none_ids() -> None: """A None life-cycle id must not reach int() before the target is configured.""" manager = StreamingKVCacheEventManager( @@ -523,6 +477,7 @@ def test_validate_endpoint_ranges_rejects_port_overflow( validate_endpoint_ranges(config, 1, dp_size) +@streaming_unsupported def test_validate_streaming_support_rejects_overflowing_port_span() -> None: """Every rank must reject the span, so none reaches the following collective.""" config = KVEventsConfig(enable_kv_cache_events=True, endpoint="tcp://*:65535") diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py b/tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py new file mode 100644 index 000000000000..5f868f674568 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py @@ -0,0 +1,88 @@ +# 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. + +"""``VirtMem`` must outlive the Python handle to the allocator it borrows. + +``VirtMem`` holds ``PooledPhysMemAllocator&`` in C++ and dereferences it whenever it +unmaps physical memory. The binding's ``nb::keep_alive`` is what ties the allocator's +Python lifetime to the ``VirtMem``, and the native disaggregated bounce buffer depends +on it: ``bounce/buffer.py`` builds the allocator as a local and lets it go out of scope +while keeping only the ``VirtMem``. Without that tie the allocator is freed at the end +of that constructor and every later unmap reads freed memory. + +Only ``test_virt_mem_holds_a_reference_to_its_allocator`` guards the contract. Dropping +``keep_alive`` from the binding and rerunning this file leaves the end-to-end case below +passing, because the freed allocator reads back intact in a release build; the reference +count is what actually changes. Keep any new case here anchored on something observable +rather than on an unmap that happens to succeed. +""" + +import gc +import sys + +import pytest +import torch + +from tensorrt_llm.runtime.kv_cache_manager_v2._introspection import PooledPhysMemAllocator, VirtMem + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + +_CHUNK = 2 << 20 + + +@pytest.fixture(autouse=True) +def _cuda_context() -> None: + # PooledPhysMemAllocator drives the driver API directly, so it needs a current + # primary context; torch only creates one lazily on first use. + torch.cuda.init() + torch.zeros(1, device="cuda") + + +def _make_virt_mem_from_temporary_allocator() -> VirtMem: + """Build a VirtMem exactly as bounce/buffer.py does: allocator is a local.""" + allocator = PooledPhysMemAllocator(_CHUNK) + return VirtMem(_CHUNK, allocator, init_num_phys_mem=1) + + +def test_virt_mem_outlives_its_allocator_handle() -> None: + vm = _make_virt_mem_from_temporary_allocator() + # The only Python reference to the allocator went out of scope on return, so a + # missing keep_alive leaves the C++ reference dangling from here on. + gc.collect() + + assert vm.address != 0 + # destroy() unmaps the backing chunk, which is the path that dereferences the + # allocator. It must still be alive. + vm.destroy() + + +def test_virt_mem_holds_a_reference_to_its_allocator() -> None: + """The retention is asserted directly, not inferred from a surviving unmap. + + Dropping ``keep_alive`` leaves freed memory that a release build usually reads back + without faulting, so a test that only called ``destroy()`` could keep passing while + the contract is broken. The reference count is exact and fails the moment it goes. + """ + unheld = PooledPhysMemAllocator(_CHUNK) + baseline = sys.getrefcount(unheld) + + allocator = PooledPhysMemAllocator(_CHUNK) + vm = VirtMem(_CHUNK, allocator, init_num_phys_mem=1) + assert sys.getrefcount(allocator) == baseline + 1 + + del allocator + gc.collect() + assert vm.address != 0 + vm.destroy() From a32196002cb1654712ff9300db2c869be148ed89 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 14 Sep 2026 07:58:04 +0000 Subject: [PATCH 2/2] [TRTLLM-15217][test] Lock SSM life cycles into KVCM2 page-movement statistics Iteration statistics are keyed by life cycle and report recurrent (SSM) page movement alongside attention movement, while the global cache-hit counters stay attention-only. That split had no test on the C++ side: the only coverage lived in the Python backend's unit tests, which went away with the backend itself, and no C++ test builds an SSM life cycle at all. Add a hybrid attention + SSM fixture and three cases over it: - offload and onboard are reported for both life cycles, with byte counts matching each life cycle's slot size - an SSM onboard leaves allocTotalBlocks / allocNewBlocks to attention - host drops are reported for both life cycles Each case was confirmed to fail when the life-cycle filters are restored in KvCache::_recordMigratedSlots and _recordDroppedPages. Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/exceptions.h | 4 +- .../kvCacheManagerV2StatsTest.cpp | 164 ++++++++++++++++++ .../batch_manager/kvCacheManagerV2TestUtils.h | 31 ++++ .../test_kv_cache_manager_v2.py | 4 +- 4 files changed, 199 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h index 7fb4faa8902d..9f76e36576a3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h @@ -111,8 +111,8 @@ class LogicError : public std::logic_error }; // Mirrors a Python `assert` failure: the binding layer translates this to a -// Python AssertionError so shared tests observe the same exception type as the -// pure-Python backend. +// Python AssertionError, so a configuration mistake surfaces to Python callers +// as the exception type they would expect. class AssertionError : public std::logic_error { public: diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index 8e47633f954a..2b625fbeaab8 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -39,6 +39,7 @@ namespace using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeConfig; +using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeHybridTieredConfig; using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeTieredConfig; TEST(KvCacheManagerV2StatsTest, StatsDeltaArithmetic) @@ -454,4 +455,167 @@ TEST(KvCacheManagerV2StatsTest, DisabledStatsSuppressSuspendResumeCounters) EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); } +// Hybrid attention + SSM page-movement statistics. +// +// Iteration statistics are keyed by life cycle and must report recurrent (SSM) page +// movement alongside attention movement, otherwise KDA recurrent-state offload, onboard +// and drop are invisible to callers. The global cache-hit counters are the deliberate +// exception: they stay attention-only. +// +// The tiers below give the attention life cycle 4 GPU and 2 host slots of 1 MiB, and the +// SSM life cycle 1 GPU and 1 host slot of 2 MiB, so a second sequence evicts the first and +// a second eviction round overflows the host pool. +namespace +{ +constexpr int kHybridBlocks = 3; +constexpr CacheLevel kHybridHostLevel{1}; + +int64_t slotBytesFor(StorageManager const& storage, CacheLevel level, LifeCycleId lifeCycle) +{ + int64_t bytes = 0; + for (size_t const size : storage.slotSize(storage.getPoolGroupIndex(level, lifeCycle))) + { + bytes += static_cast(size); + } + return bytes; +} + +std::vector makeTokens(KvCacheManager const& manager, int firstToken) +{ + std::vector tokens; + for (int offset = 0; offset < kHybridBlocks * manager.tokensPerBlock(); ++offset) + { + tokens.emplace_back(TokenId{firstToken + offset}); + } + return tokens; +} + +// Fill a sequence, park it, then start a second one that needs the same slots. Returns the +// still-open second sequence so the caller can close it to trigger the onboard. +std::pair, std::shared_ptr> evictFirstSequence( + KvCacheManager& manager, cudaStream_t stream, int firstToken) +{ + auto const tokens = makeTokens(manager, firstToken); + auto first = manager.createKvCache(); + EXPECT_TRUE(first->resume(reinterpret_cast(stream))); + EXPECT_TRUE(first->resize(static_cast(tokens.size()))); + first->commit(toSpan(tokens)); + first->suspend(); + + auto second = manager.createKvCache(); + EXPECT_TRUE(second->resume(reinterpret_cast(stream))); + EXPECT_TRUE(second->resize(static_cast(tokens.size()))); + return {std::move(first), std::move(second)}; +} +} // namespace + +TEST(KvCacheManagerV2StatsTest, OffloadAndOnboardAreRecordedForAttentionAndSsmLifeCycles) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + cudaStream_t stream{}; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + auto manager = std::make_shared(makeHybridTieredConfig()); + auto& storage = manager->storage(); + LifeCycleId const attention{0}; + LifeCycleId const ssm{1}; + ASSERT_TRUE(std::holds_alternative(manager->lifeCycles().getLifeCycle(attention))); + ASSERT_FALSE(std::holds_alternative(manager->lifeCycles().getLifeCycle(ssm))); + + manager->getAndResetIterationStats(); + auto [first, second] = evictFirstSequence(*manager, stream, 0); + + auto const offload = manager->getAndResetIterationStats(); + ASSERT_EQ(offload.size(), 2) << "both life cycles must report offload"; + for (LifeCycleId const lifeCycle : {attention, ssm}) + { + auto const& stats = offload.at(lifeCycle); + EXPECT_GT(stats.iterOffloadBlocks, 0) << "life cycle " << lifeCycle.value(); + EXPECT_EQ(stats.iterOffloadBytes, stats.iterOffloadBlocks * slotBytesFor(storage, kHotLevel, lifeCycle)); + } + + second->close(); + ASSERT_TRUE(first->resume(reinterpret_cast(stream))); + + auto const onboard = manager->getAndResetIterationStats(); + ASSERT_EQ(onboard.size(), 2) << "both life cycles must report onboard"; + for (LifeCycleId const lifeCycle : {attention, ssm}) + { + auto const& stats = onboard.at(lifeCycle); + EXPECT_GT(stats.iterOnboardBlocks, 0) << "life cycle " << lifeCycle.value(); + EXPECT_EQ(stats.iterOnboardBytes, stats.iterOnboardBlocks * slotBytesFor(storage, kHotLevel, lifeCycle)); + } + + first->close(); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST(KvCacheManagerV2StatsTest, SsmOnboardLeavesGlobalAllocCountersToAttention) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + cudaStream_t stream{}; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + auto manager = std::make_shared(makeHybridTieredConfig()); + LifeCycleId const attention{0}; + LifeCycleId const ssm{1}; + + manager->getAndResetIterationStats(); + auto [first, second] = evictFirstSequence(*manager, stream, 0); + manager->getAndResetIterationStats(); + + second->close(); + auto const allocTotalBefore = manager->getCommittedStats().allocTotalBlocks; + auto const allocNewBefore = manager->getCommittedStats().allocNewBlocks; + ASSERT_TRUE(first->resume(reinterpret_cast(stream))); + auto const allocTotalDelta = manager->getCommittedStats().allocTotalBlocks - allocTotalBefore; + auto const allocNewDelta = manager->getCommittedStats().allocNewBlocks - allocNewBefore; + + auto const onboard = manager->getAndResetIterationStats(); + ASSERT_EQ(onboard.size(), 2); + auto const attentionOnboard = onboard.at(attention).iterAllocTotalBlocks; + auto const ssmOnboard = onboard.at(ssm).iterAllocTotalBlocks; + // Both life cycles onboard, so a global delta equal to the attention share alone is + // only possible if the SSM share was excluded. + ASSERT_GT(attentionOnboard, 0); + ASSERT_GT(ssmOnboard, 0); + EXPECT_EQ(allocTotalDelta, attentionOnboard); + EXPECT_EQ(allocNewDelta, attentionOnboard); + + first->close(); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST(KvCacheManagerV2StatsTest, HostDropIsRecordedForAttentionAndSsmLifeCycles) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + cudaStream_t stream{}; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + auto manager = std::make_shared(makeHybridTieredConfig()); + auto& storage = manager->storage(); + LifeCycleId const attention{0}; + LifeCycleId const ssm{1}; + + // First round fills the host pools. + auto [first, second] = evictFirstSequence(*manager, stream, 0); + second->close(); + first->close(); + + // Second round uses disjoint tokens, so nothing is reused and the host pools overflow. + manager->getAndResetIterationStats(); + auto [third, fourth] = evictFirstSequence(*manager, stream, 1000); + + auto const dropped = manager->getAndResetIterationStats(); + ASSERT_EQ(dropped.size(), 2) << "both life cycles must report host drops"; + for (LifeCycleId const lifeCycle : {attention, ssm}) + { + auto const& stats = dropped.at(lifeCycle); + EXPECT_GT(stats.iterHostDroppedBlocks, 0) << "life cycle " << lifeCycle.value(); + EXPECT_EQ(stats.iterHostDroppedBytes, + stats.iterHostDroppedBlocks * slotBytesFor(storage, kHybridHostLevel, lifeCycle)); + } + + fourth->close(); + third->close(); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + } // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TestUtils.h b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TestUtils.h index 479de2036327..678951d851d6 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TestUtils.h +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TestUtils.h @@ -51,4 +51,35 @@ inline KVCacheManagerConfig makeTieredConfig() return config; } +//! Attention and SSM life cycles side by side, over a GPU and a host tier. +//! +//! Life cycles are registered in layer order, so the attention layer is LifeCycleId{0} and +//! the SSM layer is LifeCycleId{1}. The buffer sizes differ so per-life-cycle byte counters +//! identify which life cycle they came from. +//! +//! The quotas give attention 4 GPU and 2 host slots of 1 MiB, and SSM 1 GPU and 1 host slot +//! of 2 MiB. A three-block sequence therefore fits on the GPU but a second sequence evicts +//! it, and a second eviction round overflows the host pools. +inline KVCacheManagerConfig makeHybridTieredConfig() +{ + KVCacheManagerConfig config; + config.tokensPerBlock = 4; + config.cacheTiers.emplace_back(GpuCacheTierConfig{6UL << 20}); + config.cacheTiers.emplace_back(HostCacheTierConfig{4UL << 20}); + + AttentionLayerConfig attention; + attention.layerId = 0; + attention.buffers.push_back(BufferConfig{"key", 1UL << 20, std::nullopt}); + config.layers.emplace_back(std::move(attention)); + + SsmLayerConfig ssm; + ssm.layerId = 1; + ssm.buffers.push_back(BufferConfig{"ssm_state", 2UL << 20, std::nullopt}); + config.layers.emplace_back(std::move(ssm)); + + // KVCacheManagerConfig::validate() rejects an SSM layer without this. + config.commitMinSnapshot = true; + return config; +} + } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2::test diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index c5dff10bcc15..26997a7508a7 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -2524,8 +2524,8 @@ def test_ssm_resume_records_intra_device_copy(self) -> None: First resume of a cache reusing an SSM snapshot copies the snapshot into a private slot; the copy must appear in the SSM life cycle's - iteration stats (TRTLLM-15217). Runs against the selected backend, so - it checks the default C++ implementation and Python-backend parity. + iteration stats (TRTLLM-15217). Offload, onboard and host-drop for the + same life cycle are covered by KvCacheManagerV2StatsTest. """ tokens_per_block = 32 cfg = self._make_ssm_config(tokens_per_block=tokens_per_block)