From e8f0948fe11329401a6b03084bb1eff8e72acb56 Mon Sep 17 00:00:00 2001 From: dylan Date: Sat, 15 Aug 2026 06:15:17 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20ix=5Fops=20integration=20layer=20?= =?UTF-8?q?=E2=80=94=20wire=20ix=5Ffull=5Fbridge.so=20into=20vllm=20hot=20?= =?UTF-8?q?path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Architecture (CCCL dispatch pattern): base_image ixformer::infer → ix_full_bridge.so → ix_ops.py → vllm patches New files: ex_engine/python/ix_ops.py — Python API for all 14 ixformer::infer ops ex_engine/python/patch_vllm_ops.py — monkey-patch vllm GemmaRMSNorm, SiluAndMul ex_engine/deploy_ix_bridge.sh — build-time deployment script Modified: qwen3_6_scripts/patch_ops.sh — integrated ix_bridge deployment + startup hook Call chain: DecoderLayer.forward → GemmaRMSNorm → ix_ops.fused_add_rms_norm → ixformer::infer::residual_rms_norm (fused C++ kernel) --- ex_engine/deploy_ix_bridge.sh | 154 +++++++++++++ ex_engine/python/ix_ops.py | 343 +++++++++++++++++++++++++++++ ex_engine/python/patch_vllm_ops.py | 201 +++++++++++++++++ qwen3_6_scripts/patch_ops.sh | 59 +++++ 4 files changed, 757 insertions(+) create mode 100755 ex_engine/deploy_ix_bridge.sh create mode 100644 ex_engine/python/ix_ops.py create mode 100644 ex_engine/python/patch_vllm_ops.py diff --git a/ex_engine/deploy_ix_bridge.sh b/ex_engine/deploy_ix_bridge.sh new file mode 100755 index 00000000..19b25ae5 --- /dev/null +++ b/ex_engine/deploy_ix_bridge.sh @@ -0,0 +1,154 @@ +#!/bin/bash +# ex_engine/deploy_ix_bridge.sh — Deploy ix_full_bridge + Python ops into vllm +# +# Architecture (CCCL build pattern): +# CCCL: cmake → compile → install to site-packages +# EX: torch.utils.cpp_extension → compile bridge → deploy to vllm pkg +# +# What this does: +# 1. Find ixformer .so libraries in base image +# 2. Either use prebuilt ix_full_bridge.so or JIT-compile from source +# 3. Deploy .so + Python modules into vllm package +# 4. Verify dlopen chain works +# +# Source mapping: +# ex_engine/csrc/ix_full_bridge_v2.cpp → pybind11 bridge to ixformer::infer +# ex_engine/python/ix_ops.py → Python API layer +# ex_engine/python/patch_vllm_ops.py → vllm monkey-patches +# +# Called from: qwen3_6_scripts/patch_ops.sh + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +VLLM_ROOT="${1:-$(python3 -c 'import vllm; import os; print(os.path.dirname(vllm.__file__))' 2>/dev/null || echo '/usr/local/corex/lib/python3/dist-packages/vllm')}" + +echo "[ix_bridge] VLLM_ROOT=${VLLM_ROOT}" +echo "[ix_bridge] SCRIPT_DIR=${SCRIPT_DIR}" + +# ========================================================================= +# Step 1: Deploy prebuilt .so if available +# ========================================================================= +PREBUILT="${SCRIPT_DIR}/../qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10" +BRIDGE_SO="" + +if [[ -f "${PREBUILT}/ix_full_bridge.so" ]]; then + cp "${PREBUILT}/ix_full_bridge.so" "${VLLM_ROOT}/ix_full_bridge.so" + BRIDGE_SO="${VLLM_ROOT}/ix_full_bridge.so" + echo "[ix_bridge] deployed prebuilt ix_full_bridge.so" +fi + +# Deploy all corex_*.so and xllm_*.so +if [[ -d "$PREBUILT" ]]; then + for so_file in "${PREBUILT}"/*.so; do + base=$(basename "$so_file") + if [[ "$base" != "ix_full_bridge.so" ]]; then + cp "$so_file" "${VLLM_ROOT}/${base}" 2>/dev/null || true + echo "[ix_bridge] deployed ${base}" + fi + done +fi + +# ========================================================================= +# Step 2: Deploy Python integration modules +# ========================================================================= +# Create ex_engine package in vllm +EX_PKG="${VLLM_ROOT}/ex_engine" +mkdir -p "${EX_PKG}" + +cat > "${EX_PKG}/__init__.py" << 'PYEOF' +"""ex_engine — Algorithm factor replacement engine for BI-V100.""" +PYEOF + +# Deploy ix_ops.py +cp "${SCRIPT_DIR}/python/ix_ops.py" "${EX_PKG}/ix_ops.py" +echo "[ix_bridge] deployed ix_ops.py" + +# Deploy patch_vllm_ops.py +cp "${SCRIPT_DIR}/python/patch_vllm_ops.py" "${EX_PKG}/patch_vllm_ops.py" +echo "[ix_bridge] deployed patch_vllm_ops.py" + +# Also make ix_ops importable from vllm.ex_engine +# and from the top-level ex_engine path +SITE_EX="${SCRIPT_DIR}/python" +if [[ -d "$SITE_EX" ]]; then + # Ensure __init__.py exists + touch "${SITE_EX}/../__init__.py" 2>/dev/null || true +fi + +# ========================================================================= +# Step 3: Create auto-patch entry point +# ========================================================================= +# This script is sourced by patch_ops.sh to ensure ix_ops patches +# are applied at vllm startup +cat > "${VLLM_ROOT}/ix_startup_patch.py" << 'PYEOF' +""" +ix_startup_patch.py — Apply ix_ops patches at vllm startup. + +Import this module early in the vllm startup to replace PyTorch fallbacks +with fused C++ kernels from the base image. + +Architecture (CCCL dispatch pattern): + import vllm → vllm.__init__ → ix_startup_patch → patch_vllm_ops +""" +import logging +logger = logging.getLogger("ix_startup_patch") + +def apply(): + """Apply all available ix_ops patches.""" + try: + from vllm.ex_engine.patch_vllm_ops import apply_all_patches + n = apply_all_patches() + if n > 0: + logger.info("ix_startup_patch: %d patches applied", n) + return n + except Exception as e: + logger.warning("ix_startup_patch failed: %s", e) + return 0 + +# Auto-apply on import +_n_patches = apply() +PYEOF +echo "[ix_bridge] deployed ix_startup_patch.py" + +# ========================================================================= +# Step 4: Deploy bridge C++ source for JIT fallback +# ========================================================================= +CSRC_DEST="${VLLM_ROOT}/ex_engine/csrc" +mkdir -p "${CSRC_DEST}" +for cpp in "${SCRIPT_DIR}/csrc/ix_full_bridge_v2.cpp" \ + "${SCRIPT_DIR}/csrc/ix_full_bridge.cpp" \ + "${SCRIPT_DIR}/csrc/ix_moe_bridge.cpp"; do + if [[ -f "$cpp" ]]; then + cp "$cpp" "${CSRC_DEST}/" + echo "[ix_bridge] deployed $(basename $cpp) for JIT fallback" + fi +done + +# ========================================================================= +# Step 5: Verify deployment +# ========================================================================= +echo "" +echo "[ix_bridge] === Deployment Summary ===" +echo "[ix_bridge] Bridge .so: ${BRIDGE_SO:-'(JIT compile at runtime)'}" +echo "[ix_bridge] Python ops: ${EX_PKG}/ix_ops.py" +echo "[ix_bridge] vllm patches: ${EX_PKG}/patch_vllm_ops.py" +echo "[ix_bridge] Startup hook: ${VLLM_ROOT}/ix_startup_patch.py" + +# Quick Python import test +python3 -c " +import sys +sys.path.insert(0, '${VLLM_ROOT}') +try: + from vllm.ex_engine import ix_ops + print('[ix_bridge] ✓ ix_ops importable') +except Exception as e: + print(f'[ix_bridge] ✗ ix_ops import failed: {e}') +try: + from vllm.ex_engine import patch_vllm_ops + print('[ix_bridge] ✓ patch_vllm_ops importable') +except Exception as e: + print(f'[ix_bridge] ✗ patch_vllm_ops import failed: {e}') +" 2>&1 || true + +echo "[ix_bridge] === Done ===" diff --git a/ex_engine/python/ix_ops.py b/ex_engine/python/ix_ops.py new file mode 100644 index 00000000..f6ded5f0 --- /dev/null +++ b/ex_engine/python/ix_ops.py @@ -0,0 +1,343 @@ +""" +ix_ops.py — Drop-in operator replacements via ix_full_bridge.so + +Architecture (CCCL dispatch pattern): + CCCL: compute_capability → policy_selector → tuned_kernel + EX: base_image_so → ix_full_bridge → ixformer::infer + +This module provides torch.nn.Module-compatible replacements for: + 1. RMSNorm → residual_rms_norm / rms_norm (fused kernel) + 2. SiluAndMul → silu_and_mul (fused activation) + 3. RotaryEmbedding → xllm_rotary_embedding (fused RoPE) + 4. reshape_and_cache → xllm_reshape_and_cache (fused KV write) + 5. paged_attention → xllm_paged_attention (fused decode attn) + 6. flash_attn_prefill → ixinfer_flash_attn_unpad (fused prefill attn) + 7. linear → ixformer_linear / linear_ex (GEMM) + +Loading: tries prebuilt ix_full_bridge.so first, then JIT-compiles +ix_full_bridge_v2.cpp as fallback. + +Source mapping: + upstream_ref/xllm_latest/core/kernels/ilu/*.cpp → this file (Python side) + ex_engine/csrc/ix_full_bridge_v2.cpp → .so (C++ side) + ixformer::infer namespace (base image) → actual CUDA kernels +""" + +import os +import sys +import logging +import importlib +import importlib.util +import glob +import torch +from typing import Optional, Tuple, List + +logger = logging.getLogger("ex_engine.ix_ops") + +# ========================================================================= +# Bridge loader +# ========================================================================= +_bridge = None +_loaded = False +_available = False + + +def _try_prebuilt(): + """Load prebuilt ix_full_bridge.so.""" + search = [ + # Deployed by patch_ops.sh into vllm package + "/usr/local/corex/lib/python3/dist-packages/vllm/ix_full_bridge.so", + ] + # Also check vllm package dir + try: + import vllm + vd = os.path.dirname(vllm.__file__) + search.insert(0, os.path.join(vd, "ix_full_bridge.so")) + except ImportError: + pass + # Check prebuilt dir + here = os.path.dirname(os.path.abspath(__file__)) + search.append(os.path.join(here, "..", "..", "qwen3_6_scripts", "prebuilt", + "corex-3.2.3-ivcore10", "ix_full_bridge.so")) + + for path in search: + path = os.path.normpath(path) + if not os.path.isfile(path): + continue + try: + spec = importlib.util.spec_from_file_location("ix_full_bridge", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + fns = [x for x in dir(mod) if not x.startswith("_")] + logger.info("ix_ops: loaded prebuilt %s: %s", path, fns) + return mod + except Exception as e: + logger.debug("ix_ops: prebuilt %s failed: %s", path, e) + return None + + +def _try_jit(): + """JIT compile ix_full_bridge_v2.cpp.""" + here = os.path.dirname(os.path.abspath(__file__)) + cpp_candidates = [ + os.path.join(here, "..", "csrc", "ix_full_bridge_v2.cpp"), + os.path.join(here, "..", "csrc", "ix_full_bridge.cpp"), + "/workspace/ex_engine/csrc/ix_full_bridge_v2.cpp", + "/workspace/qwen3_6_scripts/ix_full_bridge_v2.cpp", + ] + cpp_file = None + for c in cpp_candidates: + c = os.path.normpath(c) + if os.path.isfile(c): + cpp_file = c + break + if cpp_file is None: + return None + + extra_ldflags = [] + # Link ixformer .so libraries + try: + import ixformer + ixf_dir = os.path.dirname(ixformer.__file__) + for so in glob.glob(os.path.join(ixf_dir, "*.so")): + extra_ldflags.append(so) + extra_ldflags.append(f"-Wl,-rpath,{ixf_dir}") + except ImportError: + pass + # Also link corex libraries + corex_lib = "/usr/local/corex/lib64" + if os.path.isdir(corex_lib): + for lib in ["libixattn.so", "libixformer.so", "libcublas.so"]: + p = os.path.join(corex_lib, lib) + if os.path.isfile(p): + extra_ldflags.append(p) + extra_ldflags.append(f"-Wl,-rpath,{corex_lib}") + + try: + from torch.utils.cpp_extension import load + logger.info("ix_ops: JIT compiling %s", cpp_file) + mod = load( + name="ix_full_bridge_v2", + sources=[cpp_file], + extra_cflags=["-O2", "-std=c++17"], + extra_ldflags=extra_ldflags, + verbose=False, + ) + fns = [x for x in dir(mod) if not x.startswith("_")] + logger.info("ix_ops: JIT compiled: %s", fns) + return mod + except Exception as e: + logger.warning("ix_ops: JIT compile failed: %s", e) + return None + + +def _ensure_loaded(): + global _bridge, _loaded, _available + if _loaded: + return _available + _loaded = True + _bridge = _try_prebuilt() + if _bridge is None: + _bridge = _try_jit() + _available = _bridge is not None + if _available: + logger.info("ix_ops: bridge available with %d functions", + len([x for x in dir(_bridge) if not x.startswith("_")])) + else: + logger.warning("ix_ops: bridge NOT available, all ops will be no-op") + return _available + + +def is_available() -> bool: + return _ensure_loaded() + + +def get_bridge(): + if not _ensure_loaded(): + raise RuntimeError("ix_ops bridge not available") + return _bridge + + +# ========================================================================= +# Feature probes — check what the loaded bridge supports +# ========================================================================= +def has_silu_and_mul() -> bool: + return is_available() and hasattr(_bridge, "silu_and_mul") + +def has_rms_norm() -> bool: + return is_available() and hasattr(_bridge, "rms_norm") + +def has_fused_add_rms_norm() -> bool: + return is_available() and hasattr(_bridge, "fused_add_rms_norm") + +def has_rotary_embedding() -> bool: + return is_available() and hasattr(_bridge, "rotary_embedding") + +def has_reshape_and_cache() -> bool: + return is_available() and hasattr(_bridge, "reshape_and_cache") + +def has_paged_attention() -> bool: + return is_available() and hasattr(_bridge, "paged_attention") + +def has_flash_attn_prefill() -> bool: + return is_available() and hasattr(_bridge, "flash_attn_prefill") + +def has_linear() -> bool: + return is_available() and hasattr(_bridge, "linear") + +def has_topk_softmax() -> bool: + return is_available() and hasattr(_bridge, "topk_softmax") + +def has_fused_moe_forward() -> bool: + return is_available() and hasattr(_bridge, "fused_moe_forward") + + +# ========================================================================= +# Op wrappers — match xllm upstream signatures +# Source: upstream_ref/xllm_latest/core/kernels/ilu/*.cpp +# ========================================================================= + +def silu_and_mul(input: torch.Tensor) -> torch.Tensor: + """Fused SiLU activation + element-wise multiply. + + Source: xllm/core/kernels/ilu/activation.cpp → infer::silu_and_mul + input: (T, 2*I) → output: (T, I) + """ + return _bridge.silu_and_mul(input) + + +def rms_norm(output: torch.Tensor, input: torch.Tensor, + weight: torch.Tensor, eps: float = 1e-6) -> None: + """RMSNorm: output = rms_norm(input, weight, eps). + + Source: xllm/core/kernels/ilu/norm.cpp → infer::rms_norm + """ + _bridge.rms_norm(output, input, weight, eps) + + +def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor, + weight: torch.Tensor, output: torch.Tensor, + residual_output: torch.Tensor, + eps: float = 1e-6) -> None: + """Fused residual addition + RMSNorm. + + Source: xllm/core/kernels/ilu/norm.cpp → infer::residual_rms_norm + output = rms_norm(input + residual, weight, eps) + residual_output = input + residual + """ + _bridge.fused_add_rms_norm(input, residual, weight, output, + residual_output, eps) + + +def rotary_embedding(positions: torch.Tensor, query: torch.Tensor, + key: torch.Tensor, head_size: int, + cos_sin_cache: torch.Tensor, + is_neox: bool = True) -> None: + """Fused rotary position embedding (in-place on query and key). + + Source: xllm/core/kernels/ilu/rope.cpp → infer::xllm_rotary_embedding + """ + _bridge.rotary_embedding(positions, query, key, head_size, + cos_sin_cache, is_neox) + + +def reshape_and_cache(key: torch.Tensor, value: torch.Tensor, + key_cache: torch.Tensor, value_cache: torch.Tensor, + slot_mapping: torch.Tensor) -> None: + """Write KV to paged cache. + + Source: xllm/core/kernels/ilu/attention.cpp → infer::xllm_reshape_and_cache + """ + _bridge.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping) + + +def paged_attention(output: torch.Tensor, query: torch.Tensor, + key_cache: torch.Tensor, value_cache: torch.Tensor, + num_kv_heads: int, scale: float, + block_tables: torch.Tensor, seq_lens: torch.Tensor, + block_size: int, max_context_len: int, + alibi_slopes: Optional[torch.Tensor] = None + ) -> torch.Tensor: + """Paged attention decode. + + Source: xllm/core/kernels/ilu/attention.cpp → infer::xllm_paged_attention + """ + return _bridge.paged_attention( + output, query, key_cache, value_cache, + num_kv_heads, scale, block_tables, seq_lens, + block_size, max_context_len, alibi_slopes) + + +def flash_attn_prefill(query: torch.Tensor, key_cache: torch.Tensor, + value_cache: torch.Tensor, output: torch.Tensor, + block_tables: torch.Tensor, + cu_seq_q: torch.Tensor, cu_seq_k: torch.Tensor, + max_query_len: int, max_seq_len: int, + scale: float, is_causal: bool = True, + window_left: int = -1, + window_right: int = -1) -> torch.Tensor: + """Flash attention prefill with paged KV cache. + + Source: xllm/core/kernels/ilu/attention.cpp → + infer::ixinfer_flash_attn_unpad_with_block_tables + """ + return _bridge.flash_attn_prefill( + query, key_cache, value_cache, output, block_tables, + cu_seq_q, cu_seq_k, max_query_len, max_seq_len, + scale, is_causal, window_left, window_right) + + +def linear(input: torch.Tensor, weight: torch.Tensor, + bias: Optional[torch.Tensor] = None) -> torch.Tensor: + """GEMM via ixformer (auto-selects linear vs linear_ex). + + Source: xllm/core/kernels/ilu/matmul.cpp → infer::ixformer_linear[_ex] + """ + return _bridge.linear(input, weight, bias) + + +# ========================================================================= +# MoE ops — full 7-step pipeline +# Source: xllm/core/layers/ilu/fused_moe.cpp +# ========================================================================= +def topk_softmax(gating_output: torch.Tensor, topk: int, + renormalize: bool = True): + """Fused topk + softmax routing.""" + return _bridge.topk_softmax(gating_output, topk, renormalize) + + +def moe_gen_idx(expert_id: torch.Tensor, expert_num: int): + """Build expert permutation maps.""" + return _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): + """Expand input tokens by expert assignment.""" + return _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): + """Batched expert GEMM.""" + return _bridge.group_gemm(inputs, weights, token_count, output_n) + + +def moe_combine_result(input: torch.Tensor, weight: torch.Tensor): + """Weighted scatter-back of expert outputs.""" + return _bridge.moe_combine_result(input, weight) + + +def fused_moe_forward(hidden_states: torch.Tensor, + router_logits: torch.Tensor, + w13: torch.Tensor, w2: torch.Tensor, + topk: int, num_experts: int, + renormalize: bool = True) -> torch.Tensor: + """Full fused MoE forward (7-step pipeline). + + Source: xllm/core/layers/ilu/fused_moe.cpp → FusedMoEImpl::forward_experts + Pipeline: topk → gen_idx → expand → gemm1(w13) → silu → gemm2(w2) → combine + """ + return _bridge.fused_moe_forward( + hidden_states, router_logits, w13, w2, + topk, num_experts, renormalize) diff --git a/ex_engine/python/patch_vllm_ops.py b/ex_engine/python/patch_vllm_ops.py new file mode 100644 index 00000000..6a23fcbb --- /dev/null +++ b/ex_engine/python/patch_vllm_ops.py @@ -0,0 +1,201 @@ +""" +patch_vllm_ops.py — Wire ix_full_bridge C++ kernels into vllm's hot path. + +Architecture (CCCL policy_selector pattern): + Base image provides fused C++ kernels in ixformer::infer namespace. + ix_full_bridge.so wraps these with pybind11. + This module monkey-patches vllm's Python operators to call the bridge + instead of PyTorch fallback code. + +Problem statement (683 → 8000 gap): + vllm's _custom_ops.py fails to load on BI-V100 (no vllm C++ extensions). + Without patches, EVERY norm/activation/rope/cache/attention call goes + through pure PyTorch — multiple kernel launches per op instead of 1. + + Sub168 (competitor): all ops fused via xllm C++ engine → 11.9 TPS + Sub655 (us without patches): Python fallback → 2.6 TPS + +Solution: + Patch vllm's operator dispatch points so they call our bridge .so, + which links against the SAME ixformer .so files in the base image. + +Patched modules and their vllm paths: + 1. vllm.model_executor.layers.layernorm.GemmaRMSNorm + → ix_ops.rms_norm / ix_ops.fused_add_rms_norm + 2. vllm.model_executor.layers.activation.SiluAndMul + → ix_ops.silu_and_mul + 3. vllm._custom_ops (ops fallback registry) + → ix_ops for all registered ops + +Source mapping: + upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp → rms_norm patch + upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp → silu_and_mul patch + upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp → rotary_embedding patch + upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp → cache/attention patch +""" + +import os +import sys +import logging +import torch +from typing import Optional, Tuple + +logger = logging.getLogger("ex_engine.patch_vllm_ops") + +_patched = False + + +def apply_all_patches() -> int: + """Apply all available patches. Returns count of patches applied.""" + global _patched + if _patched: + return 0 + _patched = True + + from ex_engine.python import ix_ops + if not ix_ops.is_available(): + logger.warning("ix_ops bridge not available — no patches applied") + return 0 + + n = 0 + n += _patch_layernorm() + n += _patch_silu_and_mul() + n += _patch_custom_ops() + logger.info("patch_vllm_ops: %d patches applied", n) + return n + + +# ========================================================================= +# Patch 1: GemmaRMSNorm → fused C++ kernel +# ========================================================================= +def _patch_layernorm() -> int: + """Replace GemmaRMSNorm.forward with ix_ops.rms_norm.""" + from ex_engine.python import ix_ops + if not ix_ops.has_rms_norm(): + logger.debug("ix_ops missing rms_norm, skip layernorm patch") + return 0 + + try: + from vllm.model_executor.layers.layernorm import GemmaRMSNorm + except ImportError: + logger.debug("Cannot import GemmaRMSNorm, skip") + return 0 + + _orig_forward = GemmaRMSNorm.forward + + def _patched_forward(self, x, residual=None): + if residual is not None: + # fused_add_rms_norm: norm(x + residual) → (normed, new_residual) + if ix_ops.has_fused_add_rms_norm(): + out = torch.empty_like(x) + residual_out = torch.empty_like(x) + ix_ops.fused_add_rms_norm( + x, residual, self.weight, out, residual_out, + self.variance_epsilon) + return out, residual_out + else: + # Two-step fallback using just rms_norm + new_residual = x + residual + out = torch.empty_like(x) + ix_ops.rms_norm(out, new_residual, self.weight, + self.variance_epsilon) + return out, new_residual + else: + out = torch.empty_like(x) + ix_ops.rms_norm(out, x, self.weight, self.variance_epsilon) + return out + + GemmaRMSNorm.forward = _patched_forward + logger.info("PATCHED: GemmaRMSNorm.forward → ix_ops.rms_norm") + return 1 + + +# ========================================================================= +# Patch 2: SiluAndMul → fused C++ kernel +# ========================================================================= +def _patch_silu_and_mul() -> int: + """Replace SiluAndMul.forward with ix_ops.silu_and_mul.""" + from ex_engine.python import ix_ops + if not ix_ops.has_silu_and_mul(): + logger.debug("ix_ops missing silu_and_mul, skip activation patch") + return 0 + + try: + from vllm.model_executor.layers.activation import SiluAndMul + except ImportError: + logger.debug("Cannot import SiluAndMul, skip") + return 0 + + def _patched_forward(self, x): + return ix_ops.silu_and_mul(x) + + SiluAndMul.forward = _patched_forward + logger.info("PATCHED: SiluAndMul.forward → ix_ops.silu_and_mul") + return 1 + + +# ========================================================================= +# Patch 3: _custom_ops fallback registry +# ========================================================================= +def _patch_custom_ops() -> int: + """Patch vllm's _custom_ops to use ix_ops for registered ops.""" + from ex_engine.python import ix_ops + count = 0 + + try: + import vllm._custom_ops as ops + except ImportError: + logger.debug("Cannot import vllm._custom_ops, skip") + return 0 + + # Patch silu_and_mul + if ix_ops.has_silu_and_mul() and hasattr(ops, 'silu_and_mul'): + def _silu_and_mul(out, x): + result = ix_ops.silu_and_mul(x) + out.copy_(result) + ops.silu_and_mul = _silu_and_mul + count += 1 + logger.info("PATCHED: _custom_ops.silu_and_mul → ix_ops") + + # Patch rms_norm + if ix_ops.has_rms_norm() and hasattr(ops, 'rms_norm'): + def _rms_norm(out, input, weight, eps): + ix_ops.rms_norm(out, input, weight, eps) + ops.rms_norm = _rms_norm + count += 1 + logger.info("PATCHED: _custom_ops.rms_norm → ix_ops") + + # Patch fused_add_rms_norm + if ix_ops.has_fused_add_rms_norm() and hasattr(ops, 'fused_add_rms_norm'): + def _fused_add_rms_norm(input, residual, weight, eps): + out = torch.empty_like(input) + residual_out = torch.empty_like(input) + ix_ops.fused_add_rms_norm(input, residual, weight, + out, residual_out, eps) + input.copy_(out) + residual.copy_(residual_out) + ops.fused_add_rms_norm = _fused_add_rms_norm + count += 1 + logger.info("PATCHED: _custom_ops.fused_add_rms_norm → ix_ops") + + # Patch rotary_embedding + if ix_ops.has_rotary_embedding() and hasattr(ops, 'rotary_embedding'): + def _rotary_embedding(positions, query, key, head_size, + cos_sin_cache, is_neox): + ix_ops.rotary_embedding(positions, query, key, head_size, + cos_sin_cache, is_neox) + ops.rotary_embedding = _rotary_embedding + count += 1 + logger.info("PATCHED: _custom_ops.rotary_embedding → ix_ops") + + return count + + +# ========================================================================= +# Auto-apply on import if requested +# ========================================================================= +if os.environ.get("IX_OPS_AUTO_PATCH", "0") == "1": + try: + apply_all_patches() + except Exception as e: + logger.warning("ix_ops auto-patch failed: %s", e) diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 44737c26..abac4a57 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -199,6 +199,65 @@ if [ -d "$PREBUILT_DIR" ]; then done fi +# --- Deploy ix_bridge Python integration layer -------------------------------- +build_stage "deploying ix_bridge operator replacements" +EX_ENGINE_DIR="$(cd "$(dirname "$0")/../ex_engine" 2>/dev/null && pwd || echo "")" +if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR" ]; then + EX_ENGINE_DIR="$(cd "$(dirname "$0")" && pwd)/../ex_engine" +fi + +if [ -d "$EX_ENGINE_DIR/python" ]; then + # Create ex_engine package inside vllm + mkdir -p "${VLLM_ROOT}/ex_engine/csrc" + echo '"""ex_engine — Algorithm factor replacement for BI-V100."""' > "${VLLM_ROOT}/ex_engine/__init__.py" + + # Deploy Python modules + cp "$EX_ENGINE_DIR/python/ix_ops.py" "${VLLM_ROOT}/ex_engine/ix_ops.py" + cp "$EX_ENGINE_DIR/python/patch_vllm_ops.py" "${VLLM_ROOT}/ex_engine/patch_vllm_ops.py" + echo "[patch_ops] deployed ix_ops.py + patch_vllm_ops.py → ${VLLM_ROOT}/ex_engine/" + + # Deploy bridge C++ source for JIT fallback + for cpp in "$EX_ENGINE_DIR"/csrc/ix_full_bridge*.cpp "$EX_ENGINE_DIR"/csrc/ix_moe_bridge.cpp; do + [ -f "$cpp" ] && cp "$cpp" "${VLLM_ROOT}/ex_engine/csrc/" && \ + echo "[patch_ops] deployed $(basename $cpp) for JIT fallback" + done + + # Create startup hook that patches vllm ops at import time + cat > "${VLLM_ROOT}/ix_startup_patch.py" << 'STARTUP_EOF' +"""Apply ix_ops patches at vllm startup.""" +import logging +_logger = logging.getLogger("ix_startup_patch") +def apply(): + try: + from vllm.ex_engine.patch_vllm_ops import apply_all_patches + n = apply_all_patches() + if n > 0: + _logger.info("ix_startup_patch: %d patches applied", n) + return n + except Exception as e: + _logger.warning("ix_startup_patch failed: %s", e) + return 0 +_n_patches = apply() +STARTUP_EOF + echo "[patch_ops] deployed ix_startup_patch.py" + + # Hook into vllm __init__.py to auto-apply patches on import + VLLM_INIT="${VLLM_ROOT}/__init__.py" + if [ -f "$VLLM_INIT" ]; then + if ! grep -q "ix_startup_patch" "$VLLM_INIT" 2>/dev/null; then + echo "" >> "$VLLM_INIT" + echo "# Auto-apply ix_bridge operator patches" >> "$VLLM_INIT" + echo "try:" >> "$VLLM_INIT" + echo " from vllm import ix_startup_patch" >> "$VLLM_INIT" + echo "except Exception:" >> "$VLLM_INIT" + echo " pass" >> "$VLLM_INIT" + echo "[patch_ops] hooked ix_startup_patch into vllm/__init__.py" + fi + fi +else + echo "[patch_ops] WARN: ex_engine/python not found, skip ix_bridge deployment" +fi + # --- sequence.py: fix completion_tokens inflation under chunked prefill ------ # Bug: get_output_token_ids_to_return(delta=True) with num_new_tokens=0 # returns _cached_all_token_ids[-0:] == [0:] (the ENTIRE prompt+output list).