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)