feat: ix_moe_bridge + ix_attn_bridge — dlopen bridges for full ixformer::infer API
Bridge architecture (from xllm/core/kernels/ilu/ixformer.h): ix_moe_bridge.so (MoE 7-step fused pipeline): - topk_softmax → moe_compute_token_index_api → moe_expand_input - moe_w16a16_group_gemm (x2) → silu_and_mul → moe_output_reduce_sum - fused_moe_forward(): replaces entire Python expert loop - Fix: group_gemm format NT→TN (match xllm trans_b=true) ix_attn_bridge.so (attention + linear): - ixinfer_flash_attn_unpad_with_block_tables (fused prefill) - xllm_paged_attention (fused paged decode) - ixformer_linear (matmul + activation) - residual_rms_norm (fused residual + norm) Integration: - ix_fused_moe.py: Python loader (prebuilt .so → JIT → unavailable) - qwen3_5.py: Tier 0 dispatch in _pure_pytorch_experts() - patch_ops.sh: deploys ix_fused_moe.py + all prebuilt/*.so Source: jd-opensource/xllm (fresh clone, all ILU kernels verified SAME) Sync: upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h (+32 lines) Build on real machine: bash qwen3_6_scripts/build_ix_moe_bridge.sh bash qwen3_6_scripts/build_ix_attn_bridge.sh
This commit is contained in:
84
qwen3_6_scripts/build_ix_moe_bridge.sh
Normal file
84
qwen3_6_scripts/build_ix_moe_bridge.sh
Normal file
@@ -0,0 +1,84 @@
|
||||
#!/bin/bash
|
||||
# build_ix_moe_bridge.sh — Build ix_moe_bridge.so on real BI-V100
|
||||
#
|
||||
# This compiles ex_engine/csrc/ix_moe_bridge.cpp into a prebuilt .so
|
||||
# that can be deployed without JIT compilation in Docker.
|
||||
#
|
||||
# Run on real machine: bash qwen3_6_scripts/build_ix_moe_bridge.sh
|
||||
# Output: qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ix_moe_bridge.so
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
|
||||
CPP_SOURCE="${PROJECT_DIR}/ex_engine/csrc/ix_moe_bridge.cpp"
|
||||
PREBUILT_DIR="${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10"
|
||||
|
||||
if [ ! -f "$CPP_SOURCE" ]; then
|
||||
# Also try the local copy
|
||||
CPP_SOURCE="${SCRIPT_DIR}/ix_moe_bridge.cpp"
|
||||
fi
|
||||
|
||||
if [ ! -f "$CPP_SOURCE" ]; then
|
||||
echo "ERROR: ix_moe_bridge.cpp not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "=== Building ix_moe_bridge.so ==="
|
||||
echo "Source: $CPP_SOURCE"
|
||||
echo "Output: $PREBUILT_DIR/ix_moe_bridge.so"
|
||||
|
||||
python3 -c "
|
||||
import os, sys, glob
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
cpp_source = '$CPP_SOURCE'
|
||||
extra_ldflags = []
|
||||
|
||||
# Find ixformer .so to link against
|
||||
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
|
||||
|
||||
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}')
|
||||
|
||||
print(f'Linking: {extra_ldflags}')
|
||||
|
||||
mod = load(
|
||||
name='ix_moe_bridge',
|
||||
sources=[cpp_source],
|
||||
extra_cflags=['-O2', '-std=c++17'],
|
||||
extra_ldflags=extra_ldflags,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Find the compiled .so and copy to prebuilt
|
||||
import torch.utils.cpp_extension as ext
|
||||
build_dir = ext._get_build_directory('ix_moe_bridge', verbose=False)
|
||||
print(f'Build dir: {build_dir}')
|
||||
|
||||
import shutil
|
||||
for f in glob.glob(os.path.join(build_dir, '*.so')):
|
||||
dst = '$PREBUILT_DIR/ix_moe_bridge.so'
|
||||
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
||||
shutil.copy2(f, dst)
|
||||
sz = os.path.getsize(dst)
|
||||
print(f'✓ ix_moe_bridge.so ({sz} bytes) → {dst}')
|
||||
break
|
||||
|
||||
# Verify
|
||||
fns = [x for x in dir(mod) if not x.startswith('_')]
|
||||
print(f'Functions: {fns}')
|
||||
print('=== Build SUCCESS ===')
|
||||
"
|
||||
Reference in New Issue
Block a user