Compare commits
2 Commits
3a5cc2a589
...
4702505bf9
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4702505bf9 | ||
|
|
c152bd5a89 |
24
Dockerfile
24
Dockerfile
@@ -3,7 +3,7 @@ FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.1
|
||||
RUN mkdir -p /workspace
|
||||
WORKDIR /workspace/
|
||||
|
||||
# Copy all sources (vendor_overrides pre-staged inside qwen3_6_scripts/)
|
||||
# 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
|
||||
@@ -21,28 +21,12 @@ RUN python3 /workspace/ex_engine/precompile_moe_topk.py 2>&1 | tee -a /workspace
|
||||
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 + prebuilt .so)
|
||||
# Step 4: Deploy patches (serving + engine fixes + prebuilt .so + bridge)
|
||||
RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
|
||||
cd /workspace/qwen3_6_scripts && \
|
||||
bash ./patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
||||
bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
||||
echo "[Dockerfile] patch_ops exit code: $?"
|
||||
|
||||
# Step 5: Build ix_unified_bridge.so (ixformer symbols resolved at runtime)
|
||||
RUN chmod +x /workspace/ex_engine/build_unified_bridge.sh && \
|
||||
(bash /workspace/ex_engine/build_unified_bridge.sh 2>&1 || echo "[Dockerfile] bridge build FAILED (non-fatal)") | tee -a /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] ix_unified_bridge build exit code: $?"
|
||||
|
||||
# Step 6: Deploy ex_engine Python modules to vllm path
|
||||
RUN VLLM_ROOT=$(python3 -c "import vllm; print(vllm.__path__[0])" 2>/dev/null | tail -1 || echo "/usr/local/corex/lib64/python3/dist-packages/vllm") && \
|
||||
cp /workspace/ex_engine/python/ix_unified.py "${VLLM_ROOT}/ix_unified.py" 2>/dev/null || true && \
|
||||
cp /workspace/ex_engine/python/corex_so_loader.py "${VLLM_ROOT}/corex_so_loader.py" 2>/dev/null || true && \
|
||||
cp /workspace/ex_engine/python/moe_fused_dispatch.py "${VLLM_ROOT}/moe_fused_dispatch.py" 2>/dev/null || true && \
|
||||
if ls /workspace/ex_engine/build/ix_unified_bridge*.so 1>/dev/null 2>&1; then \
|
||||
cp /workspace/ex_engine/build/ix_unified_bridge*.so "${VLLM_ROOT}/" 2>/dev/null || true ; \
|
||||
fi ; \
|
||||
echo "[Dockerfile] ex_engine Python modules deployed"
|
||||
|
||||
# Step 7: Precompile GDN kernel (needs vllm in path, so after patch_ops)
|
||||
# Step 5: 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: $?"
|
||||
|
||||
100
ex_engine/python/ex_topk_bridge.py
Normal file
100
ex_engine/python/ex_topk_bridge.py
Normal file
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -337,6 +337,41 @@ 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 "building ix_unified_bridge (optional)"
|
||||
if [[ -x /workspace/ex_engine/build_unified_bridge.sh ]]; then
|
||||
bash /workspace/ex_engine/build_unified_bridge.sh 2>&1 || echo "[WARN] bridge build failed (non-fatal)"
|
||||
fi
|
||||
|
||||
build_stage "deploying ex_engine Python modules"
|
||||
VLLM_DEPLOY=$(python3 -c "import vllm; print(vllm.__path__[0])" 2>/dev/null | tail -1 || echo "")
|
||||
if [[ -n "$VLLM_DEPLOY" && -d "$VLLM_DEPLOY" ]]; then
|
||||
for f in ix_unified.py corex_so_loader.py moe_fused_dispatch.py ex_topk_bridge.py; do
|
||||
cp "/workspace/ex_engine/python/$f" "${VLLM_DEPLOY}/$f" 2>/dev/null || true
|
||||
done
|
||||
ls /workspace/ex_engine/build/ix_unified_bridge*.so 1>/dev/null 2>&1 && \
|
||||
cp /workspace/ex_engine/build/ix_unified_bridge*.so "${VLLM_DEPLOY}/" 2>/dev/null || true
|
||||
echo "[ok] ex_engine modules deployed to ${VLLM_DEPLOY}"
|
||||
fi
|
||||
|
||||
build_stage "patch script completed"
|
||||
|
||||
Reference in New Issue
Block a user