[https://nvbugs/6255037][fix] Count DSA indexer K-cache correctly as UINT8 in KV cache size estimate

[eopXD]

Description

The log change mentioned by the bug description was caused by the following code, which I have changed from the upper to latter in my latest PR #13745 for supporting Gemma4 for AutoDeploy.

  cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:3276:
  - cacheSizeBytes += cacheVolume * BufferDataType(mDataType).getSize();      // manager-wide dtype
  + cacheSizeBytes += cacheVolume * BufferDataType(poolDataType).getSize();   // per-pool: primaryPool->getDataType()

The test cases runs with BF16 MLA KV cache and UINT8 indexer K cache:
  - Old (everything ×2 B): 78 × (1152 + 132×2) = 78 × 1416 = 110,448 B/tok → exactly your kv size per token is 110448.
  - New (indexer at 1 B): 78 × (1152 + 132) = 78 × 1284 = 100,152 B/tok → matches the bad run's 133.44 GiB ÷ 1,430,656 = 100,150. 
  - The gap: phantom second byte on the indexer = 78 × 132 = 10,296 B/tok × 1,430,656 tok = 13.7 GiB ≈ the 147.32 − 133.44 = 13.9 GiB drop. 
This explains the reporting paged KV cache creation drop.

Given the above, this merge request accounts for token size correctly inside dsa.py.

With the MR, running perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-input_output_len:2000,500-ep:8-gpus:8], I see the improved KV cache blocks 49324 compared to the older 44708 (~10% increase).

[TRT-LLM] [I] [batchmgr][RANK 0] Max KV cache blocks per sequence: 79 [window size=2500], tokens per block=32, primary blocks=49324, secondary blocks=0, max sequence length=2500

[INFO] Running command for perf/test_perf.py::test_perf_metric_kv_cache_size[glm_5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-input_output_len:2000,500-ep:8-gpus:8]
[INFO] Reusing cached logs for command index 2.
[INFO] Searching for metric perf/test_perf.py::test_perf_metric_kv_cache_size[glm_5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-input_output_len:2000,500-ep:8-gpus:8] from output log of command 2 ...
[INFO] Use max value 147.22 out of [1.76, 1.76, 1.76, 1.76, 1.76, 1.76, 1.76, 1.76, 147.22, 147.22, 147.22, 147.22, 147.22, 147.22, 147.22, 147.22] for perf metric perf/test_perf.py::test_perf_metric_kv_cache_size[glm_5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-input_output_len:2000,500-ep:8-gpus:8].
[INFO] pytorch/TRT model config: {'print_iter_log': True, 'cuda_graph_config': {'enable_padding': True}, 'enable_autotuner': False}

Test Coverage

Tested locally. The QA CI should restore its KV cache allocation size to the reported after this MR is merged.

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

DSACacheManager cache size accounting is refactored to explicitly separate MLA latent K-cache and indexer K-cache memory computations. Both get_cache_size_per_token and get_cache_bytes_per_token now compute indexer bytes independently and add them to MLA latent bytes, replacing the previous combined kv_factor multiplier approach.

Changes

DSA Cache Size Accounting Refactoring

Layer / File(s)	Summary
K-cache memory computation separation tensorrt_llm/_torch/attention_backend/sparse/dsa.py	get_cache_size_per_token computes MLA latent K-cache bytes, then separately calculates indexer_bytes_per_token (UINT8 allocation with FP4/FP8 sizing), and adds both. get_cache_bytes_per_token derives indexer_bytes_per_token from sum(num_kv_heads_per_layer) and directly increments cache_size_bytes_per_token, replacing the prior combined kv_factor-based approach in both methods.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested reviewers

• yechank-nvidia
• nvpohanh
• JunyiXu-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.
Description check	✅ Passed	The PR description covers the issue, solution with code examples, test coverage, and PR checklist items. All required sections from the template are addressed.
Title check	✅ Passed	The title accurately describes the main change: fixing DSA indexer K-cache counting as UINT8 in KV cache size estimation, which directly aligns with the changeset's core modification of separating indexer bytes calculation from MLA latent bytes.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[eopXD]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52723 [ run ] triggered by Bot. Commit: a04f718 Link to invocation

[tensorrt-cicd]

PR_Github #52723 [ run ] completed with state SUCCESS. Commit: a04f718
/LLM/main/L0_MergeRequest_PR pipeline #41989 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

[eopXD]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52811 [ run ] triggered by Bot. Commit: a04f718 Link to invocation

[tensorrt-cicd]

PR_Github #52811 [ run ] completed with state SUCCESS. Commit: a04f718
/LLM/main/L0_MergeRequest_PR pipeline #42065 completed with status: 'SUCCESS'

CI Report

Link to invocation

[QiJune]

LGTM