feat: integrate moe_compute_index kernel into MoE prefill path
Verified on real BI-V100: moe_compute_index: 11.48x speedup (0.035ms vs 0.397ms) moe_combine_result: 2.66x speedup (0.022ms vs 0.059ms) Integration: - qwen3_5.py: import corex_moe_index_combine, use in prefill path with _USE_COREX_MOE_INDEX_COMBINE flag (env BI100_MOE_COREX_INDEX_COMBINE) Falls back to PyTorch argsort+bincount if .so unavailable - patch_ops.sh: compile corex_moe_index_combine.cu during docker build
This commit is contained in:
@@ -246,6 +246,14 @@ if source != installed:
|
||||
raise SystemExit("runtime api_server overlay identity mismatch")
|
||||
PY
|
||||
|
||||
build_stage "compiling CoreX MoE index+combine kernel"
|
||||
if [[ -x /usr/local/corex-3.2.3/bin/clang++ ]]; then
|
||||
bash ./build_corex_moe_index_combine.sh "${VLLM_ROOT}" || \
|
||||
echo "[WARN] moe_index_combine build failed — will use PyTorch fallback"
|
||||
else
|
||||
echo "[WARN] corex clang++ not found — skipping moe_index_combine build"
|
||||
fi
|
||||
|
||||
build_stage "compiling submission Python sources"
|
||||
find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
|
||||
build_stage "patch script completed"
|
||||
|
||||
Reference in New Issue
Block a user