[None][fix] Restore KVCacheManagerV2 mypyc compilability - #15990
Open
thorjohnsen wants to merge 1 commit into
Open
[None][fix] Restore KVCacheManagerV2 mypyc compilability#15990thorjohnsen wants to merge 1 commit into
thorjohnsen wants to merge 1 commit into
Conversation
…g issues) The package's mypyc build (make -C kv_cache_manager_v2 all) fails on current main. This fixes every pre-existing issue needed to get a green build with the pinned mypy==1.19.1: - typing: missing return/parameter annotations, bare generic casts (TypedIndexList, array.array, rawref.ref), shadowed loop variables with conflicting narrowed types, Self on _stats.copy(), import- untyped ignores for cuda/tensorrt_llm imports, int() wraps for helpers whose defining module is outside the compiled set; - HalfOpenRange is no longer a tuple subclass: mypyc lowers len(), bool() and 'in' on tuple-typed values to raw tuple operations, silently bypassing the overrides (an empty range reported len 2 and was truthy, corrupting scratch-slot accounting); - LayerConfig.type ClassVars became properties: compiled dataclasses count a leading defaulted ClassVar as an init field; - KVCacheDesc/BatchDesc get explicit __deepcopy__: the compiled frozen dataclass's generated __setstate__ assigns via plain setattr and trips its own frozen __setattr__ (KVCacheManager.__init__ deepcopies its config); - DeltaScratchSlots moved to module level (nested class definitions are not compilable); - setup_mypyc.py excludes _core/_kv_cache.py and _block_radix_tree.py for now: compiling the former leaks gc-invisible strong page references (breaks block reuse) and the latter strands pages in eviction queues on resize_quota teardown; both need root-causing in mypyc codegen before they can be re-enabled. With this, the hybrid build (20/22 modules compiled) passes the full v2 regression suite identically to the interpreted baseline (61 passed / 12 skipped). Signed-off-by: thorjohnsen <41591019+thorjohnsen@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughThis PR applies typing-focused cleanup across kv_cache_manager_v2: adds explicit type annotations, casts, and type-ignore comments to imports; converts layer-type ClassVars to properties; adds deepcopy overrides; rewrites HalfOpenRange as a non-tuple class; and removes two modules from mypyc compilation. No functional logic changes. ChangesType annotation and typing-behavior cleanup
Estimated code review effort: 2 (Simple) | ~15 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.py (1)
152-158: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAvoid sharing the same list object here.
[[]] * self.num_life_cyclesaliases every inner list. Since this function returns nested mutable lists, that can cross-contaminate life-cycle entries if one is mutated later. Build the outer list with a comprehension instead.Suggested fix
- ret: list[list[int]] = [[]] * self.num_life_cycles + ret: list[list[int]] = [[] for _ in range(int(self.num_life_cycles))]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.py` around lines 152 - 158, The slot_to_page_indices method is creating aliased inner lists by initializing ret with [[]] * self.num_life_cycles, which can cause life_cycle entries to share mutable state. Update slot_to_page_indices in _config.py to build ret with a comprehension so each life_cycle gets its own independent list, while keeping the rest of the loop over slot_desc_list and slot.coalesced_buffers unchanged.
🧹 Nitpick comments (2)
tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py (2)
67-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse overloads for
_unwrapinstead ofAny.
_unwraphas fixed arity-based outcomes (None,T, ortuple[T, U]), soAnyleaks type information into the CUDA query and copy helpers. Overloads keep those call sites typed without changing runtime behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py` around lines 67 - 74, The _unwrap helper currently returns Any even though its behavior is fixed by tuple arity, so replace the broad annotation with overloads on _unwrap to preserve typing for the CUDA query and copy helpers. Add distinct overload signatures for the CUresult-only case, the single-value tuple case, and the two-value tuple case, while keeping the existing runtime implementation unchanged and using the _unwrap symbol to guide the edit.Source: Coding guidelines
237-250: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
ParamSpecfor these decorators.noexceptandnot_implementedcurrently erase wrapped signatures withCallable[..., T]andAny; switching toCallable[P, T]/*args: P.args, **kwargs: P.kwargskeeps the decorators type-safe with minimal churn.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py` around lines 237 - 250, Update the noexcept and not_implemented decorators to preserve wrapped function signatures by introducing ParamSpec and using Callable[P, T] instead of Callable[..., T]. Adjust the inner wrappers in _utils.py to accept *args: P.args and **kwargs: P.kwargs, and remove the broad Any-based signature so the decorators remain type-safe while keeping behavior unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py`:
- Line 17: The import in the module-level setup uses typing.Self, which is not
available on Python 3.10. Update the _stats.py import to pull Self from
typing_extensions instead, and keep the rest of the type annotations in the file
unchanged so the class and method signatures that reference Self continue to
work on all supported versions.
---
Outside diff comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.py`:
- Around line 152-158: The slot_to_page_indices method is creating aliased inner
lists by initializing ret with [[]] * self.num_life_cycles, which can cause
life_cycle entries to share mutable state. Update slot_to_page_indices in
_config.py to build ret with a comprehension so each life_cycle gets its own
independent list, while keeping the rest of the loop over slot_desc_list and
slot.coalesced_buffers unchanged.
---
Nitpick comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py`:
- Around line 67-74: The _unwrap helper currently returns Any even though its
behavior is fixed by tuple arity, so replace the broad annotation with overloads
on _unwrap to preserve typing for the CUDA query and copy helpers. Add distinct
overload signatures for the CUresult-only case, the single-value tuple case, and
the two-value tuple case, while keeping the existing runtime implementation
unchanged and using the _unwrap symbol to guide the edit.
- Around line 237-250: Update the noexcept and not_implemented decorators to
preserve wrapped function signatures by introducing ParamSpec and using
Callable[P, T] instead of Callable[..., T]. Adjust the inner wrappers in
_utils.py to accept *args: P.args and **kwargs: P.kwargs, and remove the broad
Any-based signature so the decorators remain type-safe while keeping behavior
unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 63ce790e-ca73-4925-991d-9eee91168e12
📒 Files selected for processing (17)
tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_moving_average.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.pytensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.pytensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.pytensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytensorrt_llm/runtime/kv_cache_manager_v2/_stats.pytensorrt_llm/runtime/kv_cache_manager_v2/_storage/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_utils.pytensorrt_llm/runtime/kv_cache_manager_v2/rawref/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py
| # limitations under the License. | ||
|
|
||
| from dataclasses import dataclass, fields | ||
| from typing import Self |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '## File\n'
cat -n tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py | sed -n '1,120p'
printf '\n## Python version declarations\n'
rg -n --hidden --glob '!**/.git/**' -e 'python_requires|requires-python|3\.10|3\.11|3\.12|3\.13' pyproject.toml setup.py setup.cfg tox.ini noxfile.py requirements*.txt .github/workflows AGENTS.md CLAUDE.md legacy-files.txt tensorrt_llm -g '!**/node_modules/**' || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 4992
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py')
print(p.read_text())
PY
printf '\n## Python version declarations\n'
rg -n --hidden --glob '!**/.git/**' -e 'python_requires|requires-python|3\.10|3\.11|3\.12|3\.13' pyproject.toml setup.py setup.cfg tox.ini noxfile.py requirements*.txt .github/workflows AGENTS.md CLAUDE.md legacy-files.txt tensorrt_llm || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 4467
🏁 Script executed:
#!/bin/bash
set -euo pipefail
pwd
git ls-files tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py pyproject.toml setup.py setup.cfg | sed -n '1,20p'Repository: NVIDIA/TensorRT-LLM
Length of output: 251
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
print(Path('pyproject.toml').read_text())
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- pyproject.toml ---'
sed -n '1,120p' pyproject.toml
echo
echo '--- Search for requires-python / python versions ---'
rg -n 'requires-python|python_requires|3\.10|3\.11|3\.12|3\.13' pyproject.toml setup.py setup.cfg tox.ini noxfile.py .github/workflows AGENTS.md CLAUDE.md tensorrt_llm | head -n 200 || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 7664
Import Self from typing_extensions. Python 3.10 is still a supported target, so typing.Self will fail at import time here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py` at line 17, The import in
the module-level setup uses typing.Self, which is not available on Python 3.10.
Update the _stats.py import to pull Self from typing_extensions instead, and
keep the rest of the type annotations in the file unchanged so the class and
method signatures that reference Self continue to work on all supported
versions.
Sources: Coding guidelines, Learnings
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57839 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57839 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
KVCacheManagerV2 is designed to be compilable with mypyc for production performance (see its
AGENTS.mdandsetup_mypyc.py), but the build (make -C kv_cache_manager_v2 all) fails on current main with the pinnedmypy==1.19.1. This PR fixes every pre-existing issue needed to get a green, regression-clean build:Typing fixes (annotation-only): missing return/parameter annotations, bare generic casts (
TypedIndexList,array.array,rawref.ref), loop variables reusing a name with a conflicting narrowed type,Selfon_stats.copy(),# type: ignorecomment placement (mypy only honors it as the first comment), import-untyped ignores forcuda/tensorrt_llmimports, andint()wraps for helpers whose defining module is outside the compiled set.Semantic fixes for mypyc backend behavior (behavior-neutral in interpreted mode, verified by the regression suite):
HalfOpenRangeis no longer atuplesubclass: mypyc lowerslen(),bool()andinon tuple-typed values to raw tuple operations, silently bypassing the class's overrides — an empty range reportedlen() == 2and was truthy, which corrupted scratch-slot accounting when compiled. It is now a plain slotted class implementing the same protocol (indexing, unpacking, value-equality with 2-tuples).LayerConfig.typeClassVars became properties: compiled dataclasses count a leading defaulted ClassVar as an__init__field (TypeError: non-default argument 'layer_id' follows default argumentat import).KVCacheDesc/BatchDescget explicit__deepcopy__: the compiled frozen dataclass's generated__setstate__assigns via plainsetattrand trips its own frozen__setattr__(KVCacheManager.__init__deepcopies its config).DeltaScratchSlotsmoved to module level (mypyc cannot compile class definitions nested in a class body).Known remaining mypyc issues (excluded from compilation in
setup_mypyc.pywith comments, so the default build is green):_core/_kv_cache.pyleaks strong page references invisible togcintrospection (committed pages surviveclear_reusable_blocks()with zero visible referrers, breaking block reuse and pool shrink at shutdown);_block_radix_tree.pystrands pages in eviction queues onresize_quotateardown, with an occasional segfault in the subsequent__del__cascade.Both reproduce deterministically via
TestNoBatching::test_cache_reuse_0/TestResizeQuota::test_resize_quotawhen the respective module is added back to the compile list; they appear to be mypyc codegen issues (refcount ortp_traverse) and are left for follow-up.Test Coverage
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py, run on H100:PR Checklist
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Refactor