[Bug] Hybrid Mamba KV-cache estimation clamps max_tokens to 0 on integrated/unified-memory GPUs (regression in 1.3.0rc15, commit 091ad7b0)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
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
- 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.
- 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."
- 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 091ad7b0 (#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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in tensorrt_llm/_torch/pyexecutor/_util.py at KvCacheCreator._get_token_num_for_estimation and CacheCost.tokens_for_budget, then trace the call to resource_manager.py get_max_atten_window_upper_bound. Reproduce the KV-cache estimation with a hybrid Mamba+attention model on an integrated or unified-memory GPU and compare it with a discrete GPU. Done means estimation retains a positive attention-window block count and the server starts without the assertion.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- ai-infra-agents, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100