From 1ae398eeee17b070d97d6dd5d86ab03cdf25c835 Mon Sep 17 00:00:00 2001
From: EX Engine <dev@enginex.ai>
Date: Mon, 10 Aug 2026 04:36:16 +0000
Subject: [PATCH] fix(interface): corex_moe accepts w13 merged format + no
 silent fallback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

corex_moe.py: moe_forward now accepts both formats:
  Format A: w1(E,I,H) + w2(E,H,I) + w3(E,I,H) — xllm style, separate gate/up
  Format B: w13(E,2*I,H) + w2(E,H,I) + w3=None — vllm style, merged gate_up
  Auto-detects by checking if w3 is None, splits w13 internally.

qwen3_5.py:
  - Fix corex_moe call: use keyword args (w3=None, topk=self.top_k)
    prevents topk integer going to w3 tensor position
  - Remove silent fallback on corex_moe failure — raise RuntimeError
    with full shape info for diagnosis. Zero score with no error log
    is worse than a crash.
---
 ex_engine/python/corex_moe.py | 18 ++++++++++++++++--
 qwen3_6_scripts/qwen3_5.py    | 12 ++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/ex_engine/python/corex_moe.py b/ex_engine/python/corex_moe.py
index cf88a2f2..d551aa20 100644
--- a/ex_engine/python/corex_moe.py
+++ b/ex_engine/python/corex_moe.py
@@ -127,20 +127,34 @@ def topk_softmax(
 def moe_forward(
     hidden_states: torch.Tensor,
     gate_output: torch.Tensor,
-    w1: torch.Tensor,
+    w1_or_w13: torch.Tensor,
     w2: torch.Tensor,
-    w3: torch.Tensor,
+    w3: Optional[torch.Tensor] = None,
     topk: int = 8,
     renormalize: bool = True,
     **kwargs,
 ) -> torch.Tensor:
     """
     Full MoE pipeline: CUDA topk → per-expert GEMM (cublas) → silu → GEMM → scatter-add.
+
+    Accepts two weight formats:
+      Format A (xllm style):  w1=(E,I,H), w2=(E,H,I), w3=(E,I,H) — gate and up separate
+      Format B (vllm style):  w13=(E,2*I,H), w2=(E,H,I), w3=None — gate_up merged
     """
     num_tokens = hidden_states.shape[0]
     hidden_size = hidden_states.shape[1]
     dtype = hidden_states.dtype
 
+    # Detect weight format
+    if w3 is None:
+        # Format B: w13 merged — split into w1 (gate) and w3 (up)
+        w13 = w1_or_w13
+        inter2 = w13.shape[1]
+        w1 = w13[:, :inter2 // 2, :]  # (E, I, H)
+        w3 = w13[:, inter2 // 2:, :]  # (E, I, H)
+    else:
+        w1 = w1_or_w13
+
     topk_weights, topk_ids = topk_softmax(gate_output, topk, renormalize)
 
     num_experts = w1.shape[0]
diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py
index bf3ff9e8..213b6678 100644
--- a/qwen3_6_scripts/qwen3_5.py
+++ b/qwen3_6_scripts/qwen3_5.py
@@ -1207,12 +1207,16 @@ class Qwen3_5MoeSparseBlock(nn.Module):
                 routed_out = self._corex_moe_forward(
                     hidden_states, router_logits,
                     self.experts.w13_weight, self.experts.w2_weight,
-                    self.top_k,
+                    w3=None, topk=self.top_k,
                 )
             except Exception as e:
-                logger.warning("CoreX MoE forward failed (%s), falling back permanently", e)
-                self._use_corex_moe = False
-                routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
+                # NO FALLBACK — crash with error log so we can diagnose
+                logger.error("CoreX MoE forward FAILED: %s", e)
+                raise RuntimeError(
+                    f"corex_moe.moe_forward failed: {e}. "
+                    f"Shapes: hidden={hidden_states.shape}, router={router_logits.shape}, "
+                    f"w13={self.experts.w13_weight.shape}, w2={self.experts.w2_weight.shape}"
+                ) from e
         else:
             routed_out = self._pure_pytorch_experts(hidden_states, router_logits)