[TRTLLM-5838][fix] fix max batch size and max tokens in kv cache estimations for Nemotron-H

[tomeras91]

Throughout the code there are a few places where the number of available kv cache tokens and/or maximum batch size are estimated. These estimations are based on the available free GPU memory and on the memory size of a single kv cache entry. These estimations didn't take into account hybrid models like Nemotron-H, in which not all layers are attention layers and require a kv cache.

Changes in this PR:

• consider mamba cache for max batch size estimation in trtllm-bench throughput command
• take only attention layers into account when estimating maximum number of tokens in kv cache, both in trtllm-bench throughput and in KvCacheCreator
• propagate kv_cache_gpu_mem_fraction from trtllm-bench throughput CLI arg to the function that estimates the maximum batch size
• release mamba cache memory when MambaHybridCacheManager is shut down (+ a small refactor to increase readability of MambaHybridCacheManager).

[Copilot]

Pull Request Overview

This PR fixes the estimation of maximum batch size and maximum token count for the KV cache when using hybrid models like Nemotron-H, ensuring that only attention layers are considered in the KV cache estimations and that mamba cache memory is also taken into account. Key changes include:

• Adjusting the byte-per-token calculation to count only attention layers using the hybrid override pattern.
• Propagating the kv_cache_gpu_mem_fraction CLI argument and applying a conservative adjustment for mamba hybrid models.
• Refactoring resource manager methods to better handle mamba cache blocks and releasing cache memory upon shutdown.

Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tensorrt_llm/bench/build/tuning.py	Adjustments to KV cache estimations and logging for hybrid models.
tensorrt_llm/bench/build/dataclasses.py	Adding hybrid_override_pattern and mamba_config fields to model configurations.
tensorrt_llm/bench/build/build.py	Propagation of kv_cache_gpu_mem_fraction into benchmark engine settings.
tensorrt_llm/bench/benchmark/utils/general.py	Passing the new CLI argument for kv_cache memory fraction.
tensorrt_llm/_torch/pyexecutor/resource_manager.py	Renaming and refactoring resource methods and adding a shutdown method for mamba cache release.
tensorrt_llm/_torch/pyexecutor/config_utils.py	Updating hybrid check logic using getattr.
tensorrt_llm/_torch/pyexecutor/_util.py	Adjusting the cache size calculation using attention layers in hybrid models.

Comments suppressed due to low confidence (1)

tensorrt_llm/bench/build/tuning.py:95

• Consider adding an inline comment to explain the rationale behind squaring kv_cache_gpu_mem_fraction for mamba hybrid models, as it improves clarity on why a more conservative memory fraction is applied.

        kv_cache_gpu_mem_fraction *= kv_cache_gpu_mem_fraction

[Copilot]

[nitpick] The variable 'num_attention_layers' is being reassigned to represent the number of mapped pipeline layers instead of its original meaning. Consider using a new variable name (e.g., 'mapped_attention_layers') to preserve clarity and avoid confusion.

Suggested change

	num_attention_layers = max(len(mapping.pp_layers(num_attention_layers)),
	1)
	mem_per_token *= num_attention_layers * head_dim
	mapped_attention_layers = max(len(mapping.pp_layers(num_attention_layers)),
	1)
	mem_per_token *= mapped_attention_layers * head_dim

[tomeras91]

/bot run

[tensorrt-cicd]

PR_Github #9520 [ run ] triggered by Bot

[tomeras91]

/bot run

[tensorrt-cicd]

PR_Github #9532 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #9520 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #9532 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #6994 completed with status: 'FAILURE'