From 1eda4f4e71e74e8e796579b1c84961f8e2e653e9 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:00:07 +0000 Subject: [PATCH] [https://nvbugs/6459792][fix] Teach bench ModelConfig about kv_lora_rank/qk_rope_head_dim, add is_mla() calc_engine_setting used the plain MHA formula (2 * layers * kv_heads * head_size) for every model, which overestimates KV bytes/token by ~25x for MLA models. For DeepSeek-V3.2 that caps the auto-selected max_batch_size at 64 while the runtime happily sustains 242, costing 20.47% output throughput. The runtime KV allocator (kv_cache_manager_v2.py / resource_manager.py) already special-cases MLA with head_dim = kv_lora_rank + qk_rope_head_dim and kv_factor = 1; only the bench heuristic disagreed. Teach ModelConfig to parse kv_lora_rank / qk_rope_head_dim (with a model_type-based fallback to the DeepSeek 512/64 defaults), add is_mla(), and branch calc_engine_setting on it. Rebased from PR #16707, whose original files tensorrt_llm/bench/build/{tuning,dataclasses}.py were deleted by 3c07ada3c5 (#16612). The change is ported to their successors tensorrt_llm/bench/tuning/{heuristics,dataclasses}.py; the defect was carried over unchanged by that move. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- tensorrt_llm/bench/tuning/dataclasses.py | 57 ++++++++++++++++++++++++ tensorrt_llm/bench/tuning/heuristics.py | 26 +++++++---- 2 files changed, 75 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/bench/tuning/dataclasses.py b/tensorrt_llm/bench/tuning/dataclasses.py index ca7e58a4fb51..1cfb942c7eec 100755 --- a/tensorrt_llm/bench/tuning/dataclasses.py +++ b/tensorrt_llm/bench/tuning/dataclasses.py @@ -19,6 +19,26 @@ load_pretrained_config, ) +# Model types that use Multi-Head Latent Attention (MLA). The runtime KV-cache +# formula for these is a single compressed head with +# head_dim = kv_lora_rank + qk_rope_head_dim, kv_factor = 1 -- matching the MLA +# branch in tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py. +_MLA_MODEL_TYPES = frozenset( + { + "deepseek_v2", + "deepseek_v3", + "deepseek_v32", + "kimi_k2", + "glm_moe_dsa", + } +) + +# Standard DeepSeek-family MLA head geometry (see +# tensorrt_llm/_torch/configs/deepseek_v3.py). Used only when a known-MLA +# model_type is set but the HF config we parsed did not surface these fields. +_MLA_DEFAULT_KV_LORA_RANK = 512 +_MLA_DEFAULT_QK_ROPE_HEAD_DIM = 64 + # Mapping from safetensors dtype strings to bytes per element. # Used to compute checkpoint size from per-dtype element counts. SAFETENSORS_DTYPE_BYTES = { @@ -194,6 +214,26 @@ class ModelConfig(BaseModel): dtype: Literal["float16", "bfloat16", "float32", None] = Field( default="float16", validation_alias=AliasChoices("dtype", "torch_dtype") ) + # MLA-specific attention geometry. Present on DeepSeek V2/V3/V3.2, Kimi-K2, + # and other MLA checkpoints; None for standard MHA / GQA models. When set, + # KV bytes/token uses a single compressed head: + # head_dim = kv_lora_rank + qk_rope_head_dim, kv_factor = 1. + kv_lora_rank: Optional[int] = Field( + default=None, + validation_alias=AliasChoices( + "kv_lora_rank", + AliasPath("text_config", "kv_lora_rank"), + AliasPath("language_config", "kv_lora_rank"), + ), + ) + qk_rope_head_dim: Optional[int] = Field( + default=None, + validation_alias=AliasChoices( + "qk_rope_head_dim", + AliasPath("text_config", "qk_rope_head_dim"), + AliasPath("language_config", "qk_rope_head_dim"), + ), + ) @model_validator(mode="after") def set_values_if_none(self): @@ -206,8 +246,25 @@ def set_values_if_none(self): self.head_size = self.hidden_size // self.num_attention_heads if self.num_attention_layers is None: self.num_attention_layers = self.num_hidden_layers + # For known MLA model_types whose HF configs didn't surface + # kv_lora_rank / qk_rope_head_dim, backfill the standard + # DeepSeek-family geometry so the bench heuristic can still use + # the MLA formula instead of falling back to MHA. + if self.model_type in _MLA_MODEL_TYPES: + if self.kv_lora_rank is None: + self.kv_lora_rank = _MLA_DEFAULT_KV_LORA_RANK + if self.qk_rope_head_dim is None: + self.qk_rope_head_dim = _MLA_DEFAULT_QK_ROPE_HEAD_DIM return self + def is_mla(self) -> bool: + """True when this model uses Multi-Head Latent Attention. + + `set_values_if_none` backfills the MLA head geometry for known + MLA `model_type`s, so checking the parsed fields alone is sufficient. + """ + return bool(self.kv_lora_rank and self.qk_rope_head_dim) + @classmethod def get_param_count_and_checkpoint_size(cls, model_hf_name, hf_model_path): """Read parameter count and checkpoint size from safetensors metadata. diff --git a/tensorrt_llm/bench/tuning/heuristics.py b/tensorrt_llm/bench/tuning/heuristics.py index 0b41fee9767f..1a1d4af1ad59 100755 --- a/tensorrt_llm/bench/tuning/heuristics.py +++ b/tensorrt_llm/bench/tuning/heuristics.py @@ -71,14 +71,24 @@ def calc_engine_setting( logger.info(f"Number of attention layers: {model_config.num_attention_layers}") - gb_per_token = ( - 2 - * model_config.num_attention_layers - * adjusted_num_kv_heads - * model_config.head_size - * byte_per_kv_elem - / (1024**3) - ) + if model_config.is_mla(): + # MLA stores a single compressed KV entry per token: + # head_dim = kv_lora_rank + qk_rope_head_dim, kv_factor = 1. + # Mirrors tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py and + # resource_manager.py, so the bench heuristic agrees with the runtime. + mla_head_dim = model_config.kv_lora_rank + model_config.qk_rope_head_dim + gb_per_token = ( + model_config.num_attention_layers * mla_head_dim * byte_per_kv_elem / (1024**3) + ) + else: + gb_per_token = ( + 2 + * model_config.num_attention_layers + * adjusted_num_kv_heads + * model_config.head_size + * byte_per_kv_elem + / (1024**3) + ) # Number of GPU used for this run. n_gpus = tp_size * pp_size