[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:
  • Creates Indexer-K pool rows only for “full indexer”/indexer-owning layers; shared layers are masked out.
  • Adds compact layer-to-pool-row mapping plus consumer-side validation (masked-out layers/buffers reject access).
• Threads indexer-layer masking through cache/transceiver plumbing:
  • Adds indexerLayerNumPerPP / indexerKCacheLayerMask parameters across C++ cache state, KV cache manager (pool creation + mapping), and CacheTransceiver.
  • Adds Indexer-K-specific cache transmission topology/routing (targetIRanksForIndexerKCache) and buffer offset-ratio handling for BufferKind::kKV_INDEXER.
  • Updates PP resharding / split-concat logic to use indexer-layer counts in “indexer layer space”.
• Updates serialization and Python bindings:
  • Extends kv_cache::CacheState serialization/deserialization with indexerLayerNumPerPP.
  • Nanobind Python bindings accept indexer_k_cache_layer_mask and expose layer→pool-row mapping with masking enforcement; Python NIXL path rejects masked pools (per objective).
• GLM 5.2 correctness/capacity:
  • Reduces Indexer-K cache accounting per token (55,224 → 47,700 bytes/token), improving 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; outputs remain bitwise-identical.
• Disaggregated serving / resource-extractor behavior:
  • Adds fail-fast validation in KV extractor for unsupported masked ownership patterns (raises 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
  • Modified construction to pass indexerKCacheLayerMask (uses std::nullopt).
  • Added masked layout coverage for Indexer-K pool row mapping + default dense mapping behavior.
• cpp/tests/unit_tests/executor/serializeUtilsTest.cpp
  • Modified SerializeUtilsTest.CacheStateIndexerKCache to include indexerKCacheUseFp4 and indexerLayerNumPerPP, and validate serialization round-trip equality.
• cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
  • Added AsymmetricalCacheTestWithIndexerLayerMask and updated setup/block-data logic to use indexer-layer counts when testing Indexer-K transceiver send/receive.

Coverage in tests/integration/test_lists/: these unit tests are not directly represented as list entries (they run via unit test binaries).
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 GLM52 NVFP4 NIXL QA entry for cache_mgr_v1 to use test_nvfp4_nixl[...] instead of test_nvfp4_nixl_python[...].
• tests/integration/test_lists/test-db/l0_dgx_b200.yml
  • Updated DGX B200 coverage for cache_mgr_v1 to use test_nvfp4_nixl[...] instead of test_nvfp4_nixl_python[...] (timeout unchanged).
• tests/integration/test_lists/waives.txt
  • Added a waived multi-GPU test entry for cpp/test_multi_gpu.py::test_cache_transceiver with 8proc-ucx_kvcache-90 (NVBUG 5838199 referenced).

Verdict (integration lists): needs follow-up (CI failures reported for specific commits in the PR context; additionally, CBTS coverage status for new/updated entries is not provided here).

Overall QA verdict

needs follow-up due to the reported L0 pipeline failures (commits listed in PR context) and to confirm the updated masked Indexer-K paths are stable end-to-end on all targeted configurations.

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

[Tabrizian]

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

[Tabrizian]

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

[tensorrt-cicd]

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

[tensorrt-cicd]

PR_Github #61427 [ run ] completed with state ABORTED. Commit: 64de65f

Link to invocation

[tensorrt-cicd]

PR_Github #61444 [ run ] completed with state SUCCESS. Commit: 64de65f
/LLM/main/L0_MergeRequest_PR pipeline #49668 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 --add-multi-gpu-test

[tensorrt-cicd]

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

[tensorrt-cicd]

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