From 1615a313eac5d99ab1a80aabd82c02e1ae6ee26d Mon Sep 17 00:00:00 2001 From: root Date: Tue, 25 Aug 2026 04:35:36 +0000 Subject: [PATCH] =?UTF-8?q?[fix]=20F.linear=E2=86=92ix=5Fmoe=5Fbridge.line?= =?UTF-8?q?ar?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- qwen3_6_scripts/qwen3_5.py | 40 +++++++++++++++++++++++++++++++++++--- 1 file changed, 37 insertions(+), 3 deletions(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index a0a0e06c..8392bba2 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -166,6 +166,40 @@ except ImportError: except ImportError: _xllm_moe = None +# ix_moe_bridge: direct GEMV (4.1x faster than F.linear for decode M=1) +# Benchmark: F.linear 133us vs br.linear 32us on BI-V100 +try: + from vllm import ix_moe_bridge as _ix_moe_bridge +except ImportError: + try: + import importlib.util as _ilu + for _p in ["/workspace/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ix_moe_bridge.so", + os.path.join(os.path.dirname(__file__), "prebuilt", + "corex-3.2.3-ivcore10", "ix_moe_bridge.so")]: + if os.path.isfile(_p): + _spec = _ilu.spec_from_file_location("ix_moe_bridge", _p) + _ix_moe_bridge = _ilu.module_from_spec(_spec) + _spec.loader.exec_module(_ix_moe_bridge) + break + else: + _ix_moe_bridge = None + except Exception: + _ix_moe_bridge = None +_HAS_BRIDGE_LINEAR = ( + _ix_moe_bridge is not None + and hasattr(_ix_moe_bridge, 'linear')) +if _HAS_BRIDGE_LINEAR: + logger.info("ix_moe_bridge.linear ENABLED — 4.1x GEMV speedup for decode") + + +def _fast_linear(x: torch.Tensor, weight: torch.Tensor, + bias=None) -> torch.Tensor: + """Drop-in F.linear replacement using ix_moe_bridge GEMV. + 4.1x faster than F.linear for M=1 decode on BI-V100.""" + if _HAS_BRIDGE_LINEAR and x.dtype == torch.float16 and weight.dtype == torch.float16: + return _ix_moe_bridge.linear(x, weight, bias) + return F.linear(x, weight, bias) + try: from vllm import corex_gdn_chunk_recurrent as _corex_gdn_chunk_recurrent except ImportError: @@ -1855,7 +1889,7 @@ class Qwen3_5MoeSparseBlock(nn.Module): # 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( + gate_up = _fast_linear( hidden_states, w13_sel.reshape(-1, H), # (K*2*I, H) ) # (1, K*2*I) @@ -1945,10 +1979,10 @@ class Qwen3_5MoeSparseBlock(nn.Module): continue tok_ids = sorted_tok_ids[start:end] tokens = hidden_states[tok_ids] # (n, H) - gate_up = F.linear(tokens, w13[eid]) # (n, 2*I) + gate_up = _fast_linear(tokens, w13[eid]) # (n, 2*I) gate, up = gate_up.chunk(2, dim=-1) act = F.silu(gate) * up # (n, I) - expert_out = F.linear(act, w2[eid]) # (n, H) + expert_out = _fast_linear(act, w2[eid]) # (n, H) weights = sorted_weights[start:end].unsqueeze(-1) out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype)) start = end