# Copyright 2025 The LG AI Research Team # Copyright 2023-2024 SGLang Team # 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. # ============================================================================== # Adapted from the vLLM version of EXAONE-MoE model """Inference-only ExaoneMoE model compatible with HuggingFace weights.""" import logging from collections.abc import Iterable from typing import Any, Dict, Optional, Tuple, Union import torch from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import ( get_moe_expert_parallel_world_size, get_pp_group, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.dp_attention import ( get_attention_tp_rank, get_attention_tp_size, is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput from sglang.srt.layers.moe import get_moe_a2a_backend from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.moe.utils import RoutingMethodType from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers logger = logging.getLogger(__name__) _is_cuda = is_cuda() class ExaoneMoEMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: Optional[QuantizationConfig] = None, reduce_results: bool = True, prefix: str = "", tp_rank: Optional[int] = None, tp_size: Optional[int] = None, ) -> None: super().__init__() gateup_quant_config = quant_config down_quant_config = quant_config if quant_config and hasattr(quant_config, "ignore") and quant_config.ignore: if add_prefix("gate_proj", prefix) in quant_config.ignore: gateup_quant_config = None if add_prefix("down_proj", prefix) in quant_config.ignore: down_quant_config = None self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=gateup_quant_config, prefix=add_prefix("gate_up_proj", prefix), tp_rank=tp_rank, tp_size=tp_size, ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=down_quant_config, reduce_results=reduce_results, prefix=add_prefix("down_proj", prefix), tp_rank=tp_rank, tp_size=tp_size, ) if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. " "Only silu is supported for now." ) self.act_fn = SiluAndMul() def forward( self, x, forward_batch=None, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ): gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj( x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter, ) return x class ExaoneMoESparseMoEBlock(nn.Module): def __init__( self, layer_id: int, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, alt_stream: Optional[torch.cuda.Stream] = None, prefix: str = "", ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() self.moe_ep_size = get_moe_expert_parallel_world_size() self.layer_id = layer_id self.routed_scaling_factor = config.routed_scaling_factor self.alt_stream = alt_stream self.n_routed_experts = config.num_experts if self.tp_size > config.num_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than " f"the number of experts {config.num_experts}." ) self.gate = ReplicatedLinear( config.hidden_size, config.num_experts, bias=False, quant_config=None, prefix=add_prefix("gate", prefix), ) self.e_score_correction_bias = nn.Parameter( torch.empty(config.num_experts, dtype=torch.float32) ) self.experts = get_moe_impl_class(quant_config)( num_experts=config.num_experts + get_global_server_args().ep_num_redundant_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, layer_id=self.layer_id, quant_config=quant_config, prefix=add_prefix("experts", prefix), routing_method_type=RoutingMethodType.RenormalizeNaive, ) self.topk = TopK( top_k=config.num_experts_per_tok, renormalize=config.norm_topk_prob, use_grouped_topk=True, num_expert_group=config.n_group, topk_group=config.topk_group, correction_bias=self.e_score_correction_bias, routed_scaling_factor=self.routed_scaling_factor, apply_routed_scaling_factor_on_output=True, scoring_func="sigmoid", ) if config.num_shared_experts is not None: intermediate_size = config.moe_intermediate_size * config.num_shared_experts self.shared_experts = ExaoneMoEMLP( hidden_size=config.hidden_size, intermediate_size=intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, reduce_results=False, prefix=add_prefix("shared_experts", prefix), **( dict(tp_rank=0, tp_size=1) if get_moe_a2a_backend().is_deepep() else {} ), ) if get_moe_a2a_backend().is_deepep(): self.ep_size = get_moe_expert_parallel_world_size() self.num_experts = ( config.num_experts + get_global_server_args().ep_num_redundant_experts ) self.top_k = config.num_experts_per_tok def get_moe_weights(self): return [ x.data for name, x in self.experts.named_parameters() if name not in ["correction_bias"] ] def _forward_shared_experts(self, hidden_states: torch.Tensor) -> torch.Tensor: shared_output = self.shared_experts(hidden_states) return shared_output def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): shared_output = None if hidden_states.shape[0] > 0: router_logits, _ = self.gate(hidden_states) shared_output = self._forward_shared_experts(hidden_states) topk_output = self.topk( hidden_states, router_logits, num_token_non_padded=forward_batch.num_token_non_padded, expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( layer_id=self.layer_id, ), ) else: topk_output = self.topk.empty_topk_output(hidden_states.device) final_hidden_states = self.experts( hidden_states=hidden_states, topk_output=topk_output, ) if shared_output is not None: final_hidden_states.add_(shared_output) return final_hidden_states def _forward_router_experts(self, hidden_states: torch.Tensor) -> torch.Tensor: router_logits, _ = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) return self.experts(hidden_states, topk_output) def forward_normal_dual_stream( self, hidden_states: torch.Tensor, ) -> torch.Tensor: current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) shared_output = self._forward_shared_experts(hidden_states.clone()) with torch.cuda.stream(self.alt_stream): router_output = self._forward_router_experts(hidden_states) current_stream.wait_stream(self.alt_stream) return router_output, shared_output def forward( self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, use_reduce_scatter: bool = False, ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) if get_moe_a2a_backend().is_deepep(): return self._forward_deepep(hidden_states, forward_batch) if ( self.alt_stream is not None and hidden_states.shape[0] > 0 and get_is_capture_mode() ): final_hidden_states, shared_output = self.forward_normal_dual_stream( hidden_states ) else: shared_output = self._forward_shared_experts(hidden_states) final_hidden_states = self._forward_router_experts(hidden_states) if shared_output is not None: final_hidden_states = final_hidden_states + shared_output if self.tp_size > 1 and not use_reduce_scatter: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) class ExaoneMoEAttention(nn.Module): def __init__( self, config: PretrainedConfig, hidden_size: int, num_heads: int, num_kv_heads: int, layer_id: int = 0, rope_theta: float = 1000000, rope_scaling: Optional[Dict[str, Any]] = None, rope_is_neox_style: bool = True, max_position_embeddings: int = 8192, quant_config: Optional[QuantizationConfig] = None, bias: bool = False, prefix: str = "", ) -> None: super().__init__() self.hidden_size = hidden_size attn_tp_rank = get_attention_tp_rank() attn_tp_size = get_attention_tp_size() self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= attn_tp_size: # Number of KV heads is greater than TP size, so we partition # the KV heads across multiple tensor parallel GPUs. assert self.total_num_kv_heads % attn_tp_size == 0 else: # Number of KV heads is less than TP size, so we replicate # the KV heads across multiple tensor parallel GPUs. assert attn_tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) # MistralConfig has an optional head_dim introduced by Mistral-Nemo self.head_dim = getattr( config, "head_dim", self.hidden_size // self.total_num_heads ) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.max_position_embeddings = max_position_embeddings qkv_quant_config = quant_config o_quant_config = quant_config if quant_config and hasattr(quant_config, "ignore") and quant_config.ignore: if add_prefix("q_proj", prefix) in quant_config.ignore: qkv_quant_config = None if add_prefix("o_proj", prefix) in quant_config.ignore: o_quant_config = None self.qkv_proj = QKVParallelLinear( hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=bias, quant_config=qkv_quant_config, prefix=add_prefix("qkv_proj", prefix), tp_rank=attn_tp_rank, tp_size=attn_tp_size, ) self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, hidden_size, bias=bias, quant_config=o_quant_config, prefix=add_prefix("o_proj", prefix), tp_rank=attn_tp_rank, tp_size=attn_tp_size, ) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if quant_config is not None and quant_config.get_name() == "gguf": rope_is_neox_style = False self.sliding_window = config.layer_types[layer_id] == "sliding_attention" # apply rotary embeddings to every layer in full attention models self.apply_rope_all_layers = "sliding_attention" not in config.layer_types self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.head_dim, max_position=max_position_embeddings, base=rope_theta, rope_scaling=rope_scaling, is_neox_style=rope_is_neox_style, ) self.attn = RadixAttention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, layer_id=layer_id, prefix=add_prefix("attn", prefix), sliding_window_size=( config.sliding_window if self.sliding_window else None ), ) self.layer_id = layer_id def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q = q.reshape(-1, self.head_dim) q = self.q_norm(q) q = q.reshape(-1, self.num_heads * self.head_dim) k = k.reshape(-1, self.head_dim) k = self.k_norm(k) k = k.reshape(-1, self.num_kv_heads * self.head_dim) if self.sliding_window or self.apply_rope_all_layers: q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v, forward_batch) output, _ = self.o_proj(attn_output) return output class ExaoneMoEDecoderLayer(nn.Module): def __init__( self, config: PretrainedConfig, layer_id: int = 0, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size self.config = config rope_theta = getattr(config, "rope_theta", 1000000) rope_scaling = getattr(config, "rope_scaling", None) if rope_scaling is not None and getattr( config, "original_max_position_embeddings", None ): rope_scaling["original_max_position_embeddings"] = ( config.original_max_position_embeddings ) rope_is_neox_style = getattr(config, "rope_is_neox_style", True) max_position_embeddings = getattr(config, "max_position_embeddings", 131072) attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) self.attn_tp_size = get_attention_tp_size() self.attn_tp_rank = get_attention_tp_rank() self.self_attn = ExaoneMoEAttention( config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, layer_id=layer_id, rope_theta=rope_theta, rope_scaling=rope_scaling, rope_is_neox_style=rope_is_neox_style, max_position_embeddings=max_position_embeddings, quant_config=quant_config, bias=attention_bias, prefix=add_prefix("self_attn", prefix), ) if config.is_moe_layer[layer_id]: self.mlp = ExaoneMoESparseMoEBlock( layer_id=layer_id, config=config, quant_config=quant_config, alt_stream=alt_stream, prefix=add_prefix("mlp", prefix), ) else: self.mlp = ExaoneMoEMLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: if residual is None: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) # Self Attention hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) # Fully Connected hidden_states = self.mlp(hidden_states) return hidden_states, residual class ExaoneMoEModel(nn.Module): fall_back_to_pt_during_load = False def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.config = config self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.pp_group = get_pp_group() if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, enable_tp=not is_dp_attention_enabled(), ) else: self.embed_tokens = PPMissingLayer() self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: ExaoneMoEDecoderLayer( layer_id=idx, config=config, quant_config=quant_config, prefix=prefix, alt_stream=alt_stream, ), pp_rank=self.pp_group.rank_in_group, pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: self.norm = PPMissingLayer(return_tuple=True) # for EAGLE3 support self.layers_to_capture = [] def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: if self.pp_group.is_first_rank: if input_embeds is None: hidden_states = self.embed_tokens(input_ids) else: hidden_states = input_embeds residual = None else: assert pp_proxy_tensors is not None hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] aux_hidden_states = [] for i in range(self.start_layer, self.end_layer): with get_global_expert_distribution_recorder().with_current_layer(i): if i in self.layers_to_capture: aux_hidden_states.append(hidden_states + residual) layer = self.layers[i] hidden_states, residual = layer( positions, hidden_states, forward_batch, residual ) if not self.pp_group.is_last_rank: return PPProxyTensors( { "hidden_states": hidden_states, "residual": residual, } ) else: if hidden_states.shape[0] != 0: if residual is None: hidden_states = self.norm(hidden_states) else: hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) == 0: return hidden_states return hidden_states, aux_hidden_states class ExaoneMoEForCausalLM(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config alt_stream = torch.cuda.Stream() if _is_cuda else None self.model = ExaoneMoEModel( config, quant_config=quant_config, prefix=add_prefix("model", prefix), alt_stream=alt_stream, ) if self.config.tie_word_embeddings: self.lm_head = self.model.embed_tokens else: self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=add_prefix("lm_head", prefix), use_attn_tp_group=get_global_server_args().enable_dp_lm_head, ) self.logits_processor = LogitsProcessor(config) # For EAGLE3 support self.capture_aux_hidden_states = False self._routed_experts_weights_of_layer = LazyValue( lambda: { layer_id: self.model.layers[layer_id].mlp.get_moe_weights() for layer_id in range(self.start_layer, self.end_layer) if isinstance(self.model.layers[layer_id].mlp, ExaoneMoESparseMoEBlock) } ) @property def routed_experts_weights_of_layer(self): return self._routed_experts_weights_of_layer.value @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> LogitsProcessorOutput: hidden_states = self.model( input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors=pp_proxy_tensors, ) aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states if self.pp_group.is_last_rank: return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states, ) else: return hidden_states @torch.no_grad() def forward_split_prefill( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, split_interval: Tuple[int, int], # [start, end) 0-based input_embeds: torch.Tensor = None, ): start, end = split_interval # embed if start == 0: if input_embeds is None: forward_batch.hidden_states = self.model.embed_tokens(input_ids) else: forward_batch.hidden_states = input_embeds # decoder layer for i in range(start, end): layer = self.model.layers[i] forward_batch.hidden_states, forward_batch.residual = layer( positions, forward_batch.hidden_states, forward_batch, forward_batch.residual, ) if end == self.model.config.num_hidden_layers: # norm hidden_states, _ = self.model.norm( forward_batch.hidden_states, forward_batch.residual ) forward_batch.hidden_states = hidden_states # logits process result = self.logits_processor( input_ids, forward_batch.hidden_states, self.lm_head, forward_batch ) else: result = None return result @property def start_layer(self): return self.model.start_layer @property def end_layer(self): return self.model.end_layer def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight def set_embed_and_head(self, embed, head): del self.model.embed_tokens.weight del self.lm_head.weight self.model.embed_tokens.weight = embed self.lm_head.weight = head torch.cuda.empty_cache() torch.cuda.synchronize() def load_weights( self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False ): stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), ] expert_params_mapping = FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=self.config.num_experts, ) params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: if is_mtp: if "mtp" not in name: continue if name in [ "mtp.fc.weight", "mtp.pre_fc_norm_embedding.weight", "mtp.pre_fc_norm_hidden.weight", ]: name = name.replace("mtp.", "") else: name = name.replace("mtp", "model") if not is_mtp and "mtp" in name: continue if "rotary_emb.inv_freq" in name or "projector" in name: continue if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue if name.startswith("model.vision_tower") and name not in params_dict: continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue if "mlp.experts" in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) param = params_dict[name] weight_loader = param.weight_loader weight_loader( param, loaded_weight, name, expert_id=expert_id, shard_id=shard_id, ) break else: # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue if name in params_dict.keys(): param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) else: logger.warning(f"Parameter {name} not found in params_dict") @classmethod def get_model_config_for_expert_location(cls, config): return ModelConfigForExpertLocation( num_layers=config.num_hidden_layers, num_logical_experts=config.num_experts, num_groups=None, ) def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None): if not get_pp_group().is_last_rank: return self.capture_aux_hidden_states = True if layer_ids is None: num_layers = self.config.num_hidden_layers self.model.layers_to_capture = [ 2, num_layers // 2, num_layers - 3, ] # Specific layers for EAGLE3 support else: self.model.layers_to_capture = [val + 1 for val in layer_ids] EntryClass = ExaoneMoEForCausalLM