diff --git a/ex_engine/python/ex_topk_bridge.py b/ex_engine/python/ex_topk_bridge.py new file mode 100644 index 00000000..050b1f6c --- /dev/null +++ b/ex_engine/python/ex_topk_bridge.py @@ -0,0 +1,100 @@ +"""ex_topk_bridge.py — ctypes bridge for ex_factor_0.so topk_softmax + +CCCL pattern: ex_registry → ex_dispatch → kernel +Python bridge: ctypes.CDLL → ex_dispatch_moe_topk_softmax() + +Usage: + from ex_engine.python.ex_topk_bridge import ex_topk_softmax + ex_topk_softmax(topk_weights, topk_ids, token_expert_indices, gating_output) +""" +import ctypes +import os +import glob +import logging +import torch + +logger = logging.getLogger("ex_topk_bridge") + +_lib = None +_dispatch_fn = None + + +def _load(): + global _lib, _dispatch_fn + if _dispatch_fn is not None: + return True + + # Search for ex_factor_0.so + search = [ + os.path.join(os.path.dirname(__file__), "..", "build"), + "/workspace/ex_engine/build", + os.path.join(os.path.dirname(__file__), ".."), + ] + # Also check vllm model path (where build.sh factor compile puts it) + for p in ["/usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/models/ex_engine", + "/usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models/ex_engine"]: + search.append(p) + + for d in search: + so = os.path.join(d, "ex_factor_0.so") + if os.path.isfile(so): + try: + _lib_local = ctypes.CDLL(so) + fn = _lib_local.ex_dispatch_moe_topk_softmax + fn.restype = ctypes.c_int + fn.argtypes = [ + ctypes.c_void_p, # float* topk_weights + ctypes.c_void_p, # int32_t* topk_ids + ctypes.c_void_p, # const float* logits + ctypes.c_int, # T + ctypes.c_int, # E + ctypes.c_int, # top_k + ctypes.c_void_p, # stream + ] + _lib = _lib_local + _dispatch_fn = fn + logger.info("ex_factor_0.so loaded from %s", so) + return True + except Exception as e: + logger.warning("Failed to load %s: %s", so, e) + + return False + + +def ex_topk_softmax(topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor) -> None: + """Drop-in replacement for _custom_ops.topk_softmax using ex_factor_0.so. + + Same interface as vllm._custom_ops.topk_softmax: + topk_weights: (T, K) float32, output + topk_ids: (T, K) int32, output + token_expert_indices: (T, K) int32, output (ignored by ex kernel) + gating_output: (T, E) float32, input + """ + if not _load(): + raise RuntimeError("ex_factor_0.so not available") + + T, E = gating_output.shape + K = topk_weights.shape[1] + + # Get CUDA stream + stream = torch.cuda.current_stream().cuda_stream + + ret = _dispatch_fn( + topk_weights.data_ptr(), + topk_ids.data_ptr(), + gating_output.data_ptr(), + T, E, K, + stream, + ) + if ret != 0: + raise RuntimeError(f"ex_dispatch_moe_topk_softmax returned {ret}") + + # token_expert_indices: vllm expects (T, K) with values k_idx * T + t_idx + # ex kernel doesn't write this, fill it here + if token_expert_indices is not None: + T_t = torch.arange(T, device=topk_ids.device, dtype=torch.int32) + for k in range(K): + token_expert_indices[:, k] = k * T + T_t diff --git a/qwen3_6_scripts/_custom_ops.py b/qwen3_6_scripts/_custom_ops.py index 0d4b1c77..d9c97e93 100644 --- a/qwen3_6_scripts/_custom_ops.py +++ b/qwen3_6_scripts/_custom_ops.py @@ -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 diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 6c277622..de120073 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -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"