Skip to content

[TRTLLM-14846][chore] Group KV Cache managers and reunite the Disaggregation transceiver halves #17899

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lori-ren merged 20 commits into from
Sep 5, 2026
Merged

[TRTLLM-14846][chore] Group KV Cache managers and reunite the Disaggregation transceiver halves #17899

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
cb8ec0f
[TRTLLM-14846][chore] Group KV Cache managers and reunite the Disaggr…
lori-ren Aug 18, 2026
70edebe
[TRTLLM-14846][chore] Move the KV Cache and Disaggregation tests alon…
lori-ren Aug 21, 2026
d8573f3
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Aug 21, 2026
2e03519
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Aug 24, 2026
38bd6c5
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Aug 24, 2026
63a36e8
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Aug 24, 2026
151c21c
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Aug 30, 2026
a5587c1
[TRTLLM-14846][chore] Keep the owners of the executor tests the colla…
lori-ren Aug 30, 2026
d6658d7
[TRTLLM-14846][fix] Re-anchor the two imports the move left pointing …
lori-ren Aug 30, 2026
a98373b
[TRTLLM-14846][chore] Move the KV Cache and Disaggregation tests upst…
lori-ren Aug 31, 2026
30d9256
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Sep 4, 2026
d62ab8e
[TRTLLM-14846][fix] Re-anchor the imports upstream added at the old m…
lori-ren Sep 4, 2026
a2c66a8
[TRTLLM-14846][chore] Move the Disaggregation tests upstream added af…
lori-ren Sep 4, 2026
1fe0b08
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Sep 4, 2026
10871ae
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Sep 4, 2026
dd33567
[TRTLLM-14846][chore] Merge main into the branch
lori-ren Sep 4, 2026
3ecec7f
[TRTLLM-14846][fix] Re-anchor the imports Qwen3.8-Flash-Next added at…
lori-ren Sep 4, 2026
50302ae
[TRTLLM-14846][chore] Merge main into the branch
lori-ren Sep 4, 2026
3c66d8a
[TRTLLM-14846][chore] Merge main into the branch
lori-ren Sep 5, 2026
d77955c
[TRTLLM-14846][chore] Resolve conflicts with main
lori-ren Sep 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 4 additions & 10 deletions .github/CODEOWNERS
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,15 +224,11 @@
/cpp/tests/unit_tests/batch_manager/kvCache* @NVIDIA/trt-llm-kv-cache-manager-devs
/cpp/tests/unit_tests/batch_manager/radixBlockTree* @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/kv_cache_compression @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/pyexecutor/kv_cache @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/pyexecutor/resource_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/runtime/kv_cache_manager_v2 @NVIDIA/trt-llm-kv-cache-manager-devs
/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tests/unittest/_torch/executor/test_kv_cache* @NVIDIA/trt-llm-kv-cache-manager-devs
/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tests/unittest/_torch/executor/kv_cache @NVIDIA/trt-llm-kv-cache-manager-devs
/tests/unittest/_torch/executor/test_kvcache_aware_router.py @NVIDIA/trt-llm-kv-cache-manager-devs
Comment thread
lori-ren marked this conversation as resolved.
/tests/unittest/_torch/executor/test_mamba_cache_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tests/unittest/kv_cache_manager_v2_tests @NVIDIA/trt-llm-kv-cache-manager-devs

# ===== DISAGGREGATED SERVING =====
Expand All @@ -249,7 +245,6 @@
/examples/disaggregated @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-doc-owners
/examples/disaggregated/slurm/benchmark @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-perf-devs
/tensorrt_llm/_torch/disaggregation @NVIDIA/trt-llm-disagg-devs
/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @NVIDIA/trt-llm-disagg-devs
/tensorrt_llm/disaggregated_params.py @NVIDIA/trt-llm-disagg-devs
/tensorrt_llm/serve/openai_disagg_server.py @NVIDIA/trt-llm-disagg-devs
# Disagg tests: sole-owned. Owners on one line are OR, so co-owning made the disagg
Expand All @@ -259,7 +254,7 @@
/tests/integration/defs/stress_test/disagg_cancel @NVIDIA/trt-llm-disagg-devs
/tests/scripts/perf-sanity/disaggregated @NVIDIA/trt-llm-perf-devs @NVIDIA/trt-llm-disagg-devs
/tests/scripts/perf/disaggregated @NVIDIA/trt-llm-perf-devs @NVIDIA/trt-llm-disagg-devs
/tests/unittest/_torch/executor/*disagg* @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs
/tests/unittest/_torch/disaggregation @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs
Comment thread
lori-ren marked this conversation as resolved.
/tests/unittest/_torch/multimodal/*disagg* @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-disagg-devs
/tests/unittest/disaggregated @NVIDIA/trt-llm-disagg-devs
/tests/unittest/llmapi/*disagg* @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs
Expand Down Expand Up @@ -412,7 +407,6 @@
/examples/disaggregated @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-doc-owners
/examples/disaggregated/slurm/benchmark @NVIDIA/trt-llm-disagg-devs @NVIDIA/trtllm-bench-reviewers
/tensorrt_llm/disaggregated_params.py @NVIDIA/trt-llm-disagg-devs
/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @NVIDIA/trt-llm-disagg-devs
/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @NVIDIA/trt-llm-disagg-devs
/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @NVIDIA/trt-llm-disagg-devs
/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @NVIDIA/trt-llm-disagg-devs
Expand All @@ -437,7 +431,7 @@
/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @NVIDIA/trt-llm-kv-cache-manager-devs
/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/attention/backends/sparse/*/cache_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/pyexecutor/kv_cache @NVIDIA/trt-llm-kv-cache-manager-devs
/tensorrt_llm/_torch/pyexecutor/resource_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs
/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.h @NVIDIA/trt-llm-kv-cache-manager-devs
/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @NVIDIA/trt-llm-kv-cache-manager-devs
Expand Down
8 changes: 4 additions & 4 deletions .pre-commit-config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,7 @@ common-files: &common_files |
tensorrt_llm/_torch/debug/__init__.py |
tensorrt_llm/_torch/debug/debug_hook.py |
tensorrt_llm/_torch/device_mesh.py |
tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py |
tensorrt_llm/_torch/distributed/__init__.py |
tensorrt_llm/_torch/distributed/communicator.py |
tensorrt_llm/_torch/distributed/ops.py |
Expand Down Expand Up @@ -299,10 +300,9 @@ common-files: &common_files |
tensorrt_llm/_torch/pyexecutor/guided_decoder.py |
tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py |
tensorrt_llm/_torch/pyexecutor/handle_logits.py |
tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py |
tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py |
tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py |
tensorrt_llm/_torch/pyexecutor/llm_request.py |
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py |
tensorrt_llm/_torch/pyexecutor/model_engine.py |
tensorrt_llm/_torch/pyexecutor/model_loader.py |
tensorrt_llm/_torch/pyexecutor/py_executor.py |
Expand Down Expand Up @@ -924,6 +924,7 @@ legacy-files: &legacy_files |
tensorrt_llm/_torch/debug/__init__.py |
tensorrt_llm/_torch/debug/debug_hook.py |
tensorrt_llm/_torch/device_mesh.py |
tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py |
tensorrt_llm/_torch/distributed/__init__.py |
tensorrt_llm/_torch/distributed/communicator.py |
tensorrt_llm/_torch/distributed/ops.py |
Expand Down Expand Up @@ -1058,10 +1059,9 @@ legacy-files: &legacy_files |
tensorrt_llm/_torch/pyexecutor/guided_decoder.py |
tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py |
tensorrt_llm/_torch/pyexecutor/handle_logits.py |
tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py |
tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py |
tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py |
tensorrt_llm/_torch/pyexecutor/llm_request.py |
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py |
tensorrt_llm/_torch/pyexecutor/model_engine.py |
tensorrt_llm/_torch/pyexecutor/model_loader.py |
tensorrt_llm/_torch/pyexecutor/py_executor.py |
Expand Down
8 changes: 4 additions & 4 deletions examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,13 +53,13 @@
import tensorrt_llm.bindings
import tensorrt_llm.bindings.executor as trtllm
from tensorrt_llm import DisaggregatedParams
from tensorrt_llm._torch.distributed import Distributed
from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector
from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import (
from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import (
create_kv_cache_transceiver,
maybe_enable_fabric_memory_for_python_transceiver,
)
from tensorrt_llm._torch.distributed import Distributed
from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, LlmRequestType
from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, CacheTransceiverConfig
Expand Down
4 changes: 2 additions & 2 deletions legacy-files.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,7 @@ tensorrt_llm/_torch/cute_dsl_utils.py
tensorrt_llm/_torch/debug/__init__.py
tensorrt_llm/_torch/debug/debug_hook.py
tensorrt_llm/_torch/device_mesh.py
tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py
tensorrt_llm/_torch/distributed/__init__.py
tensorrt_llm/_torch/distributed/communicator.py
tensorrt_llm/_torch/moe/fused_moe/communication/moe_alltoall.py
Expand Down Expand Up @@ -291,10 +292,9 @@ tensorrt_llm/_torch/pyexecutor/grammar_matcher.py
tensorrt_llm/_torch/pyexecutor/guided_decoder.py
tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py
tensorrt_llm/_torch/pyexecutor/handle_logits.py
tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py
tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py
tensorrt_llm/_torch/pyexecutor/llm_request.py
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/pyexecutor/model_loader.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,7 @@ exclude = [
"tensorrt_llm/_torch/debug/__init__.py",
"tensorrt_llm/_torch/debug/debug_hook.py",
"tensorrt_llm/_torch/device_mesh.py",
"tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py",
"tensorrt_llm/_torch/distributed/__init__.py",
"tensorrt_llm/_torch/distributed/communicator.py",
"tensorrt_llm/_torch/distributed/ops.py",
Expand Down Expand Up @@ -348,10 +349,9 @@ exclude = [
"tensorrt_llm/_torch/pyexecutor/guided_decoder.py",
"tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py",
"tensorrt_llm/_torch/pyexecutor/handle_logits.py",
"tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py",
"tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py",
"tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py",
"tensorrt_llm/_torch/pyexecutor/llm_request.py",
"tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py",
"tensorrt_llm/_torch/pyexecutor/model_engine.py",
"tensorrt_llm/_torch/pyexecutor/model_loader.py",
"tensorrt_llm/_torch/pyexecutor/py_executor.py",
Expand Down
4 changes: 2 additions & 2 deletions ruff-legacy.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,7 @@ include = [
"tensorrt_llm/_torch/debug/__init__.py",
"tensorrt_llm/_torch/debug/debug_hook.py",
"tensorrt_llm/_torch/device_mesh.py",
"tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py",
"tensorrt_llm/_torch/distributed/__init__.py",
"tensorrt_llm/_torch/distributed/communicator.py",
"tensorrt_llm/_torch/distributed/ops.py",
Expand Down Expand Up @@ -308,10 +309,9 @@ include = [
"tensorrt_llm/_torch/pyexecutor/guided_decoder.py",
"tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py",
"tensorrt_llm/_torch/pyexecutor/handle_logits.py",
"tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py",
"tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py",
"tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py",
"tensorrt_llm/_torch/pyexecutor/llm_request.py",
"tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py",
"tensorrt_llm/_torch/pyexecutor/model_engine.py",
"tensorrt_llm/_torch/pyexecutor/model_loader.py",
"tensorrt_llm/_torch/pyexecutor/py_executor.py",
Expand Down
4 changes: 2 additions & 2 deletions tensorrt_llm/_torch/attention/backends/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,8 @@

from ...memory_buffer_utils import Buffers
from ...metadata import KVCacheParams
from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
from ...pyexecutor.mamba_cache_manager import BaseMambaCacheManager
from ...pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
from ...pyexecutor.kv_cache.mamba_cache_manager import BaseMambaCacheManager
from ...pyexecutor.resource_manager import KVCacheManager
from ...pyexecutor.trace_log_utils import log_tensor_size
from ...utils import get_model_extra_attrs
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import torch

from tensorrt_llm._torch.pyexecutor import llm_request
from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import GPU_LEVEL, KVCacheManagerV2
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import GPU_LEVEL, KVCacheManagerV2
from tensorrt_llm._utils import (
TensorWrapper,
convert_to_torch_tensor,
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

import tensorrt_llm
import tensorrt_llm.bindings
from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role
from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager, get_pp_layers
from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes
from tensorrt_llm.bindings import DataType
Expand Down
3 changes: 1 addition & 2 deletions tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import torch

from tensorrt_llm._torch.disaggregation.resource.page import MapperKind
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role
from tensorrt_llm._utils import (
TensorWrapper,
binding_to_torch_dtype,
Expand All @@ -42,8 +43,6 @@
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 .....pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role


class MiniMaxM3SparseIndexCache:
"""Plain-tensor side cache for the M3 sparse index branch.
Expand Down
4 changes: 2 additions & 2 deletions tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@

import torch

from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManagerV2
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role
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
Expand Down
3 changes: 2 additions & 1 deletion tensorrt_llm/_torch/attention/backends/trtllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1915,7 +1915,8 @@ def forward(
else:
forward_args.fmha_scheduler_counter.zero_()
assert forward_args.latent_cache is not None
from ...pyexecutor.mamba_cache_manager import BaseMambaCacheManager
from ...pyexecutor.kv_cache.mamba_cache_manager import \
BaseMambaCacheManager

# Hybrid (mamba/masked-layer) KV managers take the graph-safe
# append; the same predicate interface.py uses to detect hybrid
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
from transformers.activations import ACT2FN
from transformers.utils import ModelOutput

from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager
from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager

from ...distributed.common import broadcast
from ...shim.interface import CachedSequenceInterface
Expand Down
10 changes: 5 additions & 5 deletions tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,15 +21,15 @@
from torch._prims_common import DeviceLikeType

from tensorrt_llm._torch.autotuner import AutoTuner
from tensorrt_llm._torch.distributed import Distributed
from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID
from tensorrt_llm._torch.pyexecutor.guided_decoder import GuidedDecoder
from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import (
from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import (
AttentionTypeCpp,
create_kv_cache_transceiver,
)
from tensorrt_llm._torch.distributed import Distributed
from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID
from tensorrt_llm._torch.pyexecutor.guided_decoder import GuidedDecoder
from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import BaseMambaCacheManager
from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length
from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import BaseMambaCacheManager
from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine, PyTorchModelEngine
from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
from tensorrt_llm._torch.pyexecutor.py_executor_creator import get_guided_decoding_config
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/auto_deploy/shim/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@

if TRTLLM_AVAILABLE:
import tensorrt_llm.bindings
from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import (
MambaHybridCacheManager,
MixedMambaHybridCacheManager,
)
Expand Down
13 changes: 6 additions & 7 deletions ..._torch/pyexecutor/kv_cache_transceiver.py → ...ch/disaggregation/kv_cache_transceiver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,12 @@
CacheTransceiverConfig)
from tensorrt_llm.mapping import Mapping

from .config_utils import resolve_cache_transceiver_config
from .llm_request import LlmRequest
from .mamba_cache_manager import (BaseMambaCacheManager,
CppMambaHybridCacheManager,
MambaHybridCacheManagerV2,
MixedMambaHybridCacheManager)
from .resource_manager import KVCacheManager
from ..pyexecutor.config_utils import resolve_cache_transceiver_config
from ..pyexecutor.kv_cache.mamba_cache_manager import (
BaseMambaCacheManager, CppMambaHybridCacheManager,
MambaHybridCacheManagerV2, MixedMambaHybridCacheManager)
from ..pyexecutor.llm_request import LlmRequest
from ..pyexecutor.resource_manager import KVCacheManager

CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver
AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

import numpy as np

from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager

Expand Down
4 changes: 2 additions & 2 deletions tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@
compute_layer_byte_ranges,
get_physical_pool,
)
from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role
from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import (
MambaHybridCacheManager,
MambaHybridCacheManagerV2,
)
Expand Down
17 changes: 10 additions & 7 deletions tensorrt_llm/_torch/disaggregation/transceiver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,11 @@
WaitResult,
get_unique_rid,
)
from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import (
CtxTransferStatus,
GenTransferStatus,
KvCacheTransceiver,
)
from tensorrt_llm._torch.disaggregation.native.bounce import (
config_from_size as bounce_config_from_size,
)
Expand All @@ -45,17 +50,15 @@
from tensorrt_llm._torch.disaggregation.resource.page import CacheKind
from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool
from tensorrt_llm._torch.distributed.communicator import Distributed
from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2
from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import (
CtxTransferStatus,
GenTransferStatus,
KvCacheTransceiver,
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import (
BlockReusePolicy,
KVCacheManagerV2,
)
from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import (
MambaHybridCacheManager,
MambaHybridCacheManagerV2,
)
from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
from tensorrt_llm._utils import nvtx_range
from tensorrt_llm.bindings import LlmRequestState
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import triton
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS

from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
from tensorrt_llm._utils import prefer_pinned
from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import (
copy_batch_block_offsets_to_device,
Expand All @@ -37,7 +38,6 @@

from ...distributed import allgather
from ...modules.top_k import TopK, TopKImplementation
from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
from ...pyexecutor.llm_request import LlmRequestState
from ...pyexecutor.resource_manager import KVCacheCompressionManager
from ...utils import next_positive_power_of_2
Expand Down
Loading
Loading