NVIDIA / NVIDIA/TensorRT-LLM

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

Open
#15,178 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

KV-Cache Management Pytorch
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
  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 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.