dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat(MoE): wire full ix_fused_moe_forward as Tier 0 dispatch

Browse Source 
ix_bridge.py: expose all 6 ixformer::infer functions + fused_moe_forward()
qwen3_5.py: 4-tier MoE dispatch (fused C++ → CUB topk → ix topk → PyTorch)
patch_ops.sh: deploy ix_moe_bridge.cpp to 4 search paths for JIT
This commit is contained in:
EX Engine
2026-08-10 03:36:39 +00:00
parent 1be9449883
commit 388f6b2d1a
3 changed files with 192 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

184
ex_engine/python/ix_bridge.py
View File

@@ -1,71 +1,179 @@
"""
ix_bridge.py — Load ix_moe_bridge C++ extension at runtime.
ix_bridge.py — Full ixformer MoE pipeline bridge.
Calls ixformer::infer::topk_softmax() via C++ torch extension,
bypassing the missing Python binding in ixformer.functions.
Loads ix_moe_bridge.so via JIT and exposes both individual ops and the full
fused MoE forward pass that replaces the Python for-loop in qwen3_5.py.
Build: JIT-compiled on first import via torch.utils.cpp_extension.load()
(same mechanism as flash_qla_sm70 GDN kernel — proven to work on BI-V100)
Pipeline (mirrors xllm/core/layers/ilu/fused_moe.cpp):
topk_softmax → moe_gen_idx → moe_expand_input → group_gemm(w13)
→ silu_and_mul → group_gemm(w2) → moe_combine_result
All 6 ixformer::infer C++ functions are called through ix_moe_bridge.cpp
which forward-declares them and links against the base image SDK.
"""
import os
import logging
import torch
from typing import Tuple, Optional, List
logger = logging.getLogger("ex_engine.ix_bridge")
_ix_bridge = None
_ix_bridge_loaded = False # True after attempt, even if failed
_ix_bridge_available = False
def _find_cpp_source():
"""Find ix_moe_bridge.cpp in multiple locations."""
candidates = []
# 1. Relative to this file: ex_engine/csrc/
here = os.path.dirname(os.path.abspath(__file__))
candidates.append(os.path.join(here, "..", "csrc", "ix_moe_bridge.cpp"))
# 2. Deployed path inside vllm model dir
candidates.append(os.path.join(here, "ix_moe_bridge.cpp"))
# 3. /workspace paths
candidates.append("/workspace/ex_engine/csrc/ix_moe_bridge.cpp")
candidates.append("/workspace/qwen3_6_scripts/ix_moe_bridge.cpp")
for c in candidates:
p = os.path.normpath(c)
if os.path.exists(p):
return p
return None
def _load_bridge():
"""JIT-compile and load ix_moe_bridge.so"""
global _ix_bridge, _ix_bridge_available
if _ix_bridge is not None:
"""JIT-compile and load ix_moe_bridge.so — called once."""
global _ix_bridge, _ix_bridge_loaded, _ix_bridge_available
if _ix_bridge_loaded:
return _ix_bridge_available
csrc_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "csrc")
cpp_file = os.path.join(csrc_dir, "ix_moe_bridge.cpp")
if not os.path.exists(cpp_file):
# Try deployed path (inside vllm model dir)
alt_dir = os.path.dirname(os.path.abspath(__file__))
cpp_file = os.path.join(alt_dir, "ix_moe_bridge.cpp")
if not os.path.exists(cpp_file):
logger.warning("ix_moe_bridge.cpp not found at %s", cpp_file)
_ix_bridge_available = False
_ix_bridge_loaded = True
cpp_file = _find_cpp_source()
if cpp_file is None:
logger.warning("ix_moe_bridge.cpp not found in any search path")
return False
try:
from torch.utils.cpp_extension import load
logger.info("JIT-compiling ix_moe_bridge.cpp ...")
logger.info("JIT-compiling ix_moe_bridge.cpp from %s ...", cpp_file)
_ix_bridge = load(
name="ix_moe_bridge",
sources=[cpp_file],
extra_cflags=["-O2"],
extra_cflags=["-O2", "-std=c++17"],
verbose=False,
)
_ix_bridge_available = True
logger.info("ix_moe_bridge loaded successfully: %s", dir(_ix_bridge))
fns = [x for x in dir(_ix_bridge) if not x.startswith("_")]
logger.info("ix_moe_bridge loaded: %s", fns)
return True
except Exception as e:
logger.warning("ix_moe_bridge JIT compile failed: %s", e)
_ix_bridge_available = False
return False
def topk_softmax(gating_output: torch.Tensor, topk: int, renormalize: bool = True):
def is_available() -> bool:
"""Check if bridge is available (lazy-load on first call)."""
if not _ix_bridge_loaded:
_load_bridge()
return _ix_bridge_available
# =========================================================================
# Individual ops (thin wrappers with type safety)
# =========================================================================
def topk_softmax(
gating_output: torch.Tensor,
topk: int,
renormalize: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Fused topk+softmax via ixformer C++ API.
FAIL FAST: if bridge not available, raises RuntimeError immediately.
No silent fallback — 0 score with no error log is worse than a crash.
Fused topk+softmax via ixformer::infer::topk_softmax.
Returns: (topk_weights [T, K] fp32, topk_ids [T, K] int32)
"""
if not _ix_bridge_available:
if not _load_bridge():
raise RuntimeError(
"ix_moe_bridge: FATAL — ixformer C++ topk_softmax not available. "
"JIT compile failed. Run probe_ixformer_symbols.py on real machine "
"to diagnose. Cannot fall back silently — would produce 0 score."
)
if not is_available():
raise RuntimeError("ix_moe_bridge not available — JIT compile failed")
return _ix_bridge.topk_softmax(gating_output, topk, renormalize)
def moe_gen_idx(
expert_id: torch.Tensor,
expert_num: int,
) -> List[torch.Tensor]:
"""
Build expert permutation maps.
Returns: [src_dst, dst_src, expert_sizes, cumsum]
"""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.moe_gen_idx(expert_id, expert_num)
def moe_expand_input(
input: torch.Tensor,
gather_index: torch.Tensor,
combine_idx: torch.Tensor,
topk: int,
) -> torch.Tensor:
"""Gather tokens by expert assignment."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.moe_expand_input(input, gather_index, combine_idx, topk)
def group_gemm(
inputs: torch.Tensor,
weights: torch.Tensor,
token_count: torch.Tensor,
output_n: int,
) -> torch.Tensor:
"""Batched expert GEMM via ixformer."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.group_gemm(inputs, weights, token_count, output_n)
def silu_and_mul(input: torch.Tensor) -> torch.Tensor:
"""Fused SiLU gate activation."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.silu_and_mul(input)
def moe_combine_result(
input: torch.Tensor,
weight: torch.Tensor,
) -> torch.Tensor:
"""Weighted reduce for MoE output."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.moe_combine_result(input, weight)
# =========================================================================
# Full fused MoE forward — replaces _pure_pytorch_experts() entirely
# =========================================================================
def fused_moe_forward(
hidden_states: torch.Tensor, # (T, H)
router_logits: torch.Tensor, # (T, E)
w13: torch.Tensor, # (E, 2*I, H) gate_up
w2: torch.Tensor, # (E, H, I) down
topk: int,
num_experts: int,
renormalize: bool = True,
) -> torch.Tensor:
"""
Full fused MoE forward via ixformer C++ pipeline.
Pipeline: topk → gen_idx → expand → gemm1(w13) → silu → gemm2(w2) → combine
Returns: (T, H) — partial output, needs all-reduce after.
"""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.fused_moe_forward(
hidden_states, router_logits, w13, w2, topk, num_experts, renormalize
)

4
qwen3_6_scripts/patch_ops.sh
View File

@@ -185,9 +185,11 @@ if [ -d "$EX_ENGINE_SRC/python" ]; then
EX_DST="$VLLM/model_executor/models/ex_engine"
mkdir -p "$EX_DST/python" "$EX_DST/csrc"
cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST/python/" 2>/dev/null || true
# ix_moe_bridge.cpp for JIT compile
# ix_moe_bridge.cpp for JIT compile — deploy to ALL search paths
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/csrc/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/python/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "/workspace/ex_engine/csrc/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "/workspace/qwen3_6_scripts/" 2>/dev/null || true
touch "$EX_DST/__init__.py"
touch "$EX_DST/python/__init__.py"
# Copy built .so files

60
qwen3_6_scripts/qwen3_5.py
View File

@@ -71,29 +71,35 @@ except ImportError:
# corex_gdn/corex_moe: these are custom modules that teams package into their
# Docker image. If present, they provide fused GDN/MoE kernels.
# ix_bridge: C++ bridge to ixformer::infer::topk_softmax (bypasses missing Python binding)
_ix_bridge_module = None
# ix_bridge: C++ bridge to ixformer::infer (full MoE pipeline)
_ix_bridge_available = False
_ix_topk_softmax = None
_ix_fused_moe_forward = None
try:
from ex_engine.python.ix_bridge import topk_softmax as _ix_topk_softmax
from ex_engine.python.ix_bridge import (
topk_softmax as _ix_topk_softmax,
fused_moe_forward as _ix_fused_moe_forward,
is_available as _ix_bridge_check,
)
_ix_bridge_available = True
logger.info("ix_bridge: ixformer C++ topk_softmax available")
logger.info("ix_bridge: full ixformer MoE pipeline available (topk + fused_moe)")
except ImportError:
try:
# Try deployed path inside vllm models dir
import importlib, sys
import sys
_ex_dir = os.path.join(os.path.dirname(__file__), "ex_engine")
if os.path.isdir(_ex_dir) and _ex_dir not in sys.path:
sys.path.insert(0, os.path.dirname(_ex_dir))
from ex_engine.python.ix_bridge import topk_softmax as _ix_topk_softmax
from ex_engine.python.ix_bridge import (
topk_softmax as _ix_topk_softmax,
fused_moe_forward as _ix_fused_moe_forward,
is_available as _ix_bridge_check,
)
_ix_bridge_available = True
logger.info("ix_bridge: ixformer C++ topk_softmax available (deployed path)")
logger.info("ix_bridge: full ixformer MoE pipeline available (deployed path)")
except ImportError as e:
# NOT silent: log the exact error so we can diagnose from docker logs
logger.warning(
"ix_bridge: IMPORT FAILED (%s). MoE will use PyTorch topk. "
"This is 3x slower. Run probe_ixformer_symbols.py to diagnose.", e)
"ix_bridge: IMPORT FAILED (%s). MoE will use PyTorch fallback. "
"This is 3-10x slower.", e)
_corex_gdn_available = False
_corex_moe_available = False
@@ -1073,13 +1079,36 @@ class Qwen3_5MoeSparseBlock(nn.Module):
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
) -> torch.Tensor:
"""Pure-PyTorch MoE (ixformer has no MoE kernels on BI-V100).
"""MoE expert computation with tiered dispatch.
Dispatch order:
Tier 0: ix_fused_moe_forward — full C++ pipeline (7 kernel launches)
Tier 1: EX Engine CUB topk kernel + PyTorch GEMM
Tier 2: ix_bridge topk_softmax + PyTorch GEMM
Tier 3: Pure PyTorch (torch.softmax + torch.topk + for-loop)
w13_weight: (num_experts, 2*inter_per_partition, hidden) [TP-sharded]
w2_weight: (num_experts, hidden, inter_per_partition) [TP-sharded]
Output is partial (pre-all-reduce), same contract as FusedMoE
with reduce_results=False.
Output is partial (pre-all-reduce), same contract as FusedMoE.
"""
w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I)
# Tier 0: Full fused MoE pipeline via ixformer C++
# 7 kernel launches vs 3*E in Python loop
if _ix_fused_moe_forward is not None and _ix_bridge_available:
try:
return _ix_fused_moe_forward(
hidden_states, router_logits,
w13, w2,
self.top_k, self.num_experts,
renormalize=True,
)
except Exception as e:
if not getattr(self, '_ix_fused_warned', False):
logger.warning("ix_fused_moe_forward failed (%s), falling back to tiered dispatch", e)
self._ix_fused_warned = True
# Routing: fused topk+softmax dispatch chain
# Tier 1: EX Engine CUB kernel → Tier 2: ix_bridge → Tier 3: PyTorch
if _ex_moe_topk_available:
@@ -1105,9 +1134,6 @@ class Qwen3_5MoeSparseBlock(nn.Module):
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
topk_weights = topk_weights.to(hidden_states.dtype)
w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I)
T = hidden_states.shape[0]
if T == 1:
# Fast path: single token (decode).