You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[TRTLLM-14855][chore] Use the canonical modules instead of the layout compatibility shims - #22086
The code-layout Epic (TRTLLM-14558) left compatibility forwarding modules (shims) at retired import paths, so that callers outside the repository keep working during the deprecation window. Code inside the repository is expected to use the canonical paths. This PR updates the following imports and code comments that name a shim path:
a sentence naming MLA.forward_impl in tensorrt_llm/_torch/modules/attention.py
removed; the caller contract above it stays
The mamba_cache_manager_v2.py import never resolved, because the attention_backend shim is a package with only an __init__.py and has no interface submodule. The import sits under TYPE_CHECKING, so only static type checkers ever evaluated it.
A follow-up PR under the same ticket will delete the shims, together with their tests, their CODEOWNERS entries and the legacy keys in the restricted-unpickle allowlist.
There is no runtime behaviour change. The shims stay in place, and each changed import names the canonical module that defines what it imports.
Test Coverage
Local, on a CPU node without Ray: pytest --collect-only on the two changed test files. Collection runs their module-level imports (tensorrt_llm.executor.ray.utils, and the tensorrt_llm.executor.ray.stub fallback when Ray is absent).
The RayWorkerWrapper import in test_mnnvl_allreduce.py runs inside a Ray actor. It runs whenever that test creates its workers, which happens in its existing pre-merge GPU entries: l0_dgx_b200 lists the file, and l0_dgx_h100 selects it with -m "gpu2".
The TYPE_CHECKING import is checked statically: locally, pyrefly reports it as a missing import before this change and reports no error after it.
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.
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
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
Six files changed: 8 additions and 9 deletions.
Imports and documentation references now use canonical TensorRT-LLM module paths.
No production logic, public API, defaults, or test definitions changed.
Import resolution is the main verification point.
QA Engineer Review
Two existing test files were modified. No test files or test-list files were added or removed.
Existing test functions and waivers remain unchanged. The affected tests are test_mnnvl_allreduce_over_process_group, test_refit_isolation_engine, test_refit_isolation_decode, and test_refit_isolation_report.
The MNNVL test is listed in tests/integration/test_lists/test-db/l0_dgx_b200.yml.
No matching CI or manual-QA entry exists for the refit isolation test.
Current-head test results are not available. Coverage verdict: needs follow-up.
Per-File QA Perspective
cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp: Comment-only path correction. No observable behavior changes.
tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py: Docstring-only change. The pre-sliced cache contract remains unchanged.
tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py: The TYPE_CHECKING import now uses the canonical path. Static type checking should resolve AttentionMetadata.
tensorrt_llm/_torch/speculative/dflash.py: Documentation-only path correction. No observable behavior changes.
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_mnnvl_allreduce.py: Covers Ray-orchestrated MNNVL all-reduce with fusion and no-fusion cases. The canonical Ray imports must resolve during collection and worker execution. The test is listed in the DGX B200 CI database.
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_refit_compile_isolation.py: Covers engine refit isolation, decode behavior, and reporting. The canonical decorator import must resolve during collection. No matching CI or manual-QA test-list entry exists.
Reviewing files that changed from the base of the PR and between 272cc2b and 65b9635.
You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.
Use the checkbox below for a quick retry:
Trigger review
No actionable comments were generated in the recent review.
Walkthrough
This change updates module references in comments and imports. It also removes implementation detail from a cache docstring while retaining its pre-slicing contract.
Changes
Path and documentation updates
Layer / File(s)
Summary
Update Ray references cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp, tests/unittest/_torch/ray_orchestrator/multi_gpu/*
The Ray executor comment and test imports now use updated module paths.
The type-checking import and documentation use updated attention references. The cache docstring no longer describes how MLA.forward_impl slices latent_cache; its pre-slicing contract remains.
This is a mechanical path-reference update with no runtime behavior change, and all new targets resolve. No actionable merge-blocking risk remains.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title clearly identifies the ticket, change type, and primary change: replacing layout compatibility shim references with canonical modules.
Docstring Coverage
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Description check
The description follows the required structure. It clearly explains the compatibility-shim cleanup, lists affected files and references, documents test coverage, and includes the PR checklist. The rep…
Full details: Docstring Coverage
Explanation
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 2 functions across 6 files.
Full details: Description check
Explanation
The description follows the required structure. It clearly explains the compatibility-shim cleanup, lists affected files and references, documents test coverage, and includes the PR checklist. The reported validation and no-runtime-change scope are sufficiently detailed.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Reviewing files that changed from the base of the PR and between b4cede8 and 272cc2b.
Files selected for processing: 1
tensorrt_llm/_torch/speculative/dflash.py
Files skipped from review as they are similar to previous changes: 1
tensorrt_llm/_torch/speculative/dflash.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
1 points — Tiny mechanical attention review: check the canonical metadata import and confirm that removing the stale caller example preserves the documented cache contract.
Affinity (0–1)
0.500 component/tag match
Outstanding points at selection
48 → 49 (+1)
Points already reserved on this PR
0
Effective weight
1
Completed review points (preceding 168 hours)
13
Combined points after assignment
62
Demand-derived target
61.00 points (soft share)
Additional fair share before assignment
0.00 points
Share explanation
Whole PR / qualification or coverage constraint exceeds a soft share; assignment is allowed.
Selection rule
balance-first maintenance
Assignment scope (farazkh80)
all
Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once.
The reason will be displayed to describe this comment to others. Learn more.
No blockers, possibly-blocking or required items; 3 nice-to-haves (2 inline, 1 below). Each changed import now names the module that defines the object the shim forwarded: RayWorkerWrapper.physical_to_local_id, control_action_decorator, AttentionMetadata, and the Ray stub's remote/__getattr__. tensorrt_llm/executor/ray/__init__.py re-exports nothing, so the Ray-less stub fallback still imports without Ray. The two corrected comments match the code: tensorrt_llm/executor/ray/executor.py:161-164 sets TLLM_DISABLE_MPI=1 for the workers, and attention/backends/utils.py:233-234 holds the V1 divisor. No other .py file imports a shim outside the shims' own tests. CI (PR_Github #84053) is still running. The in-actor RayWorkerWrapper import only runs on the DGX multi-GPU stages your Test Coverage section cites, and those stages also need the ci: full pre-merge approved label.
[Nice to have] The sweep misses the perf cookbooks under .claude/skills/: 14 lines in 9 files still name the pre-move paths. The one that will break is perf-regression-cookbook/references/cuda-graph-and-compile/piecewise-attention-torch-compile-host-overhead.md:56: its detection command runs import tensorrt_llm._torch.modules.attention. It already prints (False, False), because maybe_compiled_cat now lives in tensorrt_llm/_torch/utils.py, and it will raise ModuleNotFoundError once the follow-up deletes that shim. The other lines grep tensorrt_llm/_torch/attention_backend/sparse/dsa.py, .../attention_backend/fmha/flashinfer_trtllm_gen.py or '*/attention_backend/*', which no longer exist or match only the shim's __init__.py. Please point them at tensorrt_llm/_torch/attention/backends/sparse/dsa/, tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py, tensorrt_llm/_torch/attention/mla.py and tensorrt_llm._torch.utils. Do it here, or add them to the follow-up's list in the description. To find them: git grep -n -E "attention_backend/|modules/attention\.py|_torch\.modules\.attention" -- .claude/skills.
The reason will be displayed to describe this comment to others. Learn more.
[Nice to have] Only the path in the removed sentence was stale. MLA now lives in tensorrt_llm/_torch/attention/mla.py, and MLA.forward_impl still slices latent_cache[:num_ctx_tokens, ...] for context (mla.py:1117) and latent_cache[num_ctx_tokens:, ...] for generation (mla.py:1148) before dispatching. That pre-slicing is what the contract above requires of callers, so I'd keep the pointer and fix only the path:
The reason will be displayed to describe this comment to others. Learn more.
Thanks — the path was indeed the only stale part. I dropped the sentence rather than re-pointing it, for two reasons:
The slicing lives in the private MLA._forward_impl (mla.py:1117/1148), not forward_impl, which skips it on the sparse-hooks path; with external RoPE the context latent is rebuilt via torch.cat, and latent_cache_gen can be supplied by the caller.
The direct callers of this helper are the backends (e.g. fmha/fp4_mla.py slices latent_cache[:num_tokens] itself), so the caller contract above already states the requirement without a cross-file pointer that can drift again.
The reason will be displayed to describe this comment to others. Learn more.
[Nice to have] Consider naming the function instead of the file, so the reference greps and survives the next move. The V1 divisor is hardcoded in append_mla_latent_cache_generation_cuda_graph_safe (tensorrt_llm/_torch/attention/backends/utils.py:139); the // (num_pool_layers * kv_factor) decode is at :233-234.
PR_Github #84116 [ run ] completed with state DISABLED
The CI pipeline is currently frozen to deploy an urgent hotfix. For time-sensitive requests, please contact trt-llm-infra-devs.
… compatibility shims
Update the following imports and code comments that name a layout
compatibility shim:
- mamba_cache_manager_v2.py: the TYPE_CHECKING import of AttentionMetadata
now names tensorrt_llm._torch.attention.backends.interface. The old
tensorrt_llm._torch.attention_backend.interface never resolved: that
shim is a package with only an __init__.py.
- test_mnnvl_allreduce.py and test_refit_compile_isolation.py import the
Ray helpers from tensorrt_llm.executor.ray.{stub,gpu_worker,utils}.
- dflash.py: the docstring names
append_mla_latent_cache_generation_cuda_graph_safe, the function that
hardcodes the V1 divisor, instead of its file.
- moeAlltoAllOp.cpp: the comment names tensorrt_llm/executor/ray/executor.py
as the place that sets TLLM_DISABLE_MPI for Ray workers.
- cache_update.py: drop the docstring sentence that named a shim module as
the code slicing latent_cache; the caller contract before it stays.
Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
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
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.
sunnyqgg
left a comment
Description
The code-layout Epic (TRTLLM-14558) left compatibility forwarding modules (shims) at retired import paths, so that callers outside the repository keep working during the deprecation window. Code inside the repository is expected to use the canonical paths. This PR updates the following imports and code comments that name a shim path:
tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.pytensorrt_llm._torch.attention_backend.interface(underTYPE_CHECKING)tensorrt_llm._torch.attention.backends.interfacetests/unittest/_torch/ray_orchestrator/multi_gpu/test_mnnvl_allreduce.pytensorrt_llm.ray_stub,tensorrt_llm.executor.ray_gpu_workertensorrt_llm.executor.ray.stub,tensorrt_llm.executor.ray.gpu_workertests/unittest/_torch/ray_orchestrator/multi_gpu/test_refit_compile_isolation.pytensorrt_llm._ray_utilstensorrt_llm.executor.ray.utilstensorrt_llm/_torch/speculative/dflash.py(docstring)attention_backend/utils.pyappend_mla_latent_cache_generation_cuda_graph_safecpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp(comment)executor/ray_executor.pytensorrt_llm/executor/ray/executor.pytensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py(docstring)MLA.forward_implintensorrt_llm/_torch/modules/attention.pyThe
mamba_cache_manager_v2.pyimport never resolved, because theattention_backendshim is a package with only an__init__.pyand has nointerfacesubmodule. The import sits underTYPE_CHECKING, so only static type checkers ever evaluated it.A follow-up PR under the same ticket will delete the shims, together with their tests, their CODEOWNERS entries and the legacy keys in the restricted-unpickle allowlist.
There is no runtime behaviour change. The shims stay in place, and each changed import names the canonical module that defines what it imports.
Test Coverage
pytest --collect-onlyon the two changed test files. Collection runs their module-level imports (tensorrt_llm.executor.ray.utils, and thetensorrt_llm.executor.ray.stubfallback when Ray is absent).RayWorkerWrapperimport intest_mnnvl_allreduce.pyruns inside a Ray actor. It runs whenever that test creates its workers, which happens in its existing pre-merge GPU entries:l0_dgx_b200lists the file, andl0_dgx_h100selects it with-m "gpu2".TYPE_CHECKINGimport is checked statically: locally, pyrefly reports it as a missing import before this change and reports no error after it.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.Dev Engineer Review
QA Engineer Review
test_mnnvl_allreduce_over_process_group,test_refit_isolation_engine,test_refit_isolation_decode, andtest_refit_isolation_report.tests/integration/test_lists/test-db/l0_dgx_b200.yml.Per-File QA Perspective
cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp: Comment-only path correction. No observable behavior changes.tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py: Docstring-only change. The pre-sliced cache contract remains unchanged.tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py: TheTYPE_CHECKINGimport now uses the canonical path. Static type checking should resolveAttentionMetadata.tensorrt_llm/_torch/speculative/dflash.py: Documentation-only path correction. No observable behavior changes.tests/unittest/_torch/ray_orchestrator/multi_gpu/test_mnnvl_allreduce.py: Covers Ray-orchestrated MNNVL all-reduce with fusion and no-fusion cases. The canonical Ray imports must resolve during collection and worker execution. The test is listed in the DGX B200 CI database.tests/unittest/_torch/ray_orchestrator/multi_gpu/test_refit_compile_isolation.py: Covers engine refit isolation, decode behavior, and reporting. The canonical decorator import must resolve during collection. No matching CI or manual-QA test-list entry exists.