[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

[tensorrt-cicd]

PR_Github #62045 [ run ] triggered by Bot. Commit: d03ad87 Link to invocation

[tensorrt-cicd]

PR_Github #62045 [ run ] completed with state SUCCESS. Commit: d03ad87
/LLM/main/L0_MergeRequest_PR pipeline #50230 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 #62119 [ run ] triggered by Bot. Commit: d03ad87 Link to invocation

[tensorrt-cicd]

PR_Github #62119 [ run ] completed with state FAILURE. Commit: d03ad87
/LLM/main/L0_MergeRequest_PR pipeline #50300 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 #63604 [ run ] triggered by Bot. Commit: 2d76dde Link to invocation

[tensorrt-cicd]

PR_Github #63604 [ run ] completed with state FAILURE. Commit: 2d76dde
/LLM/main/L0_MergeRequest_PR pipeline #51565 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 #63674 [ run ] triggered by Bot. Commit: 75e27c5 Link to invocation

[tensorrt-cicd]

PR_Github #63674 [ run ] completed with state FAILURE. Commit: 75e27c5
/LLM/main/L0_MergeRequest_PR pipeline #51627 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 #63722 [ run ] triggered by Bot. Commit: 2a62a72 Link to invocation