[None][feat] Support the masked DSA indexer k-cache pool in the Python cache transceiver - #17283
Conversation
|
/bot run --disable-fail-fast |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughV1 DSA page-table construction now supports masked indexer K-cache layers and omits fully masked pools. GLM DSA models now default to the Python transceiver. Bounce transport sizing counts distinct physical pools, and fan-in validation checks replicated views on both endpoints. ChangesMasked DSA cache handling
Bounce transport safety and sizing
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant KVCacheManager
participant KVCacheExtractor
participant Receiver
participant PeerPageTable
participant ReceiverPageTable
KVCacheManager->>KVCacheExtractor: provide masked layers and packed indices
KVCacheExtractor->>ReceiverPageTable: build owning-layer views
Receiver->>PeerPageTable: inspect replicated pool views
Receiver->>ReceiverPageTable: inspect replicated pool views
Receiver-->>Receiver: allow or reject fan-in bounce
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/unittest/disaggregated/test_extractor.py (1)
211-215: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate
indexer_k_cache_layer_mask.Declare the new parameter as
list[bool] | None. This preserves the helper contract and matches the documented global mask format.As per coding guidelines, “Annotate every function” and “prefer built-in generic types and
|.”🤖 Prompt for 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. In `@tests/unittest/disaggregated/test_extractor.py` around lines 211 - 215, Update the _make_v1_dsa_manager parameter annotation for indexer_k_cache_layer_mask to list[bool] | None, preserving its existing default and behavior.Source: Coding guidelines
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py (1)
336-337: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse declared
KVCacheManagerattributes.
KVCacheManager.__init__always setsenable_indexer_k_cacheandindexer_k_cache_local_layer_mask. Replace bothgetattrcalls with direct attribute access. This keeps the manager contract type-checkable and fails fast on an invalid manager.As per coding guidelines, “Avoid reflection when ordinary explicit code is sufficient.”
🤖 Prompt for 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. In `@tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py` around lines 336 - 337, In the KV cache extraction logic, replace both getattr calls on kv_cache_manager with direct access to its declared enable_indexer_k_cache and indexer_k_cache_local_layer_mask attributes, preserving the existing conditional behavior and allowing invalid managers to fail fast.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py`:
- Around line 336-337: In the KV cache extraction logic, replace both getattr
calls on kv_cache_manager with direct access to its declared
enable_indexer_k_cache and indexer_k_cache_local_layer_mask attributes,
preserving the existing conditional behavior and allowing invalid managers to
fail fast.
In `@tests/unittest/disaggregated/test_extractor.py`:
- Around line 211-215: Update the _make_v1_dsa_manager parameter annotation for
indexer_k_cache_layer_mask to list[bool] | None, preserving its existing default
and behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 37d069cf-5d7c-4823-a88e-1fc952bcb392
📒 Files selected for processing (4)
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytests/unittest/disaggregated/test_extractor.pytests/unittest/llmapi/test_llm_args.py
|
PR_Github #63946 [ run ] triggered by Bot. Commit: |
|
PR_Github #63946 [ run ] completed with state
|
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
tensorrt_llm/inputs/prefix_token_cache.py (1)
30-30: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd complete cache API contracts.
Add return types for
_bucket_key()and_find_longest_prefix(). Typetokenizerand**kwargsinencode(), preferably with a tokenizerProtocol. Add docstrings for exportedprefix_cache_enabled()and publicencode().As per coding guidelines, “Annotate every function” and “Use docstrings rather than comments for externally usable interfaces.”
Also applies to: 67-70, 96-96
🤖 Prompt for 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. In `@tensorrt_llm/inputs/prefix_token_cache.py` at line 30, Complete the cache API contracts in _bucket_key(), _find_longest_prefix(), encode(), and prefix_cache_enabled(): add explicit return annotations, type encode()’s tokenizer and **kwargs (prefer a tokenizer Protocol), and document the exported prefix_cache_enabled() and public encode() APIs with docstrings.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
3422-3427: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the unreachable disabled implementations.
Both methods return before their former logic. Delete the unreachable code and remove now-unused admission helpers if no other caller needs them.
tensorrt_llm/_torch/pyexecutor/py_executor.py#L3422-L3427: remove the unreachable admission-controller branch.tensorrt_llm/_torch/pyexecutor/py_executor.py#L3521-L3536: remove the unreachable idle-progress branch.🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3422 - 3427, Remove the unreachable admission-controller branch at tensorrt_llm/_torch/pyexecutor/py_executor.py:3422-3427 and the unreachable idle-progress branch at tensorrt_llm/_torch/pyexecutor/py_executor.py:3521-3536 from their respective methods, preserving the active return behavior. Afterward, remove any admission helpers left unused if no other callers depend on them.tensorrt_llm/_torch/disaggregation/native/bounce/config.py (1)
121-123: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse Python 3.10 annotation syntax consistently.
Replace
Optional[int]withint | None. ReplaceList[int]withlist[int].
tensorrt_llm/_torch/disaggregation/native/bounce/config.py#L121-L123: useint | Noneformin_blocksandmin_bytes.tensorrt_llm/_torch/disaggregation/native/bounce/impl.py#L67-L72: uselist[int]forblock_bytes_per_group.tensorrt_llm/_torch/disaggregation/native/bounce/impl.py#L102-L114: uselist[int]andfloat | Nonestyle where applicable.🤖 Prompt for 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. In `@tensorrt_llm/_torch/disaggregation/native/bounce/config.py` around lines 121 - 123, Use Python 3.10 union and built-in generic annotation syntax consistently: in config.py lines 121-123, update config_from_size parameters min_blocks and min_bytes to int | None; in impl.py lines 67-72, change block_bytes_per_group to list[int]; and in impl.py lines 102-114, replace applicable List[int] and Optional[float] annotations with list[int] and float | None.Sources: Coding guidelines, Learnings
🤖 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 `@tensorrt_llm/inputs/prefix_token_cache.py`:
- Around line 42-53: Update the prefix-cache configuration flow around __init__
and the environment parsing at the referenced locations to safely parse invalid
numeric values with defaults before cache creation. Validate max_entries,
overlap, resync, min_chars, and bucket_chars constructor bounds before assigning
them, explicitly rejecting overlap <= 0, while preserving valid configured
values and preventing cache initialization from raising.
- Around line 96-151: Add direct PrefixTokenCache tests covering exact token
IDs, seam-resynchronization fallback, concurrent access, eviction, and invalid
settings. In PrefixTokenCache.encode, update the cache entry’s position in
_order whenever an existing entry is reused, moving its ID to the most-recent
end before _evict runs, so eviction follows LRU rather than FIFO.
In `@tensorrt_llm/inputs/registry.py`:
- Around line 141-147: Update the exception handling around
PrefixTokenCache.encode in the prefix-cache branch to catch only the dedicated
cache-fallback exception or explicitly enumerated recoverable exceptions.
Preserve the existing cache-disable fallback for those failures, while allowing
programming errors and unexpected tokenizer exceptions to propagate.
---
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/native/bounce/config.py`:
- Around line 121-123: Use Python 3.10 union and built-in generic annotation
syntax consistently: in config.py lines 121-123, update config_from_size
parameters min_blocks and min_bytes to int | None; in impl.py lines 67-72,
change block_bytes_per_group to list[int]; and in impl.py lines 102-114, replace
applicable List[int] and Optional[float] annotations with list[int] and float |
None.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3422-3427: Remove the unreachable admission-controller branch at
tensorrt_llm/_torch/pyexecutor/py_executor.py:3422-3427 and the unreachable
idle-progress branch at tensorrt_llm/_torch/pyexecutor/py_executor.py:3521-3536
from their respective methods, preserving the active return behavior. Afterward,
remove any admission helpers left unused if no other callers depend on them.
In `@tensorrt_llm/inputs/prefix_token_cache.py`:
- Line 30: Complete the cache API contracts in _bucket_key(),
_find_longest_prefix(), encode(), and prefix_cache_enabled(): add explicit
return annotations, type encode()’s tokenizer and **kwargs (prefer a tokenizer
Protocol), and document the exported prefix_cache_enabled() and public encode()
APIs with docstrings.
🪄 Autofix
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: 6c97242f-8dfe-4002-8997-496531318a02
📒 Files selected for processing (11)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/disaggregation/native/bounce/config.pytensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/inputs/prefix_token_cache.pytensorrt_llm/inputs/registry.pytests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yamltests/integration/defs/disaggregated/test_disaggregated.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pytests/unittest/disaggregated/test_bounce.py
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: Cannot merge as-is — mergeable_state is dirty (rebase needed), and the PR bundles several unrelated, untested behavioural changes on top of the masked-indexer feature. No proven crash, but the risk-without-test surface is large.
Concerns
-
[MAJOR]
tensorrt_llm/inputs/prefix_token_cache.py- new 177-line cache has no tests- What is wrong:
encode()/_find_longest_prefix()/_evict()implement a thread-safe LRU + BPE-seam resync splice, but no test file for this module is in the diff. - How it fails: with
TLLM_PREFIX_TOKEN_CACHE=1, a splice/resync edge case (e.g.len(pids)-reuse <= 0, non-contiguous seam offsets) would silently return wrong token IDs for any prompt ≥min_chars, corrupting generation. Off-by-default limits blast radius but the path is unexercised. - Suggested fix: add unit tests (hit/miss, resync fallback, eviction, byte-identical vs. full tokenization). This module is also unrelated to the PR title — consider splitting it out.
- What is wrong:
-
[MAJOR]
tensorrt_llm/_torch/pyexecutor/py_executor.py:3422&:3521- disagg admission/idle collectives disabled, untested- What is wrong:
_apply_disagg_transfer_admissionnow unconditionally returns(requests, False)and_check_disagg_transfer_progress_when_idlereturns immediately, removing the in-flight-transfer bound and two rank-scoped collectives. - How it fails: correctness (rank-uniformity, completion reaped elsewhere) is asserted only in comments. If any rank ever still needed the status wait, dropping the WORLD/TP collective can hang/diverge disagg ranks; unbounded admission can pressure transfer buffers under load. No test covers this, and it is unrelated to the feature.
- Suggested fix: cover with a disagg integration test or gate behind a flag; at minimum delete the now-unreachable code below each
return.
- What is wrong:
-
[MAJOR]
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py:326- masked transfer only tested at page-table level- What is wrong: the masked REPLICATED indexer view is validated by construction only; the PR's claim that the masked subset "transfers correctly, including under PP resharding" has no end-to-end test.
- How it fails: a per-pool matching mismatch for the masked subset under PP resharding would misroute/drop indexer K-cache and go uncaught.
- Suggested fix: add a masked-layout analogue of
test_v1_dsa_indexer_replicated_transfer_across_pp.
Minor notes (non-blocking)
py_executor.py:3428- remove the dead code after the early returns.config.py:118- default gate semantics change (min_blocks 96→1, new min_bytes=2 MiB); flag in release notes.kv_extractor.py:336- replacegetattron always-set manager attrs with direct access.test_extractor.py- annotateindexer_k_cache_layer_maskaslist[bool] | None.
QA view
- Test coverage: partial - dsa.py MTP-draft rebuild and bounce byte-gate are well covered;
prefix_token_cache.pyand the py_executor disabling are uncovered; masked indexer is construction-only. - SM coverage: DSA/GLM-5.2 masked path targets Hopper (sm90) and Blackwell (sm100/fp8); new DSA tests run
skip_pre_hopper+ DeepGEMM (Hopper only), no explicit Blackwell run for the masked path. - Test code: masked page-table test asserts shape only; no prefix-cache test; helper param unannotated.
- Test time: small - two Hopper-gated unit tests + CPU bounce/extractor cases; integration test only renamed an env var.
- Needs
/qa-verify: yes - disagg behavioural changes + preference flip (GLM 5.2 → PYTHON) + untested new module warrant a real GLM-5.2 disagg run (incl. PP resharding) before trusting.
Possible new issues
- Unbounded disagg admission may exhaust transfer buffers under concurrency.
- Removed status-vote collectives depend on rank-uniformity at every site; if violated, cross-rank hang.
- Byte-gate default change flips which transfers take the coalesced-bounce path for existing models.
What I could not verify
- Runtime safety of removing the disagg collectives (depends on invariants at other executor call sites not shown).
- Whether
get_indexer_k_cache_pool_layer_idx/local_layer_idsindexing holds under all PP configs. - End-to-end masked-indexer transfer correctness — only page-table construction is visible in the diff.
Automated review by NVCortex Lite, run by @fredricz-20070104.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: Bundles disagg-serving behavioural changes (admission control + collective removal) and a masked-indexer transceiver path with no end-to-end disagg test, plus a preference flip (GLM 5.2 now PYTHON) and an untested new tokenization cache; a human QA should run the GLM-5.2 disagg path (incl. PP resharding) before trusting this.
Automated review by NVCortex Lite, run by @fredricz-20070104.
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The core masked-DSA work here is solid — the kv_extractor.py rewrite is correct (I verified the mask is local-layer-keyed on both the Python manager and the C++ getIndexerKCachePoolLayerIdx binding, and the packed-row offset math checks out), and the new tests pin it well. But this PR is six logical commits under a title/description that covers exactly one of them:
- Masked DSA indexer support + GLM 5.2 preference revert (described)
min_blocks→min_bytesbounce gate (60a26e8c3, undescribed)- Multi-pool bounce reservation fix (
40ff740b2, undescribed) - Removal of the disagg transfer admission controller and per-iteration idle-progress collectives in
py_executor.py(d693ae1dc, undescribed) - DSA MTP token-to-request map rebuild fix (
00b3b4aed, undescribed) - A new prefix-tokenization cache feature:
prefix_token_cache.py+ registry hook (174778777, undescribed)
Items 4 and 6 are each PR-worthy on their own. Item 4 in particular changes disagg executor scheduling behavior for every model and deployment, and shipping it silently inside a "[feat] masked DSA indexer" PR means nobody bisecting a future disagg regression will find it from the PR title. Per the repo's own AGENTS.md ("one concern per PR"), please split this — at minimum pull out 4 and 6.
Other whole-PR items:
[None]ticket tag: a transceiver feature plus two perf changes and two bug fixes should carry a JIRA/NVBug reference.- Undocumented env vars:
TLLM_PREFIX_TOKEN_CACHE(+_ENTRIES/_OVERLAP/_RESYNC/_MIN_CHARS) andTRTLLM_KV_CACHE_BOUNCE_MIN_BYTESare user-visible knobs with no docs anywhere. The oldMIN_BLOCKSdefault also changes semantics (96-block gate → vacuous 1), which alters which transfers bounce on existing deployments. prefix_token_cache.pyships with zero tests — 177 lines of BPE-seam splicing, resync fallback, eviction, and locking, all pure Python and eminently unit-testable. The module docstring cites offline validation (192/192 matches), but nothing in-tree guards it against regression.
60a26e8 to
25b5f9b
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/disaggregated/test_extractor.py (1)
277-421: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd
-> Nonetotest_v1_dsa_masked_indexer_page_table_covers_owning_layers.Coverage is sufficient. The three changed tests are covered by
tests/unittest/disaggregated/test_extractor.py, which is registered inl0_a10.ymlandl0_h100.yml.🤖 Prompt for 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. In `@tests/unittest/disaggregated/test_extractor.py` around lines 277 - 421, Add the explicit -> None return annotation to test_v1_dsa_masked_indexer_page_table_covers_owning_layers, matching the annotations used by the other test functions. Leave the test behavior unchanged.Source: Path instructions
🤖 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_extractor.py`:
- Around line 277-284: Add the required -> None return type annotation to
test_v1_dsa_masked_indexer_page_table_covers_owning_layers, preserving the
test’s existing behavior and docstring.
---
Outside diff comments:
In `@tests/unittest/disaggregated/test_extractor.py`:
- Around line 277-421: Add the explicit -> None return annotation to
test_v1_dsa_masked_indexer_page_table_covers_owning_layers, matching the
annotations used by the other test functions. Leave the test behavior unchanged.
🪄 Autofix
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: ab0831d3-1c79-4405-97df-956d3c3a6d0c
📒 Files selected for processing (6)
tensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_extractor.pytests/unittest/llmapi/test_llm_args.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tensorrt_llm/_torch/models/modeling_deepseekv3.py
- tests/unittest/llmapi/test_llm_args.py
- tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py
- tensorrt_llm/_torch/disaggregation/native/bounce/impl.py
- tests/unittest/disaggregated/test_bounce.py
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Approving — all prior-round items check out on the current branch, verified against the code rather than the thread:
- The fan-in bounce gate hole is closed for real:
_fanin_bounce_safe(transfer.py:1590) now takes the receiver's own page table and refuses equal-split fan-in when either endpoint advertises a REPLICATED view; the gate unit test covers the masked-representative case (NHD sender / REPLICATED receiver), andtest_cache_transceiver_v1_masked_dsa_indexer_across_asymmetric_ppexercises the realKvCacheTransceiverV2CTX PP2→GEN PP1 path with the metadata rank fully masked. It's registered inl0_h100.ymland the QA list. - The branch is still 6 on-theme commits; none of the previously-dropped unrelated changes (py_executor idle check, prefix_token_cache, registry swallow) have crept back.
- I re-verified the
GlmMoeDsaForCausalLMsplit:_arch_index.py:59still maps the architecture tomodeling_deepseekv3, so lazy registration resolves the new subclass, and themodel_type == 'glm_moe_dsa'rewrite inDeepseekV3ForCausalLM.__init__is inherited.
One non-blocking ask: the PR description is now stale. It still claims "no transfer-machinery changes are needed", but the branch changes transfer.py (the fan-in gate above) and bounce/impl.py — block_bytes_per_group now sums every distinct physical pool in a layer group instead of pool 0 only. That second change fixes a bug that exists on main independent of masking: with bounce enabled, a dense DSA-indexer receiver under-reserves its bounce region (KV pool only) while the sender's coalesced write also carries the indexer fragments, overrunning into the neighboring slot. Please add both to the description; the sizing fix is exactly what someone triaging a bounce corruption will search PR text for. An NVBug reference for that fix would be ideal given it corrects shipped behavior, though I won't block on it.
|
/bot run --disable-fail-fast |
|
PR_Github #65393 [ run ] triggered by Bot. Commit: |
|
PR_Github #65393 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65496 [ run ] triggered by Bot. Commit: |
|
PR_Github #65496 [ run ] completed with state |
Signed-off-by: Simeng Liu <109828133+SimengLiu-nv@users.noreply.github.com>
|
/bot skip --comment "The pipeline passed. Resolve conflict only." |
|
PR_Github #65599 [ skip ] triggered by Bot. Commit: |
|
PR_Github #65599 [ skip ] completed with state |
…made it prefer the Python transceiver) Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…made it prefer the Python transceiver) Signed-off-by: Eric Tsai <ertsai@nvidia.com>
16 participants
Description
The Python (v2) KV-cache transceiver previously raised
NotImplementedErrorfor the per-layer masked DSA indexer k-cache pool (cross-layer indexer sharing, e.g. GLM 5.2), so those checkpoints were forced onto the C++ transceiver — #16558 added aGlmMoeDsaForCausalLM -> CPPpreference as a stop-gap while the Python path lacked this support.This PR teaches the Python transceiver to handle the masked layout, mirroring the C++ support added in #16558:
build_page_table(tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py): the indexerREPLICATEDpool view now covers only the indexer-owning layers — onebuffer_entriesrow per owning layer, each mapped to its packed pool row viaimpl.get_indexer_k_cache_pool_layer_idx(lid)— and skips the indexer pool entirely for a layer group with no owning layers (which would otherwise hit the null-pool getter). The dense/unmasked layout is byte-for-byte unchanged.pool_role+global_layer_idoverlap, so a masked subset transfers correctly, including under PP resharding (the generic analogue of the C++indexerLayerNumPerPP/targetIRanksForIndexerKCacheinterval logic).GlmMoeDsaForCausalLMprefers the Python transceiver again like the other DeepSeek-family checkpoints, reverting theCPPoverride that [None][perf] Allocate DSA indexer k-cache only for layers that own an indexer #16558 added only because Python lacked masked-pool support.Test Coverage
tests/unittest/disaggregated/test_extractor.py::test_v1_dsa_masked_indexer_page_table_covers_owning_layers(new) — builds a V1KVCacheManagerwith a per-layer indexer mask ([True, False, True, False]) and asserts the indexer view isREPLICATED, covers exactly the two owning layers, and maps them to the correct packed pool rows/offsets.test_extractor.py::test_v1_dsa_indexer_page_table_is_replicated_with_per_layer_entriesand::test_v1_dsa_indexer_replicated_transfer_across_pp(existing) — continue to guard the dense layout and the end-to-end replicated transfer path.tests/unittest/llmapi/test_llm_args.py::TestDeepseekTransceiverPreference::test_preference_per_architecture— updated to expectPYTHONforglm_moe_dsaagain.Perf
28703342633578PR Checklist
get_preferred_transceiver_runtimeis an internal preference hook).GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
build_page_tablesupports masked indexer K-cache ownership.GlmMoeDsaForCausalLMnow prefers the Python transceiver.L0_MergeRequest_PRpipeline failed. Another CI run is required after the failures are fixed.QA Engineer Review
Modified test code covers:
tests/unittest/disaggregated/test_extractor.py.tests/unittest/disaggregated/test_bounce.py.tests/unittest/llmapi/test_llm_args.py.tests/unittest/disaggregated/test_cache_transceiver_single_process.py.No corresponding
test-db/orqa/entries were modified. CI coverage data is unavailable.Verdict: needs follow-up.