[TRTLLM-14846][chore] Group KV Cache managers and reunite the Disaggregation transceiver halves - #17899
Conversation
013ad78 to
7241a76
Compare
7241a76 to
5cb994f
Compare
Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
|
PR_Github #71638 [ run ] triggered by Bot. Commit: |
|
PR_Github #71628 [ run ] completed with state |
|
PR_Github #71638 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #71658 [ run ] triggered by Bot. Commit: |
|
PR_Github #71658 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #71668 [ run ] triggered by Bot. Commit: |
Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #71685 [ run ] triggered by Bot. Commit: |
|
PR_Github #71668 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
PR_Github #71718 [ run ] triggered by Bot. Commit: |
|
PR_Github #71685 [ run ] completed with state |
|
PR_Github #71718 [ run ] completed with state |
…ssion Resolves the conflict GitHub reported as mergeable_state "dirty": main moved tests/unittest/_torch/executor/test_benchmark_disagg.py to tests/unittest/_torch/disaggregation/test_benchmark_disagg.py in NVIDIA#17899 while this branch modified the old path. Rename detection (89% similarity) carries the one-hunk MockPadDummyExecutor change to the new location; no other path conflicted, textually or semantically. Checks after the merge, since main rewrote large parts of the files this branch touches (model_engine.py, py_executor.py, _util.py, eagle3.py, mtp.py): the AST guard over get_spec_resource_manager reports no manager built without num_seq_slots; all six slot-pool managers still accept num_seq_slots with a None default; no reference to the pre-rename should_enable_disagg_adp_overlap_headroom survives; yapf and ruff format both report the touched files clean under their respective legacy-files.txt families; and ruff check produces a lint set identical to main's. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
…he moves Two directory moves landed while the PRs below were in review, so their imports point at packages that no longer exist and every executor import fails at collection time: NVIDIA#17968 moved _torch/attention_backend -> _torch/attention/backends NVIDIA#17899 moved pyexecutor/kv_cache_manager_v2 -> pyexecutor/kv_cache/ Retarget the nine stale references. No behavior change; each symbol is re-exported unchanged from its new location. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… backends NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, was in flight at the same time and still imports the old path, so main now fails at: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable and every unit test that touches an attention backend fails collection. Point the three fmha modules and the three FMHA unit tests at pyexecutor.kv_cache.kv_cache_manager_v2, matching the sibling backends (interface.py, sparse/dsa/cache_manager.py). Adds a CPU-only import guard so a future move of this module is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…ak main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. NVIDIA#17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. NVIDIA#18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…ak main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. NVIDIA#17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. NVIDIA#18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py, the visual-gen FlashInfer backend and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…ak main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. NVIDIA#17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. NVIDIA#18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py, the visual-gen FlashInfer backend and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com> (cherry picked from commit 0c221e3) Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…ak main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. NVIDIA#17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. NVIDIA#18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py, the visual-gen FlashInfer backend and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com> (cherry picked from commit 0c221e3) Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…egation transceiver halves
Dev Engineer Review
_torch/pyexecutor/kv_cache/._torch/disaggregation/.KVCacheV2Scheduler.QA Engineer Review
AsyncTransferManagertransfers and state handling.IndexMapperslot release.tests/integration/test_lists/test-db/l0_a10.yml.unittest/_torch/disaggregationto:l0_b300.ymll0_cpu.ymll0_dgx_b300.ymll0_gb300_multi_gpus.ymll0_h100.ymltests/integration/defs/.test_durationsandwaives.txt.Description
Pure relocation — no production code changed, only import paths.
_torch/pyexecutor/kv_cache_manager_v2.py_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py_torch/pyexecutor/mamba_cache_manager.py_torch/pyexecutor/kv_cache/mamba_cache_manager.py_torch/pyexecutor/kv_cache_transceiver.py_torch/disaggregation/kv_cache_transceiver.py_torch/disaggregation/transceiver.py_torch/disaggregation/transceiver_v2.pyNew package
_torch/pyexecutor/kv_cache/groups the two KV cache managers. Thetransceiver's two halves — the
KvCacheTransceiverinterface plus factory, and theKvCacheTransceiverV2implementation — now sit in one package, and the file namesdistinguish them. The factory's deferred import is kept verbatim: the two modules still
form a cycle, and package membership does not change the module graph.
CODEOWNERS: two exact-file rules for
kv_cache_transceiver.pyretire (the file now fallsunder the
_torch/disaggregationdirectory rule); the KV cache rules collapse into onepyexecutor/kv_cachedirectory rule. No file changes owner.The second commit moves the matching tests so they sit with the modules they cover — 13
to
tests/unittest/_torch/executor/kv_cache/, 4 totests/unittest/_torch/disaggregation/— along with the CODEOWNERS rules and test-list entries that name them.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.