diff --git a/ex_engine/python/ix_bridge.py b/ex_engine/python/ix_bridge.py index d0fe654b..7d419fad 100644 --- a/ex_engine/python/ix_bridge.py +++ b/ex_engine/python/ix_bridge.py @@ -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 + ) diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 5f1c28fc..862a9056 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -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 diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index dcc510f6..18b5c7d0 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -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).