fix Gemini review.

Signed-off-by: Asher Zhang <[email protected]>

Author: kzjeef

vllm/model_executor/models/hunyuan_v1_eagle3.py

@@ -25,7 +25,6 @@
 from vllm.v1.sample.metadata import SamplingMetadata
 
 from .utils import AutoWeightsLoader, maybe_prefix
-from .llama_eagle3  import LlamaModel as LlamaEagle3Model
 
 logger = init_logger(__name__)
 
@@ -87,9 +86,7 @@ def forward(
 
 
 @support_torch_compile
-class Eagle3HunYuanModel(LlamaEagle3Model):
-    # Most function are same as Llama Eagle 3 support.
-    # only different is from init layer.
+class Eagle3HunYuanModel(nn.Module):
 
     def __init__(
         self,
@@ -98,8 +95,7 @@ def __init__(
         start_layer_id: int = 0,
         prefix: str = "",
     ) -> None:
-        # llama 's init will setup layers, which cuase conflict
-        nn.Module.__init__(self)
+        super().__init__()
         self.config = vllm_config. \
             speculative_config.draft_model_config.hf_config
         self.vocab_size = self.config.vocab_size
@@ -130,7 +126,59 @@ def __init__(
             eps=self.config.rms_norm_eps,
         )
 
-class Eagle3HunYuanDenseV1ForCausalLM(HunYuanDenseV1ForCausalLM):
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_embeds = self.embed_tokens(input_ids)
+        assert hidden_states.shape[-1] == input_embeds.shape[-1]
+
+        residual = None
+        hidden_states, residual = self.layers[0](
+            positions,
+            input_embeds,
+            hidden_states,
+            residual,
+        )
+
+        hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
+        return hidden_states, hidden_prenorm
+
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
+        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),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if 'midlayer.' in name:
+                name = name.replace('midlayer.', 'layers.0.')
+            for param_name, weight_name, shard_id in stacked_params_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, shard_id)
+                break
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Eagle3HunYuanDenseV1ForCausalLM(nn.Module):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)