From 0b577a7263f6b95015c39ea280f0a24d11c434eb Mon Sep 17 00:00:00 2001
From: "Chendi.Xue" <chendi.xue@intel.com>
Date: Thu, 19 Feb 2026 04:57:11 -0600
Subject: [PATCH] [v0.16.0][BUG FIX]Fix hunyuan MOE after update to 0.16.0
 (#1401)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>
---
 .../hunyuan_image_3_transformer.py            | 14 -----
 .../pipeline_hunyuan_image_3.py               | 58 +++++++++----------
 2 files changed, 29 insertions(+), 43 deletions(-)

diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py
index adb582741..60d4ea738 100644
--- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py
@@ -1588,20 +1588,6 @@ class HunyuanFusedMoE(SharedFusedMoE):
         self._init_hook_handle.remove()
 
     def forward(self, hidden_states, router_logits):
-        from vllm.model_executor.layers.fused_moe.layer import get_forward_context
-
-        ctx = get_forward_context()
-        if not ctx.remaining_moe_layers:
-            import re
-
-            moe_names = [name for name in ctx.no_compile_layers.keys() if ".mlp.experts" in name]
-
-            def get_layer_num(name):
-                match = re.search(r"layers\.(\d+)\.mlp", name)
-                return int(match.group(1)) if match else -1
-
-            moe_names.sort(key=get_layer_num, reverse=True)
-            ctx.remaining_moe_layers.extend(moe_names)
         return super().forward(hidden_states, router_logits)
 
 
diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py
index 9a8f26d38..c9617eef4 100644
--- a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py
+++ b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py
@@ -804,17 +804,14 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
             )
             kwargs["num_image_tokens"] = num_image_tokens
             # 50 and 5.0 hard code
-            from vllm.forward_context import set_forward_context
-
-            with set_forward_context(None, self.vllm_config):
-                results = self.pipeline(
-                    batch_size=len(batch_gen_image_info),
-                    image_size=[batch_gen_image_info[0].image_height, batch_gen_image_info[0].image_width],
-                    num_inference_steps=kwargs.get("num_inference_steps", 50),
-                    guidance_scale=kwargs.get("guidance_scale", 5.0),
-                    generator=generator,
-                    model_kwargs=kwargs,
-                )
+            results = self.pipeline(
+                batch_size=len(batch_gen_image_info),
+                image_size=[batch_gen_image_info[0].image_height, batch_gen_image_info[0].image_width],
+                num_inference_steps=kwargs.get("num_inference_steps", 50),
+                guidance_scale=kwargs.get("guidance_scale", 5.0),
+                generator=generator,
+                model_kwargs=kwargs,
+            )
             samples = results[0]
             return samples
 
@@ -929,24 +926,27 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
             )
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            custom_pos_emb=custom_pos_emb,
-            mode=mode,
-            first_step=first_step,
-            query_lens=query_lens,
-            seq_lens=seq_lens,
-            num_image_tokens=num_image_tokens,
-            gen_timestep_scatter_index=gen_timestep_scatter_index,
-        )
+        from vllm.forward_context import set_forward_context
+
+        with set_forward_context(None, self.vllm_config):
+            outputs = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                custom_pos_emb=custom_pos_emb,
+                mode=mode,
+                first_step=first_step,
+                query_lens=query_lens,
+                seq_lens=seq_lens,
+                num_image_tokens=num_image_tokens,
+                gen_timestep_scatter_index=gen_timestep_scatter_index,
+            )
         hidden_states = outputs[0]
 
         if mode == "gen_text":