c8a982c4e8
Ported from ex_engine/csrc/ix_full_bridge_v2.cpp + ix_moe_bridge.cpp.
Source: upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h
Exposes 14 ixformer::infer functions as Python-callable torch extension:
Attention: paged_attention, flash_attn_prefill, reshape_and_cache
MoE: topk_softmax, moe_gen_idx, moe_expand_input, group_gemm,
moe_combine_result, fused_moe_forward
Activation: silu_and_mul
Norm: rms_norm, fused_add_rms_norm
Linear: linear
RoPE: rotary_embedding
Build: torch.utils.cpp_extension.load() in docker build (patch_ops.sh)
Links against libixformer.so from base image at runtime.
This replaces PyTorch MoE fallback (the #1 performance bottleneck).
Without bridge: MoE loops over experts in Python → ~3 TPS decode
With bridge: fused 7-step pipeline in C++ → ~16 TPS decode (sub168 level)
96 lines
3.1 KiB
Bash
Executable File
96 lines
3.1 KiB
Bash
Executable File
#!/bin/bash
|
|
# Build ix_full_bridge.so — bridges ixformer::infer C++ symbols to Python
|
|
# Compiled via torch.utils.cpp_extension at docker build time
|
|
# Runtime: dlopen links against libixformer.so in base image
|
|
|
|
set -eo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
|
VLLM_ROOT="${1:?usage: build_ix_bridge.sh VLLM_ROOT}"
|
|
BRIDGE_SRC="${SCRIPT_DIR}/ix_full_bridge.cpp"
|
|
MOE_SRC="${SCRIPT_DIR}/ix_moe_bridge.cpp"
|
|
|
|
if [ ! -f "$BRIDGE_SRC" ]; then
|
|
echo "[bridge] SKIP: $BRIDGE_SRC not found"
|
|
exit 0
|
|
fi
|
|
|
|
# Find ixformer .so directory
|
|
IX_LIB=""
|
|
for d in /usr/local/corex/lib/python3/dist-packages/ixformer \
|
|
/usr/local/corex/lib64 \
|
|
/usr/lib/python3/dist-packages/ixformer; do
|
|
if [ -d "$d" ]; then
|
|
IX_LIB="$d"
|
|
break
|
|
fi
|
|
done
|
|
|
|
# Find torch library path
|
|
TORCH_LIB=$(python3 -c "import torch; print(torch.__path__[0] + '/lib')" 2>/dev/null)
|
|
|
|
echo "[bridge] Building ix_full_bridge via torch.utils.cpp_extension..."
|
|
echo "[bridge] ixformer lib: ${IX_LIB:-not found}"
|
|
echo "[bridge] torch lib: ${TORCH_LIB:-not found}"
|
|
|
|
python3 << PYEOF
|
|
import torch
|
|
from torch.utils.cpp_extension import load
|
|
import os, shutil
|
|
|
|
# Extra link flags: find libixformer.so and link
|
|
extra_ldflags = []
|
|
ix_lib = "${IX_LIB}"
|
|
torch_lib = "${TORCH_LIB}"
|
|
|
|
# Search for libixformer.so
|
|
for d in [ix_lib, "/usr/local/corex/lib64", "/usr/local/corex/lib"]:
|
|
if d and os.path.exists(os.path.join(d, "libixformer.so")):
|
|
extra_ldflags.extend([f"-L{d}", "-lixformer"])
|
|
break
|
|
else:
|
|
# No explicit libixformer.so — symbols may be in already-loaded .so
|
|
# (ixformer is imported as Python module which loads the .so)
|
|
try:
|
|
import ixformer
|
|
print("[bridge] ixformer Python module available — symbols in process")
|
|
except:
|
|
print("[bridge] WARNING: no libixformer.so found, link may fail at runtime")
|
|
|
|
# Add torch lib to rpath
|
|
if torch_lib and os.path.isdir(torch_lib):
|
|
extra_ldflags.append(f"-Wl,-rpath,{torch_lib}")
|
|
|
|
try:
|
|
mod = load(
|
|
name="ix_full_bridge",
|
|
sources=["${BRIDGE_SRC}"],
|
|
extra_cflags=["-O2", "-std=c++17"],
|
|
extra_ldflags=extra_ldflags,
|
|
verbose=True,
|
|
)
|
|
# Copy compiled .so to VLLM_ROOT for import
|
|
import importlib
|
|
spec = importlib.util.find_spec("ix_full_bridge")
|
|
if spec and spec.origin:
|
|
dest = os.path.join("${VLLM_ROOT}", "ix_full_bridge.so")
|
|
shutil.copy2(spec.origin, dest)
|
|
print(f"[bridge] SUCCESS: {dest}")
|
|
else:
|
|
# Try finding it in torch extension build dir
|
|
build_dir = os.path.expanduser("~/.cache/torch_extensions")
|
|
for root, dirs, files in os.walk(build_dir):
|
|
for f in files:
|
|
if f.startswith("ix_full_bridge") and f.endswith(".so"):
|
|
src = os.path.join(root, f)
|
|
dest = os.path.join("${VLLM_ROOT}", "ix_full_bridge.so")
|
|
shutil.copy2(src, dest)
|
|
print(f"[bridge] SUCCESS: {src} -> {dest}")
|
|
break
|
|
except Exception as e:
|
|
print(f"[bridge] FAILED: {e}")
|
|
raise SystemExit(1)
|
|
PYEOF
|
|
|
|
echo "[bridge] Build complete"
|