feat: ex_factor_0.so ctypes桥接 + ex_engine package部署

1. ex_topk_bridge.py (100行):
   ctypes.CDLL加载ex_factor_0.so → ex_dispatch_moe_topk_softmax()
   CCCL warp-shuffle kernel, 零SMEM, 64 experts × topk=8

2. _custom_ops.py topk_softmax调用链新增Priority 1:
   P0: ix_bridge → ixformer::infer
   P1: ex_factor_0.so → CCCL warp kernel  ← NEW
   P2: _moe_C.so → vllm v0.5.5 kernel
   P3: moe_topk_softmax_v3.so → 自编译kernel

3. patch_ops.sh补齐ex_engine package部署:
   ex_engine/python/*.py + build/*.so → site-packages/ex_engine/
This commit is contained in:
Claude
2026-08-11 09:57:25 +00:00
parent 3a5cc2a589
commit c152bd5a89
3 changed files with 131 additions and 1 deletions

View File

@@ -1114,7 +1114,18 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
except Exception as e:
logger.warning("topk_softmax ix_bridge failed (%s), trying CUDA kernel", e)
# Priority 1: CUDA kernel (_moe_C or moe_topk_softmax_v3)
# Priority 1: ex_factor_0.so → CCCL warp-shuffle topk kernel (compiled for BI-V100)
try:
from ex_engine.python.ex_topk_bridge import ex_topk_softmax as _ex_topk
gating = gating_output if isinstance(gating_output, torch.Tensor) else gating_output
_ex_topk(topk_weights, topk_ids, token_expert_indicies, gating.float())
return
except Exception as e:
if not getattr(topk_softmax, '_ex_warned', False):
logger.warning("ex_factor_0 topk failed (%s), trying _moe_C", e)
topk_softmax._ex_warned = True
# Priority 2: CUDA kernel (_moe_C or moe_topk_softmax_v3)
if _moe_topk_ext is not None:
try:
gating = gating_output if isinstance(gating_output, torch.Tensor) else gating_output

View File

@@ -337,6 +337,25 @@ if [[ -n "$VLLM2" ]]; then
echo "[ok] mirrored all patches to VLLM2"
fi
build_stage "deploying ex_engine package to Python path"
_SITE=""
for _s in /usr/local/corex/lib64/python3/dist-packages \
/usr/local/corex/lib/python3/dist-packages \
/usr/local/lib/python3.10/site-packages; do
[[ -d "$_s" ]] && _SITE="$_s" && break
done
if [[ -n "$_SITE" ]]; then
_EX_DST="$_SITE/ex_engine"
mkdir -p "$_EX_DST/python" "$_EX_DST/build"
touch "$_EX_DST/__init__.py" "$_EX_DST/python/__init__.py"
cp /workspace/ex_engine/python/*.py "$_EX_DST/python/" 2>/dev/null || true
if [[ -d /workspace/ex_engine/build ]]; then
cp /workspace/ex_engine/build/*.so "$_EX_DST/build/" 2>/dev/null || true
cp /workspace/ex_engine/build/*.so "$_EX_DST/" 2>/dev/null || true
fi
echo "[ok] ex_engine deployed to $_EX_DST ($(ls "$_EX_DST/build/"*.so 2>/dev/null | wc -l) .so files)"
fi
build_stage "compiling submission Python sources"
find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile 2>&1 || echo "[WARN] some .py files failed to compile (non-fatal)"
build_stage "patch script completed"