7aa5054574
System design: algorithm factor replacement, not a connector.
All ops go through ixformer::infer C++ namespace (no Python fallback).
New files:
build_ix_bridge.sh — compile ix_full_bridge_v2.cpp on BI-V100
build_xllm_ilu_kernels.sh — compile upstream xllm ILU wrappers
deploy_ilu_pipeline.sh — wire everything into patch_ops.sh
ix_ops_dispatch.py — runtime dispatcher (12 ops via C++ bridge)
corex_fa2_dispatch.py — 3-mode attention (prefill/v1/flash paged)
fused_moe_ilu.py — 7-step MoE pipeline (no expert for-loop)
Upstream sources used (not rewritten):
xllm/core/kernels/ilu/*.cpp (ILU kernel wrappers)
xllm/core/kernels/ilu/ixformer.h (14 C++ function declarations)
ds_vllm/csrc/libtorch_stable/*.cu (kernel references)
Call chain:
patch_ops.sh → deploy_ilu_pipeline.sh → build_ix_bridge.sh
→ ix_full_bridge_v2.so → ixformer::infer::*
→ silu_and_mul, rms_norm, rotary_embedding, paged_attention,
topk_softmax, group_gemm, expand_input, combine_result
128 lines
3.9 KiB
Bash
Executable File
128 lines
3.9 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# build_xllm_ilu_kernels.sh — Compile xllm upstream ILU kernel wrappers
|
|
#
|
|
# Source: upstream_ref/xllm/xllm/core/kernels/ilu/*.cpp
|
|
# Already: ex_engine/xllm_kernels/ilu/ (copied from upstream)
|
|
# Header: upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h
|
|
#
|
|
# These .cpp files are thin wrappers that call ixformer::infer C++ functions.
|
|
# They're already proven to work on BI-V100 (xllm uses them in production).
|
|
# We compile them into xllm_ilu_ops.so with pybind11 bindings.
|
|
#
|
|
# Usage:
|
|
# bash build_xllm_ilu_kernels.sh [VLLM_ROOT]
|
|
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
|
|
|
# Source locations — prefer ex_engine copy, fall back to upstream_ref
|
|
ILU_DIR="${SCRIPT_DIR}/xllm_kernels/ilu"
|
|
if [[ ! -d "$ILU_DIR" ]]; then
|
|
ILU_DIR="${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/ilu"
|
|
fi
|
|
|
|
if [[ ! -d "$ILU_DIR" ]]; then
|
|
echo "[xllm_ilu] ERROR: ILU kernel source not found" >&2
|
|
exit 1
|
|
fi
|
|
|
|
# Header with ixformer::infer declarations
|
|
IXFORMER_H="${ILU_DIR}/ixformer.h"
|
|
if [[ ! -f "$IXFORMER_H" ]]; then
|
|
# Copy from upstream
|
|
cp "${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h" \
|
|
"${ILU_DIR}/ixformer.h" 2>/dev/null || true
|
|
cp "${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/ilu/utils.h" \
|
|
"${ILU_DIR}/utils.h" 2>/dev/null || true
|
|
fi
|
|
|
|
echo "[xllm_ilu] Source dir: ${ILU_DIR}"
|
|
echo "[xllm_ilu] Files:"
|
|
ls -la "$ILU_DIR"/*.cpp "$ILU_DIR"/*.h 2>/dev/null || true
|
|
|
|
# --- Compile via torch.utils.cpp_extension ---
|
|
VLLM_ROOT="${1:-}"
|
|
|
|
python3 << PYEOF
|
|
import os
|
|
import sys
|
|
import glob
|
|
|
|
# Set up paths
|
|
ilu_dir = "${ILU_DIR}"
|
|
script_dir = "${SCRIPT_DIR}"
|
|
vllm_root = "${VLLM_ROOT}" if "${VLLM_ROOT}" else None
|
|
|
|
# Find all .cpp files in the ILU directory
|
|
cpp_files = sorted(glob.glob(os.path.join(ilu_dir, "*.cpp")))
|
|
if not cpp_files:
|
|
print("[xllm_ilu] ERROR: No .cpp files found in", ilu_dir)
|
|
sys.exit(1)
|
|
|
|
print(f"[xllm_ilu] Found {len(cpp_files)} source files:")
|
|
for f in cpp_files:
|
|
print(f" {os.path.basename(f)}")
|
|
|
|
# Find ixformer .so files for linking
|
|
corex_root = os.environ.get("COREX_ROOT", "/usr/local/corex")
|
|
ix_so_files = []
|
|
rpath_dirs = set()
|
|
for search_dir in [
|
|
os.path.join(corex_root, "lib", "python3", "dist-packages", "ixformer"),
|
|
os.path.join(corex_root, "lib64", "python3", "dist-packages", "ixformer"),
|
|
os.path.join(corex_root, "lib64"),
|
|
]:
|
|
if os.path.isdir(search_dir):
|
|
rpath_dirs.add(search_dir)
|
|
for so in glob.glob(os.path.join(search_dir, "*.so")):
|
|
ix_so_files.append(so)
|
|
for so in glob.glob(os.path.join(search_dir, "lib*.so")):
|
|
if so not in ix_so_files:
|
|
ix_so_files.append(so)
|
|
|
|
extra_ldflags = list(ix_so_files)
|
|
for d in rpath_dirs:
|
|
extra_ldflags.append(f"-Wl,-rpath,{d}")
|
|
|
|
print(f"[xllm_ilu] Linking against {len(ix_so_files)} ixformer .so files")
|
|
|
|
try:
|
|
from torch.utils.cpp_extension import load
|
|
mod = load(
|
|
name="xllm_ilu_ops",
|
|
sources=cpp_files,
|
|
extra_include_paths=[ilu_dir],
|
|
extra_cflags=["-O2", "-std=c++17"],
|
|
extra_ldflags=extra_ldflags,
|
|
verbose=True,
|
|
)
|
|
print("[xllm_ilu] ✓ Compilation successful")
|
|
|
|
# Save the .so
|
|
import torch
|
|
so_path = os.path.join(script_dir, "prebuilt", "xllm_ilu_ops.so")
|
|
os.makedirs(os.path.dirname(so_path), exist_ok=True)
|
|
|
|
# Find the compiled .so in the torch cache
|
|
import importlib
|
|
spec = importlib.util.find_spec("xllm_ilu_ops")
|
|
if spec and spec.origin:
|
|
import shutil
|
|
shutil.copy2(spec.origin, so_path)
|
|
print(f"[xllm_ilu] ✓ Saved to {so_path}")
|
|
|
|
if vllm_root:
|
|
dst = os.path.join(vllm_root, "ex_engine", "xllm_ilu_ops.so")
|
|
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
|
shutil.copy2(spec.origin, dst)
|
|
print(f"[xllm_ilu] ✓ Deployed to {dst}")
|
|
|
|
except Exception as e:
|
|
print(f"[xllm_ilu] ERROR: {e}")
|
|
sys.exit(1)
|
|
PYEOF
|
|
|
|
echo "[xllm_ilu] Done"
|