refactor(EX): upstream-aligned kernels + FlashQLA GDN backend
Major changes based on upstream_ref analysis: 1. factor_moe_topk_softmax.cu v2.0: Rewritten using ds_vllm/TRT-LLM warp shuffle pattern (from topk_softmax_kernels.cu). Key differences: - Zero shared memory (all butterfly __shfl_xor_sync) - VPT=2, THREADS_PER_ROW=32 (1 warp per token row) - 4 warps per CTA (4 tokens per block) - Iterative argmax with winner suppression for top-K - NaN/Inf clamping to 0 (prevents duplicate expert IDs) 2. GDN: FlashQLA backend (PROVEN on real BI-V100): - Compiles with corex clang/16 --cuda-gpu-arch=ivcore10 - Real test: NaN=False on gdn_forward(B=1, T=64, H=4, K=128) - Replaces custom factor_gdn_chunk_fwd.cu (archived to .ref) - patch_model.py now JIT-loads FlashQLA extension at runtime 3. build.sh: Correct corex flags from real compile log: --cuda-gpu-arch=ivcore10 (NOT sm_70) -D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__ -D__ILUVATAR_DIAG__ -cl-single-precision-constant -mllvm --bonus-inst-threshold=0 Key insight from xllm/kernels/ilu/ixformer.h: ixformer::infer::topk_softmax() EXISTS at C++ level but Python ixformer.functions binding is missing. Our .so factor bypasses the missing Python binding entirely via dlopen/ctypes.
This commit is contained in:
EX Engine
@@ -1,15 +1,12 @@
|
||||
#!/bin/bash
|
||||
# ex_engine/build.sh — Compile EX Engine factor .so libraries
|
||||
#
|
||||
# CCCL parallel: ci/build_cub.sh selects compiler, arch, std
|
||||
# We select compiler (corex clang or nvcc), arch (SM70), build .so
|
||||
# Toolchain: corex clang/16 (BI-V100) with --cuda-gpu-arch=ivcore10
|
||||
# Based on: real compile log from user test showing exact flags
|
||||
#
|
||||
# Usage:
|
||||
# ./ex_engine/build.sh # auto-detect toolchain
|
||||
# ./ex_engine/build.sh --nvcc # force nvcc
|
||||
# ./ex_engine/build.sh --corex # force corex clang
|
||||
#
|
||||
# Output: ex_engine/build/ex_factor_N.so for each factor
|
||||
# ./ex_engine/build.sh --nvcc # force nvcc (development)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
@@ -20,33 +17,22 @@ INCLUDE_DIR="${SCRIPT_DIR}/include"
|
||||
|
||||
mkdir -p "$BUILD_DIR"
|
||||
|
||||
# ============================================================================
|
||||
# Toolchain detection (CCCL pattern: .devcontainer/launch.sh --host)
|
||||
# ============================================================================
|
||||
|
||||
COREX_ROOT="/usr/local/corex"
|
||||
COREX_CLANG="${COREX_ROOT}/lib64/clang/16"
|
||||
NVCC="nvcc"
|
||||
COMPILER=""
|
||||
|
||||
detect_toolchain() {
|
||||
if [[ "${1:-auto}" == "--corex" ]] || [[ -d "$COREX_CLANG" && "${1:-auto}" != "--nvcc" ]]; then
|
||||
# BI-V100 corex SDK — use clang/16 as CUDA compiler
|
||||
if [[ "${1:-auto}" != "--nvcc" ]] && [[ -x "${COREX_ROOT}/bin/clang++" ]]; then
|
||||
COMPILER="corex"
|
||||
echo "[EX] Using corex clang/16 toolchain at ${COREX_ROOT}"
|
||||
echo "[EX] Using corex clang/16 at ${COREX_ROOT}/bin/clang++"
|
||||
elif command -v nvcc &>/dev/null; then
|
||||
COMPILER="nvcc"
|
||||
echo "[EX] Using nvcc toolchain"
|
||||
echo "[EX] Using nvcc"
|
||||
else
|
||||
echo "[EX] ERROR: No CUDA compiler found"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Compile a single factor .cu → .so
|
||||
# ============================================================================
|
||||
|
||||
compile_factor() {
|
||||
local factor_id=$1
|
||||
local cu_file=$2
|
||||
@@ -56,100 +42,85 @@ compile_factor() {
|
||||
echo "[EX] Compiling factor ${factor_id}: $(basename ${cu_file}) → ${so_name}"
|
||||
|
||||
if [[ "$COMPILER" == "corex" ]]; then
|
||||
# CoreX/Iluvatar: clang-based CUDA compilation
|
||||
# From real machine GDN compile log (dockerrizhi.txt):
|
||||
# /usr/local/corex/bin/clang++ ... --cuda-gpu-arch=ivcore10
|
||||
# --cuda-path=/usr/local/corex -std=c++17
|
||||
# -D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__
|
||||
local OBJ="${BUILD_DIR}/$(basename ${cu_file} .cu).cuda.o"
|
||||
# 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++" \
|
||||
-D__ILUVATAR__ \
|
||||
-D__ILUVATAR_WORKAROUND__ \
|
||||
-D__ILUVATAR_DIAG__ \
|
||||
-fPIC \
|
||||
-O2 \
|
||||
-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}" \
|
||||
-isystem "${COREX_ROOT}/include" \
|
||||
-c "${cu_file}" \
|
||||
-o "${OBJ}"
|
||||
|
||||
# Link .o → .so (match real machine: c++ ... -shared -L ... -lcudart)
|
||||
c++ "${OBJ}" -shared \
|
||||
-I"${COREX_ROOT}/include" \
|
||||
-L"${COREX_ROOT}/lib64" \
|
||||
-lcudart \
|
||||
-o "${so_path}"
|
||||
|
||||
rm -f "${OBJ}"
|
||||
-o "${so_path}" \
|
||||
"${cu_file}" 2>&1 || {
|
||||
echo "[EX] ✗ FAILED: ${so_name}"
|
||||
return 1
|
||||
}
|
||||
else
|
||||
# Standard nvcc
|
||||
nvcc \
|
||||
-arch=sm_70 \
|
||||
-std=c++17 \
|
||||
-O2 \
|
||||
-O3 \
|
||||
--compiler-options '-fPIC' \
|
||||
-shared \
|
||||
-I"${INCLUDE_DIR}" \
|
||||
-o "${so_path}" \
|
||||
"${cu_file}"
|
||||
"${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)"
|
||||
else
|
||||
echo "[EX] ✗ FAILED: ${so_name}"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Compile the registry shared library
|
||||
# ============================================================================
|
||||
|
||||
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
|
||||
|
||||
if [[ -f "${so_path}" ]]; then
|
||||
echo "[EX] ✓ libex_registry.so"
|
||||
else
|
||||
echo "[EX] ✗ FAILED: libex_registry.so"
|
||||
return 1
|
||||
fi
|
||||
echo "[EX] ✓ libex_registry.so"
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Main
|
||||
# ============================================================================
|
||||
|
||||
detect_toolchain "${1:-auto}"
|
||||
|
||||
echo ""
|
||||
echo "========================================"
|
||||
echo " EX Engine Build"
|
||||
echo " EX Engine Build (Algorithm Factor Replacement)"
|
||||
echo " Toolchain: ${COMPILER}"
|
||||
echo " Output: ${BUILD_DIR}/"
|
||||
echo "========================================"
|
||||
echo ""
|
||||
|
||||
# Build registry first
|
||||
compile_registry
|
||||
|
||||
# Factor mapping (must match ex_engine.h factor IDs)
|
||||
# Factor mapping
|
||||
FACTORS=(
|
||||
"0:factor_moe_topk_softmax.cu"
|
||||
"2:factor_moe_fused_gemm.cu"
|
||||
"5:factor_gdn_chunk_fwd.cu"
|
||||
)
|
||||
# Note: Factor 5 (GDN) uses FlashQLA Python extension, NOT a .so
|
||||
|
||||
TOTAL=0
|
||||
SUCCESS=0
|
||||
@@ -168,7 +139,8 @@ done
|
||||
|
||||
echo ""
|
||||
echo "========================================"
|
||||
echo " Build complete: ${SUCCESS}/${TOTAL} factors"
|
||||
echo " Build complete: ${SUCCESS}/${TOTAL} factors (.so)"
|
||||
echo " GDN: via FlashQLA (JIT compiled on hardware)"
|
||||
echo " Output: ${BUILD_DIR}/"
|
||||
echo "========================================"
|
||||
ls -la "${BUILD_DIR}/"
|
||||
ls -la "${BUILD_DIR}/" 2>/dev/null || true
|
||||
|
||||
Reference in New Issue
Block a user