[https://nvbugs/6327149][fix] Handle EXAONE 4.5 33B memory constraints

tensorrt_llm/_torch/models/modeling_exaone4_5.py

@@ -10,6 +10,7 @@
 
 from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper
 from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg
+from tensorrt_llm.logger import logger
 
 from ...inputs import (
     ContentFormat,
@@ -36,6 +37,31 @@
 )
 from .modeling_utils import ModelConfig, register_auto_model, register_vision_encoder
 
+
+def _normalize_exaone4_5_mtp_layer_types(text_config: dict) -> None:
+    """Remove MTP-only entries from the base decoder layer layout."""
+    layer_types = text_config.get("layer_types")
+    num_hidden_layers = text_config.get("num_hidden_layers")
+    num_mtp_layers = text_config.get("_num_mtp_layers")
+    num_nextn_predict_layers = text_config.get("num_nextn_predict_layers")
+
+    if not (
+        isinstance(layer_types, list)
+        and isinstance(num_hidden_layers, int)
+        and isinstance(num_mtp_layers, int)
+        and num_mtp_layers > 0
+        and num_mtp_layers == num_nextn_predict_layers
+        and len(layer_types) == num_hidden_layers + num_mtp_layers
+    ):
+        return
+
+    logger.warning(
+        f"EXAONE 4.5 config includes {num_mtp_layers} trailing MTP layer type(s); "
+        f"excluding them from the {num_hidden_layers}-layer base decoder layout."
+    )
+    text_config["layer_types"] = layer_types[:num_hidden_layers]
+
+
 # transformers >= 5.8 ships native Exaone4.5 configs with the same
 # sub-config-instantiation logic we'd otherwise re-implement here. Prefer
 # the HF classes when available and only register local fallbacks on older
@@ -66,6 +92,7 @@ def __init__(
         ):
             if isinstance(text_config, dict):
                 text_config = copy.deepcopy(text_config)
+                _normalize_exaone4_5_mtp_layer_types(text_config)
                 model_type = text_config.get("model_type", "exaone4")
                 # BC: EXAONE 4.5 first released with the text model type
                 # as `exaone4_5_text`, later renamed to `exaone4`.

tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py

@@ -131,6 +131,7 @@ class TestExaone4_5_33B(LlmapiAccuracyTestHarness):
         ],
         ids=["full_budget", "forced_chunked_prefill"],
     )
+    @pytest.mark.skip_less_device_memory(60000)
     def test_auto_dtype(self, enable_chunked_prefill, max_num_tokens):
         with LLM(
             self.MODEL_PATH,

tests/integration/test_lists/waives.txt

@@ -289,8 +289,6 @@ full:L40S/accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[
 full:L40S/accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-1] SKIP (https://nvbugs/6322045)
 full:L40S/accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-flashinfer_ssm-False] SKIP (https://nvbugs/6327147)
 full:L40S/accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-triton_ssm-False] SKIP (https://nvbugs/6327147)
-full:L40S/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6327149)
-full:L40S/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6327149)
 full:RTXPro6000D/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/5961814)
 full:RTXPro6000D/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/5961814)
 full:RTXPro6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass-torch_compile=True] SKIP (https://nvbugs/5929339)

tests/unittest/_torch/modeling/test_modeling_exaone4_5.py

@@ -1,12 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import copy
 import os
 from dataclasses import dataclass
 from typing import List
 
 import pytest
 import torch
+from huggingface_hub.errors import StrictDataclassClassValidationError
 from test_modeling_multimodal import MultimodalScenario, TestModelingMultimodal
 from transformers import AutoProcessor
 from utils.llm_data import llm_models_root
@@ -110,6 +112,35 @@
 _EXAONE_4_5_ASSET_PATH = EXAONE_4_5_TEST_CONFIG.get("_name_or_path")
 
 
+def test_exaone4_5_config_normalizes_trailing_mtp_layer_types():
+    config = copy.deepcopy(EXAONE_4_5_TEST_CONFIG)
+    text_config = config["text_config"]
+    text_config["_num_mtp_layers"] = 1
+    text_config["num_nextn_predict_layers"] = 1
+    text_config["layer_types"].append("sliding_attention")
+
+    hf_config = Exaone4_5Config(**config)
+
+    assert hf_config.text_config.layer_types == [
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+    ]
+    assert len(text_config["layer_types"]) == 5
+
+
+def test_exaone4_5_config_preserves_unexpected_layer_type_mismatch():
+    config = copy.deepcopy(EXAONE_4_5_TEST_CONFIG)
+    text_config = config["text_config"]
+    text_config["_num_mtp_layers"] = 1
+    text_config["num_nextn_predict_layers"] = 1
+    text_config["layer_types"].extend(["sliding_attention", "sliding_attention"])
+
+    with pytest.raises(StrictDataclassClassValidationError, match="number of layer types"):
+        Exaone4_5Config(**config)
+
+
 @dataclass(repr=False)
 class TestExaone4_5Scenario(MultimodalScenario):
     """Scenario config for Exaone4.5 multimodal smoke tests."""