[#15178][fix] Fix unified-memory Mamba KV estimation

[peter941221]

Description

Refs #15178.

On integrated GPUs, the estimation dry run can start from a mem_get_info() budget that is already depressed by mmap-backed weights sharing the same physical memory pool.

When hybrid Mamba models use an affine CacheCost, _get_token_num_for_estimation() subtracts the recurrent-state intercept from that reduced budget and can clamp the provisional token cap to zero. That is enough to trip assert max_blocks_per_seq > 0 in the attention-window path even though the later affine sizing still succeeds.

This change keeps the final affine sizing unchanged. It only relaxes the provisional estimation cap on integrated GPUs by dropping the affine intercept from the dry-run budget calculation.

The regression test covers the zero-clamp case by mocking an affine CacheCost, a small mem_get_info() budget, and an integrated device.

Test Coverage

Validated in the matching TensorRT-LLM 1.3.0rc18 CUDA 13 / PyTorch 2.10 runtime with:

• python -m pytest tests/unittest/_torch/executor/test_kv_cache_estimation.py -k integrated_gpu_estimation_ignores_affine_intercept -q
• python -m pytest tests/unittest/_torch/executor/test_kv_cache_estimation.py -q

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.

Dev Engineer Review

• Integrated-GPU KV-cache estimation ignores the fixed CacheCost intercept only during provisional token-cap calculation.
• Final affine KV sizing remains unchanged.
• No public API signatures changed.
• Discrete-GPU behavior remains unchanged.
• The fix addresses the zero-token clamp caused by temporarily reduced unified-memory availability from mmap-backed weights.
• The change does not address the separate MambaHybridCacheManagerV2 quota error.

QA Engineer Review

• Modified tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py.
• Added test_integrated_gpu_estimation_ignores_affine_intercept().
• The test covers limited memory, an integrated device, and an affine CacheCost intercept that exceeds the available budget.
• No matching test-db/ or qa/ entry was found.
• Reported targeted tests, full tests, and DGX Spark hardware testing completed. Internal CI failures were also reported, so CI status needs follow-up.
• Coverage verdict: needs follow-up.

Per-File QA Perspective

• tensorrt_llm/_torch/pyexecutor/_util.py: Verify that only the provisional integrated-GPU calculation omits the intercept. Verify that final sizing and discrete-GPU behavior remain unchanged.
• tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py: Covers the integrated-GPU zero-token-clamp regression. No matching test-db/ or qa/ list entry was found.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 24bb4587-89be-4f75-be72-1cad1ef16a73

📥 Commits

Reviewing files that changed from the base of the PR and between afe52a0 and d2f6f56.

📒 Files selected for processing (1)

• tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

KV-cache token estimation now ignores the fixed-cost intercept on integrated devices during provisional capacity calculation. A regression test verifies a 64-token estimate when the intercept exceeds available memory.

Changes

Integrated GPU KV-cache Estimation

Layer / File(s)	Summary
Intercept-zeroing logic and regression test tensorrt_llm/_torch/pyexecutor/_util.py, tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py	When the device is integrated and the intercept is positive, the estimator uses a slope-only cache cost for provisional token-cap calculation. The regression test verifies a 64-token estimate with 100 bytes of available memory and an intercept of 128.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to d2f6f

The change prevents integrated-GPU KV-cache estimation from collapsing to zero while preserving final sizing behavior; the supplied regression coverage leaves no merge-blocking risk identified.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the fix for unified-memory Mamba KV estimation and uses the required issue and type format.
Description check	✅ Passed	The description explains the problem, the solution, test coverage, and checklist status. It is complete and aligned with the pull request changes.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 3 files.
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.

✨ Finishing Touches 💡 1

⚔️ Resolve merge conflicts 💡

• Resolve merge conflict in branch fix/unified-mem-mamba-kv-est-clean

🧪 Generate unit tests (beta)

• Create a new PR

---

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

[janbernloehr]

Confirmed on hardware — this fixes #15178 for us. We applied this exact change on top of 1.3.0rc18 and re-ran the workload that was failing: a ~31B hybrid Mamba+attention model (bf16, TP=1) served on a DGX Spark (GB10, unified memory).

• Before (stock rc15–rc18): engine init aborts with AssertionError: Impossible to fit in any sequence in kvCache — the KV-cache estimation dry-run clamps max_tokens to 0 on the integrated device, so the attention window gets 0 primary blocks.
• After (rc18 + this PR): the dry-run no longer clamps to 0, the server starts, and the full benchmark runs to completion (exit 0). The is_device_integrated() gate fires as intended and the final affine KV sizing is unchanged, so this only relaxes the provisional dry-run cap.

Discrete-GPU (RTX Pro 6000) runs were never affected, consistent with the integrated-memory root cause. Thanks for the fix @peter941221 — would be great to land this. Happy to test a candidate build on DGX Spark if useful. cc @joyang-nv

[peter941221]

Thanks for testing this on DGX Spark and confirming that it fixes the rc18 unified-memory KV-cache failure while leaving the final affine sizing unchanged.

[tongyuantongyu]

@VALLIS-NERIA this seems to related to your #12896. Could you take a look?

[VALLIS-NERIA]

Hi @peter941221 , thank you for the fix. It looks good to me but I'm not quite clear about the cause. My understanding is:
The weights are mmap-ed to CPU memory space, which is unified memory in this case, and causes mem_get_info to get a much smaller number. After the estimation dry run, the weights are unloaded and we get a normal free memory number, right?

[VALLIS-NERIA]

/bot run

[tensorrt-cicd]

PR_Github #55449 [ run ] triggered by Bot. Commit: 38864aa Link to invocation

[peter941221]

Hi @peter941221 , thank you for the fix. It looks good to me but I'm not quite clear about the cause. My understanding is: The weights are mmap-ed to CPU memory space, which is unified memory in this case, and causes mem_get_info to get a much smaller number. After the estimation dry run, the weights are unloaded and we get a normal free memory number, right?

@VALLIS-NERIA Yes — on integrated/unified-memory devices, the dry-run sees a temporarily depressed mem_get_info() budget from mmap-backed weights, and subtracting the affine intercept can clamp the provisional token cap to 0. The final affine sizing is unchanged.

[tensorrt-cicd]

PR_Github #55449 [ run ] completed with state FAILURE. Commit: 38864aa
/LLM/main/L0_MergeRequest_PR pipeline #44382 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

[VALLIS-NERIA]

/bot run

[tensorrt-cicd]

PR_Github #55471 [ run ] triggered by Bot. Commit: fc3b3fb Link to invocation

[tensorrt-cicd]

PR_Github #55471 [ run ] completed with state FAILURE. Commit: fc3b3fb
/LLM/main/L0_MergeRequest_PR pipeline #44399 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

[peter941221]

The public GitHub checks are green on the current head. The remaining failure is the internal blossom-ci run on fc3b3fb after main was merged into this branch, and I can't see the internal failure details from my side. Could you share the failing job or log snippet so I can check whether it comes from this patch or from the merged main changes?

[mihai-chiorean]

@peter941221 nice catch on the affine-intercept clamp. fyi #12301 is sitting in the same UMA-detection territory — different files (cpp kvCacheManager + llm_args.py) so no merge conflict, but same symptom (zero/negative caps on UMA) from a different angle. Forces host_cache_size=0 via cudaDevAttrIntegrated. Probably complementary with yours, worth looking at together if a reviewer picks one up.

[peter941221]

@mihai-chiorean
Thanks, that makes sense. This PR is scoped to the Python-side dry-run path where the affine intercept can collapse the provisional token cap to 0 on integrated / unified-memory systems. #12301 looks complementary from the KV-cache-manager / host-cache-budget side, so I agree they are related but fixing different layers. I’d prefer to keep this patch narrowly scoped unless a reviewer wants them aligned more explicitly.

[janbernloehr]

@peter941221 can you maybe merge main / rebase so that we can try to get this finally merged. THX

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[peter941221]

Merged the latest main into this branch in commit �fe52a0bd6. Resolved the _util.py import conflict by preserving both the current main imports and the integrated-device fix, and moved the regression coverage to the current ests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py path. The diff against current main remains limited to the intended two files. The merge commit is GPG-signed and DCO-signed.

[janbernloehr]

Thanks @peter941221 for refreshing this PR! One follow-up for #15178: our testing on newer releases still encounters ValueError: The V2 Mamba GPU cache quota is too small.

This PR addresses the provisional token-cap calculation in the V1 estimation path. With NemotronH using MambaHybridCacheManagerV2 from rc24 onward (via #16598), the V2 quota check needs separate attention; landing the V1 guard alone would not resolve that remaining failure.

Could a KV-cache owner advise whether to cover V2 here or track it in a linked follow-up? It would be helpful to make that scope explicit while moving this fix toward merge. Happy to help validate a candidate on Spark. Thanks!

---

Prepared with assistance from Codex (GPT-6).

[janbernloehr]

@peter941221 can you address the pre-commit ci failures? THX

[peter941221]

Applied the pre-commit fix in d2f6f56. The only failing hook on the previous head was ruff-format; the formatted test file now passes the full pre-commit suite locally. Please trigger CI again on the new head.