[TRTLLM-14558][chore] Consolidate the Attention domain: modules, backends and unit tests by YihuiLu512 · Pull Request #17968 · NVIDIA/TensorRT-LLM

YihuiLu512 · 2026-08-19T11:38:14Z

Description

Layout-only consolidation of the Attention domain into a single package tree, plus the reference
updates that requires. Epic TRTLLM-14558, 21 commits, 319 files, +1035/−881, of which 87 are
renames.

ticket	what moves
TRTLLM-14841	`_torch/modules/{attention,cross_attention,mla,qk_norm_attention,rotary_embedding}.py` + `ATTENTION_DEVELOPER_GUIDE.md` → `tensorrt_llm/_torch/attention/`
TRTLLM-14842	`_torch/attention_backend/` → `_torch/attention/backends/` (whole tree, incl. `fmha/` and `sparse/` with its algorithm sub-packages)
TRTLLM-14847	12 files under `tests/unittest/_torch/` → `tests/unittest/_torch/attention/{,kernels/,multi_gpu/}`; `_torch/custom_ops/` is emptied and retired

The remaining commits repoint references (docs, skills, CODEOWNERS, lint lists, CI test lists) and
answer review feedback.

Code moves, code does not change. No top-level definition is edited anywhere — verified by AST
comparison, not by reading the diff. The only in-file edits are the import lines the moves require
plus reference updates. The entire C++ delta is one comment line in
cpp/tensorrt_llm/common/attentionOp.h. _torch/visual_gen/attention_backend/ is a different tree
and is deliberately untouched (TRTLLM-15864).

One in-file change is not a move, and is called out because it is a behaviour fix: moving
attention_backend/ one package deeper left the relative import in backends/sparse/dsa/metadata.py
a level short. ModuleNotFoundError subclasses ImportError, so the except on the next line
swallowed it and warmup_selfsampling_topk() silently became a no-op on the opt-in
TRTLLM_GVR_SELF_SAMPLING=1 path. Fixed with the absolute form; all 228 relative imports under
_torch/attention/ now resolve.

Merge order is load-bearing. Forwarding shims for the retired import paths live in #17969
(draft), which must merge after this PR. Adding them here would leave git no delete to pair each
add with, losing the rename records and pointing git blame at this refactor rather than at the
original authors. Between the two merges, out-of-repo callers importing through the old paths break.

Ownership was replayed over the whole tree with the repo's own
.github/scripts/label_component.py: 0 files become unowned; 4 change membership, in every case
@NVIDIA/trt-llm-runtime-devs stepping off a file it held only through a directory or catch-all rule.
All four were waived or acked in review.

Why two post_merge entries gain `TIMEOUT (120)`

On main, test_mla_helix.py is a standalone test-list entry with a per-entry budget of its own,
and test_mha_helix.py is not listed anywhere — grep across tests/integration/test_lists/
returns nothing for it. After the move both are swept by the unittest/_torch/attention directory
entry, so three budgets collapse into one.

Measured on a 4-GPU B300 node: the merged entry takes 3597.44 s against the 3600 s default
per-entry timeout — 2.6 s of headroom, and an earlier run on a slightly slower node was in fact
killed by that default. The helix files account for 1322.1 s (37 %) of the total; without them
the entry finishes in roughly 2275 s.

Hence TIMEOUT (120) on the post_merge attention entry in l0_dgx_b300.yml and
l0_gb300_multi_gpus.yml. The 1-GPU stages are left alone: helix skips there on
device_count < 2 and adds no time.

Test Coverage

No test is added, removed or rewritten; 12 files are relocated. The
parallel / parallel_hw_agnostic / serial execution split those files had under thop/ is not
preserved — per review every --ignore= flag was dropped and they now run inside the single
attention entry.

Coverage delta, replayed entry by entry across every test-db list, main vs HEAD, with the rename
map applied symmetrically to both sides: 0 test files lose coverage on any list, and 4 gain it —
attention/kernels/test_deepseek_v4_block_table.py, attention/kernels/test_fused_inv_rope_fp8_quant.py,
attention/multi_gpu/test_mha_helix.py, attention/test_mla_registry.py. Confirmed on hardware by a
per-test JUnit diff: 0 cases lost, 0 outcome changes on the shared set.

Full rerun of DGX_B300-4_GPUs-PyTorch-Post-Merge on a 4×B300 node (x86, SM103) — every entry of
the stage, not only the helix files:

entries	cases	failures	errors
01 `unittest/_torch/attention`	2991	0	0
02–08	2007	0	0

test_mha_helix 24/24 and test_mla_helix 24/24 PASS, reproduced across two independent
runs. Neither guard (device_count < 2, skip_pre_blackwell) fired, so these are real
executions rather than skips.
Zero cross-file contamination: the 12 downstream neighbours scheduled after helix inside the
same entry were each accounted for individually — no mpi4py/MPI, serializer, CUDA-context or OOM
signature in any of them.
Multi-GPU helix therefore runs in CI for the first time on these two post_merge stages; this is a
net increase in coverage, not a preservation of it.

Limitations, stated rather than discovered later:

The TIMEOUT (120) on l0_gb300_multi_gpus.yml (GB300, aarch64) has no measurement behind it.
It is a symmetry argument from the x86 B300 result.
The rerun deviates from CI in one place: the CI base image could not be pulled (registry moved to
artifactory.nvidia.com, no local credentials), so the previous base was used —
NIXL 1.3.1 / UCX 1.21 against CI's 1.4.0 / 1.22. This stage is entirely attention, modeling and
accuracy tests and does not exercise the disagg-NIXL path.
Entry 01 was measured once; 3597.44 s has no repeat sample, so node-to-node variance is
unquantified.

Other checks: path-existence plus AST symbol validation over every doc and skill reference this
PR edits; check_test_list.py --validate passes; both attention waives in waives.txt still resolve
through the merged entry and were confirmed SKIPPED on hardware. .test_durations is deliberately
left byte-identical to main — jenkins/UpdateTestDurations.groovy rebuilds it from real runs and
the next pass drops the four now-stale keys and re-measures the directory entries, so the drift is
self-healing.

Follow-ups filed

TRTLLM-15863 — PEP 562 lazy re-export of the Attention module-layer symbols from
_torch/attention/; kept out here because it changes the public import surface.
TRTLLM-15864 — the mirror consolidation for VisualGen under _torch/visual_gen/attention/.
TRTLLM-15893 — repair a pre-existing dead deferred import in
_torch/pyexecutor/model_engine.py; this PR only rewrote the dead path's prefix, it did not make
it dead.
Per-entry parallelism tuning for the ex-thop/ files, confirmed with @yuxianq as a follow-up
rather than part of this move.

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.

…vel short Moving attention_backend/ under _torch/attention/backends/ put this module one package level deeper, so every relative import in the subtree needed one more dot. The sibling sparse/minimax_m3/cache_manager.py was updated; this one was not, and four dots now resolve to tensorrt_llm._torch.attention.cute_dsl_kernels, which does not exist. The bug is silent three ways: it is a function-local deferred import, so no smoke import touches it; ModuleNotFoundError subclasses ImportError, so the except clause on the next line swallows it; and a residue scan cannot see it, because the coordinate it produces is a new path that happens not to exist rather than a retired one. Effect: warmup_selfsampling_topk() became a no-op whenever TRTLLM_GVR_SELF_SAMPLING=1 on SM100/103, so the radix-topk kernels fell back to lazy JIT on first touch -- exactly what the warmup exists to prevent, notably during CUDA-graph capture. Switched to the absolute form rather than adding a fifth dot, matching the neighbouring warmup_cute_dsl_radix_topk() in this same file and backends/fmha/cute_dsl_mla.py, so a further move cannot break it again. Re-ran the audit that missed it, now resolving each level from the module's package rather than its file path: 228 relative imports under _torch/attention/, 0 unresolved. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

…tale These lines name a file plus a symbol in prose, which the path-existence scan does not see: it checks dotted module.Symbol references, so a reference whose directory was repointed but whose leaf is stale passes through. All of them described MLA yet pointed at attention/attention.py, which holds only class Attention -- class MLA has always lived in its own module. They were equally wrong before this move, when the leaf was modules/attention.py, but attention.py and mla.py are now one-word-different siblings, which makes the wrong pointer easier to follow and harder to notice. Nine of them are a leaf swap to attention/mla.py. Two needed more: - sparse-mla-topk-attention.md named forward_absorption, which is not a symbol -- the real ones are forward_absorption_context and forward_absorption_generation -- and forward_generation_dsa, which does not exist anywhere and is now forward_generation_sparse_attn in backends/sparse/dsa/module.py. - split-custom-op-for-piecewise-capture.md lumped five symbols under one wrong path. Split across backends/sparse/dsa/custom_ops.py for the two registered ops and backends/sparse/dsa/module.py for the three Python functions, and corrected forward_dsa_attn to _forward_dsa_attn and _should_use_short_mha to should_use_short_mha. The three grep examples in trtllm-codebase-exploration are executable and returned nothing; they now produce hits. Verified after the edit: every path named exists, every symbol named is found in the file it is attributed to, and the four executable examples produce output. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

…eveloper guide The guide's Scope list and Key File Map were repointed by the move but still enumerate only the members the file had when it sat under _torch/modules/. Now that it is the README of a self-contained package, naming the whole package is its job. Adds qk_norm_attention.py, cross_attention.py, rotary_embedding.py and the new kernels/ subpackage to both lists, and the two new test sub-trees to the key test files. The kernels/ row carries the ownership rule its __init__ docstring states, since that rule is the non-obvious part. AGENTS.md had the same blind spot: its "read before modifying" scope named attention.py, mla.py and backends/, which is now an incomplete subset of the package. Widened to the package rather than extended to seven paths. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

…d sub-trees The comment block above this entry is shared verbatim with the GPU lists, where "this entry is what picks them up" is true. On the CPU-only list it is not: pytest_ignore_collect in tests/unittest/conftest.py drops every test file without a pytest.mark.cpu_only marker before importing it when the stage runs -m cpu_only, and only 3 of the 53 files under tests/unittest/_torch/attention/ carry that marker. Worth stating, because the paragraph above warns that the helix tests rebind mpi4py's serializer at import on every stage that collects them -- a warning that reads as if it applies here, and does not. Comment only; no entry changes. check_test_list.py --validate still reports 1979 unique entries. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

…ouched A re-audit with a consistent rule -- in scope iff this PR modifies the line carrying the reference, regardless of whether it was also wrong on main -- turned up three the earlier pass missed. All three are the same "directory repointed, leaf stale" class already fixed elsewhere in this PR. mla-kv-cache-reuse.md:35 prepare_paged_context_mla Renamed to prepare_context_mla_with_cached_kv by NVIDIA#6538; introduced by NVIDIA#3571, which this same line cites. Line 28 named it too, plus mla_context_paged_kv, which NVIDIA#6538 deleted outright with no like-for-like successor -- used enable_context_mla_with_cached_kv (backends/interface.py), the gate the sentence is actually about. auxiliary-cache-in-kv-manager.md:90 _get_pool_block_indices Now get_pool_block_indices. Line 52 named it too. trtllm-model-onboard-multimodal/SKILL.md:105 Qwen2_5_VLVisionAttention.prepare_attn_metadata That class defines only __init__, apply_rope and forward. prepare_attn_metadata belongs to Qwen2_5_VisionModel. Companion lines 28 and 52 were not themselves modified by this PR, but leaving one half of a file naming a symbol the other half just corrected is worse than the scope rule is strict. Verified after the edit: every path named exists, every symbol named is found in the file it is attributed to, and the four retired names have zero occurrences left in docs/ and .claude/. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

…tion guard An earlier commit in this PR deleted /tests/unittest/_torch/custom_ops on the grounds that a rule matching no file is a dead entry. That reasoning does not hold for this file: it already carries 27 rules that match nothing, and main carries the same 27, so "matches nothing" is not a deletion criterion here. Restoring it costs one line and closes a real gap. Without the rule, a recreated custom_ops/ falls through to the global * fallback instead of to runtime-devs -- an ownership change nobody decided, arriving silently. The comment states that, so the next reader does not delete it again as dead. This does add one dead rule, 27 -> 28, deliberately. The two per-file rules that were under that directory are not restored: those files moved into attention/kernels/ and are covered there by /tests/unittest/_torch/attention, so restoring them would name paths that no longer exist. Measured with the repo's own label_component.py over all 10899 tracked files: the rule matches nothing, so no file's owner set changes. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

…racker IDs Two comments this PR adds cited internal Jira IDs. A reader of the public repository cannot open those; the PR number is the reference they can follow, and it is where the reasoning actually lives. See TRTLLM-14844. -> See PR NVIDIA#17968 for the discussion. -- TRTLLM-14558 moved these ... -> -- PR NVIDIA#17968 moved these ... Scope is what this PR introduces. The pre-existing "Epic TRTLLM-14558" mention in the compatibility-forwarding block came from NVIDIA#17857 and is left alone, as are the tracker IDs in the other 70 files that carry them -- including pull_request_template.md and pr-check.yml, which exist to require the ID in the title. Commit and PR titles keep the ID, per the repository's own format rule. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

… rebase The rebase onto the latest main replayed this PR's relocation commits over seven upstream lines that were added after the previous merge base and still name the retired coordinates: * `_torch/attention_backend/...` -> `_torch/attention/backends/...` - `attention/backends/sparse/deepseek_v4/footer_scale_kv.py` (TYPE_CHECKING import) - `_torch/visual_gen/attention_backend/cute_dsl/fmha.py` (skip_softmax params) - `visual_gen/sparse_attention.py` (one new lazy import; the other four were already repointed by the move commits) - `tests/.../attention/sparse/deepseek_v4/test_compressor_module.py` and `tests/.../attention/test_combined_fmha.py` (mock/monkeypatch target strings) * `_torch/modules/rotary_embedding` -> `_torch/attention/rotary_embedding` - `_torch/models/modeling_dflash.py` (new upstream file) * `_torch/modules/attention` -> `_torch/attention/attention` - `tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py` (new upstream file) Also wraps the `attention.backends.interface` import in `tests/unittest/_torch/attention/test_fmha_page_index.py`, which ruff requires now that the longer module path pushes the line past the limit. Import repointing only; no behavioural change. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

…/GB300 Consolidating the attention unit tests pulled the two helix files out of multi_gpu/ and into _torch/attention/, so the directory-level test-list entry now collects them as well. On main, test_mla_helix.py was a standalone entry with a default budget of its own and test_mha_helix.py was not listed at all, so neither cost was ever charged against unittest/_torch/attention. A full rerun of the DGX_B300-4_GPUs-PyTorch-Post-Merge stage on a B300 node measured the merged entry at 3597.44s against the 3600s default per-entry timeout, leaving 2.6s of headroom. The helix files account for 1322.1s (37%) of that total; without them the entry finishes in roughly 2275s. The previous run, on a slightly slower node, had already been killed by the 3600s timeout. The rerun is functionally clean: entries 01-08 of that stage cover about 4998 cases with 0 failures and 0 errors, and helix passed 48/48 across both repetitions with no cross-file interference. Only the time budget is short, so raise both affected post_merge entries to TIMEOUT (120). Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>

YihuiLu512 · 2026-09-04T22:31:13Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-04T22:38:39Z

PR_Github #71631 [ run ] triggered by Bot. Commit: 72737f3 Link to invocation

tensorrt-cicd · 2026-09-05T02:42:08Z

PR_Github #71631 [ run ] completed with state FAILURE. Commit: 72737f3
/LLM/main/L0_MergeRequest_PR pipeline #58716 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

YihuiLu512 · 2026-09-05T02:47:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-05T02:54:09Z

PR_Github #71647 [ run ] triggered by Bot. Commit: 72737f3 Link to invocation

tensorrt-cicd · 2026-09-05T04:21:35Z

PR_Github #71647 [ run ] completed with state FAILURE. Commit: 72737f3
/LLM/main/L0_MergeRequest_PR pipeline #58732 completed with status: 'UNSTABLE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

YihuiLu512 · 2026-09-05T04:46:33Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-05T04:53:08Z

PR_Github #71662 [ run ] triggered by Bot. Commit: 72737f3 Link to invocation

tensorrt-cicd · 2026-09-05T08:11:24Z

PR_Github #71662 [ run ] completed with state FAILURE. Commit: 72737f3
/LLM/main/L0_MergeRequest_PR pipeline #58745 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

YihuiLu512 · 2026-09-05T08:24:20Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-05T08:30:58Z

PR_Github #71679 [ run ] triggered by Bot. Commit: 72737f3 Link to invocation

tensorrt-cicd · 2026-09-05T09:47:53Z

PR_Github #71679 [ run ] completed with state SUCCESS. Commit: 72737f3
/LLM/main/L0_MergeRequest_PR pipeline #58761 completed with status: 'SUCCESS'

CI Report

Link to invocation

YihuiLu512 · 2026-10-21T06:55:39Z

Follow-up from the post-merge check of this PR: the DSv4 indexer-Q CuTe DSL prewarm imported DeepseekV4Indexer from ...sparse.deepseek_v4.deepseek_v4, a module that has never existed, so the except ImportError: return made the prewarm a no-op. This was not introduced here: the broken path came in with #17376, and this PR only renamed its attention_backend prefix to attention/backends. It is tracked in TRTLLM-15893 and fixed in #20734.

YihuiLu512 requested review from a team as code owners August 19, 2026 11:38

YihuiLu512 requested review from Wanli-Jiang, bmarimuthu-nv, chenfeiz0326, dhansen-nvidia, dpitman-nvda, luyiyun1021, niukuo, pengbowang-nv, xrq-phys, yechank-nvidia and yunruis August 19, 2026 11:38

github-actions Bot assigned YihuiLu512 Aug 19, 2026

YihuiLu512 requested review from BowenFu, brnguyen2, cascade812, chang-l, nv-guomingz and tongyuantongyu August 19, 2026 11:38

YihuiLu512 mentioned this pull request Aug 19, 2026

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

Merged

1 task done

YihuiLu512 added 9 commits September 5, 2026 06:25

yizhang-nv mentioned this pull request Oct 13, 2026

[None][test] Migrate model and executor fixtures to KV cache manager V2 #20098

Merged

1 task done

YihuiLu512 mentioned this pull request Oct 13, 2026

[TRTLLM-15864][chore] Consolidate the VisualGen attention domain under _torch/visual_gen/attention/ #20214

Merged

1 task done

chenfeiz0326 mentioned this pull request Oct 13, 2026

[None][perf] Experimental: compile and tune Triton, DeepGEMM and CuTe DSL JIT variants before serving #19843

Draft

YihuiLu512 mentioned this pull request Oct 20, 2026

[TRTLLM-15886][infra] Remove redundant attention unit-test runs from L0 lists #20683

Merged

1 task done

YihuiLu512 mentioned this pull request Nov 3, 2026

[TRTLLM-14853][chore] Update CODEOWNERS after the code-layout moves #21514

Merged

1 task done

Conversation

YihuiLu512 commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Why two post_merge entries gain TIMEOUT (120)

Test Coverage

Follow-ups filed

PR Checklist

GitHub Bot Help

Uh oh!

YihuiLu512 commented Sep 4, 2026

Uh oh!

tensorrt-cicd commented Sep 4, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

YihuiLu512 commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

YihuiLu512 commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

YihuiLu512 commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

YihuiLu512 commented Oct 21, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

13 participants

YihuiLu512 commented Aug 19, 2026 •

edited

Loading

Why two post_merge entries gain `TIMEOUT (120)`