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:
EX Engine
2026-08-10 04:36:16 +00:00
parent 44d36e6ccc
commit 1ae398eeee
2 changed files with 24 additions and 6 deletions

View File

@@ -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]