[None][fix] align KV slice token_range.end with transferred block count - #15824
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughChanges the default ChangesTokenRange calculation fix
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
3e54ed8 to
a2052de
Compare
|
Could you please simplify the comments? Thanks! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/disaggregated/test_cache_reuse_adapter.py`:
- Around line 334-348: The boundary regression test in
test_extra_tokens_do_not_cross_block_boundary currently assumes the helper still
produces a crossing case only because tokens_per_block happens to be 8. Harden
the setup by asserting the fixture conditions around
_build_transceiver_for_kv_slice and _create_kv_slice still force a
block-boundary crossing before the final equality check, so the test fails
loudly if tokens_per_block or the constructed prompt/extra-token lengths no
longer exercise the intended case.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d992714e-931e-457c-a3e3-a4efdcd95942
📒 Files selected for processing (2)
tensorrt_llm/_torch/disaggregation/transceiver.pytests/unittest/disaggregated/test_cache_reuse_adapter.py
|
PR_Github #56896 [ run ] triggered by Bot. Commit: |
|
PR_Github #56896 [ run ] completed with state
|
da73c60 to
b3bed32
Compare
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_B200-8_GPUs-PyTorch-4, GB200-4_GPUs-PyTorch-5, A30-PyTorch-1, DGX_B200-AutoDeploy-1, DGX_H100-PyTorch-1, RTX5090-PyTorch-1, DGX_H100-PyTorch-2" |
|
PR_Github #57074 [ run ] triggered by Bot. Commit: |
|
PR_Github #57074 [ run ] completed with state
|
|
/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-4, DGX_B200-4_GPUs-AutoDeploy-1" |
|
PR_Github #57314 [ run ] triggered by Bot. Commit: |
|
PR_Github #57314 [ run ] completed with state
|
|
/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-4, DGX_B200-4_GPUs-AutoDeploy-1" |
|
PR_Github #57429 [ run ] triggered by Bot. Commit: |
|
PR_Github #57429 [ run ] completed with state |
|
/bot skip --comment " all tests have passed" |
|
PR_Github #57652 [ skip ] triggered by Bot. Commit: |
|
PR_Github #57652 [ skip ] completed with state |
_create_kv_slice set the default token_range.end to req.prompt_len + num_extra_kv_tokens, while the block list is trimmed to total_blocks = ceil(prompt_len / tpb). The sender reconstructs total_blocks from token_range.end via ceil(end / tpb) to derive per-layer token starts, so the two block counts disagree whenever num_extra_kv_tokens pushes token_range.end across a block boundary. In current configurations the resulting offset is applied symmetrically to the src and dst token starts and cancels in _align_kv_blocks, so no transfer is corrupted today (dst_start_token is always None and num_extra_kv_tokens is smaller than tokens_per_block). It is nonetheless a latent violation of the KVSlice contract (total_blocks = ceil(token_range.end / tpb)) that would surface if dst_start_token becomes non-None, or if num_extra_kv_tokens exceeds tokens_per_block together with asymmetric prefix reuse. Set token_range.end to req.prompt_len so it matches the trimmed block list. Not transferring the extra speculative-decoding blocks follows the NVIDIA/TensorRT-LLM PR NVIDIA#14546 review discussion; the accuracy impact on speculative decoding is not yet confirmed. Update tests: flip test_includes_num_extra_kv_tokens to test_excludes_num_extra_kv_tokens and add a block-boundary regression. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
7bb8790 to
b308579
Compare
|
/bot skip --comment "all tests have passed" |
|
PR_Github #59407 [ skip ] triggered by Bot. Commit: |
|
PR_Github #59407 [ skip ] completed with state |
5 participants
_create_kv_slice set the default token_range.end to req.prompt_len + num_extra_kv_tokens, but the block list is trimmed to total_blocks = ceil(prompt_len / tpb). The sender re-derives per-layer token starts from token_range.end via total_blocks = ceil(end / tpb), so including num_extra_kv_tokens inflates total_blocks past the trimmed block count whenever the extra tokens cross a block boundary, mis-deriving token_start by one block and corrupting the transfer.
Set token_range.end to req.prompt_len so it matches the blocks actually sent. The extra blocks reserved for speculative decoding hold uninitialized KV and should not be transferred (see NVIDIA/TensorRT-LLM PR #14546 review discussion).
Update the unit tests: flip test_includes_num_extra_kv_tokens to test_excludes_num_extra_kv_tokens and add a block-boundary regression.
Summary by CodeRabbit
Bug Fixes
Tests
Description
Description
KvCacheTransceiverV2._create_kv_slicesets the defaulttoken_range.endtoreq.prompt_len + num_extra_kv_tokens, but the block list is independentlytrimmed to
total_blocks = (req.prompt_len + tpb - 1) // tpb. This makes theslice metadata inconsistent with the blocks it carries:
base/transfer.pydocuments the contract
total_blocks = ceil(token_range.end / tpb), and thesender reconstructs
total_blocksfromtoken_range.end(native/transfer.py)to derive per-layer token starts. When
num_extra_kv_tokenspushestoken_range.endacross a block boundary, the reconstructed count no longermatches the trimmed block list.
Scope — this is a metadata/contract-consistency fix, not a data-movement
change. The extra
num_extra_kv_tokensblocks are already excluded from thetransfer by the existing
block_ids[:total_blocks]trim, so which blocks getsent does not depend on
token_range.end. The inflatedtotal_blocksonlyshifts
src_startanddst_startby the same delta, which cancels in_align_kv_blocks(the sole asymmetric consumer,dst_start_token, is alwaysNone). So no transfer is corrupted today; this fix restores the documentedinvariant and removes a latent hazard that would surface if
dst_start_tokenbecomes non-
None, or ifnum_extra_kv_tokens > tokens_per_blockcombines withasymmetric prefix reuse.
Set
token_range.end = req.prompt_lenso the metadata matches what is actuallytransferred. Keeping the extra speculative-decoding blocks out of the transfer
is consistent with the review discussion in
https://github.com/NVIDIA/TensorRT-LLM/pull/14546/changes#r3386942085; the
accuracy impact on speculative decoding is not yet confirmed.
Test Coverage
tests/unittest/disaggregated/test_cache_reuse_adapter.pytest_includes_num_extra_kv_tokens→test_excludes_num_extra_kv_tokens(asserts
token_range.end == prompt_len).test_extra_tokens_do_not_cross_block_boundary(block-boundaryconsistency regression).
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.