Conversation
40a5e12 to
6493b80
Compare
|
/bot run --disable-fail-fast |
| 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. |
There was a problem hiding this comment.
Do we need to remove the existing mypy checks?
https://github.com/NVIDIA/TensorRT-LLM/blob/main/jenkins/Build.groovy#L493
https://github.com/NVIDIA/TensorRT-LLM/blob/main/.pre-commit-config.yaml#L1685
There was a problem hiding this comment.
No — those should stay. Different tools with confusingly similar names:
- mypyc is the ahead-of-time compiler that built the pure-Python KVCacheManagerV2 into a
.so. That is what this PR removes (--mypyc,TRTLLM_ENABLE_MYPYC,setup_mypyc.py, and thesetup.pypackaging surgery they needed). Zeromypycreferences remain in tracked files. - mypy is the static type checker. Both links you posted run it via
scripts/run_mypy.sh— Build.groovy'stypeCheckstage and thetype-checkpre-commit hook — and neither has anything to do with mypyc.
The only pyproject.toml mypy change here is dropping the [[tool.mypy.overrides]] block for tensorrt_llm.runtime.kv_cache_manager_v2.*, whose disallow_any_generics = false existed to accommodate the deleted Python implementation.
If anything the type check now covers slightly more: the PR restores runtime/kv_cache_manager_v2/__init__.pyi and fills in 12 exported names the stub was missing, so mypy --strict has more to read, not less.
| if config.algorithm == "quantization_for_cold_page": | ||
| from tensorrt_llm.runtime.kv_cache_manager_v2 import _BACKEND | ||
|
|
||
| if _BACKEND == "python": |
There was a problem hiding this comment.
Removing _BACKEND and this admission check also requires migrating test_quantization_for_cold_page.py. Four tests still start with monkeypatch.setattr(runtime_v2_mod, "_BACKEND", ...), which now raises AttributeError before their assertions. The first also expects the deleted Python-backend rejection, so raising=False alone would still fail. Please remove those obsolete patches/assertions while retaining the SM100, speculative-mode, and estimation checks.
There was a problem hiding this comment.
Valid — confirmed by running them, all four failed as you described. Fixed.
One correction for anyone following along: the file is tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py (not _torch/executor/kv_cache/). Everything else matched, including the detail that raising=False alone would not have been enough for the first test, since it asserts the deleted Python-backend rejection.
Changes: dropped that rejection block, removed the six now-meaningless monkeypatch.setattr(runtime_v2_mod, "_BACKEND", "cpp") pins, and cleaned up what they orphaned — the runtime_v2_mod import and the monkeypatch fixture parameter on the two tests that no longer use it. The SM100, speculative-mode and estimation checks are all retained as you asked.
47 passed in that file (was 4 failing).
This one is squarely my miss: my post-removal sweep greps module paths rather than symbol names, so a _BACKEND attribute reference in a directory I was not running never surfaced. Thanks for catching it.
ce691ea to
aad3c39
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #75485 [ run ] triggered by Bot. Commit: |
|
PR_Github #75485 [ run ] completed with state
|
aad3c39 to
c26cbfd
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #75514 [ run ] triggered by Bot. Commit: |
|
PR_Github #75514 [ run ] completed with state
|
c26cbfd to
83cb8a1
Compare
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #75528 [ run ] triggered by Bot. Commit: |
|
PR_Github #75529 [ run ] triggered by Bot. Commit: |
|
PR_Github/19154-83cb8a1 #75528 was force-killed by a newer pipeline run. |
|
PR_Github #75529 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #75610 [ run ] triggered by Bot. Commit: |
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 <lowsfer@users.noreply.github.com>
…atistics
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 <lowsfer@users.noreply.github.com>
83cb8a1 to
4acacda
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #75626 [ run ] triggered by Bot. Commit: |
|
PR_Github #75610 [ run ] completed with state |
This branch has not been deployed
9 participants
Summary
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, andis 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 +
rawrefbuild pipeline that existed only to make it fast enough tomatter.
tensorrt_llm/runtime/kv_cache_manager_v2/is now a re-export shim over the nanobindmodule plus the backend-agnostic
_introspectiondispatcher: 36 tracked files down to 4,net −13k lines.
Commits
behaviour that only the deleted Python tests guarded.
New bindings
Only one consumer needed native support. The KV-aware router's V2 hashing is expressed with
the existing
sequence_to_blockchain_keys(every caller chains from a reuse-scope root, andits first yielded pair is that root), so
v2_sha256_block_hasheris gone and no hashingbinding was added. Equivalence was verified against a reference implementation of the old
per-chunk chaining across 3 block sizes x 2 algorithms x salted/unsalted -- byte-identical --
and against golden digests captured from the Python hasher before deleting it.
The native disaggregated bounce buffer did need something:
PooledPhysMemAllocatorandVirtMemwrap the existingcudaVirtMem.{h,cpp}. They register on the_introspectionsubmodule, not the package surface, because they carry no stability promise -- the package
namespace is the stable surface. Three members total (
device_id,address,destroy); thepackage's exported API shrinks from 74 to 72 symbols.
BREAKING
TLLM_KV_CACHE_MANAGER_V2_BACKEND=pythonno longer exists.build_wheel.py --mypycandTRTLLM_ENABLE_MYPYCare gone, along withsetup_mypyc.py,the
rawrefC extension and thesetup.pypackaging surgery they required.kv_cache_config.kv_events_config) are dropped. The sink isduck-typed Python and the C++ radix tree calls its sink natively, so there is no live
path.
validate_streaming_supportnow rejects the config and points at the buffered pathvia
event_buffer_max_size. The interface is kept as a stub and its tests are skippedrather than deleted, so a native sink can restore it later.
Test coverage
Deleting the Python implementation orphaned
test_kv_cache_stats_life_cycles.py, whichdrove the Python page-movement recorders through a duck-typed stand-in. That behaviour —
SSM/recurrent life cycles appearing in iteration stats, with global cache-hit counters
staying attention-only — was fixed in both backends by #17447, but only ever tested in
Python, and no C++ test builds an SSM life cycle at all.
The second commit closes that with a hybrid attention + SSM fixture (the first in the C++
suite) and three cases: offload/onboard per life cycle, the attention-only alloc-counter
guard, and host drops. Each was confirmed to fail when the life-cycle filters are restored
in
KvCache::_recordMigratedSlots/_recordDroppedPages— a regression lock that cannotdetect the regression is worthless.
Host-drop coverage is new for attention too; nothing asserted it against a live manager
before.
Verification
kvCacheManagerV2StatsTest: 13/13 (10 pre-existing + 3 new)tests/unittest/kv_cache_manager_v2_tests/: 255 passed, 20 skippedtests/unittest/_torch/executor/kv_cache/: 788 passedexecutor/test_stats_serializer.py,disaggregated/test_router.py: greenAll run on a B200. Note the local dev box is an H100 while
cpp/buildwas configuredCUDA_ARCHITECTURES=100-real, which aborts on kernel launch — an artefact of the buildconfig, not a code defect, and it reproduces identically on unmodified
main.Dev Engineer Review
KVCacheManagerV2backend, backend selection,rawref, mypyc build support, and Python implementation modules.PooledPhysMemAllocatorandVirtMem.event_buffer_max_size.sequence_to_blockchain_keys. Verify downstream hash compatibility.QA Engineer Review
l0_h100.yml,l0_b200.yml,l0_cpu.yml, andl0_a10.yml. Coverage verdict: needs follow-up.Per-File QA Perspective
.gitignore: Removes Python and mypyc artifact exclusions. Verify obsolete artifacts cannot enter source or packaging workflows.AssertionError.PooledPhysMemAllocatorandVirtMem. Verify construction, lifetime retention, address access, device ID access, and destruction._introspectionexports and remove backend-specific validation. Verify imports, disaggregation allocation, and NVFP4 validation.test-dbsuites.