feat(moe): wire silu_and_mul through C++ bridge in corex_moe.py

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)
This commit is contained in:
project6-dev
2026-08-10 06:36:40 +00:00
parent 7127d18491
commit 905bf4db2c

View File

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