diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index c9acc9b61ece..5eb2c47c505d 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -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 -/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 ===== @@ -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 @@ -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 /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 @@ -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 @@ -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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0f7bbe4b3a2e..f6f71a361e6a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -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 | @@ -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 | @@ -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 | @@ -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 | diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index c42433190658..d54be2f24667 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -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 diff --git a/legacy-files.txt b/legacy-files.txt index 3c1369c08750..37fd59a51b0e 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -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 @@ -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 diff --git a/pyproject.toml b/pyproject.toml index c15831866268..27d66cd7f84e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", @@ -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", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index ab08f675c7b0..3a566963ac9e 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -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", @@ -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", diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 67de4b723580..7d11cdb28f8d 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -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 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 a838717cadd6..b1dba9188127 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 @@ -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, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py index 681791a9d986..b3bdfef69a2a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py @@ -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 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 3cf626bb9046..043acaaaa73b 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 @@ -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, @@ -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. 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 52fdcccd6930..7029f111cfdf 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -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 diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 21ca568ab81d..5a5715b1d273 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -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 diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py index 346cbcf0b490..2c1a5e9fce28 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py @@ -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 diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index bbd4009b6d7e..b781799df9d5 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -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 diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index d3250137e532..ed8eb3c7170c 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -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, ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py similarity index 98% rename from tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py rename to tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py index 607e512a75eb..b83758bd4c53 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py @@ -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 diff --git a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py index 142161093f12..850b9dd401e1 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py +++ b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py @@ -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 diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 77a35a904e07..7d7490146a6a 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -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, ) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 62107d31e447..8bd406d879dd 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -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, ) @@ -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 diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 468c26eff699..e0f32806ddd6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -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, @@ -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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 82876b40773b..cd56117042ff 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -42,6 +42,9 @@ from tensorrt_llm.quantization import QuantAlgo from ..attention.backends import get_sparse_attn_kv_cache_manager +from ..disaggregation.kv_cache_transceiver import ( + AttentionTypeCpp, create_kv_cache_transceiver, + maybe_enable_fabric_memory_for_python_transceiver) from ..hostfunc import set_low_latency_dispatch from ..model_config import ModelConfig from ..models.modeling_multimodal_mixin import MultimodalModelMixin @@ -57,16 +60,13 @@ from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder -from .kv_cache_manager_v2 import KVCacheManagerV2 -from .kv_cache_transceiver import ( - AttentionTypeCpp, create_kv_cache_transceiver, - maybe_enable_fabric_memory_for_python_transceiver) +from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager, + CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, + MixedMambaHybridCacheManager, + use_py_mamba_cache_manager) from .llm_request import ExecutorResponse, LlmRequestState -from .mamba_cache_manager import (BaseMambaCacheManager, - CppMambaHybridCacheManager, - MambaHybridCacheManagerV2, - MixedMambaHybridCacheManager, - use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor from .resource_manager import (KVCacheCompressionManager, KVCacheManager, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/__init__.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/__init__.py @@ -0,0 +1,14 @@ +# 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. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py similarity index 99% rename from tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py rename to tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 7b5dcd955d30..d4a93d8de9e5 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -27,11 +27,16 @@ from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp +from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import ( TensorWrapper, + binding_to_torch_dtype, convert_to_torch_tensor, get_size_in_bytes, + mpi_rank, + nvtx_range, prefer_pinned, + str_dtype_to_torch, ) from tensorrt_llm.bindings.internal.batch_manager import KvCacheIterationStats, KvCacheStats from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( @@ -39,6 +44,8 @@ copy_batch_block_offsets_to_device, ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_hash import get_effective_kv_cache_event_hash_algo from tensorrt_llm.runtime.kv_cache_manager_v2 import ( _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, @@ -79,21 +86,17 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfMemoryError as KVCacheOutOfMemoryError from tensorrt_llm.sampling_params import SamplingParams -from ..._utils import binding_to_torch_dtype, mpi_rank, nvtx_range, str_dtype_to_torch -from ...logger import logger -from ...mapping import Mapping -from ..utils import maybe_compile -from .config_utils import uses_vswa_kv_cache_layout -from .connectors.kv_cache_connector import KvCacheConnectorManager -from .kv_cache_stats import ( +from ..config_utils import uses_vswa_kv_cache_layout +from ..connectors.kv_cache_connector import KvCacheConnectorManager +from ..kv_cache_stats import ( KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, KVCacheV2PoolGroupIterationStats, KVCacheV2SsmLifeCycleIterationStats, KVCacheV2SsmSnapshotIterationStats, ) -from .llm_request import LlmRequest, LlmRequestState, get_draft_token_length -from .resource_manager import ( +from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length +from ..resource_manager import ( BaseResourceManager, CacheTypeCpp, DataType, @@ -105,7 +108,7 @@ get_pp_layers, request_context, ) -from .scheduler import ScheduledRequests +from ..scheduler import ScheduledRequests if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata @@ -362,12 +365,12 @@ def _estimate_swa_cache_size( def _get_dflash_generation_kv_capacity_headroom(spec_config) -> Optional[int]: """DFlash KV capacity reserved ahead of committed history.""" - from ..speculative.interface import SpeculativeDecodingMode + from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode if spec_config is None or spec_config.spec_dec_mode != SpeculativeDecodingMode.DFLASH: return None - from ..speculative import get_num_extra_kv_tokens + from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens return get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step @@ -971,7 +974,7 @@ def __init__( self.max_batch_size = max_batch_size self.max_num_tokens = max_num_tokens self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 - from ..speculative import get_num_extra_kv_tokens + from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) # Mirror V1: expose max_draft_len so the native disagg AuxBuffer diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py similarity index 99% rename from tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py rename to tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 6348db3d78cd..71ed3d3f3fb5 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -30,7 +30,7 @@ from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig from tensorrt_llm.sampling_params import SamplingParams -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, KVCacheManagerV2, Role) from tensorrt_llm._torch.pyexecutor.kv_cache_stats import \ KVCacheV2IterationStatsReport @@ -442,7 +442,7 @@ def has_kda_replay_caches(self) -> bool: def commit_conv_window(self, slot_indices: torch.Tensor, conv_pool: torch.Tensor) -> None: """Seed the KDA replay caches from the live ``W - 1`` conv pool.""" - from ..modules.kimi_kda._kda_kernels import \ + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ copy_kda_replay_conv_window copy_kda_replay_conv_window( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 09792b233567..f23325af2101 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -90,11 +90,12 @@ mm_encoder_cache_enabled, setup_mm_encoder_attn_metadata) from .guided_decoder import CapturableGuidedDecoder -from .kv_cache_manager_v2 import KVCacheManagerV2 +from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager, + MambaHybridCacheManager) from .layerwise_nvtx_marker import LayerwiseNvtxMarker from .llm_request import (LlmRequest, LlmRequestState, get_draft_token_length, get_multimodal_embedding_lengths) -from .mamba_cache_manager import BaseMambaCacheManager, MambaHybridCacheManager from .model_loader import ModelLoader, _construct_checkpoint_loader from .resource_manager import (BaseResourceManager, KVCacheManager, PeftCacheManager, ResourceManager, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cebdbc942c27..339f7ecb26bc 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -60,6 +60,8 @@ NoopDisaggCoordinator) from ..disaggregation.executor.pp_termination import DisaggPPTerminationHandler from ..disaggregation.executor.transfer_manager import AsyncTransferManager +from ..disaggregation.kv_cache_transceiver import ( + KvCacheTransceiver, is_disagg_inflight_cancel_enabled) from ..distributed import Distributed from ..distributed.communicator import ReduceOp from ..models.modeling_multimodal_mixin import \ @@ -81,18 +83,16 @@ from .handle_logits import HandleLogits from .hang_detector import (HangDetector, hard_kill_on_rank_crash, propagate_hard_kill, start_rank_crash_kill_watchdog) -from .kv_cache_manager_v2 import KVCacheManagerV2 +from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager, + MixedMambaHybridCacheManager) from .kv_cache_stats import append_kv_cache_iteration_stats -from .kv_cache_transceiver import (KvCacheTransceiver, - is_disagg_inflight_cancel_enabled) from .llm_request import (ATTENTION_DP_DUMMY_REQUEST_ID, MAX_SPEC_DECODE_POSITIONS, ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, MultimodalEncoderRequestError, get_draft_token_length, initialize_multimodal_encoder_request, is_multimodal_encoder_ready) -from .mamba_cache_manager import (BaseMambaCacheManager, - MixedMambaHybridCacheManager) from .model_engine import ModelEngine from .perf_metrics_manager import PerfMetricsManager from .pp_utils import PPCommTag diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index e72974db9b82..6077f62d1ae8 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -70,7 +70,7 @@ from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, KvCacheCompressionConfig) - from .kv_cache_manager_v2 import KVCacheManagerV2 + from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 BlocksPerWindow = Dict[int, Tuple[ int, @@ -1447,7 +1447,7 @@ def update_resources(self, # applies to it. if self.kv_cache_type != CacheTypeCpp.CROSS: if not self.is_draft: - from .kv_cache_manager_v2 import \ + from .kv_cache.kv_cache_manager_v2 import \ _update_kv_cache_draft_token_location _update_kv_cache_draft_token_location(self, scheduled_batch, @@ -2819,7 +2819,7 @@ def bind_kv_cache_managers( draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None, ) -> None: """Bind the target and optional draft KVCMs after their construction.""" - from .kv_cache_manager_v2 import KVCacheManagerV2 + from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 if not isinstance(kv_cache_manager, KVCacheManagerV2): raise TypeError("KV-cache compression requires KVCacheManagerV2") diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index cf12b4b91a7e..b78f57694c8c 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -168,7 +168,7 @@ def __init__( self.max_num_requests = ( scheduler_capacity if scheduler_capacity is not None else max_batch_size ) - from ..kv_cache_manager_v2 import KVCacheManagerV2 + from ..kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 assert isinstance(kv_cache_manager, KVCacheManagerV2), ( f"KVCacheV2Scheduler requires KVCacheManagerV2, got {type(kv_cache_manager).__name__}" @@ -906,7 +906,7 @@ def _try_schedule_cross_context(self, req: LlmRequest) -> ScheduleAction: req_tokens = self._get_optional_encoder_output_len(req) if req_tokens is None: return ScheduleAction.SCHEDULED - from ..kv_cache_manager_v2 import KVCacheManagerV2 + from ..kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 if isinstance(self.cross_kv_cache_manager, KVCacheManagerV2): if not self._try_schedule_cross_context_v2( diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 39f68365509f..ddc358a5ef39 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -25,7 +25,7 @@ from tensorrt_llm.mapping import Mapping from ..attention.backends import AttentionMetadata -from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.resource_manager import BaseResourceManager from .accept_stats import maybe_create_recorder from .dflash_attention import get_dflash_trtllm_gen_ops, validate_dflash_trtllm_gen_runtime diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 44c890fe3d54..17fd8efbc962 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -15,8 +15,8 @@ from ..attention.backends import AttentionMetadata from ..attention.backends.flashinfer import FlashInferAttentionMetadata from ..model_config import ModelConfig +from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest -from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.scheduler import ScheduledRequests from .interface import (INVALID_PROMPT_LOOKAHEAD_TOKEN, SpecMetadata, diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 87258b133e17..685143e2d975 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -9,8 +9,8 @@ from tensorrt_llm._utils import prefer_pinned from ..attention.backends import AttentionMetadata +from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest -from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.scheduler import ScheduledRequests from .interface import SpecMetadata, SpecWorkerBase diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 11adc3dfafb6..86895f4c3558 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -20,7 +20,7 @@ import torch import triton -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from tensorrt_llm._utils import get_sm_version, nvtx_range from tensorrt_llm.mapping import Mapping diff --git a/tensorrt_llm/_torch/speculative/sa_worker.py b/tensorrt_llm/_torch/speculative/sa_worker.py index d2c6f0ec0400..148808298f7c 100644 --- a/tensorrt_llm/_torch/speculative/sa_worker.py +++ b/tensorrt_llm/_torch/speculative/sa_worker.py @@ -30,7 +30,7 @@ from tensorrt_llm._utils import prefer_pinned -from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from .interface import SpecMetadata, SpecWorkerBase from .suffix_automaton import SuffixAutomatonManager diff --git a/tensorrt_llm/bench/tuning/dataclasses.py b/tensorrt_llm/bench/tuning/dataclasses.py index 1cfb942c7eec..fc0283e1d835 100755 --- a/tensorrt_llm/bench/tuning/dataclasses.py +++ b/tensorrt_llm/bench/tuning/dataclasses.py @@ -22,7 +22,7 @@ # Model types that use Multi-Head Latent Attention (MLA). The runtime KV-cache # formula for these is a single compressed head with # head_dim = kv_lora_rank + qk_rope_head_dim, kv_factor = 1 -- matching the MLA -# branch in tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py. +# branch in tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py. _MLA_MODEL_TYPES = frozenset( { "deepseek_v2", diff --git a/tensorrt_llm/bench/tuning/heuristics.py b/tensorrt_llm/bench/tuning/heuristics.py index 1a1d4af1ad59..5d5391679ae5 100755 --- a/tensorrt_llm/bench/tuning/heuristics.py +++ b/tensorrt_llm/bench/tuning/heuristics.py @@ -74,7 +74,7 @@ def calc_engine_setting( if model_config.is_mla(): # MLA stores a single compressed KV entry per token: # head_dim = kv_lora_rank + qk_rope_head_dim, kv_factor = 1. - # Mirrors tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py and + # Mirrors tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py and # resource_manager.py, so the bench heuristic agrees with the runtime. mla_head_dim = model_config.kv_lora_rank + model_config.qk_rope_head_dim gb_per_token = ( diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 2e7237d41ce3..2a04d17d9c2d 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -53,7 +53,7 @@ is_llm_response) if TYPE_CHECKING: - from .._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver + from .._torch.disaggregation.kv_cache_transceiver import KvCacheTransceiver from ..disaggregated_params import DisaggregatedParams __all__ = [ diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 90a3e09c4ad7..8024f8e572c8 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -42,7 +42,7 @@ from .utils import IntraProcessQueue, ProcessPoolExecutorSession, RequestError if TYPE_CHECKING: - from .._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver + from .._torch.disaggregation.kv_cache_transceiver import KvCacheTransceiver from .proxy import GenerationExecutorProxy from .worker import GenerationExecutorWorker diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e385b17ee9c2..243047f8113c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4513,7 +4513,7 @@ def _to_pybind(self): # Legacy env-var selectors for the "DEFAULT" cache-transceiver backend, # ordered by priority. Single source of truth shared with -# _torch/pyexecutor/kv_cache_transceiver.py. +# _torch/disaggregation/kv_cache_transceiver.py. _CACHE_TRANSCEIVER_BACKEND_ENV_VARS = ( ("TRTLLM_USE_NIXL_KVCACHE", "NIXL"), ("TRTLLM_USE_UCX_KVCACHE", "UCX"), diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index ce9332bc00bb..0e448e1115a1 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -33,7 +33,7 @@ load_pretrained_config, unwrap_kimi_text_config, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MixedMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.model_loader import ( ModelLoader, _construct_checkpoint_loader, diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index ec972db3e704..f2d74e33949f 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -1019,11 +1019,11 @@ "unittest/_torch/debugger": 18.826091299677767, "unittest/_torch/distributed": 22.260023017902814, "unittest/_torch/executor": 187.15932672285572, - "unittest/_torch/executor/test_disagg_index_mapper_early_release.py": 23.67072709551657, + "unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py": 23.67072709551657, "unittest/_torch/executor/test_error_classification.py": 22.753568665377177, "unittest/_torch/executor/test_hang_detector_kill.py": 82.3269842791411, - "unittest/_torch/executor/test_kv_cache_compression_manager.py": 23.403027079303676, - "unittest/_torch/executor/test_kv_cache_v2_capacity_only.py": 22.63832879377432, + "unittest/_torch/executor/kv_cache/test_kv_cache_compression_manager.py": 23.403027079303676, + "unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py": 22.63832879377432, "unittest/_torch/executor/test_model_loader_mx.py": 19.900650019208605, "unittest/_torch/executor/test_overlap_scheduler.py": 562.8957540983606, "unittest/_torch/executor/test_resource_manager.py": 31.168673469387755, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 7893aeba6ef9..4224796b6b6a 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -30,9 +30,9 @@ l0_a10: - unittest/_torch/modeling/test_modeling_minicpmv4_6.py - unittest/_torch/modeling/test_multimodal_encoder_mixin.py - unittest/_torch/sampler/test_token_ban.py - - unittest/_torch/executor/test_disagg_index_mapper_early_release.py - - unittest/_torch/executor/test_kv_cache_compression_manager.py - - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py + - unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py + - unittest/_torch/executor/kv_cache/test_kv_cache_compression_manager.py + - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/executor/test_resource_manager.py - unittest/_torch/moe/test_communication_factory.py diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 7a6cf5709cb9..8e90364f6f20 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -35,6 +35,7 @@ l0_b300: - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/thop/serial - unittest/_torch/executor # 250s + - unittest/_torch/disaggregation # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_fused_add_rms_norm_quant.py - unittest/_torch/modules/test_fused_activation_quant.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 901403843578..14b83766f9b0 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -40,6 +40,7 @@ l0_cpu: - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py - unittest/_torch/distributed - unittest/_torch/executor + - unittest/_torch/disaggregation - unittest/_torch/peft - unittest/_torch/memory - unittest/_torch/modeling diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300.yml b/tests/integration/test_lists/test-db/l0_dgx_b300.yml index e288753f4b96..6b61ff6506bc 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b300.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b300.yml @@ -29,6 +29,7 @@ l0_dgx_b300: - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120) - unittest/_torch/executor + - unittest/_torch/disaggregation - unittest/_torch/modeling -k "modeling_llama" - unittest/_torch/modeling -k "modeling_gpt_oss" - unittest/_torch/multi_gpu_modeling -k "deepseek" diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index ea51564d4639..10146986412d 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -28,6 +28,7 @@ l0_gb300_multi_gpus: - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/thop/serial - unittest/_torch/executor + - unittest/_torch/disaggregation # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_fused_add_rms_norm_quant.py - unittest/_torch/modules/test_fused_activation_quant.py diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index fa749a2bea97..8fdaf6ab1d29 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -26,6 +26,7 @@ l0_h100: - unittest/_torch/compilation - unittest/_torch/debugger - unittest/_torch/executor + - unittest/_torch/disaggregation - unittest/_torch/peft - unittest/others/test_lora_manager.py - unittest/_torch/misc diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index b1b5f63c4bff..7be8b61e8243 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -235,14 +235,14 @@ def load_internal_apis(): import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm_executor from tensorrt_llm import DisaggregatedParams - from tensorrt_llm._torch.distributed import Distributed - from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class - 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.models.modeling_utils import get_registered_model_class + 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, diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 0553f6272191..b23358200999 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -31,7 +31,7 @@ from tensorrt_llm._torch.attention.backends.utils import create_attention, get_attention_backend from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._torch.metadata import KVCacheParams -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.resource_manager import KVCacheManager from tensorrt_llm._utils import str_dtype_to_torch, torch_dtype_to_binding from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 2eb78fe65d00..b3d26621cd76 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -61,7 +61,7 @@ from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.speculative.interface import ( prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index cec23ea0b17a..043e0b5fe010 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -20,7 +20,7 @@ ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import msa_paged_kv from tensorrt_llm._torch.attention.backends.sparse.registry import _resolve_minimax_m3_backend_cls -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.bindings import DataType from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig diff --git a/tests/unittest/_torch/attention/test_attention.py b/tests/unittest/_torch/attention/test_attention.py index 67ca6f83b4f4..3bf4c7db6c3b 100644 --- a/tests/unittest/_torch/attention/test_attention.py +++ b/tests/unittest/_torch/attention/test_attention.py @@ -14,7 +14,8 @@ from tensorrt_llm._torch.attention.backends.interface import \ PredefinedAttentionMask from tensorrt_llm._torch.metadata import KVCacheParams -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.resource_manager import KVCacheManager from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 1d9c41064013..1c4838e36d26 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -13,7 +13,8 @@ TrtllmAttentionMetadata from tensorrt_llm._torch.attention.backends.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams -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, LlmRequestState, SamplingConfig) diff --git a/tests/unittest/_torch/attention/test_vanilla_attention.py b/tests/unittest/_torch/attention/test_vanilla_attention.py index 1ef5f9101d1d..001ce71db2d2 100644 --- a/tests/unittest/_torch/attention/test_vanilla_attention.py +++ b/tests/unittest/_torch/attention/test_vanilla_attention.py @@ -13,7 +13,8 @@ from tensorrt_llm._torch.attention.backends.interface import \ PredefinedAttentionMask from tensorrt_llm._torch.metadata import KVCacheParams -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.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig diff --git a/tests/unittest/_torch/executor/fake_kv_cache_transceiver.py b/tests/unittest/_torch/disaggregation/fake_kv_cache_transceiver.py similarity index 99% rename from tests/unittest/_torch/executor/fake_kv_cache_transceiver.py rename to tests/unittest/_torch/disaggregation/fake_kv_cache_transceiver.py index 253960c2917b..529bb6b07d0a 100644 --- a/tests/unittest/_torch/executor/fake_kv_cache_transceiver.py +++ b/tests/unittest/_torch/disaggregation/fake_kv_cache_transceiver.py @@ -3,7 +3,7 @@ from typing import Dict, List, Optional -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import ( +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import ( CtxTransferStatus, GenTransferStatus, KvCacheTransceiver, diff --git a/tests/unittest/_torch/executor/test_async_transfer_manager.py b/tests/unittest/_torch/disaggregation/test_async_transfer_manager.py similarity index 100% rename from tests/unittest/_torch/executor/test_async_transfer_manager.py rename to tests/unittest/_torch/disaggregation/test_async_transfer_manager.py diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py similarity index 99% rename from tests/unittest/_torch/executor/test_benchmark_disagg.py rename to tests/unittest/_torch/disaggregation/test_benchmark_disagg.py index ca89f11502bc..d9fc28274587 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py @@ -30,7 +30,7 @@ import pytest -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import GenTransferStatus +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import GenTransferStatus from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import RequestScheduler, ScheduledRequests diff --git a/tests/unittest/_torch/executor/test_disagg_coordinator.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator.py similarity index 100% rename from tests/unittest/_torch/executor/test_disagg_coordinator.py rename to tests/unittest/_torch/disaggregation/test_disagg_coordinator.py diff --git a/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py b/tests/unittest/_torch/disaggregation/test_disagg_fill_gate_stall_bound.py similarity index 100% rename from tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py rename to tests/unittest/_torch/disaggregation/test_disagg_fill_gate_stall_bound.py diff --git a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py b/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py similarity index 100% rename from tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py rename to tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/disaggregation/test_disagg_inflight_cancel_gate.py similarity index 99% rename from tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py rename to tests/unittest/_torch/disaggregation/test_disagg_inflight_cancel_gate.py index 65119d242bcf..4612ca3bc126 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_inflight_cancel_gate.py @@ -20,17 +20,17 @@ import numpy as np import pytest -from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as transceiver_module -from tensorrt_llm._torch.pyexecutor import py_executor as executor_module -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import ( +from tensorrt_llm._torch.disaggregation import kv_cache_transceiver as transceiver_module +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import ( BindKvCacheTransceiver, CtxTransferStatus, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( +from tensorrt_llm._torch.pyexecutor import py_executor as executor_module +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( CppMambaHybridCacheManager, MambaHybridCacheManagerV2, ) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig diff --git a/tests/unittest/_torch/executor/test_disagg_loop_transcript.py b/tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py similarity index 100% rename from tests/unittest/_torch/executor/test_disagg_loop_transcript.py rename to tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py diff --git a/tests/unittest/_torch/executor/test_kv_cache_transceiver_contract.py b/tests/unittest/_torch/disaggregation/test_kv_cache_transceiver_contract.py similarity index 99% rename from tests/unittest/_torch/executor/test_kv_cache_transceiver_contract.py rename to tests/unittest/_torch/disaggregation/test_kv_cache_transceiver_contract.py index 5fe1df91ebb2..f506b7f9ac05 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_transceiver_contract.py +++ b/tests/unittest/_torch/disaggregation/test_kv_cache_transceiver_contract.py @@ -17,7 +17,7 @@ import pytest from fake_kv_cache_transceiver import FakeKvCacheTransceiver -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import ( +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import ( BindKvCacheTransceiver, CtxTransferStatus, GenTransferStatus, diff --git a/tests/unittest/_torch/executor/test_kv_cache_transceiver_creation.py b/tests/unittest/_torch/disaggregation/test_kv_cache_transceiver_creation.py similarity index 96% rename from tests/unittest/_torch/executor/test_kv_cache_transceiver_creation.py rename to tests/unittest/_torch/disaggregation/test_kv_cache_transceiver_creation.py index 4d2ffde9d30d..0b3e89f6f009 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_transceiver_creation.py +++ b/tests/unittest/_torch/disaggregation/test_kv_cache_transceiver_creation.py @@ -25,7 +25,7 @@ import pytest -from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as transceiver_module +from tensorrt_llm._torch.disaggregation import kv_cache_transceiver as transceiver_module from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py similarity index 99% rename from tests/unittest/_torch/executor/test_dual_pool_kv_cache.py rename to tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py index 059cb3b79738..d2d661653106 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py @@ -30,7 +30,7 @@ import tensorrt_llm from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator -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.resource_manager import KVCacheManager, ResourceManagerType from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, KvCacheConfig, TorchLlmArgs from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py b/tests/unittest/_torch/executor/kv_cache/test_kv_block_offset_overlap_race.py similarity index 100% rename from tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_block_offset_overlap_race.py diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py similarity index 99% rename from tests/unittest/_torch/executor/test_kv_cache_budget_split.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 79eca335918e..8ddf7ebc26ed 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -21,7 +21,7 @@ from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator from tensorrt_llm._torch.pyexecutor.config_utils import uses_vswa_kv_cache_layout -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.llmapi.llm_args import KvCacheConfig pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_compression_manager.py similarity index 99% rename from tests/unittest/_torch/executor/test_kv_cache_compression_manager.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_compression_manager.py index 0d6d619d23c7..3421104a5e73 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_compression_manager.py @@ -92,7 +92,7 @@ def _factory_model_engine( def _v2_manager(*, is_draft: bool): - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 manager = KVCacheManagerV2.__new__(KVCacheManagerV2) manager.enable_block_reuse = False @@ -233,7 +233,9 @@ def test_target_update_receives_metadata_before_final_compression(self): ] def test_real_v2_target_receives_relocation_metadata(self): - from tensorrt_llm._torch.pyexecutor import kv_cache_manager_v2 as kv_cache_v2_module + from tensorrt_llm._torch.pyexecutor.kv_cache import ( + kv_cache_manager_v2 as kv_cache_v2_module, + ) target = _v2_manager(is_draft=False) target.kv_cache_map = {} diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py similarity index 99% rename from tests/unittest/_torch/executor/test_kv_cache_estimation.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 76fba30b31ce..0d3a47329a7e 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -22,7 +22,7 @@ from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window -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.speculative.interface import SpeculativeDecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MultimodalConfig, TorchLlmArgs diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py similarity index 98% rename from tests/unittest/_torch/executor/test_kv_cache_manager_v2.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 54abfb345976..c359d2eed1c1 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py @@ -22,8 +22,8 @@ import torch from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp -from tensorrt_llm._torch.pyexecutor import kv_cache_manager_v2 as kv_cache_v2_module -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( +from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, KVCacheManagerV2, _KVCacheManagerInitStatus, @@ -170,7 +170,7 @@ def build_cache_config( fake_impl.pool_group_descs = [] fake_impl.get_layer_group_id.side_effect = lambda _: 0 - module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" + module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" with ( patch(f"{module}.CuError", _CacheTierInitError), patch(f"{module}.IndexMapper"), @@ -521,7 +521,7 @@ def test_world_ranks_converge_on_hostless_fallback() -> None: def test_local_fallback_failure_is_shared_before_raising() -> None: - module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" + module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" with ( patch( @@ -553,7 +553,7 @@ def test_local_fallback_failure_is_shared_before_raising() -> None: def test_peer_fallback_failure_discards_local_candidate() -> None: initial_impl = Mock() fallback_impl = Mock() - module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" + module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" with ( patch( @@ -934,7 +934,7 @@ def test_per_conversation_policy_releases_cancelled_request( batch_b = _prepare_context_resources(manager, request_b) with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2.logger.warning" + "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2.logger.warning" ) as mock_warning: assert manager.prepare_context(request_b) mock_warning.assert_not_called() @@ -1025,7 +1025,7 @@ def test_per_conversation_policy_ignores_overlapping_request( batch_b = _prepare_context_resources(manager, request_b) with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2.logger.warning" + "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2.logger.warning" ) as mock_warning: assert manager.prepare_context(request_b) mock_warning.assert_called_once_with( @@ -1128,7 +1128,7 @@ def test_cold_pool_group_iteration_stats_sum_all_cold_levels() -> None: def test_disagg_role_mapper_kinds_default_to_indexed(): from tensorrt_llm._torch.disaggregation.resource.page import MapperKind - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role manager = object.__new__(KVCacheManagerV2) diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py similarity index 99% rename from tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index e95682373fcd..e75361c70e50 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -27,7 +27,7 @@ import pytest -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 LlmRequestState diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py similarity index 97% rename from tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py index 1d86b5d1129a..8c88bbba654d 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py @@ -8,8 +8,8 @@ import tensorrt_llm import tensorrt_llm.bindings -from tensorrt_llm._torch.pyexecutor import kv_cache_manager_v2 as kv_cache_v2_module -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig DataType = tensorrt_llm.bindings.DataType diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py similarity index 99% rename from tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py index 069c47841287..d518d288581e 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py @@ -18,7 +18,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.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py similarity index 97% rename from tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py index 6995d556398c..aed1ddfe94fb 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py @@ -4,8 +4,8 @@ import pytest import torch -import tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 as resource_manager -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +import tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 as resource_manager +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 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, diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py similarity index 99% rename from tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index e89bc6518405..2e2e77c50b0e 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -209,7 +209,7 @@ def make_scheduler( from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2.KVCacheManagerV2", + "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2.KVCacheManagerV2", new=type(kv_cache_manager), ): kwargs = {} @@ -1288,7 +1288,7 @@ def test_encoder_without_cross_manager_raises(self): mgr = make_kv_cache_manager() # Build a scheduler with ENCODER_INIT gating but no cross manager. with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2.KVCacheManagerV2", + "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2.KVCacheManagerV2", new=type(mgr), ): sched = KVCacheV2Scheduler( diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py similarity index 100% rename from tests/unittest/_torch/executor/test_kv_pool_rebalance.py rename to tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py diff --git a/tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py b/tests/unittest/_torch/executor/kv_cache/test_kvv2_host_tier_sizing.py similarity index 97% rename from tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py rename to tests/unittest/_torch/executor/kv_cache/test_kvv2_host_tier_sizing.py index 9604ee7b797f..6d677ab41c4d 100644 --- a/tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvv2_host_tier_sizing.py @@ -11,7 +11,7 @@ import pytest -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _compute_auto_host_tier_quota, _sync_host_tier_quota, ) @@ -110,7 +110,7 @@ def _host_tier_sync_worker(per_rank_mem_gib): cross-rank sync. Returns this rank's pre- and post-sync quota so the parent process can assert on convergence. """ - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _compute_auto_host_tier_quota, _sync_host_tier_quota, ) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py similarity index 99% rename from tests/unittest/_torch/executor/test_mamba_cache_manager.py rename to tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 06378da704c2..4ad86cc13ba0 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -23,13 +23,11 @@ extract_mamba_kv_cache_params, ) from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.llm_request import ( - ATTENTION_DP_DUMMY_REQUEST_ID, - LlmRequest, - SamplingConfig, +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + BlockReusePolicy, + KVCacheManagerV2, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, MambaHybridCacheManagerV2, @@ -44,6 +42,11 @@ _mamba_snapshot_rule_counts, _promote_mamba_state_triton, ) +from tensorrt_llm._torch.pyexecutor.llm_request import ( + ATTENTION_DP_DUMMY_REQUEST_ID, + LlmRequest, + SamplingConfig, +) from tensorrt_llm._torch.pyexecutor.resource_manager import ( CacheTypeCpp, DataType, @@ -1234,11 +1237,11 @@ def test_cpp_hybrid_replay_bookkeeping_is_fused_into_conv_promotion( promote_calls = [] monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", lambda *args, **kwargs: promote_calls.append((args, kwargs)), ) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._advance_replay_state", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._advance_replay_state", lambda *args, **kwargs: pytest.fail( "Cpp replay bookkeeping must not advance before conv promotion" ), @@ -1619,7 +1622,7 @@ def test_v2_hybrid_estimator_counts_dummy_states_without_attention_capacity( lambda *args, **kwargs: 11, ) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._get_local_mamba_cache_layout", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._get_local_mamba_cache_layout", lambda *args, **kwargs: ( SimpleNamespace(get_states_bytes_per_layer=lambda mapping: 64), 1, @@ -1755,7 +1758,7 @@ def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): mgr.max_seq_len = 4096 warnings_seen = [] monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager.logger.warning", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.logger.warning", lambda message: warnings_seen.append(message), ) @@ -2543,7 +2546,7 @@ def test_v2_hybrid_debug_logs_prefix_reuse_only_on_rank_zero( log_debug = MagicMock() monkeypatch.setattr(KVCacheManagerV2, "_create_kv_cache", create_kv_cache) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager.logger.debug", log_debug + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.logger.debug", log_debug ) mgr = object.__new__(MambaHybridCacheManagerV2) @@ -2611,7 +2614,7 @@ def test_v2_hybrid_logs_aggregated_recurrent_cache_status_only_on_rank_zero( monkeypatch.setattr(KVCacheManagerV2, "get_iteration_stats", MagicMock(return_value=report)) log_debug = MagicMock() monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager.logger.debug", log_debug + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.logger.debug", log_debug ) mgr = object.__new__(MambaHybridCacheManagerV2) @@ -3496,11 +3499,11 @@ def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): events = [] mgr._commit_gdn_cached_replay_history_layers = lambda *_args, **_kwargs: events.append("commit") monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._advance_replay_state", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._advance_replay_state", lambda *_args, **_kwargs: events.append("advance"), ) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", lambda *_args, **_kwargs: None, ) @@ -3540,7 +3543,7 @@ def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): mgr.all_conv_states = [torch.empty(0)] mgr.intermediate_conv_states = torch.empty(0) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", lambda *args, **kwargs: None, ) @@ -3577,7 +3580,7 @@ def capture_promoted_position(_dst, _src, _src_indices, positions, _dst_indices) promoted_positions.append(positions.clone()) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", capture_promoted_position, ) diff --git a/tests/unittest/_torch/executor/test_per_layer_head_dim.py b/tests/unittest/_torch/executor/test_per_layer_head_dim.py index ba39d7520220..851f56377081 100644 --- a/tests/unittest/_torch/executor/test_per_layer_head_dim.py +++ b/tests/unittest/_torch/executor/test_per_layer_head_dim.py @@ -20,7 +20,7 @@ import tensorrt_llm import tensorrt_llm.bindings -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, KVCacheManagerV2, Role, diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index c90b6180b33a..a040a371c93e 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -17,7 +17,8 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.peft.lora.config import LoraConfig -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.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ( diff --git a/tests/unittest/_torch/executor/test_send_kv_async_split.py b/tests/unittest/_torch/executor/test_send_kv_async_split.py index 2ca96e872de7..7603fd1405ac 100644 --- a/tests/unittest/_torch/executor/test_send_kv_async_split.py +++ b/tests/unittest/_torch/executor/test_send_kv_async_split.py @@ -17,7 +17,7 @@ import pytest from tensorrt_llm._torch.disaggregation.executor.transfer_manager import AsyncTransferManager -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import CtxTransferStatus +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import CtxTransferStatus from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index a823dfd02386..7f6b59c1adef 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -257,7 +257,7 @@ def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): def make_fake_v2(enable_block_reuse=False, *, is_draft=False): """Build an unallocated V2 double with TriAttention's production contract.""" - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 fake_v2 = KVCacheManagerV2.__new__(KVCacheManagerV2) fake_v2.enable_block_reuse = enable_block_reuse diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py index a17a1b4c6731..3be0e9861a3e 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -400,7 +400,7 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): pytest.importorskip("cutlass") import tensorrt_llm import tensorrt_llm.bindings - 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.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 9eef8be1840a..84bedef0c009 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -53,7 +53,7 @@ from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX if TYPE_CHECKING: - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 _FLASHINFER_WORKSPACE_BYTES = 320 * 1024 * 1024 _TRTLLM_GEN_TOKENS_PER_BLOCK = 32 @@ -936,7 +936,7 @@ def _build_gemma4_kv_cache_manager( sizes line up with what the model actually requests at runtime. """ import tensorrt_llm - 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.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 if quant_config is not None and quant_config.layer_quant_mode.has_fp8_kv_cache(): diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index ef2f3821c03d..c3fa5f427a5f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -16,7 +16,7 @@ from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_llama import LlamaForCausalLM -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import \ +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ _update_kv_cache_draft_token_location from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager diff --git a/tests/unittest/_torch/modeling/test_modeling_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_multimodal.py index 00856c7f35e7..60d12664c15b 100644 --- a/tests/unittest/_torch/modeling/test_modeling_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_multimodal.py @@ -26,7 +26,7 @@ is_nemotron_hybrid, is_qwen3_hybrid, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import CppMambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import CppMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import str_dtype_to_torch from tensorrt_llm.bindings.executor import KvCacheConfig diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index 6a682a12c4d9..be6c1a3a29f6 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -19,7 +19,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import MambaKVCacheParams -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import KvCacheConfig, QSASparseAttentionConfig diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 81fb704c3a05..5da980b03c4c 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -433,7 +433,9 @@ def test_ple_cache_layout_excludes_separate_mtp_draft() -> None: def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: from tensorrt_llm._torch.configs import Qwen4ExpTextConfig from tensorrt_llm._torch.pyexecutor.config_utils import extract_qwen4_exp_ple_cache_params - from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping @@ -463,7 +465,7 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: def test_ple_states_use_v2_lifecycle_buffers(monkeypatch) -> None: - from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( MambaHybridCacheManagerV2, MambaRole, ) diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 9c7169e36d73..56fb548b8793 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -31,7 +31,7 @@ cu_seqlens_to_chunk_indices_offsets, cu_seqlens_to_chunk_indices_offsets_triton, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, ReplayStateUpdateMetadata, ) diff --git a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py index 04f1ff00bcf3..48ef92fb8d72 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py @@ -11,7 +11,7 @@ import pytest import torch -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( PythonMambaCacheManager, _allocate_mamba_seed_buffer, _compute_deterministic_mamba_seed, @@ -160,7 +160,9 @@ def _build_cpp_hybrid(*, spec_config, use_replay: bool, sr: bool, max_batch_size layer. Mirrors test_mamba_cache_manager._build_hybrid_with_mamba_layer but parameterizes the replay / SR flags so we can exercise the non-replay MTP SR layer-cache hand-off path.""" - from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import CppMambaHybridCacheManager + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + CppMambaHybridCacheManager, + ) from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp from tensorrt_llm.llmapi.llm_args import KvCacheConfig diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index 0289b2656d7d..31ce6b825d39 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -14,7 +14,7 @@ # limitations under the License. """Multi-rank TP tests for the KVCacheManagerV2 rebalance trigger. -``tests/unittest/_torch/executor/test_kv_pool_rebalance.py`` covers the same +``tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py`` covers the same logic against a mocked ``dist``; these tests run it across real MPI ranks with a real ``MPIDist``, which is the only way to show the collective actually agrees (and that ranks reach it in lockstep rather than deadlocking). diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py index 918de736c56f..83ca7ef66205 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py @@ -445,7 +445,9 @@ class TestKdaReplaySeedOnDisaggTransfer(unittest.TestCase): L, SLOTS, D, W, M, NHEADS = 2, 4, 6, 4, 2, 3 def _make_manager(self, use_kda_replay=True): - from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import PythonMambaCacheManager + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + PythonMambaCacheManager, + ) L, SLOTS, D, W, M, NH = (self.L, self.SLOTS, self.D, self.W, self.M, self.NHEADS) committed = W - 1 diff --git a/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py b/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py index 17e964d45fbf..b40016f8fb22 100644 --- a/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py +++ b/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py @@ -15,7 +15,7 @@ import pytest -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MixedMambaHybridCacheManager from tensorrt_llm._torch.speculative.sa_worker import SAWorker diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py index 57ee76d03d10..ed12b2706d2a 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py @@ -46,7 +46,7 @@ UnpagedResourceHandler, ) from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager # ============================================================================= diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py index afeef5471baf..76482a7140d0 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py @@ -24,7 +24,7 @@ from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionType from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs from tensorrt_llm._torch.auto_deploy.shim.ad_executor import create_autodeploy_executor -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import AttentionTypeCpp +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import AttentionTypeCpp from tensorrt_llm.llmapi import CacheTransceiverConfig pytestmark = pytest.mark.cpu_only @@ -354,7 +354,7 @@ def test_create_executor_requires_attention_type(): def test_create_executor_rejects_mamba_cache_manager_for_transceiver(): """Test create_autodeploy_executor rejects Mamba/hybrid cache transfer.""" - from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import BaseMambaCacheManager + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import BaseMambaCacheManager mock_tokenizer = MockTokenizer() diff --git a/tests/unittest/disaggregated/kv_transfer_harness.py b/tests/unittest/disaggregated/kv_transfer_harness.py index d03a1cfeab6a..305938ed40c8 100644 --- a/tests/unittest/disaggregated/kv_transfer_harness.py +++ b/tests/unittest/disaggregated/kv_transfer_harness.py @@ -33,7 +33,7 @@ import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -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, LlmRequestType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 418ea4352634..9a53343fd2ad 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -49,7 +49,7 @@ from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.utils import get_global_layer_ids from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -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, LlmRequestType from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 89002470f95f..6a1958041636 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -558,7 +558,7 @@ def test_pipelined_transfer_rejects_bounce_buffer(): def test_pipelined_transfer_rejects_mamba_cache_manager(): """Pipelined transfer does not support recurrent-state cache managers.""" from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 - from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager transceiver = object.__new__(KvCacheTransceiverV2) transceiver._mapping = SimpleNamespace(pp_size=1) @@ -577,8 +577,10 @@ def test_pipelined_transfer_rejects_mamba_cache_manager(): def test_python_transceiver_rejects_cpp_mamba_cache_manager(): """Python transceiver requires separate Python-managed Mamba state.""" - from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver - from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import CppMambaHybridCacheManager + from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import create_kv_cache_transceiver + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + CppMambaHybridCacheManager, + ) kv_cache_manager = object.__new__(CppMambaHybridCacheManager) cache_transceiver_config = CacheTransceiverConfig( diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 4db4fc943904..8c4948d0f523 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -34,7 +34,7 @@ get_physical_pool, get_unique_layers, ) -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.resource_manager import ( CacheTypeCpp, DataType, diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index fc854ed1e3d3..6250f9574ec0 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -67,8 +67,8 @@ ) from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup from tensorrt_llm._torch.pyexecutor.config_utils import extract_mamba_kv_cache_params +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MixedMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 81ced03755da..86d2bbdec42a 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -50,7 +50,7 @@ from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import CacheKind from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -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, LlmRequestType from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 19a9eeed5e9c..44b7b8040982 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -34,15 +34,15 @@ import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + MixedMambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, LlmRequestType, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( - MambaHybridCacheManagerV2, - MixedMambaHybridCacheManager, -) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp diff --git a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py index cd2fd27e8bd8..024deba0c22f 100644 --- a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py +++ b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py @@ -38,7 +38,7 @@ from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool, get_pool_bytes -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._utils import TensorWrapper, convert_to_torch_tensor from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index a5b8b280dfae..abb41c6d12a4 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -27,6 +27,10 @@ import pytest from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import ( + CtxTransferStatus, + GenTransferStatus, +) from tensorrt_llm._torch.disaggregation.native.transfer import ( TaskStatus, TransferWorker, @@ -34,7 +38,6 @@ TxSession, ) from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import CtxTransferStatus, GenTransferStatus from tensorrt_llm.bindings import LlmRequestState diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index 0c51f53d09ce..831be729ba1b 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -20,7 +20,7 @@ import pytest import torch -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, LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager as KVCacheManagerV1 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 8737091bbd2d..1179365af005 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -20,15 +20,16 @@ import tensorrt_llm import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm -from tensorrt_llm._torch.distributed import Distributed -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.kv_cache.kv_cache_manager_v2 import \ + KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import \ + MixedMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest, LlmRequestState) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import \ - MixedMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, KvCacheConfig