From b7149f810a8dfc633c4b77fb8ff6b7b2d94c2c7b Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 14:55:01 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20decode=20MoE=E8=B7=AF=E5=BE=84=E5=AF=B9?= =?UTF-8?q?=E9=BD=90base=20=E2=80=94=20F.linear+bmm=E6=9B=BF=E6=8D=A2pre-t?= =?UTF-8?q?ranspose+bmm?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit base qwen3_5.py的decode路径（已验证可跑通竞赛）: F.linear(hidden, w13_sel.reshape(-1,H)) → view → act → bmm(w2_sel, act) 我们之前的路径（未验证，probe显示更慢）: pre-transpose(w13全量) → w13_t[eids] → bmm(x_expand, w13_t_sel) → act → bmm(act, w2_t_sel) probe真机数据: loop matmul 19ms < torch.bmm 24ms 说明F.linear路径在BI-V100单token场景下更优 保持的corex加速: ✓ corex_moe_topk_softmax (topk+softmax fused) ✓ corex_moe_weight_gather (gather fused) ✓ corex_moe_exact_reduce (weighted sum fused) ✓ corex_moe_index_combine (prefill token routing fused) --- qwen3_6_scripts/qwen3_5.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index a05c80c6..130894f7 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -1824,19 +1824,14 @@ class Qwen3_5MoeSparseBlock(nn.Module): H = hidden_states.shape[-1] - # --- Pre-transpose weights for bmm (cached after first call) --- - if not hasattr(self, '_w13_t') or self._w13_t is None: - # (E, 2*I, H) → (E, H, 2*I) — one-time cost at first decode - self._w13_t = self.experts.w13_weight.transpose(1, 2).contiguous() - self._w2_t = self.experts.w2_weight.transpose(1, 2).contiguous() - # (E, H, I) → (E, I, H) - - w13_t_sel = self._w13_t[eids] # (K, H, 2*I) - w2_t_sel = self._w2_t[eids] # (K, I, H) - - # FC1: bmm (K,1,H) @ (K,H,2I) → (K,1,2I) - x_expand = hidden_states.unsqueeze(0).expand(self.top_k, -1, -1) # (K, 1, H) - gate_up = torch.bmm(x_expand, w13_t_sel).squeeze(1) # (K, 2*I) + # FC1: single large GEMM via F.linear + # (1, H) @ (K*2*I, H)^T → (1, K*2*I) + # Source: base qwen3_5.py — verified on BI-V100 (sub 655 = 683) + gate_up = F.linear( + hidden_states, + w13_sel.reshape(-1, H), # (K*2*I, H) + ) # (1, K*2*I) + gate_up = gate_up.view(self.top_k, -1) # (K, 2*I) if _USE_FUSED_MOE_ACTIVATION: act = self.act_fn(gate_up) # (K, I) @@ -1844,8 +1839,9 @@ class Qwen3_5MoeSparseBlock(nn.Module): gate, up = gate_up.chunk(2, dim=-1) act = F.silu(gate) * up - # FC2: bmm (K,1,I) @ (K,I,H) → (K,1,H) - expert_out = torch.bmm(act.unsqueeze(1), w2_t_sel).squeeze(1) # (K, H) + # FC2: bmm (K, H, I) @ (K, I, 1) → (K, H) + # w2_sel is (K, H, I), act is (K, I) + expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) if (_USE_COREX_MOE_EXACT_REDUCE and expert_out.dtype == torch.float16