[None][perf] Allocate DSA indexer k-cache only for layers that own an indexer

[Tabrizian]

Dev Engineer Review

• Implements per-layer masked Indexer-K cache pool allocation for cross-layer sharing:
  • Builds Indexer-K pool rows only for “full indexer”/indexer-owning layers; shared layers are masked out.
  • Adds a compact global layer → masked pool-row mapping (getIndexerKCachePoolLayerIdx), and enforces rejection of masked-out layers/buffers on the consumer side (C++ and Python).
• Threads indexer-layer masking through cache/transceiver plumbing:
  • Adds indexerKCacheLayerMask/indexerLayerNumPerPP through KV cache manager construction, CacheTransceiver, and kv_cache::CacheState (including equality + toString() updates).
  • Updates Indexer-K specific cache transmission topology/routing (targetIRanksForIndexerKCache) and buffer-kind offset-ratio handling for BufferKind::kKV_INDEXER.
  • Updates PP resharding/split-concat logic to operate in “indexer layer space” using indexer-layer counts.
  • Adjusts Indexer-K cache transfer buffer sizing to derive sizes directly from the Indexer-K pool (avoiding attention-window iteration when Indexer-K transfer is enabled).
• Updates serialization and language bindings:
  • Extends kv_cache::CacheState serialization/deserialization with indexerLayerNumPerPP.
  • Nanobind Python bindings add indexer_k_cache_layer_mask and expose layer→pool-row mapping with masking enforcement.
  • Python NIXL/DSA path enforces ownership: shared (non-full) indexer layers have None pools and access is blocked.
• GLM 5.2 correctness/capacity:
  • Reduces Indexer-K cache accounting per token (55,224 → 47,700 bytes/token), increasing effective KV capacity by ~15.8%.
  • Dense layouts remain unchanged when cross-layer sharing is not used; shared layers do not allocate Indexer-K pools; model outputs are intended to remain bitwise-identical.
• Disaggregated serving / resource-extractor behavior:
  • Adds fail-fast validation in the disaggregated KV extractor for unsupported masked ownership patterns, with guidance to disable disaggregated serving for cross-layer indexer sharing.

QA Engineer Review

Test-code changes (files outside tests/integration/test_lists/)

• cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
  • Added/updated unit coverage for masked Indexer-K pool row mapping and dense default behavior.
• cpp/tests/unit_tests/executor/serializeUtilsTest.cpp
  • Extended CacheStateIndexerKCache serialization/deserialization round-trip to include indexerKCacheUseFp4 and indexerLayerNumPerPP.
• cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
  • Added AsymmetricalCacheTestWithIndexerLayerMask and updated Indexer-K block-data generation/verification to use indexer-layer counts.

Coverage in tests/integration/test_lists/: these are unit-test binary changes (not directly represented in the integration test list files).
Verdict (unit tests): sufficient.

Test-list / integration changes (files under tests/integration/test_lists/)

• tests/integration/defs/accuracy/test_disaggregated_serving.py
  • Renamed TestGLM52NVFP4.test_nvfp4_nixl_python[...] → TestGLM52NVFP4.test_nvfp4_nixl[...].
• tests/integration/test_lists/qa/llm_function_core.txt
  • Updated the GLM52 NVFP4 NIXL QA entry for cache_mgr_v1 to point to test_nvfp4_nixl[...] instead of test_nvfp4_nixl_python[...].
• tests/integration/test_lists/test-db/l0_dgx_b200.yml
  • Updated the GLM52 NVFP4 coverage for cache_mgr_v1 similarly to use test_nvfp4_nixl[...] (timeout unchanged).
• tests/integration/test_lists/waives.txt
  • Added a skipped multi-GPU test entry for cpp/test_multi_gpu.py::test_cache_transceiver with 8proc-ucx_kvcache-90, referencing NVBUG 5838199.

Verdict (integration lists): needs follow-up (CI failures were reported for specific helper commit runs in the PR context; additionally, CBTS coverage sufficiency for the updated list entries is not confirmed here).

Overall QA verdict

needs follow-up due to the reported L0 pipeline failures and to validate end-to-end stability across the affected masked Indexer-K transceiver/PP-resharding and Python NIXL paths.

Description

GLM 5.2 (glm_moe_dsa, served via the DeepSeek-V3 path) uses DeepSeek Sparse Attention with cross-layer indexer sharing: only 21 of 78 layers own an indexer ("full" layers); the remaining 57 "shared" layers reuse the previous full layer's top-k indices and never touch the indexer k-cache. The V1 KVCacheManager allocated the indexer k-cache stride for every layer, so the shared layers' indexer pools were dead memory — 13.6% of the whole KV cache footprint (55,224 B/token instead of 47,700 B/token, i.e. +15.8% effective KV token capacity per rank once removed). On long-context agentic workloads with block reuse, KV capacity directly drives cache hit rate and prefill throughput.

Commit 1 — per-layer masked indexer k-cache pool (KVCacheManager V1):

• New optional per-local-layer indexerKCacheLayerMask through the KVCacheManager/BlockManager/WindowBlockManager ctors; createIndexerKCachePools now allocates one pool row per masked-in layer (rows follow the KV pool's layer order) and records the layer → pool-row map (getIndexerKCachePoolLayerIdx).
• nanobind: indexer_k_cache_layer_mask ctor kwarg; get_indexer_k_cache_pool_data translates local layer → pool row and rejects masked-out layers.
• DSACacheManager derives the mask from to_sparse_params(pretrained_config, layer_idx).is_full_indexer_layer — the exact source of truth MLA uses to decide whether a layer constructs an Indexer module (trailing spec/MTP layers always resolve to full). Shared layers get None pool entries and assert if asked for indexer buffers.
• Both memory estimators (get_cache_size_per_token / get_cache_bytes_per_token) charge indexer bytes for full-indexer layers only, PP-sliced consistently with the pool.
• Models without cross-layer sharing (e.g. DeepSeek V3.2) resolve to an all-full mask and keep the dense legacy layout bit-for-bit. DeepSeek-V4 uses its own cache manager and is unaffected. kv_cache_manager_v2.py is untouched.
• The Python KV-transceiver runtime (NIXL native path) rejects masked pools with a clear NotImplementedError (its FLAT pool-view contract assumes one row per layer); TestGLM52NVFP4::test_nvfp4_nixl_python is switched to the C++ NIXL transceiver and renamed test_nvfp4_nixl (CI list entries updated).

Commit 2 — C++ cache transceiver support for the masked pool:

• CacheState carries per-PP indexer layer counts (indexerLayerNumPerPP; empty = dense fallback to the attention counts), wired through the ctors, operator==, toString, and serialization.
• New targetIRanksForIndexerKCache keeps the attention pass's rank topology (same connection set) but recomputes per-peer layer counts in indexer layer space (interval intersections of the per-PP indexer prefix sums). splitKVCache/concatKVCache use it — plus the self indexer count as the kernel layer count — for the indexer pass.
• MLACacheFormatter sizes indexer-pass buffers with the indexer counts; inquireSupport validates indexer layout equality and equal totals, and rejects PP partitions that would produce zero-sized indexer exchanges (a rank or layer overlap without a full-indexer layer).
• The NIXL/Mooncake agent path computes a dedicated indexer-pass write-offset ratio in indexer layer space (it previously shared the attention-space KV ratio, which would corrupt masked-pool transfers under PP resharding).
• CacheTransceiver takes indexerLayerNumPerPP, validates it against the local pool row count; the Python transceiver gathers it over PP from the manager's indexer layer mask. CacheTransBufferManager sizes the indexer transfer buffer from the pool's actual (masked) layer count.

Expected GLM 5.2 server-start accounting: kv size per token is 47700 bytes/token (was 55,224) and ~+15.8% max tokens in paged KV cache at the same kv_cache_free_gpu_memory_fraction. Shared layers never used the pools, so model outputs are bitwise-identical.

Test Coverage

New tests (all passing on H100):

• cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp — IndexerKCachePoolLayerMaskTest (masked pool shape, row mapping, block size) and IndexerKCachePoolDenseDefaultTest (dense layout unchanged without a mask).
• cpp/tests/unit_tests/executor/serializeUtilsTest.cpp — CacheState round-trip including indexerLayerNumPerPP.
• cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp — a new AsymmetricalCacheTestWithIndexerLayerMask end-to-end masked-indexer transfer with ctx/gen PP resharding (also fixes a positional off-by-one in the existing CacheState construction where isIndexerKCache landed on enablePartialReuse).

Regression: full kvCacheManagerTest and serializeUtilsTest suites pass; existing DSA indexer scatter/roundtrip GPU tests pass (dense pool → identity mapping). End-to-end guard: TestGLM52::test_nvfp4 (GSM8K) exercises the masked pool implicitly; TestDeepSeekV32 covers the dense path.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[Tabrizian]

/bot run --disable-fail-fast

[Tabrizian]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #61600 [ run ] triggered by Bot. Commit: 64de65f Link to invocation

[tensorrt-cicd]

PR_Github #61600 [ run ] completed with state SUCCESS. Commit: 64de65f
/LLM/main/L0_MergeRequest_PR pipeline #49808 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61650 [ run ] triggered by Bot. Commit: 64de65f Link to invocation

[tensorrt-cicd]

PR_Github #61650 [ run ] completed with state SUCCESS. Commit: 64de65f
/LLM/main/L0_MergeRequest_PR pipeline #49856 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61660 [ run ] triggered by Bot. Commit: 643f133 Link to invocation

[tensorrt-cicd]

PR_Github #61660 [ run ] completed with state FAILURE. Commit: 643f133
/LLM/main/L0_MergeRequest_PR pipeline #49866 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61676 [ run ] triggered by Bot. Commit: 643f133 Link to invocation

[tensorrt-cicd]

PR_Github #61676 [ run ] completed with state SUCCESS. Commit: 643f133
/LLM/main/L0_MergeRequest_PR pipeline #49882 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

[Tabrizian]

/bot run --disable-fail-fast

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61989 [ run ] triggered by Bot. Commit: 81c0712 Link to invocation

[coderabbitai]

♻️ Duplicate comments (1)

tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)

3370-3379: ⚠️ Potential issue | 🟠 Major

Previously reported indexer accounting bugs remain.

The num_layers branch still counts every layer, and get_cache_bytes_per_token still sums KV heads instead of one row per owning layer. Masked DSA layouts therefore overestimate indexer bytes and understate capacity. This is the same unresolved issue as the previous review comments.

Also applies to: 3414-3424

🤖 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/attention_backend/sparse/dsa.py` around lines 3370 -
3379, The indexer accounting still overestimates masked DSA memory. Update the
num_layers handling near num_indexer_layers and get_cache_bytes_per_token so
only indexer-owning layers are counted, with one indexer row per owning layer
rather than summing KV heads; preserve the existing layer mapping and
full-indexer detection logic for both branches.

🧹 Nitpick comments (1)

tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)

3193-3205: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate the new helper interface.

pretrained_config is untyped and the return annotation uses List[bool]; use the repository’s precise config type and list[bool], and add a short Google-style docstring describing the per-layer mask and dense fallback. As per coding guidelines, every function must be annotated and built-in generics are preferred; repository learnings confirm Python 3.10+ support.

🤖 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/attention_backend/sparse/dsa.py` around lines 3193 -
3205, Update derive_indexer_k_cache_layer_mask to type pretrained_config with
the repository’s precise configuration type and change the return annotation to
list[bool]. Add a concise Google-style docstring describing that it returns a
per-layer indexer mask and uses dense layers as the fallback.

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.

Duplicate comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 3370-3379: The indexer accounting still overestimates masked DSA
memory. Update the num_layers handling near num_indexer_layers and
get_cache_bytes_per_token so only indexer-owning layers are counted, with one
indexer row per owning layer rather than summing KV heads; preserve the existing
layer mapping and full-indexer detection logic for both branches.

---

Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 3193-3205: Update derive_indexer_k_cache_layer_mask to type
pretrained_config with the repository’s precise configuration type and change
the return annotation to list[bool]. Add a concise Google-style docstring
describing that it returns a per-layer indexer mask and uses dense layers as the
fallback.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 9ecbb79c-929f-42c7-ac62-ccbe0ff9f94c

📥 Commits

Reviewing files that changed from the base of the PR and between 643f133 and 81c0712.

📒 Files selected for processing (18)

• cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
• cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
• cpp/include/tensorrt_llm/executor/dataTransceiverState.h
• cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp
• cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
• cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp
• cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp
• cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu
• cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h
• cpp/tensorrt_llm/executor/serialization.cpp
• cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp
• cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
• cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
• cpp/tests/unit_tests/executor/serializeUtilsTest.cpp
• cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
• tensorrt_llm/_torch/attention_backend/sparse/dsa.py
• tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py

🚧 Files skipped from review as they are similar to previous changes (15)

• cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
• cpp/tests/unit_tests/executor/serializeUtilsTest.cpp
• cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h
• cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp
• cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp
• cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu
• cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp
• cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
• cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp
• cpp/include/tensorrt_llm/executor/dataTransceiverState.h
• cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
• cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
• cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
• cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

[tensorrt-cicd]

PR_Github #61989 [ run ] completed with state FAILURE. Commit: 81c0712
/LLM/main/L0_MergeRequest_PR pipeline #50178 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