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:
project6-dev
2026-08-13 03:52:35 +00:00
parent 71d39a1c7e
commit 796b09952c
2 changed files with 36 additions and 9 deletions

View File

@@ -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"