3 changes that close the MoE performance gap:
1. _custom_ops.py: Add ix_bridge as Priority 0 for topk_softmax
- Before: tries our .cu kernel (fails) → PyTorch fallback (1-3 TPS)
- After: tries ix_bridge → ixformer::infer::topk_softmax() → FAST
- Call chain: _custom_ops.topk_softmax() → ix_bridge.topk_softmax()
→ ix_moe_bridge.so → ixformer::infer::topk_softmax()
2. Dockerfile: Add ix_moe_bridge.cpp precompile step
- This was the missing link: code existed but was never compiled
- Uses torch.utils.cpp_extension.load() to link against libixformer.so
3. upstream_ref sync from GitHub (cloned, not rewritten):
- xLLM-AI/xllm: ILU kernels + CUDA MoE + GDN fp32 state mgmt
- Deep-Spark/vllm: latest MoE kernel sources
39 lines
1.8 KiB
Docker
39 lines
1.8 KiB
Docker
FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
|
|
|
|
RUN mkdir -p /workspace
|
|
WORKDIR /workspace/
|
|
|
|
# Copy all sources
|
|
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
|
COPY ./computility-run.yaml /workspace/computility-run.yaml
|
|
COPY ./ex_engine /workspace/ex_engine
|
|
|
|
# Step 1: Build EX Engine .so libraries
|
|
RUN chmod +x /workspace/ex_engine/build.sh && \
|
|
bash /workspace/ex_engine/build.sh --corex 2>&1 | tee /workspace/ex_build.log ; \
|
|
echo "[Dockerfile] ex_engine build exit code: $?"
|
|
|
|
# Step 2: Precompile MoE CUDA kernels
|
|
RUN python3 /workspace/ex_engine/precompile_moe_topk.py 2>&1 | tee -a /workspace/ex_build.log ; \
|
|
echo "[Dockerfile] moe_topk precompile exit code: $?"
|
|
|
|
# Step 3: Precompile vllm v0.5.5 MoE kernels
|
|
RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee -a /workspace/ex_build.log ; \
|
|
echo "[Dockerfile] moe_v055 precompile exit code: $?"
|
|
|
|
# Step 4: Deploy patches (serving + engine fixes)
|
|
RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
|
|
bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
|
echo "[Dockerfile] patch_ops exit code: $?"
|
|
|
|
# Step 5: Precompile ix_moe_bridge.cpp → links to ixformer::infer::topk_softmax()
|
|
# This is the C++ pybind bridge that makes ixformer SDK callable from Python.
|
|
# Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp call pattern
|
|
RUN python3 /workspace/ex_engine/precompile_ix_bridge.py 2>&1 | tee -a /workspace/ex_build.log ; \
|
|
echo "[Dockerfile] ix_bridge precompile exit code: $?"
|
|
|
|
# Step 6: Precompile GDN kernel (needs vllm in path, so after patch_ops)
|
|
RUN python3 /workspace/qwen3_6_scripts/precompile_gdn.py \
|
|
/workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/ex_build.log ; \
|
|
echo "[Dockerfile] gdn precompile exit code: $?"
|