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-17229][fix] Resolve the KV-head count of mixed-precision high layers in disaggregated transfer - #21597
#20709 (per-view KV head transfer) made the disaggregated page-table builder resolve a KV-head count for every internal cache layer, either as an index into the manager's per-layer head list or through the DeepSeek-V4 virtual layer mapping.
However, #21021 (mixed-precision KV cache for speculative decoding and the sparse MLA backends) merged after it and records its mixed-precision high layers in the manager's mixed_precision_high_layers record instead of a virtual layer mapping, so building the page table of a mixed-precision KV cache manager raises (Cannot resolve KV-head count for internal layer ...). Each PR passed CI on its own; main fails test_disagg_exports_both_pools_and_pairs_sink_tail and test_deepseek_v4_mixed_disagg_names_every_layer since #21021 merged.
The builder now takes a mixed-precision high layer's head count from the layer it belongs to, read from that record. Other layers resolve as before; no manager change.
Test Coverage
tests/unittest/_torch/pyexecutor/kv_cache/test_mixed_precision_kv_cache.py::test_disagg_exports_both_pools_and_pairs_sink_tail and ::test_deepseek_v4_mixed_disagg_names_every_layer pass again (B200).
tests/unittest/disaggregated/test_extractor.py::test_v2_builder_resolves_mixed_precision_high_layer_heads (CPU) replaces the test that modelled the high layers through the earlier virtual-layer encoding; it builds the page table from a manager stand-in with the mixed_precision_high_layers record, with and without a virtual layer mapping of the manager's own layers.
Dev Engineer Review
_build_page_table_v2 now resolves mixed-precision high-layer head counts from their associated local layers before using virtual-layer mapping. Other layer-resolution paths remain unchanged. The context-first BF16 check now uses mixed_precision_low_dtype, matching the renamed manager field. No public declarations change.
QA Engineer Review
test_extractor.py adds parameterized coverage for mixed-precision high-layer resolution with and without virtual-layer mapping. It checks layer IDs, pool views, and head counts. The file is listed in the A10 and H100 CI test lists. The context-first profile test updates the NVFP4 case to use mixed_precision_low_dtype; it was not found in the inspected test lists. Coverage verdict: sufficient for the reported resolution scenarios. The author reports two B200 tests pass, but those results were not independently verified.
Per-File QA Perspective
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py: Verify that mixed-precision high layers use the head count of their associated local layer, including when virtual-layer metadata exists. Verify that other layers preserve their prior resolution behavior.
tests/unittest/disaggregated/test_extractor.py: The parameterized test covers high-layer resolution with and without virtual-layer mapping. The file is listed in the A10 and H100 CI test lists.
tensorrt_llm/_torch/disaggregation/native/context_first_profile.py: Verify that context-first BF16 eligibility reads mixed_precision_low_dtype and preserves the other eligibility conditions.
tests/unittest/disaggregated/test_context_first_profile.py: The NVFP4 rejection test now sets mixed_precision_low_dtype. It was not found in the inspected CI or manual-QA test lists.
… layers in disaggregated transfer
The per-view transfer layout resolves a KV-head count for every internal layer
id. A mixed-precision high layer has an id beyond the manager's per-layer head
list and no entry in the DeepSeek-V4 virtual layer mapping, so building the page
table of a mixed-precision KV cache manager raised. The high layer now takes the
head count of the layer it belongs to, read from the manager's
mixed_precision_high_layers record. The extractor test models the high layers
through that record and the manager's declared global layer ids instead of the
earlier virtual-layer encoding.
Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com>
No actionable comments were generated in the recent review.
Walkthrough
_build_page_table_v2 now resolves mixed-precision high-layer KV-head counts through associated local layers. Context-first BF16 eligibility now checks mixed_precision_low_dtype. Tests cover virtual and non-virtual layer lookup and the NVFP4 eligibility case.
_build_page_table_v2 maps mixed-precision high-layer IDs to local-layer IDs before virtual-layer resolution. For non-virtual managers, it uses the internal layer ID as the local offset. Tests cover virtual and non-virtual managers, including head counts, global IDs, pool views, and sliding-window size.
No merge-blocking issue was identified in the changed head resolution or BF16 eligibility behavior.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title names the main change: the disaggregated page-table builder now resolves the KV-head count of mixed-precision high layers. It follows the repository template with a JIRA ticket and the [fix]…
Description check
The description explains the problem and the fix. It cites the prior PR that introduced the per-view KV head transfer and the later PR that added mixed-precision high layers. It also lists the test co…
Docstring Coverage
Docstring coverage is 80.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. Documentati…
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.
Full details: Title check
Explanation
The title names the main change: the disaggregated page-table builder now resolves the KV-head count of mixed-precision high layers. It follows the repository template with a JIRA ticket and the [fix] type. The changeset matches this summary.
Full details: Description check
Explanation
The description explains the problem and the fix. It cites the prior PR that introduced the per-view KV head transfer and the later PR that added mixed-precision high layers. It also lists the test coverage with the failing tests and the new CPU test. The description does not follow the template's section headings exactly, but it covers the required information. The PR checklist is not filled in, which is expected for a checklist that the author completes after review.
Full details: Docstring Coverage
Explanation
Docstring coverage is 80.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 5 functions across 4 files.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
context_first_profile.py:116 still reads manager.low_precision_dtype, which #21021 renamed to mixed_precision_low_dtype. With TRTLLM_ENABLE_KV_TRANSFER_LIFECYCLE_ENFORCEMENT=1 and a BF16 KV cache, every rank now fails transceiver init with an AttributeError, and test_context_first_profile.py::test_context_first_profile_accepts_symmetric_tp (CPU) hits the same error on main.
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.
… context-first profile
NVIDIA#21021 renamed KVCacheManagerV2.low_precision_dtype to mixed_precision_low_dtype after
NVIDIA#21272 added this read. Every BF16 V2 manager reaches it, so the context-first eligibility
check raised AttributeError on main.
Signed-off-by: tianruih <tianruih@nvidia.com>
…s own methods
The extractor test installs the manager's global layer ids and view partition instead of a
stub, keeps both high layers in one group as the manager appends them, and checks each one
against the layer it belongs to. The head lookup states its three cases as if/elif/else.
Signed-off-by: tianruih <tianruih@nvidia.com>
…tests that fail on main
Both fail on main in B300's pre-merge stage: the disaggregated page
table (NVIDIA#20709) cannot resolve the KV-head count of a mixed-precision
high layer. Main waives the second on DGX_B200 only (nvbugs/7043616);
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
NVIDIA#21597 fixes the extractor. Drop these lines when it lands.
…tests that fail on main
Both fail on main in B300's pre-merge stage: the disaggregated page
table (NVIDIA#20709) cannot resolve the KV-head count of a mixed-precision
high layer. Main waives the second on DGX_B200 only (nvbugs/7043616),
and PR NVIDIA#21597 fixes the extractor. Drop these lines when it lands.
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…st that fails on main
test_deepseek_v4_mixed_disagg_names_every_layer fails on main in B300's
pre-merge stage: the disaggregated page table (NVIDIA#20709) cannot resolve
the KV-head count of a mixed-precision high layer. Main waives it on
DGX_B200 only, and waives the other test with this cause on B300
(nvbugs/7044232). PR NVIDIA#21597 fixes the extractor; drop this line when
it lands.
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…gregated directory
The CPU-Generic stages skip the whole `unittest/disaggregated` entry
through waives.txt (https://nvbugs/7042835, a main break in
test_context_first_profile.py and test_lifecycle_activation.py that
NVIDIA#21597 fixes). This PR removed the A10 file-level entries for ten
disaggregated files it marked cpu_only, so until the waive goes those
files ran nowhere. Waives match the wrapper's nodeid exactly, so
per-file entries are collected and run; the ten files passed on both
x86 and aarch64 in the TRTLLM-15162 gate at the same base.
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Round-4 multi-GPU CI failed [DGX_H100-4_GPUs-PyTorch-Others-1] on
'KVCacheManagerV2' object has no attribute 'low_precision_dtype'.
That regression window came from main (db98553 introduced it,
7d7b098 of NVIDIA#21597 fixed it) and hit every PR in the window; it has
no overlap with this PR's files. Update the base instead of
carrying an out-of-scope fix.
Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Brings the branch up to date with main. main removed the waive of
unittest/disaggregated (NVIDIA#21810) after NVIDIA#21597 fixed those tests, so the
pre-merge CPU stage now runs the tests in tests/unittest/disaggregated,
including the ones this branch adds. No conflicts.
Signed-off-by: berryz <berryz@nvidia.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.
chuangz0
left a comment
Description
#20709 (per-view KV head transfer) made the disaggregated page-table builder resolve a KV-head count for every internal cache layer, either as an index into the manager's per-layer head list or through the DeepSeek-V4 virtual layer mapping.
However, #21021 (mixed-precision KV cache for speculative decoding and the sparse MLA backends) merged after it and records its mixed-precision high layers in the manager's
mixed_precision_high_layersrecord instead of a virtual layer mapping, so building the page table of a mixed-precision KV cache manager raises (Cannot resolve KV-head count for internal layer ...). Each PR passed CI on its own; main failstest_disagg_exports_both_pools_and_pairs_sink_tailandtest_deepseek_v4_mixed_disagg_names_every_layersince #21021 merged.The builder now takes a mixed-precision high layer's head count from the layer it belongs to, read from that record. Other layers resolve as before; no manager change.
Test Coverage
tests/unittest/_torch/pyexecutor/kv_cache/test_mixed_precision_kv_cache.py::test_disagg_exports_both_pools_and_pairs_sink_tailand::test_deepseek_v4_mixed_disagg_names_every_layerpass again (B200).tests/unittest/disaggregated/test_extractor.py::test_v2_builder_resolves_mixed_precision_high_layer_heads(CPU) replaces the test that modelled the high layers through the earlier virtual-layer encoding; it builds the page table from a manager stand-in with themixed_precision_high_layersrecord, with and without a virtual layer mapping of the manager's own layers.Dev Engineer Review
_build_page_table_v2now resolves mixed-precision high-layer head counts from their associated local layers before using virtual-layer mapping. Other layer-resolution paths remain unchanged. The context-first BF16 check now usesmixed_precision_low_dtype, matching the renamed manager field. No public declarations change.QA Engineer Review
test_extractor.pyadds parameterized coverage for mixed-precision high-layer resolution with and without virtual-layer mapping. It checks layer IDs, pool views, and head counts. The file is listed in the A10 and H100 CI test lists. The context-first profile test updates the NVFP4 case to usemixed_precision_low_dtype; it was not found in the inspected test lists. Coverage verdict: sufficient for the reported resolution scenarios. The author reports two B200 tests pass, but those results were not independently verified.Per-File QA Perspective
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py: Verify that mixed-precision high layers use the head count of their associated local layer, including when virtual-layer metadata exists. Verify that other layers preserve their prior resolution behavior.tests/unittest/disaggregated/test_extractor.py: The parameterized test covers high-layer resolution with and without virtual-layer mapping. The file is listed in the A10 and H100 CI test lists.tensorrt_llm/_torch/disaggregation/native/context_first_profile.py: Verify that context-first BF16 eligibility readsmixed_precision_low_dtypeand preserves the other eligibility conditions.tests/unittest/disaggregated/test_context_first_profile.py: The NVFP4 rejection test now setsmixed_precision_low_dtype. It was not found in the inspected CI or manual-QA test lists.