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