[None][feat] Support K-EXAONE-2 in the PyTorch backend

docs/source/developer-guide/telemetry.md

@@ -228,7 +228,7 @@ unset or when the safety sanitizer rejects the runtime value.
 | `ray_placement_config.defer_workers_init` | `bool` | `value` |  |
 | `ray_placement_config.per_worker_gpu_share` | `float\|none` | `value` |  |
 | `ray_placement_config.placement_bundle_indices` | `list[list[int]]\|none` | `value` |  |
-| `reasoning_parser` | `allowlist\|none` | `categorical` | `auto`, `deepseek-r1`, `poolside_v1`, `laguna`, `qwen3`, `qwen3_5`, `minimax_m2`, `minimax_m2_append_think`, `nano-v3`, `gemma4`, `kimi_k2`, `kimi_k25` |
+| `reasoning_parser` | `allowlist\|none` | `categorical` | `auto`, `deepseek-r1`, `poolside_v1`, `laguna`, `qwen3`, `qwen3_5`, `minimax_m2`, `minimax_m2_append_think`, `nano-v3`, `gemma4`, `kimi_k2`, `kimi_k25`, `k-exaone` |
 | `reorder_policy_config.policy_args.agent_inflight_seq_num` | `int` | `value` |  |
 | `reorder_policy_config.policy_args.agent_percentage` | `float` | `value` |  |
 | `reorder_policy_config.policy_name` | `literal\|none` | `categorical` | `AgentTree` |

docs/source/models/supported-models.md

@@ -14,7 +14,8 @@ The following is a table of supported models for the PyTorch backend:
 | `DeepseekV32ForCausalLM`             | DeepSeek-V3.2                      | `deepseek-ai/DeepSeek-V3.2`                  |
 | `DeepseekV4ForCausalLM` [^11]        | DeepSeek-V4                        | `deepseek-ai/DeepSeek-V4-Pro`                |
 | `Exaone4ForCausalLM`                 | EXAONE 4.0                         | `LGAI-EXAONE/EXAONE-4.0-32B`                 |
-| `ExaoneMoEForCausalLM`               | K-EXAONE                           | `LGAI-EXAONE/K-EXAONE-236B-A23B`             |
+| `ExaoneMoeForCausalLM`               | K-EXAONE                           | `LGAI-EXAONE/K-EXAONE-236B-A23B`             |
+| `ExaoneMoeForCausalLM`               | K-EXAONE 2.0                       | `LGAI-EXAONE/K-EXAONE-2.0-750B-A37B`         |
 | `Gemma3ForCausalLM`                  | Gemma 3                            | `google/gemma-3-1b-it`                       |
 | `Gemma4ForConditionalGeneration`     | Gemma 4                            | `google/gemma-4-E2B-it`, `google/gemma-4-E4B-it`, `google/gemma-4-26B-A4B-it`, `google/gemma-4-31B-it` |
 | `Gemma4UnifiedForConditionalGeneration` | Gemma 4 12B Unified (encoder-free) | `google/gemma-4-12B`, `google/gemma-4-12B-it`              |

tensorrt_llm/_torch/models/_arch_index.py

@@ -50,7 +50,7 @@ def is_builtin_zoo_module(module_name: str) -> bool:
     "Eagle3DeepSeekV3ForCausalLM": "modeling_speculative",
     "Exaone4ForCausalLM": "modeling_exaone4",
     "Exaone4_5_ForConditionalGeneration": "modeling_exaone4_5",
-    "ExaoneMoEForCausalLM": "modeling_exaone_moe",
+    "ExaoneMoeForCausalLM": "modeling_exaone_moe",
     "Gemma3ForCausalLM": "modeling_gemma3",
     "Gemma3ForConditionalGeneration": "modeling_gemma3vl",
     "Gemma4AssistantForCausalLM": "modeling_gemma4",

tensorrt_llm/_torch/models/checkpoints/hf/exaone_moe_weight_mapper.py

@@ -1,11 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from torch import nn
 
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
 from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner
 
 
-@register_mapper("HF", "ExaoneMoEForCausalLM")
+@register_mapper("HF", "ExaoneMoeForCausalLM")
 class ExaoneMoeWeightMapper(HfWeightMapper):
     def __init__(self):
         super().__init__()
@@ -30,26 +45,51 @@ def __init__(self):
     def preprocess_weights(self, weights: dict):
         mtp_layer_offset = self.config.pretrained_config.num_hidden_layers
 
+        def rename(old_name: str, new_name: str) -> None:
+            # `weights` is a ConsumableWeightsDict, which deliberately exposes
+            # only __getitem__/__setitem__/__delitem__ (each lock-guarded) and
+            # has no `pop`. Reading does not consume, so get-then-delete is the
+            # exact equivalent of pop here.
+            weights[new_name] = weights[old_name]
+            del weights[old_name]
+
         for name in list(weights.keys()):
             if name.startswith("mtp.layers."):
                 # mtp.layers.{idx}.* -> model.layers.{offset + idx}.*
                 _, _, mtp_layer_idx, module_name = name.split(".", 3)
-                new_name = f"model.layers.{mtp_layer_offset + int(mtp_layer_idx)}.{module_name}"
-                weights[new_name] = weights.pop(name)
+                rename(
+                    name,
+                    f"model.layers.{mtp_layer_offset + int(mtp_layer_idx)}.{module_name}",
+                )
             elif name.startswith("mtp."):
                 # mtp.fc.* -> model.layers.{offset}.eh_proj.*
                 # mtp.norm.* -> model.layers.{offset}.shared_head.norm.*
                 # etc.
                 for mtp_prefix, trtllm_name in self.mtp_mapping.items():
                     if name.startswith(mtp_prefix):
                         suffix = name[len(mtp_prefix) :]
-                        new_name = f"model.layers.{mtp_layer_offset}.{trtllm_name}{suffix}"
-                        weights[new_name] = weights.pop(name)
+                        rename(
+                            name,
+                            f"model.layers.{mtp_layer_offset}.{trtllm_name}{suffix}",
+                        )
                         break
 
     def is_special_instance_module(self, module: nn.Module) -> bool:
         return is_moe_weight_owner(module)
 
+    def _renames_scales_to_inv(self) -> bool:
+        """Whether this checkpoint's MoE scales must be exposed as ``weight_scale_inv``.
+
+        True for everything except NVFP4, whose loader
+        (:meth:`NVFP4FusedMoEMethod.load_quant_scales`) indexes ``weight_scale``
+        and ``weight_scale_2`` directly. Unknown/absent quant config keeps the
+        historical behaviour so non-quantised checkpoints are unaffected.
+        """
+        from tensorrt_llm.quantization.mode import QuantAlgo
+
+        quant_config = getattr(self.config, "quant_config", None)
+        return getattr(quant_config, "quant_algo", None) != QuantAlgo.NVFP4
+
     def handle_special_instance_module(
         self,
         module: nn.Module,
@@ -81,8 +121,19 @@ def handle_special_instance_module(
                         for i in range(weight_value.shape[0]):
                             updated_module_weights[f"{i}.w2.weight"] = weight_value[i]
                         continue
-                new_weight_name = weight_name.replace("weight_scale", "weight_scale_inv")
-                if new_weight_name.endswith(".weight_scale_inv"):
+                # Only the FP8-block-scale and W4A16-MXFP4 MoE loaders read
+                # `weight_scale_inv`. The NVFP4 loader reads `weight_scale` (the
+                # per-block scales) and `weight_scale_2` (the per-tensor global
+                # scale) verbatim, so renaming here hides both from it.
+                #
+                # The old unconditional `replace("weight_scale", ...)` also
+                # rewrote `weight_scale_2` into `weight_scale_inv_2`, which no
+                # loader has ever looked for -- wrong under any quantisation.
+                new_weight_name = weight_name
+                if self._renames_scales_to_inv() and (
+                    weight_name == "weight_scale" or weight_name.endswith(".weight_scale")
+                ):
+                    new_weight_name = weight_name[: -len("weight_scale")] + "weight_scale_inv"
                     weight_value = weight_value.squeeze()
                 updated_module_weights[new_weight_name] = weight_value
             module.load_weights(

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -61,8 +61,9 @@
 from ..model_config import ModelConfig
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
-from ..moe.fused_moe import (DeepSeekV3MoeRoutingMethod, MoEWeightLoadingMode,
-                             create_moe, is_moe_weight_owner)
+from ..moe.fused_moe import (DEFAULT_MOE_ACTIVATION, DeepSeekV3MoeRoutingMethod,
+                             MoEWeightLoadingMode, SwigluActivation, create_moe,
+                             is_moe_weight_owner)
 
 # isort: off
 from ..moe.fused_moe.routing import Deepseekv3RoutingImpl
@@ -968,7 +969,8 @@ def __init__(self,
                  dtype: Optional[torch.dtype] = None,
                  model_config: ModelConfig = ModelConfig(),
                  override_quant_config: Optional[QuantConfig] = None,
-                 layer_idx: Optional[int] = None):
+                 layer_idx: Optional[int] = None,
+                 swiglu_limit_scalar: Optional[float] = None):
         from ..distributed import AllReduce
 
         super().__init__()
@@ -1001,6 +1003,9 @@ def __init__(self,
                 model_config, layer_idx)
         else:
             expert_quant_config = override_quant_config
+        moe_activation = DEFAULT_MOE_ACTIVATION
+        if swiglu_limit_scalar is not None:
+            moe_activation = SwigluActivation(clamp=swiglu_limit_scalar)
         self.experts = create_moe(
             num_experts=num_experts,
             routing_method=self.gate.routing_method,
@@ -1017,6 +1022,7 @@ def __init__(self,
             # examples/quantization/quantize_mixed_precision_moe.py
             weight_loading_mode=self._expert_weight_loading_mode(
                 expert_quant_config),
+            activation=moe_activation,
         )
 
         self.mapping = model_config.mapping
@@ -1047,6 +1053,7 @@ def __init__(self,
             overridden_tp_size=self.shared_tp_size,
             reduce_output=False,
             use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm,
+            swiglu_limit=swiglu_limit_scalar,
         )
         self.shared_experts_use_fp4 = (
             shared_quant_config is not None

tensorrt_llm/_torch/models/modeling_exaone_moe.py

@@ -57,6 +57,52 @@
 from .modeling_utils import DecoderModel, EagerFusionConfig, register_auto_model
 
 
+def _per_layer_list(
+    config: ExaoneMoeConfig, name: str, layer_idx: int
+) -> list[int | float | str] | None:
+    """Return `config.<name>` when it is a list long enough to cover `layer_idx`.
+
+    K-EXAONE2 encodes its per-layer contracts as plain lists in config.json.
+    When such a list is present it is authoritative for every layer it covers;
+    a 0 entry is the config's spelling of "not enabled on this layer".
+    """
+    values = getattr(config, name, None)
+    if values is None or layer_idx >= len(values):
+        return None
+    return values
+
+
+def get_exaone_swiglu_limit(config: ExaoneMoeConfig, layer_idx: int) -> Optional[float]:
+    """Return the configured per-layer SwiGLU clamp."""
+    limits = _per_layer_list(config, "swiglu_limits", layer_idx)
+    if limits is not None:
+        return float(limits[layer_idx]) or None
+    return None
+
+
+def get_exaone_attention_window(
+    config: ExaoneMoeConfig, layer_idx: int, is_mtp_layer: bool
+) -> Optional[int]:
+    """Return the effective attention window for an EXAONE-MoE layer."""
+    if is_mtp_layer:
+        mtp_layer_idx = layer_idx - config.num_hidden_layers
+        if mtp_layer_idx < 0:
+            return None
+        mtp_layer_types = _per_layer_list(config, "mtp_layer_types", mtp_layer_idx)
+        if mtp_layer_types is None or mtp_layer_types[mtp_layer_idx] != "sliding_attention":
+            return None
+        mtp_windows = _per_layer_list(config, "mtp_sliding_windows", mtp_layer_idx)
+        if mtp_windows is None:
+            return None
+        return int(mtp_windows[mtp_layer_idx]) or None
+    if config.layer_types[layer_idx] != "sliding_attention":
+        return None
+    windows = _per_layer_list(config, "sliding_windows", layer_idx)
+    if windows is not None:
+        return int(windows[layer_idx]) or None
+    return config.sliding_window
+
+
 def check_is_moe(config: ExaoneMoeConfig, layer_idx: int, is_mtp_layer: bool = False) -> bool:
     """
     Check if the current layer is a MoE layer.
@@ -93,14 +139,12 @@ def __init__(
     ):
         config = model_config.pretrained_config
 
-        self.attention_window_size = None
-        # A MTP layer uses the global attention.
-        self.is_sliding = not is_mtp_layer and config.layer_types[layer_idx] == "sliding_attention"
+        self.attention_window_size = get_exaone_attention_window(config, layer_idx, is_mtp_layer)
+        self.is_sliding = self.attention_window_size is not None
 
         # NOTE: In ExaoneMoe, only sliding layers apply rope.
         pos_embd_params = None
         if self.is_sliding:
-            self.attention_window_size = config.sliding_window
             pos_embd_params = PositionalEmbeddingParams(
                 type=PositionEmbeddingType.rope_gpt_neox,
                 rope=RopeParams.from_config(config),
@@ -207,6 +251,9 @@ def __init__(
         is_mtp_layer = False
         if layer_idx >= config.num_hidden_layers:
             is_mtp_layer = True
+        # The MTP block is dense but follows its own head contract; the K2
+        # late-layer clamp applies only to main decoder layers.
+        swiglu_limit = None if is_mtp_layer else get_exaone_swiglu_limit(config, layer_idx)
         self.self_attn = ExaoneMoeAttention(
             model_config, layer_idx=layer_idx, is_mtp_layer=is_mtp_layer
         )
@@ -228,6 +275,7 @@ def __init__(
                 override_quant_config=quant_config,
                 aux_stream_dict=aux_stream_dict,
                 layer_idx=layer_idx,
+                swiglu_limit_scalar=swiglu_limit,
             )
         else:
             block_size = 1
@@ -250,6 +298,7 @@ def __init__(
                 overridden_tp_size=self.mlp_tp_size,
                 layer_idx=layer_idx,
                 reduce_output=has_mlp_tp,
+                swiglu_limit=swiglu_limit,
             )
 
         self.disable_attn_allreduce = (
@@ -649,7 +698,10 @@ def forward(
         return hidden_states
 
 
-@register_auto_model("ExaoneMoEForCausalLM")
+# K-EXAONE and K-EXAONE-2 both declare `ExaoneMoeForCausalLM`. K-EXAONE configs
+# downloaded before the 2026-08 rename declare `ExaoneMoEForCausalLM` and need a
+# refreshed config.json.
+@register_auto_model("ExaoneMoeForCausalLM")
 class ExaoneMoeForCausalLM(SpecDecOneEngineForCausalLM[ExaoneMoeModel, ExaoneMoeConfig]):
     def __init__(
         self,

tensorrt_llm/llmapi/llm_args.py

@@ -5161,7 +5161,7 @@ class BaseLlmArgs(StrictBaseModel):
                                              'qwen3_5', 'minimax_m2',
                                              'minimax_m2_append_think',
                                              'nano-v3', 'gemma4', 'kimi_k2',
-                                             'kimi_k25'))
+                                             'kimi_k25', 'k-exaone'))
 
     # TODO[Superjomn]: To deprecate this config.
     decoding_config: Optional[object] = Field(

tensorrt_llm/llmapi/reasoning_parser.py

@@ -447,6 +447,7 @@ def parse(self, text: str) -> ReasoningParserResult:
     "nemotron_h": "nemotron-v3",
     "nemotron_h_puzzle": "nemotron-v3",
     "nemotron_h_omni": "nemotron-v3",
+    "exaone_moe": "k-exaone",
     "gemma4": "gemma4",
     "kimi_k2": "kimi_k2",
     "kimi_k25": "kimi_k25",
@@ -544,6 +545,7 @@ def resolve_auto_reasoning_parser(model: str) -> Optional[str]:
 
 @register_reasoning_parser("nemotron-v3")
 @register_reasoning_parser("nano-v3")
+@register_reasoning_parser("k-exaone")
 class NemotronV3ReasoningParser(DeepSeekR1Parser):
     """Reasoning parser for Nemotron Nano v3.
 

tensorrt_llm/usage/architecture_allowlist.py

@@ -39,7 +39,7 @@
         "Exaone4ForCausalLM",
         "Exaone4_5_ForConditionalGeneration",
         "ExaoneForCausalLM",
-        "ExaoneMoEForCausalLM",
+        "ExaoneMoeForCausalLM",
         "Gemma3ForCausalLM",
         "Gemma3ForConditionalGeneration",
         "Gemma3nForConditionalGeneration",

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -1172,7 +1172,8 @@
         "nano-v3",
         "gemma4",
         "kimi_k2",
-        "kimi_k25"
+        "kimi_k25",
+        "k-exaone"
       ],
       "capture_policy": "allowlist|none",
       "kind": "categorical",