From 056b7a31dbda9452d37314a3b5a47bf58bbc4f24 Mon Sep 17 00:00:00 2001 From: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com> Date: Mon, 7 Sep 2026 11:09:11 +0800 Subject: [PATCH] [TRTLLM-14841][chore] Add the forwarding modules for the old Attention paths Second half of the Attention relocation. PR #17968, merged as a56ec20396, moved the Attention module layer to tensorrt_llm/_torch/attention/ and the attention backends to tensorrt_llm/_torch/attention/backends/; this change puts definition-free forwarding modules back at the two retired paths that evidence shows callers outside this repository still use, so they keep working through the deprecation window. tensorrt_llm/_torch/modules/attention.py forwards Attention tensorrt_llm/_torch/attention_backend/__init__ forwards the old package root's whole __all__ Neither path is a guess. Before the move, both were named by imports in examples/llm-api/out_of_tree_example/modeling_opt.py -- lines 9 and 15 of the file the "adding a new model" walkthrough tells users to copy into their own tree. #17968 repointed that example and the walkthrough in the same change, so nothing in this repository is broken today; what that change cannot reach is every copy users already made from the earlier text. Those copies are what these modules serve. A forwarding module for one path and a hard cut for the other would leave such a copy broken anyway, so the two paths are treated the same way. Every other retired coordinate in the move is a hard cut: private, and clean on every public-surface criterion that can be scanned in-tree. The two export sets differ because the two retired paths are different kinds of thing. modules/attention.py is a module, and the only name any evidence source shows reaching through it is Attention; its other public names appear in no __all__ anywhere and are not forwarded. attention_backend/__init__.py is a package root, and a package root's __all__ IS its declared public surface, so this shim reproduces that list exactly -- including the two names the original adds under `if IS_FLASHINFER_AVAILABLE`, behind the same condition. Forwarding only the one name the example happens to use would give a caller who imported any of the other ten an ImportError with no migration hint, from a module that had just imported successfully and warned -- strictly worse than the clean ModuleNotFoundError a hard cut would have produced. That list is re-derived from the canonical __all__ rather than copied from an older revision: #18025 removed star attention after these modules were first written, so StarAttention and StarAttentionMetadata are not forwarded and the star_flashinfer import is gone. Keeping them would have made the shim raise ModuleNotFoundError on precisely the FlashInfer-enabled configuration it exists to serve, taking the other eleven names down with it. Four checks in tests/unittest/_torch/attention/test_backends_importable.py keep both modules honest: each retired path hands back the canonical object itself rather than a copy, the package shim re-exports every name the canonical package does, its __all__ matches the canonical __all__, and both paths warn with FutureWarning when imported. The __all__ check is the one that earns its keep -- it is what would have caught the star attention drift above, which on today's main is a ModuleNotFoundError raised only for a caller who has FlashInfer installed, the one configuration those two names ever served. It compares sets, not ordered lists: __all__ order binds nothing at import time, so freezing it would fail CI on a cosmetic reorder of the canonical list while catching no drift a caller could observe. The warning check drops the module from sys.modules and puts it back, because a module-level warnings.warn fires only on the first import in a process and an earlier test in the same one may already have spent it. The file is already collected by the directory-level unittest/_torch/attention test-list entries, so no test list changes. Custom-op registrations are unaffected: importing either file imports the canonical module, so its registration side effect runs as before. Both re-export the object rather than copying it -- `old.X is canonical.X` -- so isinstance and existing pickles still work. Both warn with FutureWarning, not DeprecationWarning: the latter is on Python's stock ignore list outside `__main__` and would never reach the callers these modules exist for. No in-tree caller routes through either of them: a scan of the tree for both retired paths reports only the forwarding modules themselves, the guard test added here that imports them on purpose, and two blog permalinks pinned to a historical commit that must stay as they are. CODEOWNERS does carry a COMPATIBILITY FORWARDING MODULES block for this Epic, and its rule is that each such module keeps the owner its pre-move path had; both paths were NVIDIA/trt-llm-torch-attention-devs before the move, and with these files present they resolve instead to NVIDIA/trt-llm-models-devs and NVIDIA/trt-llm-runtime-devs -- two different teams, neither the pre-move one. The two lines that would satisfy the block are held out of this change because /.github/CODEOWNERS is itself owned by NVIDIA/trt-llm-infra-devs and NVIDIA/trt-llm-oss-compliance, so adding them pulls two more required approvals onto a four-file PR; whether that trade is worth making is a question for review rather than something this commit should settle. Either way, add the Attention team as a reviewer by hand. This change also carries the copyright header for the relocated attention/backends/__init__.py. The header belongs on that file, but adding it during the relocation itself would have cut git's rename similarity for that short module to a few points above the detection threshold. Holding it back worked: the rename is recorded in a56ec20396 at R092, so `git blame` for those lines is anchored in history already and the header can land without risking it. The forwarding modules were kept out of #17968 because git pairs renames from adds and deletes, and a squash merge collapses any in-PR split. Landing them together with the move would have made git see new files plus rewritten old ones, and `git blame` for the moved lines would have pointed at the relocation instead of at their authors. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com> --- .../_torch/attention/backends/__init__.py | 14 ++++ .../_torch/attention_backend/__init__.py | 67 +++++++++++++++++ tensorrt_llm/_torch/modules/attention.py | 37 +++++++++ .../attention/test_backends_importable.py | 75 +++++++++++++++++++ 4 files changed, 193 insertions(+) create mode 100644 tensorrt_llm/_torch/attention_backend/__init__.py create mode 100644 tensorrt_llm/_torch/modules/attention.py diff --git a/tensorrt_llm/_torch/attention/backends/__init__.py b/tensorrt_llm/_torch/attention/backends/__init__.py index a43464c6a689..d1ea00e084d4 100644 --- a/tensorrt_llm/_torch/attention/backends/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/__init__.py @@ -1,3 +1,17 @@ +# 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. from ...flashinfer_utils import IS_FLASHINFER_AVAILABLE from .interface import AttentionBackend, AttentionForwardArgs, AttentionMetadata from .sparse import get_sparse_attn_kv_cache_manager diff --git a/tensorrt_llm/_torch/attention_backend/__init__.py b/tensorrt_llm/_torch/attention_backend/__init__.py new file mode 100644 index 000000000000..9993e532caf5 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/__init__.py @@ -0,0 +1,67 @@ +# 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. +"""Compatibility shim for ``tensorrt_llm._torch.attention_backend``. + +Will be removed once all usages are migrated to +``tensorrt_llm._torch.attention.backends``. + +DO NOT ADD ANYTHING TO THIS FILE. +""" + +import warnings + +from tensorrt_llm._torch.attention.backends import ( # noqa: F401 + AttentionBackend, + AttentionForwardArgs, + AttentionInputType, + AttentionMetadata, + TrtllmAttention, + TrtllmAttentionMetadata, + VanillaAttention, + VanillaAttentionMetadata, + get_sparse_attn_kv_cache_manager, +) +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + +warnings.warn( + "tensorrt_llm._torch.attention_backend has moved to " + "tensorrt_llm._torch.attention.backends. The old path still works for " + "now and will be removed in a future release.", + FutureWarning, + stacklevel=2, +) + +__all__ = [ + "AttentionMetadata", + "AttentionBackend", + "AttentionForwardArgs", + "AttentionInputType", + "TrtllmAttention", + "TrtllmAttentionMetadata", + "VanillaAttention", + "VanillaAttentionMetadata", + "get_sparse_attn_kv_cache_manager", +] + +if IS_FLASHINFER_AVAILABLE: + from tensorrt_llm._torch.attention.backends.flashinfer import ( # noqa: F401 + FlashInferAttention, + FlashInferAttentionMetadata, + ) + + __all__ += [ + "FlashInferAttention", + "FlashInferAttentionMetadata", + ] diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py new file mode 100644 index 000000000000..ca88ec3014ce --- /dev/null +++ b/tensorrt_llm/_torch/modules/attention.py @@ -0,0 +1,37 @@ +# 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. +"""Compatibility shim for ``tensorrt_llm._torch.modules.attention``. + +Will be removed once all usages are migrated to +``tensorrt_llm._torch.attention.attention``. + +DO NOT ADD ANYTHING TO THIS FILE. +""" + +import warnings + +from tensorrt_llm._torch.attention.attention import Attention # noqa: F401 + +warnings.warn( + "tensorrt_llm._torch.modules.attention has moved to " + "tensorrt_llm._torch.attention.attention. The old path still works " + "for now and will be removed in a future release.", + FutureWarning, + stacklevel=2, +) + +__all__ = [ + "Attention", +] diff --git a/tests/unittest/_torch/attention/test_backends_importable.py b/tests/unittest/_torch/attention/test_backends_importable.py index fa3ede085a4c..4ac1dbf9dd28 100644 --- a/tests/unittest/_torch/attention/test_backends_importable.py +++ b/tests/unittest/_torch/attention/test_backends_importable.py @@ -20,9 +20,15 @@ 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. + +The same file guards the two forwarding modules left behind at the retired +``_torch.modules.attention`` and ``_torch.attention_backend`` paths: they must +hand back the canonical objects themselves, and their re-export list must not +drift away from the canonical one. """ import importlib +import sys import pytest @@ -32,6 +38,13 @@ "tensorrt_llm._torch.attention.backends.fmha.utils", ] +# Import paths retired by the Attention consolidation and kept alive by a +# definition-free forwarding module for the deprecation window. +SHIM_MODULES = [ + "tensorrt_llm._torch.modules.attention", + "tensorrt_llm._torch.attention_backend", +] + def test_attention_backends_package_imports(): importlib.import_module("tensorrt_llm._torch.attention.backends") @@ -53,3 +66,65 @@ def test_kv_cache_manager_v2_names_resolve(): module = importlib.import_module("tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2") assert hasattr(module, "KVCacheManagerV2") assert hasattr(module, "Role") + + +def test_modules_attention_shim_forwards_canonical_class(): + """The retired module path must hand back the canonical class, not a copy of it.""" + shim = importlib.import_module("tensorrt_llm._torch.modules.attention") + canonical = importlib.import_module("tensorrt_llm._torch.attention.attention") + assert shim.Attention is canonical.Attention + + +def test_attention_backend_shim_forwards_canonical_objects(): + """Every name the retired package root re-exports must be the canonical object. + + Identity rather than equality: callers keep isinstance checks and pickles + against these classes, and a re-exported copy would silently fail both. + """ + shim = importlib.import_module("tensorrt_llm._torch.attention_backend") + canonical = importlib.import_module("tensorrt_llm._torch.attention.backends") + for name in shim.__all__: + assert getattr(shim, name) is getattr(canonical, name), name + + +def test_attention_backend_shim_exports_match_canonical(): + """The shim's ``__all__`` must track the canonical package's. + + The shim was first written while the canonical ``__all__`` still carried + ``StarAttention``/``StarAttentionMetadata``; those were later deleted, and + a stale copy of the list only surfaces as an import error on a machine that + has FlashInfer -- which is exactly the configuration the extra names serve. + """ + shim = importlib.import_module("tensorrt_llm._torch.attention_backend") + canonical = importlib.import_module("tensorrt_llm._torch.attention.backends") + # Compared as sets, not as ordered lists. ``__all__`` ordering binds nothing + # -- ``from ... import *`` is order-insensitive -- so freezing it would fail + # CI on a cosmetic reorder of the canonical list while catching no drift a + # caller could observe. Membership is the contract: a name present on only + # one side is what actually breaks, or silently stops serving, a caller. + assert canonical.__all__, "an empty canonical __all__ would make this guard vacuous" + assert set(shim.__all__) == set(canonical.__all__) + + +@pytest.mark.parametrize("module_name", SHIM_MODULES) +def test_shim_warns_on_import(module_name): + """Importing a retired path must warn, and specifically with FutureWarning. + + DeprecationWarning is on Python's stock ignore list outside ``__main__``, so + it would never reach the out-of-tree callers these modules exist for. + + The warning fires while the module body runs, and an earlier test in this + process may already have imported the shim, so the entry is dropped from + ``sys.modules`` to force exactly one fresh execution and put back afterwards. + Re-running the body is cheap: everything it imports is already loaded. + """ + parent_name, _, attribute = module_name.rpartition(".") + parent = importlib.import_module(parent_name) + previous = sys.modules.pop(module_name, None) + try: + with pytest.warns(FutureWarning, match="has moved to"): + importlib.import_module(module_name) + finally: + if previous is not None: + sys.modules[module_name] = previous + setattr(parent, attribute, previous)