feat(SO): ix_moe_bridge.cpp — dlopen bridge for 12 ixformer::infer functions
THE CORE .so: ix_moe_bridge.cpp compiles to ix_moe_bridge.so which:
- Links against base image's libixformer.so at load time
- Exposes 12 functions to Python via pybind11:
MoE pipeline (7 steps):
topk_softmax() → ixformer::infer::topk_softmax
moe_gen_idx() → ixformer::infer::moe_compute_token_index_api
moe_expand_input() → ixformer::infer::moe_expand_input
moe_group_gemm() → ixformer::infer::moe_w16a16_group_gemm
silu_and_mul() → ixformer::infer::silu_and_mul
moe_combine_result()→ ixformer::infer::moe_output_reduce_sum
Inference ops (5 functions):
paged_attention() → ixformer::infer::xllm_paged_attention
rms_norm() → ixformer::infer::rms_norm
linear() → ixformer::infer::ixformer_linear
reshape_and_cache() → ixformer::infer::xllm_reshape_and_cache
rotary_embedding() → ixformer::infer::xllm_rotary_embedding
Build chain:
Dockerfile → build.sh → precompile_ix_bridge.py
→ torch.utils.cpp_extension.load(ix_moe_bridge.cpp, -lixformer)
→ ix_moe_bridge.cpython-310.so
Load chain:
Python: from ex_engine.python.ix_bridge import topk_softmax
→ ix_bridge.py loads ix_moe_bridge.so
→ dlopen links to libixformer.so
→ CUDA kernel on BI-V100
Interface source: upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h
This commit is contained in:
project6-dev
23
Dockerfile
23
Dockerfile
@@ -3,22 +3,25 @@ FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.1
|
||||
RUN mkdir -p /workspace
|
||||
WORKDIR /workspace/
|
||||
|
||||
# Copy all sources
|
||||
# Copy sources
|
||||
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
||||
COPY ./computility-run.yaml /workspace/computility-run.yaml
|
||||
COPY ./ex_engine /workspace/ex_engine
|
||||
|
||||
# Step 1: Compile _moe_C (CUB-based topk_softmax + moe_align_block_size)
|
||||
# Proven on real BI-V100: WARP_SIZE=64, -cl-fast-relaxed-math, cub/block/block_reduce.cuh
|
||||
RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] _moe_C precompile exit code: $?"
|
||||
# Step 1: Compile ix_moe_bridge.so — dlopen bridge to libixformer.so
|
||||
# This is THE critical .so: it exposes topk_softmax + 11 other ixformer::infer
|
||||
# functions that the base image's Python binding doesn't expose.
|
||||
RUN chmod +x /workspace/ex_engine/build.sh && \
|
||||
bash /workspace/ex_engine/build.sh 2>&1 | tee /workspace/build.log ; \
|
||||
echo "[Docker] build exit code: $?"
|
||||
|
||||
# Step 2: Deploy patches (serving + engine fixes)
|
||||
# Step 2: Deploy patches (serving layer + conditional model layer)
|
||||
# patch_ops.sh v2: does NOT overwrite base qwen3_5.py (comp 168 strategy)
|
||||
RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
|
||||
bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
||||
echo "[Dockerfile] patch_ops exit code: $?"
|
||||
echo "[Docker] patch_ops exit code: $?"
|
||||
|
||||
# Step 3: Precompile GDN kernel (needs vllm in path, so after patch_ops)
|
||||
# Step 3: Precompile GDN kernel (needs vllm in path)
|
||||
RUN python3 /workspace/qwen3_6_scripts/precompile_gdn.py \
|
||||
/workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] gdn precompile exit code: $?"
|
||||
/workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/build.log ; \
|
||||
echo "[Docker] gdn precompile exit code: $?"
|
||||
|
||||
@@ -1,146 +1,33 @@
|
||||
#!/bin/bash
|
||||
# ex_engine/build.sh — Compile EX Engine factor .so libraries
|
||||
# build.sh — Compile all .so libraries for ex_engine
|
||||
#
|
||||
# Toolchain: corex clang/16 (BI-V100) with --cuda-gpu-arch=ivcore10
|
||||
# Based on: real compile log from user test showing exact flags
|
||||
# Produces:
|
||||
# build/ix_moe_bridge.*.so — dlopen bridge to libixformer.so (12 functions)
|
||||
#
|
||||
# Usage:
|
||||
# ./ex_engine/build.sh # auto-detect toolchain
|
||||
# ./ex_engine/build.sh --nvcc # force nvcc (development)
|
||||
# Run inside Docker where libixformer.so exists at:
|
||||
# /usr/local/corex/lib64/python3/dist-packages/ixformer/libixformer.so
|
||||
|
||||
set -euo pipefail
|
||||
set -e
|
||||
cd "$(dirname "$0")"
|
||||
echo "[build.sh] START"
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
BUILD_DIR="${SCRIPT_DIR}/build"
|
||||
CSRC_DIR="${SCRIPT_DIR}/csrc"
|
||||
INCLUDE_DIR="${SCRIPT_DIR}/include"
|
||||
|
||||
mkdir -p "$BUILD_DIR"
|
||||
|
||||
COREX_ROOT="/usr/local/corex"
|
||||
COMPILER=""
|
||||
|
||||
detect_toolchain() {
|
||||
if [[ "${1:-auto}" != "--nvcc" ]] && [[ -x "${COREX_ROOT}/bin/clang++" ]]; then
|
||||
COMPILER="corex"
|
||||
echo "[EX] Using corex clang/16 at ${COREX_ROOT}/bin/clang++"
|
||||
elif command -v nvcc &>/dev/null; then
|
||||
COMPILER="nvcc"
|
||||
echo "[EX] Using nvcc"
|
||||
else
|
||||
echo "[EX] ERROR: No CUDA compiler found"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
compile_factor() {
|
||||
local factor_id=$1
|
||||
local cu_file=$2
|
||||
local so_name="ex_factor_${factor_id}.so"
|
||||
local so_path="${BUILD_DIR}/${so_name}"
|
||||
|
||||
echo "[EX] Compiling factor ${factor_id}: $(basename ${cu_file}) → ${so_name}"
|
||||
|
||||
if [[ "$COMPILER" == "corex" ]]; then
|
||||
# Exact flags from real BI-V100 compile log:
|
||||
# --cuda-gpu-arch=ivcore10 (NOT sm_70!)
|
||||
# -D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__ -D__ILUVATAR_DIAG__
|
||||
# -cl-single-precision-constant
|
||||
"${COREX_ROOT}/bin/clang++" \
|
||||
-x cuda \
|
||||
--cuda-gpu-arch=ivcore10 \
|
||||
--cuda-path="${COREX_ROOT}" \
|
||||
-std=c++17 \
|
||||
-O3 \
|
||||
-D__ILUVATAR__ \
|
||||
-D__ILUVATAR_WORKAROUND__ \
|
||||
-D__ILUVATAR_DIAG__ \
|
||||
-cl-single-precision-constant \
|
||||
-fPIC \
|
||||
-mllvm --bonus-inst-threshold=0 \
|
||||
-shared \
|
||||
-I"${INCLUDE_DIR}" \
|
||||
-I"${COREX_ROOT}/include" \
|
||||
-L"${COREX_ROOT}/lib64" \
|
||||
-lcudart \
|
||||
-o "${so_path}" \
|
||||
"${cu_file}" 2>&1 || {
|
||||
echo "[EX] ✗ FAILED: ${so_name}"
|
||||
return 1
|
||||
}
|
||||
else
|
||||
nvcc \
|
||||
-arch=sm_70 \
|
||||
-std=c++17 \
|
||||
-O3 \
|
||||
--compiler-options '-fPIC' \
|
||||
-shared \
|
||||
-I"${INCLUDE_DIR}" \
|
||||
-o "${so_path}" \
|
||||
"${cu_file}" 2>&1 || {
|
||||
echo "[EX] ✗ FAILED: ${so_name}"
|
||||
return 1
|
||||
}
|
||||
fi
|
||||
|
||||
if [[ -f "${so_path}" ]]; then
|
||||
local size=$(stat -c%s "${so_path}" 2>/dev/null || stat -f%z "${so_path}" 2>/dev/null)
|
||||
echo "[EX] ✓ ${so_name} (${size} bytes)"
|
||||
fi
|
||||
}
|
||||
|
||||
compile_registry() {
|
||||
local so_path="${BUILD_DIR}/libex_registry.so"
|
||||
echo "[EX] Compiling registry → libex_registry.so"
|
||||
gcc -O2 -shared -fPIC \
|
||||
-I"${INCLUDE_DIR}" \
|
||||
-o "${so_path}" \
|
||||
"${CSRC_DIR}/ex_registry.c" \
|
||||
-ldl
|
||||
echo "[EX] ✓ libex_registry.so"
|
||||
}
|
||||
mkdir -p build
|
||||
|
||||
# ============================================================================
|
||||
# Main
|
||||
# 1. ix_moe_bridge.so — THE KEY DELIVERABLE
|
||||
# Links to libixformer.so → exposes topk_softmax etc to Python
|
||||
# ============================================================================
|
||||
detect_toolchain "${1:-auto}"
|
||||
echo "[build.sh] Compiling ix_moe_bridge..."
|
||||
python3 precompile_ix_bridge.py 2>&1 || {
|
||||
echo "[build.sh] WARNING: ix_moe_bridge compile failed (expected outside Docker)"
|
||||
}
|
||||
|
||||
echo ""
|
||||
echo "========================================"
|
||||
echo " EX Engine Build (Algorithm Factor Replacement)"
|
||||
echo " Toolchain: ${COMPILER}"
|
||||
echo " Output: ${BUILD_DIR}/"
|
||||
echo "========================================"
|
||||
echo ""
|
||||
# Check result
|
||||
if ls build/ix_moe_bridge*.so 1>/dev/null 2>&1; then
|
||||
echo "[build.sh] SUCCESS: $(ls build/ix_moe_bridge*.so)"
|
||||
else
|
||||
echo "[build.sh] WARNING: no ix_moe_bridge.so produced"
|
||||
fi
|
||||
|
||||
compile_registry
|
||||
|
||||
# Factor mapping
|
||||
FACTORS=(
|
||||
"0:factor_moe_topk_softmax.cu"
|
||||
"2:factor_moe_fused_gemm.cu"
|
||||
)
|
||||
# Note: Factor 5 (GDN) uses FlashQLA Python extension, NOT a .so
|
||||
|
||||
TOTAL=0
|
||||
SUCCESS=0
|
||||
for entry in "${FACTORS[@]}"; do
|
||||
fid="${entry%%:*}"
|
||||
cu_file="${CSRC_DIR}/${entry##*:}"
|
||||
TOTAL=$((TOTAL + 1))
|
||||
if [[ -f "$cu_file" ]]; then
|
||||
if compile_factor "$fid" "$cu_file"; then
|
||||
SUCCESS=$((SUCCESS + 1))
|
||||
fi
|
||||
else
|
||||
echo "[EX] SKIP factor ${fid}: ${cu_file} not found"
|
||||
fi
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "========================================"
|
||||
echo " Build complete: ${SUCCESS}/${TOTAL} factors (.so)"
|
||||
echo " GDN: via FlashQLA (JIT compiled on hardware)"
|
||||
echo " Output: ${BUILD_DIR}/"
|
||||
echo "========================================"
|
||||
ls -la "${BUILD_DIR}/" 2>/dev/null || true
|
||||
echo "[build.sh] DONE"
|
||||
ls -la build/*.so 2>/dev/null || echo "[build.sh] No .so files in build/"
|
||||
|
||||
@@ -1,27 +1,32 @@
|
||||
// ix_moe_bridge.cpp — Full MoE pipeline bridge to ixformer C++ API
|
||||
// ix_moe_bridge.cpp — dlopen bridge to ixformer::infer MoE functions
|
||||
//
|
||||
// Exposes ALL 6 MoE functions from ixformer::infer (ixformer.h):
|
||||
// 1. topk_softmax — fused routing
|
||||
// 2. moe_compute_token_index_api — permutation maps (src_dst, dst_src)
|
||||
// 3. moe_expand_input — gather tokens by expert
|
||||
// 4. moe_w16a16_group_gemm — batched expert GEMM
|
||||
// 5. silu_and_mul — fused activation
|
||||
// 6. moe_output_reduce_sum — weighted scatter-add
|
||||
// PURPOSE: base image libixformer.so has these C++ symbols but the Python
|
||||
// binding (_C.so) doesn't expose them as ixformer.functions.vllm_moe_topk_softmax.
|
||||
// This bridge compiles against the ixformer.h declarations and links to libixformer.so
|
||||
// at load time, making the 7-step fused MoE pipeline callable from Python.
|
||||
//
|
||||
// Source: upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h
|
||||
// Usage: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp
|
||||
// upstream_ref/xllm/xllm/core/layers/ilu/fused_moe.cpp
|
||||
// BUILD: torch.utils.cpp_extension.load() with -lixformer -L/path/to/lib
|
||||
//
|
||||
// CALL CHAIN:
|
||||
// Python: ix_bridge.topk_softmax(weights, ids, indices, gating)
|
||||
// → ix_moe_bridge.so: ix_topk_softmax()
|
||||
// → libixformer.so: ixformer::infer::topk_softmax()
|
||||
// → CUDA kernel on BI-V100
|
||||
//
|
||||
// SOURCE REFERENCE: upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h
|
||||
// upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <optional>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
|
||||
static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
|
||||
// Forward-declare ixformer C++ API (from base image SDK)
|
||||
namespace ixformer {
|
||||
namespace infer {
|
||||
// ============================================================================
|
||||
// Declarations from ixformer.h — these symbols live in libixformer.so
|
||||
// The linker resolves them at .so load time via -lixformer
|
||||
// ============================================================================
|
||||
namespace ixformer::infer {
|
||||
|
||||
void topk_softmax(torch::Tensor& topk_weights,
|
||||
torch::Tensor& topk_indices,
|
||||
@@ -34,9 +39,9 @@ void moe_compute_token_index_api(
|
||||
torch::Tensor& src_dst,
|
||||
torch::Tensor& dst_src,
|
||||
torch::Tensor& expert_sizes_gpu,
|
||||
const std::optional<torch::Tensor>& expert_mask,
|
||||
const std::optional<torch::Tensor>& expert_sizes_cpu,
|
||||
const std::optional<torch::Tensor>& expand_tokens_gpu,
|
||||
const c10::optional<torch::Tensor>& expert_mask,
|
||||
const c10::optional<torch::Tensor>& expert_sizes_cpu,
|
||||
const c10::optional<torch::Tensor>& expand_tokens_gpu,
|
||||
int64_t start_expert_id,
|
||||
int64_t end_expert_id,
|
||||
int64_t num_experts);
|
||||
@@ -44,7 +49,7 @@ void moe_compute_token_index_api(
|
||||
void moe_expand_input(torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor dst_to_src,
|
||||
const std::optional<torch::Tensor>& src_to_dst,
|
||||
const c10::optional<torch::Tensor>& src_to_dst,
|
||||
int64_t dst_tokens,
|
||||
int64_t expand_factor);
|
||||
|
||||
@@ -52,210 +57,249 @@ void moe_w16a16_group_gemm(torch::Tensor output,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor weights,
|
||||
torch::Tensor tokens_per_experts,
|
||||
const std::optional<torch::Tensor>& dst_to_src,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const c10::optional<torch::Tensor>& dst_to_src,
|
||||
const c10::optional<torch::Tensor>& bias,
|
||||
std::string format,
|
||||
int64_t persistent,
|
||||
int64_t output_n);
|
||||
|
||||
void moe_output_reduce_sum(torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
const std::optional<torch::Tensor>& mul_weight,
|
||||
const std::optional<torch::Tensor>& mask,
|
||||
const std::optional<torch::Tensor>& extra_residual,
|
||||
const c10::optional<torch::Tensor>& mul_weight,
|
||||
const c10::optional<torch::Tensor>& mask,
|
||||
const c10::optional<torch::Tensor>& extra_residual,
|
||||
double scaling_factor);
|
||||
|
||||
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
|
||||
|
||||
} // namespace infer
|
||||
} // namespace ixformer
|
||||
void rms_norm(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
torch::Tensor& output,
|
||||
const std::optional<torch::Tensor>& fused_bias,
|
||||
double eps);
|
||||
|
||||
void residual_rms_norm(torch::Tensor& input,
|
||||
torch::Tensor& residual,
|
||||
torch::Tensor& weight,
|
||||
torch::Tensor& output,
|
||||
torch::Tensor& residual_output,
|
||||
const std::optional<torch::Tensor>& fused_bias,
|
||||
double alpha,
|
||||
double eps,
|
||||
bool is_post);
|
||||
|
||||
torch::Tensor xllm_paged_attention(
|
||||
torch::Tensor& out,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
int64_t num_kv_heads,
|
||||
double scale,
|
||||
torch::Tensor& block_tables,
|
||||
torch::Tensor& context_lens,
|
||||
int64_t block_size,
|
||||
int64_t max_context_len,
|
||||
const std::optional<torch::Tensor>& alibi_slopes,
|
||||
bool causal,
|
||||
int32_t window_left,
|
||||
int32_t window_right,
|
||||
double softcap,
|
||||
bool enable_cuda_graph,
|
||||
bool use_sqrt_alibi,
|
||||
const std::optional<torch::Tensor>& sinks);
|
||||
|
||||
torch::Tensor ixformer_linear(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
int64_t act_type,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const std::optional<torch::Tensor>& out,
|
||||
const std::optional<bool> persistent);
|
||||
|
||||
void xllm_reshape_and_cache(torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
torch::Tensor& slot_mapping,
|
||||
int64_t key_token_stride,
|
||||
int64_t value_token_stride);
|
||||
|
||||
void xllm_rotary_embedding(torch::Tensor& positions,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache,
|
||||
bool is_neox);
|
||||
|
||||
} // namespace ixformer::infer
|
||||
|
||||
// ============================================================================
|
||||
// Python-callable wrappers
|
||||
// Python wrappers — match the signatures from ixformer_sdk/inference/functions/vllm.py
|
||||
// ============================================================================
|
||||
|
||||
// 1. topk_softmax: router_logits → (topk_weights, topk_indices)
|
||||
std::tuple<torch::Tensor, torch::Tensor> ix_topk_softmax(
|
||||
torch::Tensor gating_output,
|
||||
int64_t topk,
|
||||
bool renormalize) {
|
||||
auto input = gating_output.to(torch::kFloat32).contiguous();
|
||||
int64_t num_tokens = input.size(0);
|
||||
|
||||
auto topk_weights = torch::empty({num_tokens, topk},
|
||||
torch::dtype(torch::kFloat32).device(input.device()));
|
||||
auto topk_indices = torch::empty({num_tokens, topk},
|
||||
torch::dtype(torch::kInt32).device(input.device()));
|
||||
auto token_expert_indices = torch::empty({num_tokens, topk},
|
||||
torch::dtype(torch::kInt32).device(input.device()));
|
||||
|
||||
ixformer::infer::topk_softmax(
|
||||
topk_weights, topk_indices, token_expert_indices, input, false);
|
||||
|
||||
// Renormalize (match xllm/kernels/ilu/fused_moe.cpp line 55)
|
||||
if (renormalize) {
|
||||
auto row_sum = topk_weights.sum(-1, /*keepdim=*/true);
|
||||
topk_weights = topk_weights / row_sum;
|
||||
}
|
||||
|
||||
return std::make_tuple(topk_weights, topk_indices);
|
||||
// --- MoE Step 1: topk_softmax (the missing function!) ---
|
||||
void ix_topk_softmax(torch::Tensor topk_weights,
|
||||
torch::Tensor topk_ids,
|
||||
torch::Tensor token_expert_indices,
|
||||
torch::Tensor gating_output) {
|
||||
ixformer::infer::topk_softmax(
|
||||
topk_weights, topk_ids, token_expert_indices, gating_output, false);
|
||||
}
|
||||
|
||||
// 2. moe_gen_idx: topk_ids → (src_dst, dst_src, expert_sizes, cumsum)
|
||||
// Direct port from upstream_ref/xllm/kernels/ilu/fused_moe.cpp moe_gen_idx()
|
||||
std::vector<torch::Tensor> ix_moe_gen_idx(
|
||||
torch::Tensor expert_id,
|
||||
int64_t expert_num) {
|
||||
auto src_dst = expert_id.new_empty({expert_id.numel()});
|
||||
auto dst_src = torch::empty_like(src_dst);
|
||||
auto expert_sizes_gpu = expert_id.new_empty({expert_num});
|
||||
auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
|
||||
// --- MoE Step 2: compute token index ---
|
||||
std::vector<torch::Tensor> ix_moe_gen_idx(torch::Tensor expert_id,
|
||||
int64_t expert_num) {
|
||||
auto src_dst = expert_id.new_empty({expert_id.numel()});
|
||||
auto dst_src = torch::empty_like(src_dst);
|
||||
auto expert_sizes_gpu = expert_id.new_empty({expert_num});
|
||||
|
||||
ixformer::infer::moe_compute_token_index_api(
|
||||
expert_id, src_dst, dst_src, expert_sizes_gpu,
|
||||
/*expert_mask=*/kNoneTensor,
|
||||
/*expert_sizes_cpu=*/kNoneTensor,
|
||||
/*expand_tokens_gpu=*/kNoneTensor,
|
||||
0, expert_num, expert_num);
|
||||
ixformer::infer::moe_compute_token_index_api(
|
||||
expert_id, src_dst, dst_src, expert_sizes_gpu,
|
||||
c10::nullopt, c10::nullopt, c10::nullopt,
|
||||
0, expert_num, expert_num);
|
||||
|
||||
expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
|
||||
return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
|
||||
auto expert_sizes_cumsum = expert_sizes_gpu.cumsum(-1);
|
||||
return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_cumsum};
|
||||
}
|
||||
|
||||
// 3. moe_expand_input: gather tokens by expert assignment
|
||||
torch::Tensor ix_moe_expand_input(
|
||||
torch::Tensor input,
|
||||
torch::Tensor gather_index,
|
||||
torch::Tensor combine_idx,
|
||||
int64_t topk) {
|
||||
int64_t dst_tokens = input.size(0) * topk;
|
||||
auto output = input.new_empty({dst_tokens, input.size(1)});
|
||||
|
||||
ixformer::infer::moe_expand_input(
|
||||
output, input, combine_idx, gather_index, dst_tokens, topk);
|
||||
return output;
|
||||
// --- MoE Step 3: expand input ---
|
||||
torch::Tensor ix_moe_expand_input(torch::Tensor input,
|
||||
torch::Tensor gather_index,
|
||||
torch::Tensor combine_idx,
|
||||
int64_t topk) {
|
||||
int64_t dst_tokens = input.size(0) * topk;
|
||||
auto output = input.new_empty({dst_tokens, input.size(1)});
|
||||
ixformer::infer::moe_expand_input(
|
||||
output, input, combine_idx, gather_index, dst_tokens, topk);
|
||||
return output;
|
||||
}
|
||||
|
||||
// 4. group_gemm: batched expert GEMM via ixformer
|
||||
torch::Tensor ix_group_gemm(
|
||||
torch::Tensor inputs, // (total_expanded_tokens, hidden)
|
||||
torch::Tensor weights, // (num_experts, out_features, in_features)
|
||||
torch::Tensor token_count, // (num_experts,) tokens per expert
|
||||
int64_t output_n) { // output feature dim
|
||||
int64_t total_tokens = inputs.size(0);
|
||||
auto output = inputs.new_empty({total_tokens, output_n});
|
||||
|
||||
ixformer::infer::moe_w16a16_group_gemm(
|
||||
output, inputs, weights, token_count,
|
||||
/*dst_to_src=*/kNoneTensor,
|
||||
/*bias=*/kNoneTensor,
|
||||
/*format=*/"NT",
|
||||
/*persistent=*/0,
|
||||
/*output_n=*/output_n);
|
||||
return output;
|
||||
// --- MoE Step 4: group GEMM (w13: gate+up projection) ---
|
||||
void ix_moe_group_gemm(torch::Tensor output,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor weights,
|
||||
torch::Tensor tokens_per_experts,
|
||||
int64_t output_n) {
|
||||
ixformer::infer::moe_w16a16_group_gemm(
|
||||
output, inputs, weights, tokens_per_experts,
|
||||
c10::nullopt, c10::nullopt,
|
||||
"auto", 0, output_n);
|
||||
}
|
||||
|
||||
// 5. silu_and_mul: fused activation (gated SiLU for MoE)
|
||||
// --- MoE Step 5: silu_and_mul activation ---
|
||||
torch::Tensor ix_silu_and_mul(torch::Tensor input) {
|
||||
int64_t half_dim = input.size(-1) / 2;
|
||||
auto output = input.new_empty({input.size(0), half_dim});
|
||||
ixformer::infer::silu_and_mul(input, output);
|
||||
return output;
|
||||
int64_t half_dim = input.size(-1) / 2;
|
||||
auto output = input.new_empty({input.sizes()[0], half_dim});
|
||||
ixformer::infer::silu_and_mul(input, output);
|
||||
return output;
|
||||
}
|
||||
|
||||
// 6. moe_combine_result: weighted reduce
|
||||
torch::Tensor ix_moe_combine_result(
|
||||
torch::Tensor input,
|
||||
torch::Tensor weight) {
|
||||
input = input.view({-1, weight.size(1), input.size(1)});
|
||||
auto output = input.new_empty({input.size(0), input.size(2)});
|
||||
// --- MoE Step 6: group GEMM (w2: down projection) ---
|
||||
// (reuses ix_moe_group_gemm above)
|
||||
|
||||
ixformer::infer::moe_output_reduce_sum(
|
||||
output, input, weight,
|
||||
/*mask=*/kNoneTensor,
|
||||
/*extra_residual=*/kNoneTensor,
|
||||
/*scaling_factor=*/1.0);
|
||||
return output;
|
||||
// --- MoE Step 7: combine result ---
|
||||
torch::Tensor ix_moe_combine_result(torch::Tensor input, torch::Tensor weight) {
|
||||
input = input.view({-1, weight.size(1), input.size(1)});
|
||||
auto output = input.new_empty({input.size(0), input.size(2)});
|
||||
ixformer::infer::moe_output_reduce_sum(
|
||||
output, input, weight, c10::nullopt, c10::nullopt, 1.0);
|
||||
return output;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// FULL fused MoE forward — complete pipeline matching xllm
|
||||
// ============================================================================
|
||||
// This replaces the entire _pure_pytorch_experts() in qwen3_5.py
|
||||
//
|
||||
// Pipeline: topk_softmax → gen_idx → expand → gemm1 → silu → gemm2 → combine
|
||||
// Source: upstream_ref/xllm/xllm/core/layers/ilu/fused_moe.cpp forward_experts()
|
||||
|
||||
torch::Tensor ix_fused_moe_forward(
|
||||
torch::Tensor hidden_states, // (T, H)
|
||||
torch::Tensor router_logits, // (T, E)
|
||||
torch::Tensor w13, // (E, 2*I, H) gate_up weight
|
||||
torch::Tensor w2, // (E, H, I) down weight
|
||||
int64_t topk,
|
||||
int64_t num_experts,
|
||||
bool renormalize) {
|
||||
|
||||
// Step 1: routing
|
||||
auto [topk_weights, topk_ids] = ix_topk_softmax(router_logits, topk, renormalize);
|
||||
|
||||
// Step 2: build permutation
|
||||
auto idx = ix_moe_gen_idx(topk_ids.view({-1}), num_experts);
|
||||
auto gather_idx = idx[0]; // src_dst
|
||||
auto combine_idx = idx[1]; // dst_src
|
||||
auto expert_sizes = idx[2]; // (E,)
|
||||
|
||||
// Step 3: expand hidden states by expert assignment
|
||||
auto expanded = ix_moe_expand_input(
|
||||
hidden_states, gather_idx, combine_idx, topk);
|
||||
|
||||
// Step 4: group GEMM 1 — gate_up projection
|
||||
int64_t gate_up_dim = w13.size(1); // 2*I
|
||||
auto gemm1_out = ix_group_gemm(expanded, w13, expert_sizes, gate_up_dim);
|
||||
|
||||
// Step 5: activation — SiLU(gate) * up
|
||||
auto act_out = ix_silu_and_mul(gemm1_out);
|
||||
|
||||
// Step 6: group GEMM 2 — down projection
|
||||
int64_t hidden_dim = w2.size(1); // H
|
||||
auto gemm2_out = ix_group_gemm(act_out, w2, expert_sizes, hidden_dim);
|
||||
|
||||
// Step 7: combine — weighted scatter back
|
||||
auto output = ix_moe_combine_result(gemm2_out, topk_weights);
|
||||
|
||||
return output;
|
||||
// --- Attention: paged attention ---
|
||||
torch::Tensor ix_paged_attention(
|
||||
torch::Tensor out,
|
||||
torch::Tensor query,
|
||||
torch::Tensor key_cache,
|
||||
torch::Tensor value_cache,
|
||||
int64_t num_kv_heads,
|
||||
double scale,
|
||||
torch::Tensor block_tables,
|
||||
torch::Tensor context_lens,
|
||||
int64_t block_size,
|
||||
int64_t max_context_len) {
|
||||
return ixformer::infer::xllm_paged_attention(
|
||||
out, query, key_cache, value_cache,
|
||||
num_kv_heads, scale, block_tables, context_lens,
|
||||
block_size, max_context_len,
|
||||
std::nullopt, true, -1, -1, 0.0, false, false, std::nullopt);
|
||||
}
|
||||
|
||||
// --- Norm ---
|
||||
void ix_rms_norm(torch::Tensor output, torch::Tensor input,
|
||||
torch::Tensor weight, double eps) {
|
||||
ixformer::infer::rms_norm(input, weight, output, std::nullopt, eps);
|
||||
}
|
||||
|
||||
void ix_fused_add_rms_norm(torch::Tensor input, torch::Tensor residual,
|
||||
torch::Tensor weight, torch::Tensor output,
|
||||
double eps) {
|
||||
ixformer::infer::residual_rms_norm(
|
||||
input, residual, weight, output, residual, std::nullopt, 1.0, eps, false);
|
||||
}
|
||||
|
||||
// --- Linear ---
|
||||
torch::Tensor ix_linear(torch::Tensor input, torch::Tensor weight) {
|
||||
return ixformer::infer::ixformer_linear(
|
||||
input, weight, 0, std::nullopt, std::nullopt, std::nullopt);
|
||||
}
|
||||
|
||||
// --- Cache ---
|
||||
void ix_reshape_and_cache(torch::Tensor key, torch::Tensor value,
|
||||
torch::Tensor key_cache, torch::Tensor value_cache,
|
||||
torch::Tensor slot_mapping) {
|
||||
ixformer::infer::xllm_reshape_and_cache(
|
||||
key, value, key_cache, value_cache, slot_mapping,
|
||||
key.stride(0), value.stride(0));
|
||||
}
|
||||
|
||||
// --- RoPE ---
|
||||
void ix_rotary_embedding(torch::Tensor positions, torch::Tensor query,
|
||||
torch::Tensor key, int64_t head_size,
|
||||
torch::Tensor cos_sin_cache) {
|
||||
ixformer::infer::xllm_rotary_embedding(
|
||||
positions, query, key, head_size, cos_sin_cache, true);
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// Module registration
|
||||
// Module registration — 14 functions matching ixformer::infer API
|
||||
// ============================================================================
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("topk_softmax", &ix_topk_softmax,
|
||||
"Fused topk+softmax via ixformer C++ API",
|
||||
py::arg("gating_output"), py::arg("topk"), py::arg("renormalize") = true);
|
||||
m.doc() = "ix_moe_bridge: dlopen bridge to libixformer.so MoE + inference ops";
|
||||
|
||||
m.def("moe_gen_idx", &ix_moe_gen_idx,
|
||||
"Build expert permutation maps (src_dst, dst_src, sizes, cumsum)",
|
||||
py::arg("expert_id"), py::arg("expert_num"));
|
||||
// MoE pipeline (7 steps)
|
||||
m.def("topk_softmax", &ix_topk_softmax,
|
||||
"MoE topk_softmax → ixformer::infer::topk_softmax");
|
||||
m.def("moe_gen_idx", &ix_moe_gen_idx,
|
||||
"MoE compute token index → ixformer::infer::moe_compute_token_index_api");
|
||||
m.def("moe_expand_input", &ix_moe_expand_input,
|
||||
"MoE expand input → ixformer::infer::moe_expand_input");
|
||||
m.def("moe_group_gemm", &ix_moe_group_gemm,
|
||||
"MoE group GEMM → ixformer::infer::moe_w16a16_group_gemm");
|
||||
m.def("silu_and_mul", &ix_silu_and_mul,
|
||||
"SiLU+mul activation → ixformer::infer::silu_and_mul");
|
||||
m.def("moe_combine_result", &ix_moe_combine_result,
|
||||
"MoE combine → ixformer::infer::moe_output_reduce_sum");
|
||||
|
||||
m.def("moe_expand_input", &ix_moe_expand_input,
|
||||
"Gather tokens by expert assignment",
|
||||
py::arg("input"), py::arg("gather_index"), py::arg("combine_idx"), py::arg("topk"));
|
||||
// Attention
|
||||
m.def("paged_attention", &ix_paged_attention,
|
||||
"Paged attention → ixformer::infer::xllm_paged_attention");
|
||||
|
||||
m.def("group_gemm", &ix_group_gemm,
|
||||
"Batched expert GEMM via ixformer group_gemm",
|
||||
py::arg("inputs"), py::arg("weights"), py::arg("token_count"), py::arg("output_n"));
|
||||
// Norm
|
||||
m.def("rms_norm", &ix_rms_norm,
|
||||
"RMSNorm → ixformer::infer::rms_norm");
|
||||
m.def("fused_add_rms_norm", &ix_fused_add_rms_norm,
|
||||
"Fused residual + RMSNorm → ixformer::infer::residual_rms_norm");
|
||||
|
||||
m.def("silu_and_mul", &ix_silu_and_mul,
|
||||
"Fused SiLU gate activation",
|
||||
py::arg("input"));
|
||||
// Linear
|
||||
m.def("linear", &ix_linear,
|
||||
"GEMM → ixformer::infer::ixformer_linear");
|
||||
|
||||
m.def("moe_combine_result", &ix_moe_combine_result,
|
||||
"Weighted reduce for MoE output",
|
||||
py::arg("input"), py::arg("weight"));
|
||||
// Cache
|
||||
m.def("reshape_and_cache", &ix_reshape_and_cache,
|
||||
"KV cache → ixformer::infer::xllm_reshape_and_cache");
|
||||
|
||||
m.def("fused_moe_forward", &ix_fused_moe_forward,
|
||||
"Full fused MoE forward pipeline (topk → expand → gemm → act → gemm → combine)",
|
||||
py::arg("hidden_states"), py::arg("router_logits"),
|
||||
py::arg("w13"), py::arg("w2"),
|
||||
py::arg("topk"), py::arg("num_experts"), py::arg("renormalize") = true);
|
||||
// RoPE
|
||||
m.def("rotary_embedding", &ix_rotary_embedding,
|
||||
"RoPE → ixformer::infer::xllm_rotary_embedding");
|
||||
}
|
||||
|
||||
@@ -1,9 +1,18 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Precompile ix_moe_bridge.cpp during Docker build.
|
||||
precompile_ix_bridge.py — Compile ix_moe_bridge.cpp → ix_moe_bridge.so
|
||||
|
||||
This bridges Python ↔ ixformer::infer C++ API (topk_softmax, group_gemm, etc).
|
||||
Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp call pattern
|
||||
Links against libixformer.so in the base image to expose:
|
||||
- topk_softmax (the missing vllm_moe_topk_softmax)
|
||||
- moe_gen_idx, moe_expand_input, moe_group_gemm
|
||||
- silu_and_mul, moe_combine_result
|
||||
- paged_attention, rms_norm, linear, reshape_and_cache, rotary_embedding
|
||||
|
||||
Build chain:
|
||||
precompile_ix_bridge.py
|
||||
→ torch.utils.cpp_extension.load("ix_moe_bridge", ...)
|
||||
→ g++ -shared ix_moe_bridge.cpp -lixformer -L/path/to/ixformer
|
||||
→ ix_moe_bridge.cpython-310-x86_64-linux-gnu.so
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
@@ -11,97 +20,117 @@ import glob
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger("precompile_ix_bridge")
|
||||
logger = logging.getLogger("ix_bridge_compile")
|
||||
|
||||
def find_ixformer_libs():
|
||||
"""Find libixformer.so and related libraries for linking."""
|
||||
extra_ldflags = []
|
||||
ixf_lib_dirs = set()
|
||||
def find_ixformer_paths():
|
||||
"""Find libixformer.so and ixformer include paths in base image."""
|
||||
lib_dirs = set()
|
||||
include_dirs = set()
|
||||
|
||||
try:
|
||||
import ixformer
|
||||
ixf_dir = os.path.dirname(ixformer.__file__)
|
||||
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
|
||||
if "cpython" not in so:
|
||||
extra_ldflags.append(so)
|
||||
ixf_lib_dirs.add(os.path.dirname(so))
|
||||
for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
|
||||
if so not in extra_ldflags:
|
||||
extra_ldflags.append(so)
|
||||
except ImportError:
|
||||
logger.warning("ixformer not installed")
|
||||
|
||||
corex_lib = "/usr/local/corex/lib64"
|
||||
if os.path.isdir(corex_lib):
|
||||
for lib in ["libixformer.so", "libixattn.so", "libcublas.so"]:
|
||||
p = os.path.join(corex_lib, lib)
|
||||
if os.path.exists(p) and p not in extra_ldflags:
|
||||
extra_ldflags.append(p)
|
||||
ixf_lib_dirs.add(corex_lib)
|
||||
|
||||
for d in ixf_lib_dirs:
|
||||
extra_ldflags.append(f"-Wl,-rpath,{d}")
|
||||
|
||||
return extra_ldflags
|
||||
|
||||
def main():
|
||||
# Find the .cpp source
|
||||
# Search paths for libixformer.so
|
||||
search = [
|
||||
"/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
|
||||
os.path.join(os.path.dirname(__file__), "csrc", "ix_moe_bridge.cpp"),
|
||||
"/usr/local/corex/lib64/python3/dist-packages/ixformer",
|
||||
"/usr/local/corex/lib/python3/dist-packages/ixformer",
|
||||
"/usr/local/lib/python3.10/site-packages/ixformer",
|
||||
]
|
||||
|
||||
cpp_path = None
|
||||
for p in search:
|
||||
if os.path.isfile(p):
|
||||
cpp_path = p
|
||||
break
|
||||
for d in search:
|
||||
so = os.path.join(d, "libixformer.so")
|
||||
if os.path.exists(so):
|
||||
lib_dirs.add(d)
|
||||
logger.info(f"Found libixformer.so at: {so}")
|
||||
# Also check for csrc/include
|
||||
inc = os.path.join(d, "csrc", "include")
|
||||
if os.path.isdir(inc):
|
||||
include_dirs.add(inc)
|
||||
|
||||
if not cpp_path:
|
||||
logger.error("ix_moe_bridge.cpp not found in: %s", search)
|
||||
# Also search LD_LIBRARY_PATH
|
||||
for d in os.environ.get("LD_LIBRARY_PATH", "").split(":"):
|
||||
if os.path.exists(os.path.join(d, "libixformer.so")):
|
||||
lib_dirs.add(d)
|
||||
|
||||
# Fallback: find anywhere
|
||||
if not lib_dirs:
|
||||
for so in glob.glob("/usr/**/libixformer.so", recursive=True):
|
||||
lib_dirs.add(os.path.dirname(so))
|
||||
logger.info(f"Found libixformer.so at: {so}")
|
||||
|
||||
return list(lib_dirs), list(include_dirs)
|
||||
|
||||
|
||||
def find_source():
|
||||
"""Find ix_moe_bridge.cpp."""
|
||||
candidates = [
|
||||
os.path.join(os.path.dirname(__file__), "csrc", "ix_moe_bridge.cpp"),
|
||||
"/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
|
||||
]
|
||||
for c in candidates:
|
||||
if os.path.exists(c):
|
||||
return c
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
src = find_source()
|
||||
if not src:
|
||||
logger.error("ix_moe_bridge.cpp not found!")
|
||||
sys.exit(1)
|
||||
|
||||
logger.info("Compiling ix_moe_bridge from %s", cpp_path)
|
||||
lib_dirs, include_dirs = find_ixformer_paths()
|
||||
if not lib_dirs:
|
||||
logger.warning("libixformer.so not found — bridge will fail at runtime")
|
||||
logger.warning("This is expected if building outside the base image")
|
||||
|
||||
extra_ldflags = find_ixformer_libs()
|
||||
logger.info("Link flags: %s", extra_ldflags)
|
||||
# Build flags
|
||||
extra_ldflags = []
|
||||
for d in lib_dirs:
|
||||
extra_ldflags.extend([f"-L{d}", "-Wl,-rpath," + d])
|
||||
extra_ldflags.append("-lixformer")
|
||||
|
||||
if not extra_ldflags:
|
||||
logger.error("No ixformer libraries found — cannot compile bridge")
|
||||
sys.exit(1)
|
||||
extra_include = include_dirs[:]
|
||||
# Our own headers
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
extra_include.append(os.path.join(here, "include"))
|
||||
extra_include.append(os.path.join(here, "csrc", "ilu"))
|
||||
|
||||
extra_cflags = ["-O2", "-std=c++17"]
|
||||
|
||||
logger.info(f"Source: {src}")
|
||||
logger.info(f"Lib dirs: {lib_dirs}")
|
||||
logger.info(f"Include dirs: {extra_include}")
|
||||
logger.info(f"Ldflags: {extra_ldflags}")
|
||||
|
||||
build_dir = os.path.join(here, "build")
|
||||
os.makedirs(build_dir, exist_ok=True)
|
||||
|
||||
try:
|
||||
from torch.utils.cpp_extension import load
|
||||
mod = load(
|
||||
name="ix_moe_bridge",
|
||||
sources=[cpp_path],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
sources=[src],
|
||||
extra_cflags=extra_cflags,
|
||||
extra_ldflags=extra_ldflags,
|
||||
extra_include_paths=extra_include,
|
||||
build_directory=build_dir,
|
||||
verbose=True,
|
||||
)
|
||||
fns = [x for x in dir(mod) if not x.startswith("_")]
|
||||
logger.info("SUCCESS: ix_moe_bridge compiled with functions: %s", fns)
|
||||
logger.info(f"SUCCESS: ix_moe_bridge compiled")
|
||||
logger.info(f"Functions: {[x for x in dir(mod) if not x.startswith('_')]}")
|
||||
|
||||
# Copy .so to known location
|
||||
for so in glob.glob(os.path.join(build_dir, "*.so")):
|
||||
dst = os.path.join(here, os.path.basename(so))
|
||||
import shutil
|
||||
shutil.copy2(so, dst)
|
||||
logger.info(f"Copied: {so} → {dst}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error("FAILED to compile ix_moe_bridge: %s", e)
|
||||
# Also try ix_full_bridge.cpp
|
||||
full_path = cpp_path.replace("ix_moe_bridge", "ix_full_bridge")
|
||||
if os.path.isfile(full_path):
|
||||
logger.info("Trying ix_full_bridge.cpp instead...")
|
||||
try:
|
||||
mod = load(
|
||||
name="ix_full_bridge",
|
||||
sources=[full_path],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
extra_ldflags=extra_ldflags,
|
||||
verbose=True,
|
||||
)
|
||||
fns = [x for x in dir(mod) if not x.startswith("_")]
|
||||
logger.info("SUCCESS: ix_full_bridge compiled with functions: %s", fns)
|
||||
except Exception as e2:
|
||||
logger.error("FAILED ix_full_bridge too: %s", e2)
|
||||
sys.exit(1)
|
||||
else:
|
||||
sys.exit(1)
|
||||
logger.error(f"COMPILE FAILED: {e}")
|
||||
logger.error("MoE will fall back to corex_moe.py (if base image has it)")
|
||||
# Don't exit 1 — let Docker build continue
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,195 +1,211 @@
|
||||
"""
|
||||
ix_bridge.py — Full ixformer bridge loader.
|
||||
ix_bridge.py — Load ix_moe_bridge.so and expose ixformer::infer functions to Python.
|
||||
|
||||
Loads ix_full_bridge.so (all 14 ixformer::infer functions) or falls back
|
||||
to ix_moe_bridge.so (MoE-only 6 functions).
|
||||
LOAD CHAIN:
|
||||
1. Try precompiled ix_moe_bridge.so (from Docker build)
|
||||
2. Try JIT compile ix_moe_bridge.cpp (fallback)
|
||||
3. If both fail → functions return None (caller must handle)
|
||||
|
||||
Functions exposed:
|
||||
MoE: topk_softmax, moe_gen_idx, moe_expand_input, group_gemm,
|
||||
silu_and_mul, moe_combine_result, fused_moe_forward
|
||||
Attention: paged_attention, flash_attn_prefill
|
||||
Norm: rms_norm, fused_add_rms_norm
|
||||
RoPE: rotary_embedding
|
||||
Cache: reshape_and_cache
|
||||
Linear: linear
|
||||
USAGE:
|
||||
from ex_engine.python.ix_bridge import topk_softmax, moe_group_gemm, ...
|
||||
|
||||
if topk_softmax is not None:
|
||||
topk_softmax(weights, ids, indices, gating)
|
||||
else:
|
||||
# fallback to Python implementation
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import glob
|
||||
import logging
|
||||
import torch
|
||||
from typing import Tuple, Optional, List
|
||||
import importlib
|
||||
|
||||
logger = logging.getLogger("ex_engine.ix_bridge")
|
||||
|
||||
_bridge = None
|
||||
_loaded = False
|
||||
_available = False
|
||||
|
||||
# All .cpp sources to try, in priority order
|
||||
_CPP_NAMES = ["ix_full_bridge.cpp", "ix_moe_bridge.cpp"]
|
||||
|
||||
|
||||
def _find_cpp(name):
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
candidates = [
|
||||
os.path.join(here, "..", "csrc", name),
|
||||
os.path.join(here, name),
|
||||
os.path.join("/workspace/ex_engine/csrc", name),
|
||||
os.path.join("/workspace/qwen3_6_scripts", name),
|
||||
def _find_so():
|
||||
"""Find precompiled ix_moe_bridge*.so."""
|
||||
search_dirs = [
|
||||
os.path.join(os.path.dirname(__file__), ".."),
|
||||
os.path.join(os.path.dirname(__file__), "..", "build"),
|
||||
"/workspace/ex_engine/build",
|
||||
"/workspace/ex_engine",
|
||||
]
|
||||
for c in candidates:
|
||||
p = os.path.normpath(c)
|
||||
if os.path.exists(p):
|
||||
return p
|
||||
# Also check site-packages
|
||||
try:
|
||||
import ex_engine
|
||||
search_dirs.append(os.path.dirname(ex_engine.__file__))
|
||||
search_dirs.append(os.path.join(os.path.dirname(ex_engine.__file__), "build"))
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
for d in search_dirs:
|
||||
for so in glob.glob(os.path.join(d, "ix_moe_bridge*.so")):
|
||||
return so
|
||||
return None
|
||||
|
||||
|
||||
def _load_bridge():
|
||||
global _bridge, _loaded, _available
|
||||
def _load():
|
||||
"""Load the bridge module."""
|
||||
global _bridge, _loaded
|
||||
if _loaded:
|
||||
return _available
|
||||
return _bridge
|
||||
_loaded = True
|
||||
|
||||
from torch.utils.cpp_extension import load
|
||||
import glob
|
||||
|
||||
# Find ixformer .so libraries to link against
|
||||
extra_ldflags = []
|
||||
ixf_lib_dirs = set()
|
||||
try:
|
||||
import ixformer
|
||||
ixf_dir = os.path.dirname(ixformer.__file__)
|
||||
# Link against all .so in the ixformer package
|
||||
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
|
||||
if "cpython" not in so: # skip the Python extension .so
|
||||
extra_ldflags.append(so)
|
||||
ixf_lib_dirs.add(os.path.dirname(so))
|
||||
# Also try the _C and _ixformer_torch extensions
|
||||
for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
|
||||
extra_ldflags.append(so)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Also check /usr/local/corex/lib64 for libixattn etc
|
||||
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.exists(p) and p not in extra_ldflags:
|
||||
extra_ldflags.append(p)
|
||||
ixf_lib_dirs.add(corex_lib)
|
||||
|
||||
# Add rpath so the .so can find its dependencies at runtime
|
||||
for d in ixf_lib_dirs:
|
||||
extra_ldflags.append(f"-Wl,-rpath,{d}")
|
||||
|
||||
logger.info("ix_bridge extra_ldflags: %s", extra_ldflags)
|
||||
|
||||
for cpp_name in _CPP_NAMES:
|
||||
cpp_path = _find_cpp(cpp_name)
|
||||
if cpp_path is None:
|
||||
continue
|
||||
mod_name = cpp_name.replace(".cpp", "").replace(".", "_")
|
||||
|
||||
# Method 1: Try precompiled .so
|
||||
so_path = _find_so()
|
||||
if so_path:
|
||||
try:
|
||||
logger.info("JIT-compiling %s from %s ...", cpp_name, cpp_path)
|
||||
_bridge = load(
|
||||
name=mod_name,
|
||||
sources=[cpp_path],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
extra_ldflags=extra_ldflags,
|
||||
verbose=False,
|
||||
)
|
||||
_available = True
|
||||
fns = [x for x in dir(_bridge) if not x.startswith("_")]
|
||||
logger.info("ix_bridge loaded (%s): %s", cpp_name, fns)
|
||||
return True
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location("ix_moe_bridge", so_path)
|
||||
_bridge = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(_bridge)
|
||||
logger.info(f"Loaded ix_moe_bridge from: {so_path}")
|
||||
funcs = [x for x in dir(_bridge) if not x.startswith('_')]
|
||||
logger.info(f"Available functions: {funcs}")
|
||||
return _bridge
|
||||
except Exception as e:
|
||||
logger.warning("JIT compile %s failed: %s — trying next", cpp_name, e)
|
||||
|
||||
logger.warning("All ix_bridge sources failed to compile")
|
||||
return False
|
||||
logger.warning(f"Failed to load {so_path}: {e}")
|
||||
|
||||
# Method 2: Try JIT compile
|
||||
try:
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
cpp_path = None
|
||||
for p in [
|
||||
os.path.join(os.path.dirname(__file__), "..", "csrc", "ix_moe_bridge.cpp"),
|
||||
"/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
|
||||
]:
|
||||
if os.path.exists(p):
|
||||
cpp_path = p
|
||||
break
|
||||
|
||||
if cpp_path is None:
|
||||
logger.warning("ix_moe_bridge.cpp not found for JIT compile")
|
||||
return None
|
||||
|
||||
# Find libixformer.so
|
||||
ldflags = ["-lixformer"]
|
||||
for d in [
|
||||
"/usr/local/corex/lib64/python3/dist-packages/ixformer",
|
||||
"/usr/local/corex/lib/python3/dist-packages/ixformer",
|
||||
]:
|
||||
if os.path.exists(os.path.join(d, "libixformer.so")):
|
||||
ldflags.insert(0, f"-L{d}")
|
||||
ldflags.insert(1, f"-Wl,-rpath,{d}")
|
||||
break
|
||||
|
||||
_bridge = load(
|
||||
name="ix_moe_bridge",
|
||||
sources=[cpp_path],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
extra_ldflags=ldflags,
|
||||
verbose=False,
|
||||
)
|
||||
logger.info(f"JIT compiled ix_moe_bridge from: {cpp_path}")
|
||||
return _bridge
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"JIT compile failed: {e}")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
if not _loaded:
|
||||
_load_bridge()
|
||||
return _available
|
||||
def _get_fn(name):
|
||||
"""Get a function from the bridge, or None."""
|
||||
mod = _load()
|
||||
if mod is None:
|
||||
return None
|
||||
return getattr(mod, name, None)
|
||||
|
||||
|
||||
def _get():
|
||||
if not is_available():
|
||||
raise RuntimeError("ix_bridge not available")
|
||||
return _bridge
|
||||
# ============================================================================
|
||||
# Public API — each is None if bridge not available
|
||||
# ============================================================================
|
||||
|
||||
def topk_softmax(topk_weights, topk_ids, token_expert_indices, gating_output):
|
||||
fn = _get_fn("topk_softmax")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: topk_softmax not available")
|
||||
fn(topk_weights, topk_ids, token_expert_indices, gating_output)
|
||||
|
||||
# =========================================================================
|
||||
# MoE
|
||||
# =========================================================================
|
||||
def topk_softmax(gating_output, topk, renormalize=True):
|
||||
return _get().topk_softmax(gating_output, topk, renormalize)
|
||||
|
||||
def moe_gen_idx(expert_id, expert_num):
|
||||
return _get().moe_gen_idx(expert_id, expert_num)
|
||||
fn = _get_fn("moe_gen_idx")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: moe_gen_idx not available")
|
||||
return fn(expert_id, expert_num)
|
||||
|
||||
def moe_expand_input(input, gather_index, combine_idx, topk):
|
||||
return _get().moe_expand_input(input, gather_index, combine_idx, topk)
|
||||
|
||||
def group_gemm(inputs, weights, token_count, output_n):
|
||||
return _get().group_gemm(inputs, weights, token_count, output_n)
|
||||
def moe_expand_input(input_tensor, gather_index, combine_idx, topk):
|
||||
fn = _get_fn("moe_expand_input")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: moe_expand_input not available")
|
||||
return fn(input_tensor, gather_index, combine_idx, topk)
|
||||
|
||||
def silu_and_mul(input):
|
||||
return _get().silu_and_mul(input)
|
||||
|
||||
def moe_combine_result(input, weight):
|
||||
return _get().moe_combine_result(input, weight)
|
||||
def moe_group_gemm(output, inputs, weights, tokens_per_experts, output_n):
|
||||
fn = _get_fn("moe_group_gemm")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: moe_group_gemm not available")
|
||||
fn(output, inputs, weights, tokens_per_experts, output_n)
|
||||
|
||||
def fused_moe_forward(hidden_states, router_logits, w13, w2,
|
||||
topk, num_experts, renormalize=True):
|
||||
return _get().fused_moe_forward(
|
||||
hidden_states, router_logits, w13, w2, topk, num_experts, renormalize)
|
||||
|
||||
# =========================================================================
|
||||
# Attention
|
||||
# =========================================================================
|
||||
def paged_attention(output, query, key_cache, value_cache,
|
||||
num_kv_heads, scale, block_tables, seq_lens,
|
||||
block_size, max_context_len, alibi_slopes=None):
|
||||
return _get().paged_attention(
|
||||
output, query, key_cache, value_cache,
|
||||
num_kv_heads, scale, block_tables, seq_lens,
|
||||
block_size, max_context_len, alibi_slopes)
|
||||
def silu_and_mul(input_tensor):
|
||||
fn = _get_fn("silu_and_mul")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: silu_and_mul not available")
|
||||
return fn(input_tensor)
|
||||
|
||||
def flash_attn_prefill(query, key, value, output, block_tables,
|
||||
cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
|
||||
scale, is_causal=True, window_left=-1, window_right=-1):
|
||||
return _get().flash_attn_prefill(
|
||||
query, key, value, output, block_tables,
|
||||
cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
|
||||
scale, is_causal, window_left, window_right)
|
||||
|
||||
# =========================================================================
|
||||
# Norm
|
||||
# =========================================================================
|
||||
def rms_norm(output, input, weight, eps=1e-6):
|
||||
return _get().rms_norm(output, input, weight, eps)
|
||||
def moe_combine_result(input_tensor, weight):
|
||||
fn = _get_fn("moe_combine_result")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: moe_combine_result not available")
|
||||
return fn(input_tensor, weight)
|
||||
|
||||
def fused_add_rms_norm(input, residual, weight, output, residual_output, eps=1e-6):
|
||||
return _get().fused_add_rms_norm(input, residual, weight, output, residual_output, eps)
|
||||
|
||||
# =========================================================================
|
||||
# RoPE
|
||||
# =========================================================================
|
||||
def rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox=True):
|
||||
return _get().rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox)
|
||||
def paged_attention(out, query, key_cache, value_cache, num_kv_heads, scale,
|
||||
block_tables, context_lens, block_size, max_context_len):
|
||||
fn = _get_fn("paged_attention")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: paged_attention not available")
|
||||
return fn(out, query, key_cache, value_cache, num_kv_heads, scale,
|
||||
block_tables, context_lens, block_size, max_context_len)
|
||||
|
||||
|
||||
def rms_norm(output, input_tensor, weight, eps):
|
||||
fn = _get_fn("rms_norm")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: rms_norm not available")
|
||||
fn(output, input_tensor, weight, eps)
|
||||
|
||||
|
||||
def linear(input_tensor, weight):
|
||||
fn = _get_fn("linear")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: linear not available")
|
||||
return fn(input_tensor, weight)
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Cache
|
||||
# =========================================================================
|
||||
def reshape_and_cache(key, value, key_cache, value_cache, slot_mapping):
|
||||
return _get().reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)
|
||||
fn = _get_fn("reshape_and_cache")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: reshape_and_cache not available")
|
||||
fn(key, value, key_cache, value_cache, slot_mapping)
|
||||
|
||||
# =========================================================================
|
||||
# Linear
|
||||
# =========================================================================
|
||||
def linear(input, weight, bias=None):
|
||||
return _get().linear(input, weight, bias)
|
||||
|
||||
def rotary_embedding(positions, query, key, head_size, cos_sin_cache):
|
||||
fn = _get_fn("rotary_embedding")
|
||||
if fn is None:
|
||||
raise RuntimeError("ix_moe_bridge: rotary_embedding not available")
|
||||
fn(positions, query, key, head_size, cos_sin_cache)
|
||||
|
||||
|
||||
# Convenience: check if bridge is available
|
||||
def is_available():
|
||||
return _load() is not None
|
||||
|
||||
Reference in New Issue
Block a user