diff --git a/tensorrt_llm/_torch/attention/backends/fmha/phased.py b/tensorrt_llm/_torch/attention/backends/fmha/phased.py index 397728c19a77..08b09177e330 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/phased.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/phased.py @@ -24,7 +24,7 @@ CustomAttentionMask, PredefinedAttentionMask, ) -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 from .interface import Fmha diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 715602ffcbe5..afad6479067d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -29,7 +29,7 @@ AttentionForwardArgs, AttentionInputType, ) -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 binding_to_torch_dtype, get_sm_version from tensorrt_llm.bindings.internal import thop diff --git a/tensorrt_llm/_torch/attention/backends/fmha/utils.py b/tensorrt_llm/_torch/attention/backends/fmha/utils.py index 3be5f1a7a16a..28a69db913d3 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/utils.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/utils.py @@ -20,7 +20,7 @@ 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.resource_manager import KVCacheManager from tensorrt_llm.bindings.internal import thop diff --git a/tensorrt_llm/_torch/pyexecutor/engine/lora.py b/tensorrt_llm/_torch/pyexecutor/engine/lora.py index 02fb9638d655..87cc83d67c0c 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/lora.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/lora.py @@ -5,7 +5,7 @@ import torch from torch import nn -from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.peft.lora.cuda_graph_lora_manager import CudaGraphLoraManager from tensorrt_llm._torch.peft.lora.manager import LoraModelConfig diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py index 1de6323acb6c..caf3ddb71f65 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py @@ -10,7 +10,7 @@ from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ...attention_backend.interface import PredefinedAttentionMask +from ...attention.backends.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout _WORKSPACE_BYTES = 128 * 1024 * 1024 diff --git a/tests/unittest/_torch/attention/test_backends_importable.py b/tests/unittest/_torch/attention/test_backends_importable.py new file mode 100644 index 000000000000..fa3ede085a4c --- /dev/null +++ b/tests/unittest/_torch/attention/test_backends_importable.py @@ -0,0 +1,55 @@ +# 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. +"""Guards the attention backend import graph. + +``tensorrt_llm._torch.attention.backends`` pulls in the whole ``fmha`` +package, so a stale module path in any of its members breaks collection of +every test that touches an attention backend. The executor's model engine +imports the attention backends too (directly and through ``engine.lora``), +so a stale path there takes the whole PyTorch runtime down with it. These +checks are import-only and run on CPU. +""" + +import importlib + +import pytest + +FMHA_MODULES = [ + "tensorrt_llm._torch.attention.backends.fmha.phased", + "tensorrt_llm._torch.attention.backends.fmha.prims_ts", + "tensorrt_llm._torch.attention.backends.fmha.utils", +] + + +def test_attention_backends_package_imports(): + importlib.import_module("tensorrt_llm._torch.attention.backends") + + +@pytest.mark.parametrize("module_name", FMHA_MODULES) +def test_fmha_module_imports(module_name): + importlib.import_module(module_name) + + +def test_engine_lora_imports(): + """``engine.lora`` sits on the model engine's import path.""" + module = importlib.import_module("tensorrt_llm._torch.pyexecutor.engine.lora") + assert hasattr(module, "AttentionMetadata") + + +def test_kv_cache_manager_v2_names_resolve(): + """The names the fmha modules import must exist at their source.""" + module = importlib.import_module("tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2") + assert hasattr(module, "KVCacheManagerV2") + assert hasattr(module, "Role") diff --git a/tests/unittest/_torch/attention/test_combined_fmha.py b/tests/unittest/_torch/attention/test_combined_fmha.py index 65c20fb60ed2..96c3ead3c9be 100644 --- a/tests/unittest/_torch/attention/test_combined_fmha.py +++ b/tests/unittest/_torch/attention/test_combined_fmha.py @@ -29,7 +29,7 @@ AttentionForwardArgs, AttentionInputType, ) -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.bindings import DataType from tensorrt_llm.quantization.mode import QuantMode diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 763345829fed..e008f060d9e6 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -27,7 +27,7 @@ ) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.autotuner import AutoTuner -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 diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index dccc8ec04bde..87dbd7364f4b 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -38,7 +38,7 @@ AttentionInputType, PredefinedAttentionMask, ) -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 import DataType diff --git a/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py b/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py index b93791a4eb02..a6dc0d2f4736 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py +++ b/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py @@ -10,7 +10,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.visual_gen.attention_backend.flashinfer import FlashInferAttention from tensorrt_llm._torch.visual_gen.attention_backend.utils import get_visual_gen_attention_backend from tensorrt_llm.visual_gen.args import QuantAttentionConfig