[None][fix] Update stale attention and KV-cache import paths that break main

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

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

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
 

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

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

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")

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
 

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
 
 

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
 

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