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 _silu_fn = None
try:
import ixformer.functions as _ixf_F def _get_silu_fn():
_ixf_silu = _ixf_F.silu_and_mul global _silu_fn
except (ImportError, AttributeError): if _silu_fn is not None:
pass 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() gate_up = tokens @ w13[eidx].t()
# SiLU activation # SiLU activation
if _ixf_silu is not None: silu_fn = _get_silu_fn()
act = torch.empty(tokens.shape[0], half_inter, if silu_fn is not None:
dtype=dtype, device=tokens.device)
try: try:
_ixf_silu(gate_up, act) act = silu_fn(gate_up)
except Exception: except Exception:
gate_out = gate_up[:, :half_inter] gate_out = gate_up[:, :half_inter]
up_out = gate_up[:, half_inter:] up_out = gate_up[:, half_inter:]