[Bug] Hybrid Mamba KV-cache estimation clamps max_tokens to 0 on integrated/unified-memory GPUs (regression in 1.3.0rc15, commit 091ad7b0)

[janbernloehr]

System Info

• TensorRT-LLM 1.3.0rc15 through 1.3.0rc17 (fails); 1.3.0rc14 is the last good version.
• GPU: integrated / unified-memory device, compute capability 12.1 (GB10-class, e.g. DGX Spark), arch linux/arm64. CUDA 13.1, driver 580.126.x, PyTorch 2.11.0a0, container nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc17.
• Contrast device that passes: discrete RTX Pro 6000 Blackwell (dedicated VRAM), same TRT-LLM 1.3.0rc17.

Who can help?

Maintainers of the PyTorch-flow KV-cache memory estimator (tensorrt_llm/_torch/pyexecutor/_util.py: KvCacheCreator._get_token_num_for_estimation / CacheCost).

Information

• The official example scripts (standard trtllm PyTorch-flow serving)
• My own modified scripts

Tasks

• My own task: serve a hybrid Mamba+attention model on an integrated GPU.

Reproduction

1. On an integrated/unified-memory GPU (compute capability 12.x, GB10/DGX Spark), serve a ~31B hybrid Mamba+attention model (attention window 3000 + mamba layers) in bf16, TP=1, PyTorch flow, with kv_cache_config.enable_block_reuse=false.
2. The ~58.8 GB bf16 weights are mmap'd from safetensors. On integrated GPUs TRT-LLM keeps them in shared host memory and logs: "Skipping device transfer from cpu to cuda on integrated GPU to conserve shared memory."
3. KV-cache estimation runs.

Expected behavior

On a discrete GPU with the same rc17 binary and identical config, the estimation dry run yields {3000: (512, 0), -2147483647: (512, 0)} and the server starts successfully (exit 0).

actual behavior

On the integrated GPU, torch.cuda.mem_get_info().free is depressed because the mmap'd model weights are resident in the same physical pool as torch's weight copy. In _get_token_num_for_estimation the estimation budget free_gpu_memory_fraction * free is then reduced by the affine CacheCost intercept (the mamba recurrent-state fixed cost — observed at ~12.78 GB), and CacheCost.tokens_for_budget clamps the result to 0. The estimation dry run gives the attention window 0 primary blocks ({3000: (0, 0), -2147483647: (512, 0)}), and assert max_blocks_per_seq > 0 ("Impossible to fit in any sequence in kvCache") fires at resource_manager.py get_max_atten_window_upper_bound.

additional notes

• Bisected to commit 091ad7b ([TRTLLM-12026][feat] Support MTP with block reuse enabled for hybrid models #12896, "Support MTP with block reuse enabled for hybrid models"), which replaced the proportional estimation cost (max_memory // bytes_per_token) with the affine CacheCost.tokens_for_budget(max_memory) = max((budget - intercept) // slope, 0).
• Validated experimentally: dropping the intercept in the estimation dry-run sizing only (restoring rc14 proportional behavior) makes the same rc17 binary pass on the integrated GPU (dry-run blocks recover to {3000: (512, 0), ...}, exit 0). The real post-profiling allocation is unaffected by the change.
• Suggested fix direction: in the estimation dry run on integrated/unified-memory devices, either (a) size from the true device budget rather than the mmap-depressed mem_get_info().free, or (b) treat the recurrent-state fixed cost as a reservation rather than subtracting it before clamping max_tokens to 0, or (c) account for the integrated-GPU double-resident weights when computing the estimation free budget.

Before submitting a new issue...

• Searched for relevant issues and checked documentation/examples.

[peter941221]

I opened a draft fix in #15215.

The patch only relaxes the provisional dry-run token cap on integrated GPUs. It leaves the final affine sizing path unchanged and adds a regression test for the zero-clamp case from this report.