[None][feat] Support the masked DSA indexer k-cache pool in the Python cache transceiver

[Tabrizian]

Description

The Python (v2) KV-cache transceiver previously raised NotImplementedError for 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 a GlmMoeDsaForCausalLM -> CPP preference 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 indexer REPLICATED pool view now covers only the indexer-owning layers — one buffer_entries row per owning layer, each mapped to its packed pool row via impl.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.
• No transfer-machinery changes are needed: it already matches peers per-pool by pool_role + global_layer_id overlap, so a masked subset transfers correctly, including under PP resharding (the generic analogue of the C++ indexerLayerNumPerPP / targetIRanksForIndexerKCache interval logic).
• With Python support in place, GlmMoeDsaForCausalLM prefers the Python transceiver again like the other DeepSeek-family checkpoints, reverting the CPP override 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 V1 KVCacheManager with a per-layer indexer mask ([True, False, True, False]) and asserts the indexer view is REPLICATED, 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_entries and ::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 expect PYTHON for glm_moe_dsa again.

PR Checklist

• PR description clearly explains what and why.
• PR follows TRT-LLM coding guidelines.
• Test cases are provided for new code paths.
• No API changes (the affected get_preferred_transceiver_runtime is 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_table now supports masked DSA indexer K-cache pools.
• It creates REPLICATED views only for indexer-owning layers.
• It maps each layer to its packed pool row.
• It omits fully masked layer groups.
• Dense and unmasked layouts remain unchanged.
• Existing transfer logic supports masked subsets and pipeline-parallel resharding.
• GlmMoeDsaForCausalLM now prefers the Python transceiver.
• No public API, configuration, or test-list changes were identified.
• The implementation should be checked against CODING_GUIDELINES.md, with focus on offset mapping, empty-group handling, and regression coverage for dense layouts.

QA Engineer Review

• Modified test coverage includes:
  • _make_v1_dsa_manager, updated to accept an indexer K-cache mask.
  • A CUDA test for masked replicated indexer page-table construction.
  • DeepSeek transceiver preference tests, updated for Python runtime selection.
• The tests cover masked page-table entries, packed offsets, dense transfer behavior, cross-pipeline transfer, and architecture preference.
• No corresponding test-db/ or qa/ coverage was identified in the provided changes.
• Verdict: needs follow-up.

[Tabrizian]

/bot run --disable-fail-fast

[coderabbitai]

Walkthrough

V1 page-table construction now supports partially masked indexer K-cache layers. Fully masked groups omit the indexer pool. DeepSeek and GLM 5.2 configurations now use the Python KV-cache transceiver.

Changes

Masked DSA indexer cache support

Layer / File(s)	Summary
Masked indexer page-table construction tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py, tests/unittest/disaggregated/test_extractor.py	The extractor creates entries only for owning layers, uses packed layer indices, skips fully masked groups, and tests contiguous offsets.
Python transceiver selection tensorrt_llm/_torch/models/modeling_deepseekv3.py, tests/unittest/llmapi/test_llm_args.py	GLM 5.2 masked DSA indexer caches use the Python transceiver. Tests now expect Python for all covered architectures.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant KVCacheManager
  participant KVCacheExtractor
  participant IndexerPoolViews
  KVCacheManager->>KVCacheExtractor: provide layer mask and packed indices
  KVCacheExtractor->>IndexerPoolViews: append owning-layer pools and views
  KVCacheExtractor->>IndexerPoolViews: build masked page-table entries

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16558: Adds the indexer-layer masking that this change extends to partial masks.

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the Python transceiver support for masked DSA indexer K-cache pools.
Description check	✅ Passed	The description explains the problem, solution, test coverage, and checklist status with relevant implementation details.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (2)

tests/unittest/disaggregated/test_extractor.py (1)

211-215: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate 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 win

Use declared KVCacheManager attributes.

KVCacheManager.__init__ always sets enable_indexer_k_cache and indexer_k_cache_local_layer_mask. Replace both getattr calls 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

📥 Commits

Reviewing files that changed from the base of the PR and between 89bba4c and 83eb699.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py
• tensorrt_llm/_torch/models/modeling_deepseekv3.py
• tests/unittest/disaggregated/test_extractor.py
• tests/unittest/llmapi/test_llm_args.py

[tensorrt-cicd]

PR_Github #63946 [ run ] triggered by Bot. Commit: 83eb699 Link to invocation

[tensorrt-cicd]

PR_Github #63946 [ run ] completed with state SUCCESS. Commit: 83eb699
/LLM/main/L0_MergeRequest_PR pipeline #51880 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation