[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

[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

[tensorrt-cicd]

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

[tensorrt-cicd]

PR_Github #63792 [ run ] completed with state FAILURE. Commit: 2a62a72
/LLM/main/L0_MergeRequest_PR pipeline #51738 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

[fredricz-20070104]

Review: CONCERNS (approving with a follow-up request)

The dense/legacy path is byte-for-byte unchanged, the C++ masked-indexer logic is well covered (row-mapping, dense default, serialization round-trip, new PP-resharding test), and the transceiver owner has approved. Approving so this isn't blocked — but please address the following.

🔥 MAJOR — waiving a multi-GPU transceiver test this PR itself rewrites

tests/integration/test_lists/waives.txt:106 adds a SKIP for cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90] (nvbugs/5838199), in the same PR that heavily changes cacheTransceiverTest.cpp and the C++ transceiver / PP-resharding paths. The PR's own QA note says it "needs follow-up due to the reported L0 pipeline failures."

• Failure path: if the masked-indexer transceiver changes regress the UCX/sm90 path, this skip hides it. The previous waive only covered the mooncake variant; extending it to ucx while rewriting that exact code is the classic pattern for masking a real regression.
• Ask: attach a root-cause analysis showing the ucx/sm90 failure is pre-existing and unrelated to this PR, or fix it — rather than waiving your own multi-GPU test in the PR that touches it.

👀 MINOR (non-blocking)

• dsa.py:3480 — in get_cache_size_per_token, the num_layers is not None branch uses max(num_layers, 1) (counts all layers), while the else branch correctly counts only is_full_indexer_layer. For masked models this over-estimates indexer bytes/token and partially offsets the capacity win this PR is after. Direction is conservative/safe but inconsistent; CodeRabbit raised this repeatedly.
• mlaCacheFormatter.cpp:615 — the zero-size skip uses != nullptr && getSize()==0, while the sender side (~line 322) uses == nullptr ||. inquireSupport currently rejects all zero-overlap configs so this is effectively unreachable, but for consistency prefer == nullptr || ...->getSize()==0.

🔍 QA view

• Coverage: partial. C++ side is well covered. The Python masked path has no direct unit test — derive_indexer_k_cache_layer_mask, the None-guard in get_indexer_k_cache_buffers, and the two NotImplementedError branches in kv_extractor.py are untested.
• SM gap: code touches sm90 (fp8 indexer) and sm100 (GLM5.2 nvfp4). sm100 has B200 integration coverage, but the sm90 ucx multi-GPU e2e coverage is exactly what this PR waives.
• Needs /qa-verify: yes.

What I could not verify

Whether the ucx/sm90 L0 failure is genuinely pre-existing vs introduced here — that's the crux of the MAJOR and needs the RCA above.

[fredricz-20070104]

Approving to unblock. Correctness on the dense/legacy path is unchanged and the C++ masked-indexer logic is well tested. Please follow up on the MAJOR concern (the ucx/sm90 transceiver waive) with an RCA or a fix in a follow-up.

[pengbowang-nv]

Attention part LGTM. Left a small question.

[Tabrizian]

/bot skip --comment "unrelated CI failure"

[Tabrizian]

@fredricz-20070104 I was able to reproduce the UCX failure on main too.

[tensorrt-cicd]

PR_Github #63927 [ skip ] triggered by Bot. Commit: 2a62a72 Link to invocation

[Tabrizian]

/bot skip --comment "unrelated CI failure"

[tensorrt-cicd]

PR_Github #63933 [ skip ] triggered by Bot. Commit: 2a62a72 Link to invocation

[tensorrt-cicd]

PR_Github #63927 [ skip ] completed with state ABORTED. Commit: 2a62a72

Link to invocation

[tensorrt-cicd]

PR_Github #63933 [ skip ] completed with state SUCCESS. Commit: 2a62a72
Skipping testing for commit 2a62a72

Link to invocation