[#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.

Summary by CodeRabbit

Bug Fixes

• Improved KV-cache token estimation accuracy on integrated and unified-memory devices for more precise token capacity calculations.

Tests

• Added regression test for KV-cache token estimation behavior on integrated GPU systems.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c0f2d7ee-abac-48f6-b139-6752436bfaf6

📥 Commits

Reviewing files that changed from the base of the PR and between 90cb7ff and 38864aa.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/_util.py
• tests/unittest/_torch/executor/test_kv_cache_estimation.py

---

📝 Walkthrough

Walkthrough

The PR fixes KV-cache token estimation on integrated/unified-memory GPUs by importing a device detection utility and conditionally zeroing the affine cost intercept during dry-run sizing, preventing token-cap collapse. A regression test validates the fix.

Changes

KV-cache Estimation for Integrated GPUs

Layer / File(s)

Summary

Intercept-zeroing logic and regression test tensorrt_llm/_torch/pyexecutor/_util.py, tests/unittest/_torch/executor/test_kv_cache_estimation.py

Imports is_device_integrated and detects positive intercepts in the KV-size-per-token affine model; on integrated devices, the intercept is zeroed during token-budget calculation to prevent the provisional token cap from collapsing. A new regression test stubs GPU memory and verifies that token estimation still rounds correctly to the expected block-based total when is_device_integrated=True and an affine intercept is present.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Possibly related issues

• [Bug] Hybrid Mamba KV-cache estimation clamps max_tokens to 0 on integrated/unified-memory GPUs (regression in 1.3.0rc15, commit 091ad7b0) #15178: Directly addresses the integrated GPU KV-cache estimation bug by ignoring the affine intercept during dry-run sizing.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: fixing unified-memory (integrated GPU) Mamba KV cache estimation, with proper GitHub issue reference and fix type.
Description check	✅ Passed	The description includes all key required sections: a clear explanation of the issue and solution, test coverage details with specific pytest commands, and a completed PR checklist confirming guideline compliance.
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.

_{✏️ 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.}

[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.