dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat(EX): wire xllm CUB topk_softmax kernel into MoE routing

Browse Source 
Upstream: xllm/kernels/cuda/moe/moe_topk_softmax_kernels.cuh (Apache 2.0)
Adapted: CHECK→TORCH_CHECK, include path fix, cuda/functional guard, pybind11

Call chain now:
  qwen3_5.py:_pure_pytorch_experts()
    → _ex_moe_topk_softmax (fused CUB kernel, 1 launch)
    → fallback: torch.softmax + torch.topk (3 launches)

Files:
  ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh — xllm kernel (adapted)
  ex_engine/csrc/moe/device_utils.cuh — xllm device utils
  ex_engine/csrc/moe/moe_topk_softmax_ext.cu — pybind11 wrapper
  ex_engine/python/moe_topk.py — JIT loader (same pattern as flash_qla_sm70)
  qwen3_5.py — import + use in _pure_pytorch_experts()
  patch_ops.sh — deploy kernel sources for JIT
This commit is contained in:
EngineX
2026-08-10 03:02:54 +00:00
parent d00daa62f6
commit 7839982707
6 changed files with 1119 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
qwen3_6_scripts/patch_ops.sh
View File

@@ -183,19 +183,24 @@ EX_ENGINE_SRC="/workspace/ex_engine"
if [ -d "$EX_ENGINE_SRC/python" ]; then
# Deploy into vllm's model dir so qwen3_5.py can import it
EX_DST="$VLLM/model_executor/models/ex_engine"
mkdir -p "$EX_DST/python"
mkdir -p "$EX_DST/csrc"
mkdir -p "$EX_DST/python" "$EX_DST/csrc"
cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST/python/" 2>/dev/null || true
# ix_moe_bridge.cpp needs to be next to the python module for JIT compile
# ix_moe_bridge.cpp for JIT compile
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/csrc/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/python/" 2>/dev/null || true
# Also make ex_engine importable from Python path
touch "$EX_DST/__init__.py"
touch "$EX_DST/python/__init__.py"
# Copy built .so files if they exist
# Copy built .so files
if [ -d "$EX_ENGINE_SRC/build" ]; then
cp "$EX_ENGINE_SRC/build/"*.so "$EX_DST/" 2>/dev/null || true
fi
# Deploy MoE CUDA kernel sources for JIT compilation
if [ -d "$EX_ENGINE_SRC/csrc/moe" ]; then
mkdir -p "$EX_DST/csrc/moe"
cp "$EX_ENGINE_SRC/csrc/moe/"*.cu "$EX_DST/csrc/moe/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/moe/"*.cuh "$EX_DST/csrc/moe/" 2>/dev/null || true
echo "[patch_ops] MoE CUDA kernel sources deployed for JIT"
fi
echo "[patch_ops] EX Engine deployed to $EX_DST"
ls -la "$EX_DST/csrc/" 2>/dev/null || true
if [ -n "$VLLM2" ]; then
@@ -207,6 +212,19 @@ else
echo "[patch_ops] WARNING: EX Engine not found — MoE uses slow PyTorch fallback"
fi
# Also deploy ex_engine Python package to system path for direct import
EX_PY_DST="/usr/local/corex/lib/python3/dist-packages/ex_engine"
if [ -d "$EX_ENGINE_SRC/python" ]; then
mkdir -p "$EX_PY_DST"
cp "$EX_ENGINE_SRC/python/"*.py "$EX_PY_DST/" 2>/dev/null || true
if [ -d "$EX_ENGINE_SRC/csrc/moe" ]; then
mkdir -p "$EX_PY_DST/../ex_engine/csrc/moe"
cp "$EX_ENGINE_SRC/csrc/moe/"*.cu "$EX_PY_DST/../ex_engine/csrc/moe/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/moe/"*.cuh "$EX_PY_DST/../ex_engine/csrc/moe/" 2>/dev/null || true
fi
echo "[patch_ops] EX Engine Python package deployed to $EX_PY_DST"
fi
echo "[patch_ops] DONE — EX Engine + SM70 GDN kernel + serving layer deployed"
echo "[patch_ops] Deployed: qwen3_5.py, flash_qla_sm70, ex_engine factors, paged_attn.py, mamba_cache.py, sequence.py, scheduler.py, xformers patches, serving layer"
echo "[patch_ops] EX factors replace: vllm_moe_topk_softmax (2304 calls/token), gdn_chunk_fwd (NaN fix)"