[TRTLLM-14558][chore] Add the forwarding modules for the retired Attention paths

[YihuiLu512]

Description

Second half of the Attention consolidation — Epic TRTLLM-14558, ticket TRTLLM-14841 (the PR title carries the Epic key, the commit subject carries the ticket key). #17968 retired two import paths; this PR puts a definition-free forwarding module back at each of them so callers outside this repository keep working through the deprecation window, and lands the NVIDIA copyright header that #17968 had to hold back.

file	what it is	+lines
tensorrt_llm/_torch/modules/attention.py	forwarding module → _torch/attention/attention.py; re-exports Attention	37
tensorrt_llm/_torch/attention_backend/__init__.py	forwarding module → _torch/attention/backends/; re-exports all 11 names of the canonical __all__	67
tensorrt_llm/_torch/attention/backends/__init__.py	NVIDIA copyright header, nothing else	14
tests/unittest/_torch/attention/test_backends_importable.py	4 checks appended to the file #18771 created — 2 identity, 1 __all__, 1 warning	75

git diff --stat against the branch point 3a3871d8b1: 4 files, +193 / −0, a single commit.

Both forwarding modules define nothing — a header, a docstring, an import list, one warnings.warn, one __all__, and in the package one a conditional FlashInfer import block with an __all__ +=. They re-export the canonical objects rather than copies (old.X is canonical.X), so isinstance and pre-migration pickles keep working; importing either one imports the canonical module, so custom-op registration side effects run exactly as before. Both warn with FutureWarning, matching 7 of the 9 compatibility shims already in the tree.

Still Draft on purpose. No Jenkins pipeline has ever run on this branch. The blocking dependency is discharged, but the question below — whether out-of-repo evidence warrants a shim at all — is worth settling before spending pipeline time. Remaining steps, in order: settle the CODEOWNERS question, flip out of Draft, /bot run --disable-fail-fast. @NVIDIA/trt-llm-torch-attention-devs is already a requested reviewer. The VisualGen label was applied at PR creation, to a head that still touched _torch/visual_gen/; after the rebase onto these four files it no longer applies — replaying .github/scripts/label_component.py over them yields no component label at all.

Three claims the earlier body made are no longer true. It said 13 forwarded names: #18025 (55548ee861, 2026-08-24, BREAKING: remove star attention) deleted star attention, and the list is now re-derived from the canonical __all__ — 11 names, 9 unconditional plus 2 under IS_FLASHINFER_AVAILABLE, name-for-name and order-for-order identical to the canonical list (AST-checked). It said git apply of this patch against main fails outright: true then, and it is why the work was split into two PRs, but the patch applies cleanly now — that argument is history, not a live gate. And it claimed live in-tree evidence, which is the next section.

"Where is the evidence? Nothing in-tree imports these paths."

Correct, and it is the weak point, so it goes first rather than last.

At fc8969ec59 (the parent of the #17968 squash) the retired coordinates appeared 8 times across the example and the two copies of the add-a-model walkthrough: 6 import lines (examples/llm-api/out_of_tree_example/modeling_opt.py:9,15; docs/source/torch/adding_new_model.md:52,55; docs/source/models/adding-new-model.md:52,55) plus 2 prose mentions (:107 in each walkthrough). #17968 repointed all of them. On today's main the example reads:

from tensorrt_llm._torch.attention.attention import Attention
from tensorrt_llm._torch.attention.backends import AttentionMetadata

and a grep of main for either retired coordinate, outside the separate _torch/visual_gen/attention_backend/ tree, returns 2 hits — both frozen blob/d6b741ddfe permalinks in an archived tech blog (blog01_…B200_GPUs.md:129,134). Zero live in-tree callers. The in-tree evidence base went from 6 lines to 0, by this Epic's own hand.

What survives is out of tree and no grep here can see it. One checkable fact stands in for it: of the 93 v* tags in this repository, 61 contain that example, all 61 ship it with the old imports, and 0 carry the post-move form — from v1.0.0rc3 (2025-07-14) through v1.3.0rc25 (2026-08-27). Anyone who followed the documented recipe before 2026-09-05 holds a file whose lines 9 and 15 are exactly the two coordinates these modules restore. #17968's merged body states that between the two merges out-of-repo callers on the old paths break; that window has been open since 2026-09-05T09:48:09Z.

That is the whole argument: publication history, not a measurement. I have no telemetry on out-of-tree importers and cannot put a number on them. If the team's read is that a _torch path carries no compatibility promise whatever the walkthrough printed, that is a coherent position — and the right outcome is to drop the two modules and land only the copyright header, which has no other home (below). I would rather have that decided than assumed.

"Why shim a private _torch/ path at all? These would be the first."

They would be. git grep -l "Compatibility shim for" origin/main -- tensorrt_llm/ returns exactly 9 files — _common.py, serialization.py, _ray_utils.py, ray_stub.py, lora_helper.py, lora_manager.py, prompt_adapter_manager.py, executor/ray_executor.py, executor/ray_gpu_worker.py — and not one of them is under _torch/.

The counter-argument is narrow, and it is not that _torch is public. It is that this repository published these two particular _torch paths as copy-paste code in the add-a-model walkthrough and shipped that text in 61 releases, which makes this pair de facto documented regardless of the underscore. Containment: both modules are definition-free, both warn, and removal joins the Epic's existing batch — the written policy at tensorrt_llm/serialization.py:24-25 is that the compatibility window "spans at least one release (Epic decision D4 (b))", tracked by T25 (TRTLLM-14855). Two of those nine — _common.py and serialization.py — name that ticket in their docstrings; the other seven name no ticket at all. These two currently say only "will be removed in a future release" — one line each, added on request.

"Why didn't the MoE sibling need this?"

It didn't, and the asymmetry is conceded rather than explained away. Sibling Epic PR #17952 (c5c985c4a1, 2026-08-30) retired _torch/modules/fused_moe/** and _torch/expert_statistic.py with no forwarding module and no follow-up shim PR; on main the directory is gone and expert_statistic.py exists only at _torch/moe/. The criterion applied here is not "is it private" but "did a document tell users to copy it": the pre-cut fused_moe references under docs/ were six tech-blog mentions — five source links plus one prose filename — never import lines in a template the reader is instructed to reproduce. If the team prefers one rule for the whole Epic, the consistent rule is a hard cut everywhere and this PR shrinks to the header alone.

"This only covers the package-root import form."

Correct, and the root form is the minority. AST census over the pre-move tree at fc8969ec59, counting import statements, excluding everything inside the 79-file retired attention_backend/ package plus modules/attention.py:

form	statements	served
…attention_backend.<sub> (from/import, absolute or relative)	321	no
from …attention_backend import <name in __all__>	83	yes
from …attention_backend import <submodule object> (interface, utils, trtllm, sparse, flashinfer)	10	no
from …modules.attention import Attention	41	yes
from …modules.attention import <other> (helix helpers, extract_extra_attrs)	6	no
total	461	124 = 26.9 %

For the package coordinate alone — the figure worth scrutinising — 83 of 414 = 20.0 %. import tensorrt_llm._torch.attention_backend.trtllm and from …attention_backend.interface import AttentionMetadata still raise ModuleNotFoundError; the retired package now holds one file.

Not widened, deliberately. The shim's warrant is the walkthrough, and the walkthrough only ever taught the root form — those two published lines are exactly it. Serving the rest means re-creating a stub skeleton of the retired package's 8 submodules (79 .py files pre-move) and keeping it in sync, for a surface no document ever taught.

The limit points the same way and is worth stating rather than discovering: the three of #18771's fixes (75f521ddac) that #17968 caused — the stale imports in _torch/pyexecutor/engine/lora.py, _torch/visual_gen/attention_backend/flashinfer.py and tests/unittest/_torch/visual_gen/test_attention_flashinfer.py — were all …attention_backend.interface, submodule form. These modules would not have prevented any of them. That is the mechanism's limit, not an argument for it.

CODEOWNERS: the block's rule applies, and this revision does not follow it

.github/CODEOWNERS:455-463 carries a # ===== COMPATIBILITY FORWARDING MODULES ===== block naming this Epic, whose stated rule is that each shim "keeps the owner its pre-move path had". The earlier body's claim that "no other shim in the tree has one" was simply wrong. Replayed with the repo's own .github/scripts/label_component.py:

path	pre-move owner (a56ec20396^:419,420)	owner today
_torch/modules/attention.py	@NVIDIA/trt-llm-torch-attention-devs	@NVIDIA/trt-llm-models-devs (:177)
_torch/attention_backend/__init__.py	@NVIDIA/trt-llm-torch-attention-devs	@NVIDIA/trt-llm-runtime-devs (:117)

Straight concession: the rule is violated, not satisfied, and neither file keeps the owner its pre-move path had. The fix is two lines appended to that block, which I verified restore both paths to @NVIDIA/trt-llm-torch-attention-devs:

/tensorrt_llm/_torch/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs
/tensorrt_llm/_torch/attention_backend/__init__.py @NVIDIA/trt-llm-torch-attention-devs

They are held out of this revision for one reason worth weighing rather than deciding alone: /.github/CODEOWNERS is itself owned by @NVIDIA/trt-llm-infra-devs and @NVIDIA/trt-llm-oss-compliance (:466), so adding them pulls two more required approvals onto a 4-file PR that needs Attention-team eyes. My position is that the block wins and the two lines belong here — say the word and they go in. Meanwhile the other two changed files already route to that team by rule (:317, :321), and the team is a requested reviewer.

Why the copyright header rides along

_torch/attention/backends/__init__.py carries no NVIDIA header: 21 lines on main, 35 with the 14-line header. Adding it during the relocation would have destroyed the rename record. Both scenarios replayed in a scratch repository:

scenario	recorded
the move as actually merged in a56ec20396	R092 (and R098 for modules/attention.py)
counterfactual: header added in the same commit as the move	R049 — below git's default 50 % threshold, so it is recorded as A + D, no rename detected at all

git blame for the backends package root would then have pointed at the refactor instead of at the original authors. Hence: rename first, header second, here.

If this PR is dropped, the header silently never arrives. The repo has no SPDX/licence pre-commit hook to notice, every hand-written sibling __init__.py under _torch/attention/ already carries it — 10 of the other 21, with the 11 that do not being the whole vendored backends/prims_ts/** subtree, which carries a FlashInfer Apache-2.0 header instead, backends/prims_ts/__init__.py included — and nothing has touched this file since a56ec20396.

Residual risk

• The FutureWarning is unconditional at import. A downstream CI running -W error::FutureWarning turns the old path from "works with a warning" into a hard failure. Silent-by-default was rejected deliberately — a warning nobody sees migrates nobody — but the cost is real and it lands on the same out-of-tree users the shim exists for. tests/unittest/pytest.ini:14 sets -W ignore::DeprecationWarning and no -W error, so the warning is visible and non-fatal here; that is our config, not theirs.
• The number of out-of-tree callers is unknown and unknowable from this repository.
• _torch/visual_gen/attention_backend/ is a different tree and is untouched.
• No Jenkins pipeline has run on this branch.

Test Coverage

Four checks appended to tests/unittest/_torch/attention/test_backends_importable.py (55 → 130 lines, pure insertion — no existing test edited; 8 functions total, 4 pre-existing from #18771 plus 4 new, one of them parametrized over both retired paths for 5 new cases):

check	what breaks it
test_modules_attention_shim_forwards_canonical_class	the shim binding a copy or subclass instead of the canonical Attention, which silently breaks isinstance and pickles
test_attention_backend_shim_forwards_canonical_objects	the same, by identity (is), for every name the package shim re-exports
test_attention_backend_shim_exports_match_canonical	drift between the shim's __all__ and the canonical one, in either direction
test_shim_warns_on_import (×2)	the warning disappearing, or being downgraded to DeprecationWarning — which Python ignores by default outside __main__, so it would never reach the callers these modules exist for

The __all__ guard is the one that earns its keep: it is exactly the #18025 drift described above, which on today's main surfaces 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), with an assert canonical.__all__ so two empty sets cannot pass vacuously. The warning check pops the module from sys.modules and restores both the entry and the parent-package attribute in a finally, because a module-level warnings.warn fires once per process and the identity checks above it have already spent it.

These tests were not executed against a real tensorrt_llm. No interpreter here can import the package — system python3 is 3.6.8 and the local python3.11 has no torch — and no usable container runtime is available. Nothing below is a CI result. What was run this round:

• py_compile on all four changed files: clean. Longest lines 98 / 98 / 98 / 99 against ruff line-length = 100.
• AST comparison on the real tree: the shim's __all__ and the canonical __all__ are identical order-for-order, 11 names each.
• A toy-package replica under pytest 9.1.1: the three real source files and the new test block copied verbatim, with only the string tensorrt_llm._torch → toypkg rewritten, leaf symbols stubbed and IS_FLASHINFER_AVAILABLE = True so the conditional branch runs. Five seeded variants, results below.

variant	result
faithful	5 passed
canonical loses a name and its source module, shim still forwards it (the #18025 near-miss)	3 failed, 2 passed — ImportError: cannot import name … inside the shim's import list
canonical gains a name, shim does not follow	1 failed, 4 passed — only the __all__ guard fires
shim rebinds a subclass instead of re-exporting	1 failed, 4 passed — only the identity guard fires
canonical __all__ cosmetically reordered	5 passed — the set comparison raises no false alarm

The warning check was confirmed order-independent three ways: inside the full run, after the identity checks have already imported both shims; alone via -k warns; and over 5 consecutive full runs. This shows the assertions discriminate; it does not show that the real modules import.

pre-commit run --files over the four paths returned rc=0 with no hook modifying anything when the branch was built. It was not re-run for this description.

No test-list change, and none is needed. grep -rn backends_importable tests/integration/test_lists/ returns 0 hits; the file is swept by the directory-level unittest/_torch/attention entries at l0_cpu.yml:39, l0_h100.yml:25, l0_b200.yml:108, l0_b300.yml:27, l0_dgx_b300.yml:27 and l0_gb300_multi_gpus.yml:26, and adding functions to an already-swept file needs no edit. Note the file is not collected on l0_cpu: tests/unittest/conftest.py:224-240 drops any test_*.py lacking a literal pytest.mark.cpu_only when a stage runs -m cpu_only, and this file carries none — nor did #18771's four checks in it. These cases therefore run on the GPU L0 attention stages. Adding the marker was left out of scope.

Since no in-tree caller goes through either retired path, there is nothing else in this repository these modules could regress.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Dev Engineer Review

• Added compatibility shims for the retired Attention import paths.
• Shims forward canonical objects, preserve identity, define __all__, and emit FutureWarning.
• Backend exports remain synchronized with the canonical package.
• Added the required NVIDIA copyright header.
• No performance or configuration issues identified.
• Static checks and pre-commit checks passed. Full package tests and Jenkins CI did not run.

QA Engineer Review

Added tests:

• test_modules_attention_shim_forwards_canonical_class
• test_attention_backend_shim_forwards_canonical_objects
• test_attention_backend_shim_exports_match_canonical
• test_shim_warns_on_import

These tests are not covered by entries in tests/integration/test_lists/, test-db/, or qa/.

Verdict: needs follow-up. Add appropriate CI or QA test-list coverage, or document why list coverage is not required.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2aefd03d-b3e7-42a9-8585-bcbe026292cf

📥 Commits

Reviewing files that changed from the base of the PR and between 374df3c and 056b7a3.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/attention/backends/__init__.py
• tensorrt_llm/_torch/attention_backend/__init__.py
• tensorrt_llm/_torch/modules/attention.py
• tests/unittest/_torch/attention/test_backends_importable.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

Added compatibility shims for retired attention import paths. The shims re-export canonical objects, emit FutureWarning, conditionally expose FlashInfer classes, and include import behavior tests.

Changes

Attention import compatibility

Layer / File(s)	Summary
Compatibility shim implementation tensorrt_llm/_torch/attention/backends/__init__.py, tensorrt_llm/_torch/attention_backend/__init__.py, tensorrt_llm/_torch/modules/attention.py	Added the backend package license header. Added shims that re-export canonical attention objects, define __all__, emit FutureWarning, and conditionally expose FlashInfer classes.
Shim forwarding and warning validation tests/unittest/_torch/attention/test_backends_importable.py	Added tests for canonical object identity, export parity, and warnings from both retired import paths.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 056b7

This change restores retired Attention import paths by forwarding them to the canonical implementations and warning callers to migrate. The compatibility behavior is covered without an identified current-head merge risk.

Suggested reviewers: juney-nvidia

Sequence Diagram(s)

sequenceDiagram
  participant LegacyImport
  participant AttentionShim
  participant CanonicalAttention
  LegacyImport->>AttentionShim: import retired attention path
  AttentionShim-->>LegacyImport: emit FutureWarning
  AttentionShim->>CanonicalAttention: resolve canonical exports
  CanonicalAttention-->>LegacyImport: provide attention objects

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the chore and the primary change: adding forwarding modules for retired Attention import paths.
Description check	✅ Passed	The description is complete and directly addresses the change, rationale, affected files, test coverage, validation performed, limitations, and checklist status. It also discloses that Jenkins has not…
Docstring Coverage	✅ Passed	Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 4 files.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[YihuiLu512]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #71890 [ run ] triggered by Bot. Commit: 056b7a3 Link to invocation

[tensorrt-cicd]

PR_Github #71890 [ run ] completed with state SUCCESS. Commit: 056b7a3
/LLM/main/L0_MergeRequest_PR pipeline #58962 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation