fix(interface): corex_moe accepts w13 merged format + no silent fallback
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.
This commit is contained in:
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user