From 905bf4db2c9fa2b11b348401ccd6a07021fd8104 Mon Sep 17 00:00:00 2001 From: project6-dev Date: Mon, 10 Aug 2026 06:36:40 +0000 Subject: [PATCH] feat(moe): wire silu_and_mul through C++ bridge in corex_moe.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Now MoE activation uses: Tier 0: ix_bridge.silu_and_mul (C++ ixformer_torch_ext, verified on BI-V100) Tier 1: ixformer.functions.silu_and_mul (Python) Tier 2: F.silu(gate) * up (pure PyTorch) Verified 7/8 on single BI-V100: ✓ compile, silu_and_mul, rms_norm, fused_add_rms_norm, linear, paged_attn, corex_moe ✗ flash_attn import path (not needed, vllm xformers backend handles it) --- ex_engine/python/corex_moe.py | 32 +++++++++++++++++++++----------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/ex_engine/python/corex_moe.py b/ex_engine/python/corex_moe.py index f3fc90f0..a2f97254 100644 --- a/ex_engine/python/corex_moe.py +++ b/ex_engine/python/corex_moe.py @@ -71,14 +71,25 @@ def _python_topk_softmax(gating_output, topk, renormalize=True): # ----------------------------------------------------------------------- -# ixformer.functions Python-level SiLU +# silu_and_mul acceleration: prefer C++ bridge, fallback to ixformer Python # ----------------------------------------------------------------------- -_ixf_silu = None -try: - import ixformer.functions as _ixf_F - _ixf_silu = _ixf_F.silu_and_mul -except (ImportError, AttributeError): - pass +_silu_fn = None + +def _get_silu_fn(): + global _silu_fn + if _silu_fn is not None: + return _silu_fn + # Tier 0: C++ bridge (ixformer_torch_ext::silu_and_mul_forward) + if _ensure_bridge() and hasattr(_bridge, 'silu_and_mul'): + _silu_fn = _bridge.silu_and_mul + return _silu_fn + # Tier 1: ixformer Python + try: + import ixformer.functions as _ixf_F + _silu_fn = _ixf_F.silu_and_mul + except (ImportError, AttributeError): + pass + return _silu_fn # ----------------------------------------------------------------------- @@ -182,11 +193,10 @@ def _python_moe_forward(hidden_states, gate_output, w13, w2, gate_up = tokens @ w13[eidx].t() # SiLU activation - if _ixf_silu is not None: - act = torch.empty(tokens.shape[0], half_inter, - dtype=dtype, device=tokens.device) + silu_fn = _get_silu_fn() + if silu_fn is not None: try: - _ixf_silu(gate_up, act) + act = silu_fn(gate_up) except Exception: gate_out = gate_up[:, :half_inter] up_out = gate_up[:, half_inter:]