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:
@@ -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:]
|
||||||
|
|||||||
Reference in New Issue
Block a user