[TRTLLM-14558][chore] Consolidate the Attention domain: modules, backends and unit tests - #17968
Merged
Merged
Conversation
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
YihuiLu512
requested review from
BowenFu,
brnguyen2,
cascade812,
chang-l,
nv-guomingz and
tongyuantongyu
August 19, 2026 11:38
1 task
Consolidating the Attention unit tests dropped the four standalone entries for
this file (l0_b200.yml:111, l0_dgx_b300.yml:24, l0_h100.yml:26,
l0_gb300_multi_gpus.yml:23) and let the `unittest/_torch/attention` directory
entry pick it up instead. The two files that moved alongside it,
test_mha_helix.py and test_mla_registry.py, each got an `--ignore=`; this one
did not.
That is not a neutral change, because the file mutates process-global state at
*import* time:
# tests/unittest/_torch/attention/test_mla_helix.py:54-60
MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads,
pickle.HIGHEST_PROTOCOL)
`MPI.pickle.__init__` re-points mpi4py's serializer for the whole process, and
it runs even where the test itself is skipped: the skip is a collection-time
marker (`:562`, device_count() < 2), and collection has already imported the
module. So the 1-GPU pre_merge stages -- l0_b200 and l0_h100, the latter
running with `orchestrator: mpi` -- would carry that mutation into every other
attention test sharing the process. On the two 4-GPU stages where the test
really runs it also spawns MPI ranks through MPIPoolExecutor that each claim a
device, which is a different proposition from a process that has already
executed ~40 GPU-heavy attention tests.
Restore the isolation the file had on main: `--ignore=` it from the four
directory entries and list it again as its own entry, exactly the way
test_mha_helix.py is already handled. Record *why* next to the entry, since the
omission happened precisely because that reason was never written down.
l0_b300.yml is left alone -- it had no entry for this file on main, and its
existing `--ignore=` already reproduces that.
Collected-file parity verified: expanding every test-db block into the set of
files it collects, honouring `--ignore=`, is identical before and after, and
the file is still collected by the same five lists as on main. No duplicate
entry strings within a block; `scripts/check_test_list.py --validate` passes.
Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
…dated
test_unittests_v2 keys the xdist parallel-factor lookup on the *entire*
test-list entry string, `--ignore=` flags included:
# tests/integration/defs/test_unittests.py:206-220
key = (row['gpu'], row['unittest_case_name'])
cur_key = (gpu_name, case) # `case` == the full yml entry text
Rewriting the attention entries therefore invalidated the four
`unittest/_torch/attention` rows, and the three entries split out under
`attention/kernels/` never had a row at all. A miss is only a
`warnings.warn`, so nothing goes red -- the cases quietly drop to serial.
The slowdown is the visible half. The other half is that `num_workers == 1`
takes a branch with no retry (`test_unittests.py:381-413`): only the parallel
branch re-runs failures serially with `--lf`. So the mismatch would also have
taken away the automatic re-run from a suite full of CUDA-graph capture,
multi-stream and FP8 numerical comparisons -- a behavior regression, not just
a wall-clock one, in a change that is supposed to preserve behavior.
Repoint the keys rather than invent numbers, which is what this table already
does when a reorg changes an entry: NVIDIA#13756 renamed a key in place when
`_torch/speculative` was split, and NVIDIA#13751 seeded the sibling it split out of
`thop/parallel` with the parent's factor.
- four `unittest/_torch/attention` rows renamed in place, factor 4 unchanged
(the three lists carrying this entry hold byte-identical text, so one key
covers all four rows)
- seven rows for `attention/kernels/{parallel,parallel_hw_agnostic,serial}`,
each inheriting its `thop/*` parent's factor, and only on the GPUs where
that parent already had a row
No row is added where main ran serially: `kernels/serial` off GB300, the
l0_b300 and l0_dgx_b300 attention entries -- already serial before this PR
because they already carried `--ignore=` flags -- and test_mla_helix.py, which
has to stay serial in its own process. The `thop/*` rows are untouched; those
entries still exist and still carry the tests that did not move.
Replaying the lookup the way the runner does it -- (gpu_name, full entry text),
min(factor, 8), serial on miss -- over every test-list entry against every GPU
in the table gives 0 differences in worker count versus main (1839 entries, 13
GPUs). Every case that took the retry branch before takes it again.
The table is infra-generated, so these keys should be folded into the next
regeneration pass rather than left as inherited values. `.test_durations` is
deliberately left alone for the same reason and needs that pass too.
Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
…hip audit This move claimed `tests/unittest/_torch/thop/parallel/test_indexer_topk.py` on one basis: the file's exact-file CODEOWNERS rule names the Attention team. The same move deletes that rule, folding it into the `/tests/unittest/_torch/attention` directory rule -- so afterwards the only thing saying the file is Attention's would be the directory it had just been moved into. That is the circularity the classification rules single out as the failure mode to avoid. What the file covers has meanwhile stopped being Attention-only. NVIDIA#17622 landed `_torch/modules/top_k.py`, which now has two production consumers in two domains -- `attention_backend/sparse/dsa/indexer.py` and `kv_cache_compression/triattention/triattention.py` -- both selecting `TopKImplementation.CUTE_DSL_RADIX`, i.e. sharing the neutral radix tier. The op they share, `cute_dsl_indexer_topk_decode`, encodes no attention concept in its signature. The tier that does -- GVR, with `pre_idx`, `compress_ratio` and the `next_n` `preIdxOffset` shift -- is covered by the two tests under `attention/sparse/`, which this move leaves in place. So keep the file where it is, restore the exact-file rule so Attention keeps ownership explicitly rather than by inheritance, and let TRTLLM-14844 place it together with the kernel directory it exercises -- one move instead of two. The reasoning is recorded on that ticket and in a comment above the restored rule, so a later reader does not read this as a file the move forgot. Coverage is unchanged: every stage that carried the split-out `attention/kernels/parallel` entry already carries `unittest/_torch/thop/parallel`, which is where the file runs again. Dropping the entry also drops the `TIMEOUT (90)` it had inherited from the line above it -- a cap this entry never needed (3.7 min measured on DGX_B200) and one that predates the file being in `thop/parallel/` at all. The five `attention/kernels/parallel` rows added to `agg_unit_mem_df.csv` in 744de5b go with it. The two rows for `parallel_hw_agnostic` and `serial` stay, as do the `thop/*` rows the file falls back under. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
… docs Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
… move left stale
A review audit of this branch turned up one real regression and three references
whose directory was repointed but whose leaf was left stale. All four sit on
lines this PR already touches.
docs/source/features/helix.md documented
python3 -m pytest tests/unittest/_torch/attention/test_mla_helix.py -v
which does not exist. The helix tests moved a second time, into
attention/multi_gpu/, after this doc had already been repointed once. The
command resolved on the base commit and does not resolve here, so this is a
regression this move introduced, not a pre-existing error. A repo-wide sweep
confirms it is the only reference the second move left behind.
The other three had the right directory and the wrong file. This is the one
failure mode a "search for the old path" sweep structurally cannot see, since
the string it would look for is already gone:
trtllm-model-onboard-multimodal/SKILL.md attention.MLA -> mla.MLA; class MLA
is defined in attention/mla.py and is
re-exported nowhere
fuse-qk-norm-rope.md apply_qk_norm_rope is defined in
attention/qk_norm_attention.py:321,
not in attention.py
trtllm-codebase-exploration/SKILL.md backends/trtllm_attention.py ->
backends/trtllm.py
The attention test-list comment added in the previous commit made three claims
that do not hold. "Everything under attention/ runs from this single entry ...
no longer split out" is false on l0_b200.yml, which still lists eight
attention/sparse/** files separately seven lines below it, and on l0_b300.yml
and l0_dgx_b300.yml, whose entries still carry the pre-existing
sparse/test_cute_dsl_* ignores. "Used to be listed on one stage only" is false:
the base listed thop/parallel_hw_agnostic on two (l0_cpu.yml:43 and
l0_h100.yml:66). And the stated meaning of the parallel_hw_agnostic suffix is
contradicted by five of the twelve files left in that directory, which do carry
skipif or pytestmark gating; the suffix actually tracks the cpu_only markexpr
split. The helix note was incomplete too: those tests gate on Blackwell as well
as on device_count() < 2, and on every non-Blackwell stage it is the arch gate
that fires, not the GPU-count one.
Correcting those per list would have produced six texts that each go stale
separately, so the block is rewritten to state only what holds on all six, with
no line numbers, case counts or per-stage claims in it.
The test-list change is comment-only: the set of entries parsed out of every
file under tests/integration/test_lists/test-db/ is identical to the previous
commit, and check_test_list.py --validate and --check-duplicate-waives both pass.
Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
…he line limit `tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py` arrived from upstream during the rebase and had its deferred import repointed from `attention_backend.sparse.minimax_m3` to `attention.backends.sparse.minimax_m3` by the move commit. The new module path is six characters longer, which pushes the single-line form past the line limit, so ruff wraps it into parenthesised form. Pure formatting; the imported symbol is unchanged. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
…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>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #71631 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #71631 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #71647 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #71647 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #71662 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #71662 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #71679 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #71679 [ run ] completed with state |
This was referenced Sep 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
13 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
_torch/modules/{attention,cross_attention,mla,qk_norm_attention,rotary_embedding}.py+ATTENTION_DEVELOPER_GUIDE.md→tensorrt_llm/_torch/attention/_torch/attention_backend/→_torch/attention/backends/(whole tree, incl.fmha/andsparse/with its algorithm sub-packages)tests/unittest/_torch/→tests/unittest/_torch/attention/{,kernels/,multi_gpu/};_torch/custom_ops/is emptied and retiredThe 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 treeand 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 inbackends/sparse/dsa/metadata.pya level short.
ModuleNotFoundErrorsubclassesImportError, so theexcepton the next lineswallowed it and
warmup_selfsampling_topk()silently became a no-op on the opt-inTRTLLM_GVR_SELF_SAMPLING=1path. 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 blameat this refactor rather than at theoriginal 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.pyis a standalone test-list entry with a per-entry budget of its own,and
test_mha_helix.pyis not listed anywhere —grepacrosstests/integration/test_lists/returns nothing for it. After the move both are swept by the
unittest/_torch/attentiondirectoryentry, 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 inl0_dgx_b300.ymlandl0_gb300_multi_gpus.yml. The 1-GPU stages are left alone: helix skips there ondevice_count < 2and adds no time.Test Coverage
No test is added, removed or rewritten; 12 files are relocated. The
parallel/parallel_hw_agnostic/serialexecution split those files had underthop/is notpreserved — per review every
--ignore=flag was dropped and they now run inside the singleattention 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 aper-test JUnit diff: 0 cases lost, 0 outcome changes on the shared set.
Full rerun of
DGX_B300-4_GPUs-PyTorch-Post-Mergeon a 4×B300 node (x86, SM103) — every entry ofthe stage, not only the helix files:
unittest/_torch/attentiontest_mha_helix24/24 andtest_mla_helix24/24 PASS, reproduced across two independentruns. Neither guard (
device_count < 2,skip_pre_blackwell) fired, so these are realexecutions rather than skips.
same entry were each accounted for individually — no
mpi4py/MPI, serializer, CUDA-context or OOMsignature in any of them.
net increase in coverage, not a preservation of it.
Limitations, stated rather than discovered later:
TIMEOUT (120)onl0_gb300_multi_gpus.yml(GB300, aarch64) has no measurement behind it.It is a symmetry argument from the x86 B300 result.
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.
unquantified.
Other checks: path-existence plus AST symbol validation over every doc and skill reference this
PR edits;
check_test_list.py --validatepasses; both attention waives inwaives.txtstill resolvethrough the merged entry and were confirmed SKIPPED on hardware.
.test_durationsis deliberatelyleft byte-identical to
main—jenkins/UpdateTestDurations.groovyrebuilds it from real runs andthe next pass drops the four now-stale keys and re-measures the directory entries, so the drift is
self-healing.
Follow-ups filed
_torch/attention/; kept out here because it changes the public import surface._torch/visual_gen/attention/._torch/pyexecutor/model_engine.py; this PR only rewrote the dead path's prefix, it did not makeit dead.
thop/files, confirmed with @yuxianq as a follow-uprather 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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.