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