From 1f543b80931aa23845573241bcace977d723d72b Mon Sep 17 00:00:00 2001
From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Date: Thu, 28 May 2026 13:36:01 -0700
Subject: [PATCH] [nvbugs/6209806][fix] Honor explicit head_dim in LlamaConfig
 validator

transformers 5.5.x LlamaConfig.validate_architecture rejects configs where
hidden_size % num_attention_heads != 0, ignoring an explicit head_dim that
decouples q/k/v projection size from hidden_size (e.g. kanana-1.5-2.1b:
hidden_size=1792, num_attention_heads=24, head_dim=128). Replace the
validator (and its captured entry in __class_validators__) with a permissive
version that allows positive explicit head_dim.

Apply the patch at the two TRT-LLM entry points that load HF configs and
tokenizers, and remove the now-stale test waiver.

Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
---
 .../_torch/pyexecutor/config_utils.py         | 43 +++++++++++++++++++
 tensorrt_llm/tokenizer/tokenizer.py           |  6 +++
 2 files changed, 49 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py
index 383bc3834e5e..0461d40f3844 100644
--- a/tensorrt_llm/_torch/pyexecutor/config_utils.py
+++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py
@@ -532,6 +532,48 @@ def __getitem__(self, key):
 )  # NOTE: HF config.json uses deepseek_v32 as model_type but with same DSV3 config class
 
 
+def _relax_llama_validate_architecture() -> None:
+    """Honor explicit `head_dim` in transformers 5.x LlamaConfig validation.
+
+    transformers 5.5.x LlamaConfig.validate_architecture rejects any config
+    where `hidden_size % num_attention_heads != 0`, ignoring an explicit
+    `head_dim` that decouples q/k/v projection size from `hidden_size`
+    (e.g. kakaocorp/kanana-1.5-2.1b: hidden_size=1792, num_attention_heads=24,
+    head_dim=128). The check is too strict; downstream TRT-LLM code already
+    honors `head_dim`. Idempotent.
+
+    The huggingface_hub @strict decorator captures validators into
+    ``cls.__class_validators__`` at class-definition time, so we must replace
+    the entry in that list — overwriting ``LlamaConfig.validate_architecture``
+    alone is not enough.
+    """
+    try:
+        from transformers.models.llama.configuration_llama import LlamaConfig
+    except ImportError:
+        return
+
+    if getattr(LlamaConfig.validate_architecture, "_trtllm_patched", False):
+        return
+
+    def validate_architecture(self):
+        head_dim = getattr(self, "head_dim", None)
+        if isinstance(head_dim, int) and head_dim > 0:
+            return
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({self.hidden_size}) is not a multiple of "
+                f"the number of attention heads ({self.num_attention_heads}).")
+
+    validate_architecture._trtllm_patched = True
+    LlamaConfig.validate_architecture = validate_architecture
+
+    validators = getattr(LlamaConfig, "__class_validators__", None) or []
+    for i, v in enumerate(validators):
+        if getattr(v, "__name__", None) == "validate_architecture":
+            validators[i] = validate_architecture
+            break
+
+
 def load_pretrained_config(model_name_or_path: str,
                            trust_remote_code: bool = False,
                            checkpoint_format: Optional[str] = None,
@@ -541,6 +583,7 @@ def load_pretrained_config(model_name_or_path: str,
             MistralConfigLoader
         return MistralConfigLoader().load(model_name_or_path).pretrained_config
 
+    _relax_llama_validate_architecture()
     config_dict, _ = transformers.PretrainedConfig.get_config_dict(
         model_name_or_path, **kwargs)
     model_type = config_dict.get("model_type")
diff --git a/tensorrt_llm/tokenizer/tokenizer.py b/tensorrt_llm/tokenizer/tokenizer.py
index 9715865c65a1..db2457d61d84 100644
--- a/tensorrt_llm/tokenizer/tokenizer.py
+++ b/tensorrt_llm/tokenizer/tokenizer.py
@@ -675,6 +675,12 @@ def load_hf_tokenizer(model_dir: str,
     '''
 
     _ensure_gpt2_bytes_to_unicode_compat()
+    # AutoTokenizer.from_pretrained internally calls AutoConfig.from_pretrained,
+    # which triggers transformers 5.x LlamaConfig.validate_architecture.
+    # Relax it so models with explicit `head_dim` (e.g. kanana-1.5-2.1b) load.
+    from tensorrt_llm._torch.pyexecutor.config_utils import \
+        _relax_llama_validate_architecture
+    _relax_llama_validate_architecture()
 
     try:
         tokenizer = TransformersTokenizer.from_pretrained(