From 0c221e3462ab3fe65a31c84156228d653b946652 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 6 Sep 2026 09:10:15 -0700 Subject: [PATCH] [None][fix] Update stale attention and KV-cache import paths that break main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. #17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. #17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. #17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. #18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py, the visual-gen FlashInfer backend and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen --- .../_torch/attention/backends/fmha/phased.py | 2 +- .../attention/backends/fmha/prims_ts.py | 2 +- .../_torch/attention/backends/fmha/utils.py | 2 +- tensorrt_llm/_torch/pyexecutor/engine/lora.py | 2 +- .../attention_backend/flashinfer.py | 2 +- .../attention/test_backends_importable.py | 55 +++++++++++++++++++ .../_torch/attention/test_combined_fmha.py | 2 +- .../_torch/attention/test_fmha_page_index.py | 2 +- .../_torch/attention/test_prims_ts_fmha.py | 2 +- .../visual_gen/test_attention_flashinfer.py | 2 +- 10 files changed, 64 insertions(+), 9 deletions(-) create mode 100644 tests/unittest/_torch/attention/test_backends_importable.py 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