From c17fd301446ed931385577bfe79263fb0778393d Mon Sep 17 00:00:00 2001 From: claude Date: Tue, 11 Aug 2026 07:16:19 +0000 Subject: [PATCH] =?UTF-8?q?feat(CRITICAL):=20=E6=8E=A5=E5=85=A5ix=5Funifie?= =?UTF-8?q?d=E5=88=B0qwen3=5F5.py=20MoE=20prefill=E8=B7=AF=E5=BE=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - import ix_unified bridge到qwen3_5.py - MoE prefill: ix_bridge.moe_group_gemm替代per-expert for-loop - 保留fallback: ix_bridge失败自动回退到PyTorch for-loop - 新feature flag: BI100_MOE_IX_BRIDGE (default=True when bridge available) --- qwen3_6_scripts/qwen3_5.py | 35 +++++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 1823934d..dc184a22 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -143,7 +143,15 @@ from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA, logger = init_logger(__name__) -_bi100_model_trace("qwen3_5 runtime imports complete") +# --- ix_unified: bridge to ixformer::infer C++ APIs ------------------------- +try: + from ix_unified import ix as _ix_bridge + _HAS_IX_BRIDGE = (_ix_bridge._bridge is not None) +except ImportError: + _ix_bridge = None + _HAS_IX_BRIDGE = False + +_bi100_model_trace(f"qwen3_5 runtime imports complete (ix_bridge={_HAS_IX_BRIDGE})") _ALLOW_GDN_NAN_ZERO = env_bool("BI100_GDN_ALLOW_NAN_ZERO", False) _GDN_FINITE_CHECK = (env_bool("BI100_GDN_FINITE_CHECK", False) @@ -180,6 +188,7 @@ _USE_COREX_MOE_DIRECT_ROUTED = ( _corex_moe_direct_routed is not None and env_bool("BI100_MOE_COREX_DIRECT_ROUTED", False)) _USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True) +_USE_IX_BRIDGE_MOE = (_HAS_IX_BRIDGE and env_bool("BI100_MOE_IX_BRIDGE", True)) # --------------------------------------------------------------------------- @@ -1688,18 +1697,32 @@ class Qwen3_5MoeSparseBlock(nn.Module): out = (expert_out * ws.unsqueeze(-1)).sum( 0, keepdim=True).to(hidden_states.dtype) # (1, H) else: - # General path (prefill / multi-seq): group assignments once. The - # previous implementation scanned the full (T, top_k) routing - # matrix and ran nonzero() for every active expert. + # General path (prefill / multi-seq) out = torch.zeros_like(hidden_states) flat_eids = topk_ids.reshape(-1) order = torch.argsort(flat_eids, stable=True) sorted_tok_ids = torch.arange( T, device=topk_ids.device).repeat_interleave(self.top_k)[order] sorted_weights = topk_weights.reshape(-1)[order] - expert_counts = torch.bincount( - flat_eids, minlength=w13.shape[0]).tolist() + expert_counts_t = torch.bincount( + flat_eids, minlength=w13.shape[0]) + if _USE_IX_BRIDGE_MOE: + # ix_bridge path: batched group_gemm instead of per-expert loop + try: + sorted_hidden = hidden_states[sorted_tok_ids] + gate_up = _ix_bridge.moe_group_gemm( + sorted_hidden, w13, expert_counts_t) + act = self.act_fn(gate_up) + down = _ix_bridge.moe_group_gemm( + act, w2, expert_counts_t) + down_weighted = down * sorted_weights.unsqueeze(-1) + out.index_add_(0, sorted_tok_ids, down_weighted.to(out.dtype)) + return out + except Exception as e: + logger.warning("ix_bridge MoE failed (%s), falling back", e) + + expert_counts = expert_counts_t.tolist() start = 0 for eid, count in enumerate(expert_counts): end = start + count