fix(CRITICAL): merge 26e6cb40 build pipeline + HEAD features — fix docker build
Key changes: 1. Dockerfile: restore ex_engine COPY + build steps from26e6cb40(working), add vendor_overrides staging, add ix_unified_bridge build step 2. computility-run.yaml: restore Sub168 proven params (max-model-len=80000, gpu-util=0.95, max-num-seqs=2, enforce-eager, dtype=half) + corex env vars 3. patch_ops.sh: make vendor_overrides missing non-fatal (skip instead of exit 2) 4. New: corex_so_loader.py — unified loader for 12 prebuilt .so 5. New: moe_fused_dispatch.py — 3-tier MoE dispatch (CCCL policy_selector) Docker build was failing because: - HEAD removed ex_engine COPY and all build steps - patch_ops.sh exit 2 on missing vendor_overrides killed build - computility-run.yaml had max-model-len=262144 causing OOM26e6cb40scored on competition platform. This commit restores that build pipeline while adding the new HEAD features (prebuilt .so, vllm_overrides, corex dispatch env vars).
This commit is contained in:
Claude
84
Dockerfile
84
Dockerfile
@@ -1,22 +1,72 @@
|
||||
FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
|
||||
|
||||
ENV PATH=/usr/local/corex/bin:/usr/local/corex-3.2.3/bin:/usr/local/openmpi/bin:${PATH}
|
||||
ENV PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages:/usr/local/corex/lib/python3/dist-packages
|
||||
ENV LD_LIBRARY_PATH=/usr/local/corex/lib:/usr/local/corex/lib64:/usr/local/corex-3.2.3/lib:/usr/local/corex-3.2.3/lib64:/usr/local/openmpi/lib
|
||||
ENV VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 PYTHONUNBUFFERED=1 PYTHONFAULTHANDLER=1 BI100_EXECUTOR_STARTUP_DEBUG=1 ENABLE_CUSTOM_IPC=1
|
||||
ENV BI100_PREFIX_MODEL_FINGERPRINT=Qwen3.6-35B-A3B BI100_PREFIX_DTYPE=float16 BI100_PREFIX_TP_SIZE=4
|
||||
|
||||
RUN mkdir /workspace
|
||||
RUN mkdir -p /workspace
|
||||
WORKDIR /workspace/
|
||||
|
||||
# Copy all sources — ex_engine for build-time .so compilation,
|
||||
# qwen3_6_scripts for patches+prebuilt, vllm_overrides for core fixes
|
||||
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
||||
COPY ./vllm_overrides/core/evictor_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py
|
||||
COPY ./vllm_overrides/core/block/cpu_kv_content_cache.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py
|
||||
COPY ./vllm_overrides/core/block/cpu_gpu_block_allocator.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py
|
||||
COPY ./vllm_overrides/core/block/prefix_caching_block.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py
|
||||
COPY ./vllm_overrides/core/block/block_table.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py
|
||||
COPY ./vllm_overrides/core/block_manager_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py
|
||||
COPY ./vllm_overrides/sampling_params.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py
|
||||
COPY ./vllm_overrides/model_executor/sampling_metadata.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py
|
||||
COPY ./vllm_overrides/model_executor/layers/sampler.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
|
||||
RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
||||
COPY ./computility-run.yaml /workspace/computility-run.yaml
|
||||
COPY ./ex_engine /workspace/ex_engine
|
||||
COPY ./vllm_overrides /workspace/vllm_overrides
|
||||
|
||||
# Step 1: Build EX Engine .so libraries (tolerant of compile failures)
|
||||
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 (tolerant)
|
||||
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 (tolerant)
|
||||
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: Stage vendor_overrides into qwen3_6_scripts/ so patch_ops.sh finds them
|
||||
RUN mkdir -p /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block && \
|
||||
mkdir -p /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers && \
|
||||
cp /workspace/vllm_overrides/core/evictor_v2.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py && \
|
||||
cp /workspace/vllm_overrides/core/block/cpu_kv_content_cache.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py && \
|
||||
cp /workspace/vllm_overrides/core/block/cpu_gpu_block_allocator.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py && \
|
||||
cp /workspace/vllm_overrides/core/block/prefix_caching_block.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py && \
|
||||
cp /workspace/vllm_overrides/core/block/block_table.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py && \
|
||||
cp /workspace/vllm_overrides/core/block_manager_v2.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py && \
|
||||
cp /workspace/vllm_overrides/sampling_params.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py && \
|
||||
cp /workspace/vllm_overrides/model_executor/sampling_metadata.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py && \
|
||||
cp /workspace/vllm_overrides/model_executor/layers/sampler.py \
|
||||
/workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
|
||||
|
||||
# Step 5: Deploy patches (serving + engine fixes + prebuilt .so)
|
||||
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 ; \
|
||||
echo "[Dockerfile] patch_ops exit code: $?"
|
||||
|
||||
# Step 6: Build ix_unified_bridge.so (links against base image ixformer at runtime)
|
||||
RUN chmod +x /workspace/ex_engine/build_unified_bridge.sh && \
|
||||
bash /workspace/ex_engine/build_unified_bridge.sh 2>&1 | tee -a /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] ix_unified_bridge build exit code: $?"
|
||||
|
||||
# Step 7: Deploy ix_unified and ex_engine Python modules to vllm path
|
||||
RUN VLLM_ROOT=$(python3 -c "import vllm; print(vllm.__path__[0])" 2>/dev/null || echo "/usr/local/corex/lib/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 [ -f /workspace/ex_engine/build/ix_unified_bridge*.so ]; 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 8: 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: $?"
|
||||
|
||||
22
Dockerfile.broken_head
Normal file
22
Dockerfile.broken_head
Normal file
@@ -0,0 +1,22 @@
|
||||
FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
|
||||
|
||||
ENV PATH=/usr/local/corex/bin:/usr/local/corex-3.2.3/bin:/usr/local/openmpi/bin:${PATH}
|
||||
ENV PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages:/usr/local/corex/lib/python3/dist-packages
|
||||
ENV LD_LIBRARY_PATH=/usr/local/corex/lib:/usr/local/corex/lib64:/usr/local/corex-3.2.3/lib:/usr/local/corex-3.2.3/lib64:/usr/local/openmpi/lib
|
||||
ENV VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 PYTHONUNBUFFERED=1 PYTHONFAULTHANDLER=1 BI100_EXECUTOR_STARTUP_DEBUG=1 ENABLE_CUSTOM_IPC=1
|
||||
ENV BI100_PREFIX_MODEL_FINGERPRINT=Qwen3.6-35B-A3B BI100_PREFIX_DTYPE=float16 BI100_PREFIX_TP_SIZE=4
|
||||
|
||||
RUN mkdir /workspace
|
||||
WORKDIR /workspace/
|
||||
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
||||
COPY ./vllm_overrides/core/evictor_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py
|
||||
COPY ./vllm_overrides/core/block/cpu_kv_content_cache.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py
|
||||
COPY ./vllm_overrides/core/block/cpu_gpu_block_allocator.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py
|
||||
COPY ./vllm_overrides/core/block/prefix_caching_block.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py
|
||||
COPY ./vllm_overrides/core/block/block_table.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py
|
||||
COPY ./vllm_overrides/core/block_manager_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py
|
||||
COPY ./vllm_overrides/sampling_params.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py
|
||||
COPY ./vllm_overrides/model_executor/sampling_metadata.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py
|
||||
COPY ./vllm_overrides/model_executor/layers/sampler.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
|
||||
RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
||||
echo "[Dockerfile] patch_ops exit code: $?"
|
||||
@@ -8,37 +8,46 @@ command:
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '262144'
|
||||
- '80000'
|
||||
- --gpu-memory-utilization
|
||||
- '0.9'
|
||||
- '0.95'
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '4'
|
||||
- --max-num-seqs
|
||||
- '1'
|
||||
- '2'
|
||||
- --max-num-batched-tokens
|
||||
- '4096'
|
||||
- --enable-chunked-prefill
|
||||
- --disable-log-requests
|
||||
- --disable-frontend-multiprocessing
|
||||
- --max-num-batched-tokens
|
||||
- '8192'
|
||||
- --enable-chunked-prefill
|
||||
- --max-seq-len-to-capture
|
||||
- '32768'
|
||||
- --enforce-eager
|
||||
- --enable-auto-tool-choice
|
||||
- --tool-call-parser
|
||||
- qwen3_coder
|
||||
- --reasoning-parser
|
||||
- qwen3
|
||||
- --enable-prefix-caching
|
||||
- --max-seq-len-to-capture
|
||||
- '8192'
|
||||
- --dtype
|
||||
- half
|
||||
env:
|
||||
- name: VLLM_ENGINE_ITERATION_TIMEOUT_S
|
||||
value: 3600
|
||||
value: '3600'
|
||||
- name: VLLM_ATTENTION_BACKEND
|
||||
value: XFORMERS
|
||||
- name: ENABLE_CUSTOM_IPC
|
||||
value: '1'
|
||||
- name: PYTHONPATH
|
||||
value: /usr/local/corex/lib/python3/dist-packages:/usr/local/corex/lib64/python3/dist-packages
|
||||
- name: LD_LIBRARY_PATH
|
||||
value: /usr/local/corex/lib64:/usr/local/openmpi/lib:/usr/local/corex/lib64/python3/dist-packages/ixformer
|
||||
- name: PYTORCH_CUDA_ALLOC_CONF
|
||||
value: max_split_size_mb:512
|
||||
- name: OMP_NUM_THREADS
|
||||
value: '1'
|
||||
- name: BI100_MOE_COREX_DIRECT_ROUTED
|
||||
value: 1
|
||||
value: '1'
|
||||
- name: BI100_GDN_COREX_PACKED_DECODE
|
||||
value: 1
|
||||
- name: BI100_HYBRID_KV_ACCOUNTING
|
||||
value: full_attention
|
||||
- name: BI100_GDN_CACHE_POLICY
|
||||
value: admission64
|
||||
- name: BI100_GDN_RESTORE_MODE
|
||||
value: hybrid64
|
||||
value: '1'
|
||||
|
||||
178
ex_engine/python/corex_so_loader.py
Normal file
178
ex_engine/python/corex_so_loader.py
Normal file
@@ -0,0 +1,178 @@
|
||||
"""corex_so_loader.py — Unified loader for all 12 prebuilt CoreX .so modules.
|
||||
|
||||
CCCL pattern: device_reduce policy_selector — enumerate available kernels at
|
||||
init, expose a stable Python API, fall back gracefully when .so unavailable.
|
||||
|
||||
The 12 prebuilt .so files expose these operator families:
|
||||
|
||||
GDN decode pipeline (5 .so):
|
||||
corex_gdn_causal_conv → .causal_conv_update(conv_state, mixed_qkv, weight)
|
||||
corex_gdn_packed_decode → .packed_decode(temporal_state, packed_qkv, b, a, A_log, dt_bias)
|
||||
corex_gdn_beta_decay → .beta_decay(b, a, A_log, dt_bias)
|
||||
corex_gdn_qk_map → .qk_map(q, k, num_v_heads)
|
||||
corex_gdn_gated_norm → .apply_inverse(x, z)
|
||||
|
||||
Attention pipeline (3 .so):
|
||||
corex_attn_head_rms_norm → .prepare(x, eps) + .apply_inverse(x, z)
|
||||
corex_paged_kv_gather → .gather(key_cache, val_cache, block_tables, context_lens)
|
||||
corex_fused_paged_prefill → .forward(q, k_cache, v_cache, ...)
|
||||
|
||||
KV cache transfer (1 .so):
|
||||
corex_block_major_kv_transfer → .transfer(src, dst, mapping)
|
||||
|
||||
MoE pipeline (3 .so):
|
||||
corex_moe_direct_routed → .w13(hidden, w13, expert_ids)
|
||||
+ .w2_reduce(act, w2, expert_ids, weights)
|
||||
corex_moe_weight_gather → .gather(w13, w2, expert_ids)
|
||||
corex_moe_exact_reduce → .serial_float(expert_out, weights)
|
||||
|
||||
Usage:
|
||||
from ex_engine.python.corex_so_loader import corex
|
||||
if corex.gdn_causal_conv is not None:
|
||||
out = corex.gdn_causal_conv.causal_conv_update(...)
|
||||
|
||||
# Or import from vllm install root (patch_ops.sh deploys there):
|
||||
from corex_so_loader import corex
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger("corex_so_loader")
|
||||
|
||||
# All 12 .so modules in load order
|
||||
_SO_MANIFEST = [
|
||||
"corex_gdn_causal_conv",
|
||||
"corex_gdn_packed_decode",
|
||||
"corex_gdn_beta_decay",
|
||||
"corex_gdn_qk_map",
|
||||
"corex_gdn_gated_norm",
|
||||
"corex_attn_head_rms_norm",
|
||||
"corex_paged_kv_gather",
|
||||
"corex_fused_paged_prefill",
|
||||
"corex_block_major_kv_transfer",
|
||||
"corex_moe_direct_routed",
|
||||
"corex_moe_weight_gather",
|
||||
"corex_moe_exact_reduce",
|
||||
]
|
||||
|
||||
|
||||
def _find_so_dir() -> Optional[str]:
|
||||
"""Find the directory containing prebuilt CoreX .so files.
|
||||
|
||||
Search order:
|
||||
1. COREX_SO_DIR env var
|
||||
2. vllm install roots (where patch_ops.sh installs them)
|
||||
3. Bundled prebuilt directory (repo-relative)
|
||||
4. /usr/local/corex/lib64/
|
||||
"""
|
||||
candidates = []
|
||||
|
||||
env = os.getenv("COREX_SO_DIR")
|
||||
if env:
|
||||
candidates.append(env)
|
||||
|
||||
# vllm install roots (patch_ops.sh copies .so here)
|
||||
for p in sys.path:
|
||||
if "vllm" in p or "dist-packages" in p:
|
||||
candidates.append(p)
|
||||
# Also check parent/vllm/model_executor/models/
|
||||
candidates.append(os.path.join(p, "vllm", "model_executor", "models"))
|
||||
|
||||
# Repo-relative prebuilt bundle
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
candidates.append(os.path.join(here, "..", "..", "qwen3_6_scripts",
|
||||
"prebuilt", "corex-3.2.3-ivcore10"))
|
||||
candidates.append(os.path.join(here, "..", "..", "qwen3_6_scripts"))
|
||||
|
||||
# System CoreX
|
||||
candidates.append("/usr/local/corex/lib64/")
|
||||
|
||||
for d in candidates:
|
||||
d = os.path.normpath(d)
|
||||
if os.path.isdir(d):
|
||||
test_so = os.path.join(d, "corex_gdn_causal_conv.so")
|
||||
if os.path.isfile(test_so):
|
||||
return d
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _load_so(name: str, so_dir: str):
|
||||
"""Load a single .so by name from so_dir via importlib."""
|
||||
so_path = os.path.join(so_dir, f"{name}.so")
|
||||
if not os.path.isfile(so_path):
|
||||
return None
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(name, so_path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load %s: %s", so_path, e)
|
||||
return None
|
||||
|
||||
|
||||
class CoreXModules:
|
||||
"""Container for all loaded CoreX .so modules.
|
||||
|
||||
Each attribute is either the loaded module or None.
|
||||
Attribute names drop the 'corex_' prefix for brevity.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._loaded = {}
|
||||
self._so_dir = None
|
||||
|
||||
so_dir = _find_so_dir()
|
||||
if so_dir is None:
|
||||
logger.info("CoreX prebuilt .so directory not found — all modules disabled")
|
||||
for name in _SO_MANIFEST:
|
||||
short = name.replace("corex_", "", 1)
|
||||
setattr(self, short, None)
|
||||
self._loaded[name] = False
|
||||
return
|
||||
|
||||
self._so_dir = so_dir
|
||||
logger.info("CoreX .so directory: %s", so_dir)
|
||||
|
||||
loaded_count = 0
|
||||
for name in _SO_MANIFEST:
|
||||
mod = _load_so(name, so_dir)
|
||||
short = name.replace("corex_", "", 1)
|
||||
setattr(self, short, mod)
|
||||
self._loaded[name] = mod is not None
|
||||
if mod is not None:
|
||||
loaded_count += 1
|
||||
|
||||
logger.info("CoreX: %d/%d .so loaded from %s",
|
||||
loaded_count, len(_SO_MANIFEST), so_dir)
|
||||
|
||||
def summary(self) -> str:
|
||||
"""Return a human-readable summary of loaded modules."""
|
||||
lines = [f"CoreX .so loader ({self._so_dir or 'NOT FOUND'})"]
|
||||
for name in _SO_MANIFEST:
|
||||
status = "✓" if self._loaded.get(name) else "✗"
|
||||
short = name.replace("corex_", "", 1)
|
||||
mod = getattr(self, short, None)
|
||||
if mod is not None:
|
||||
funcs = [f for f in dir(mod) if not f.startswith("_")]
|
||||
lines.append(f" {status} {name} → .{', .'.join(funcs)}")
|
||||
else:
|
||||
lines.append(f" {status} {name}")
|
||||
return "\n".join(lines)
|
||||
|
||||
@property
|
||||
def all_loaded(self) -> bool:
|
||||
return all(self._loaded.values())
|
||||
|
||||
@property
|
||||
def loaded_count(self) -> int:
|
||||
return sum(1 for v in self._loaded.values() if v)
|
||||
|
||||
|
||||
# Singleton — initialized on first import
|
||||
corex = CoreXModules()
|
||||
236
ex_engine/python/moe_fused_dispatch.py
Normal file
236
ex_engine/python/moe_fused_dispatch.py
Normal file
@@ -0,0 +1,236 @@
|
||||
"""moe_fused_dispatch.py — Three-tier MoE dispatch (CCCL policy_selector pattern).
|
||||
|
||||
Port of upstream_ref/xllm/core/layers/ilu/fused_moe.cpp 7-step pipeline.
|
||||
|
||||
Dispatch hierarchy:
|
||||
Tier 0: ix_unified_bridge.so → ixformer::infer 7-step C++ pipeline
|
||||
topk_softmax → gen_idx → expand_input → group_gemm(w13) →
|
||||
silu_and_mul → group_gemm(w2) → combine_result
|
||||
Tier 1: corex prebuilt .so → direct_routed.w13/.w2_reduce (decode T=1 only)
|
||||
Tier 2: PyTorch fallback → per-expert F.linear loop
|
||||
|
||||
Usage in qwen3_5.py:
|
||||
from ex_engine.python.moe_fused_dispatch import fused_moe_forward
|
||||
out = fused_moe_forward(hidden_states, router_logits, w13, w2,
|
||||
top_k=8, num_experts=256, act_fn=silu_and_mul)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
logger = logging.getLogger("moe_fused_dispatch")
|
||||
|
||||
# Lazy imports — set at first call
|
||||
_ix = None
|
||||
_corex = None
|
||||
_init_done = False
|
||||
|
||||
|
||||
def _lazy_init():
|
||||
global _ix, _corex, _init_done
|
||||
if _init_done:
|
||||
return
|
||||
_init_done = True
|
||||
|
||||
# Tier 0: ix_unified
|
||||
try:
|
||||
from ex_engine.python.ix_unified import ix
|
||||
if ix._bridge is not None:
|
||||
_ix = ix
|
||||
logger.info("moe_fused_dispatch: Tier0 ix_unified_bridge.so available")
|
||||
else:
|
||||
logger.info("moe_fused_dispatch: Tier0 unavailable (bridge=None)")
|
||||
except Exception as e:
|
||||
logger.info("moe_fused_dispatch: Tier0 unavailable (%s)", e)
|
||||
|
||||
# Try import path used on real hardware
|
||||
if _ix is None:
|
||||
try:
|
||||
from ix_unified import ix
|
||||
if ix._bridge is not None:
|
||||
_ix = ix
|
||||
logger.info("moe_fused_dispatch: Tier0 ix_unified (direct) available")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Tier 1: corex prebuilt .so
|
||||
try:
|
||||
from ex_engine.python.corex_so_loader import corex
|
||||
if corex.moe_direct_routed is not None:
|
||||
_corex = corex
|
||||
logger.info("moe_fused_dispatch: Tier1 corex prebuilt .so available")
|
||||
except Exception as e:
|
||||
logger.info("moe_fused_dispatch: Tier1 unavailable (%s)", e)
|
||||
|
||||
|
||||
def _tier0_fused_moe(
|
||||
hidden_states: torch.Tensor, # [T, H]
|
||||
router_logits: torch.Tensor, # [T, E]
|
||||
w13: torch.Tensor, # [E, 2*I, H]
|
||||
w2: torch.Tensor, # [E, H, I]
|
||||
top_k: int,
|
||||
num_experts: int,
|
||||
act_fn: Callable,
|
||||
) -> torch.Tensor:
|
||||
"""Tier 0: Full 7-step ixformer::infer pipeline via ix_unified_bridge.so.
|
||||
|
||||
Maps 1:1 to xllm/core/layers/ilu/fused_moe.cpp::forward().
|
||||
"""
|
||||
T, H = hidden_states.shape
|
||||
|
||||
# Step 1: topk_softmax — fused softmax + topk selection
|
||||
topk_weights, topk_ids = _ix.moe_topk_softmax(router_logits, top_k,
|
||||
renormalize=True)
|
||||
|
||||
# Step 2: gen_idx — compute scatter/gather indices for expert routing
|
||||
idx_result = _ix.moe_gen_idx(topk_ids, num_experts)
|
||||
src_dst, dst_src, expert_sizes, cumsum = idx_result
|
||||
|
||||
# Step 3: expand_input — scatter tokens to expert order
|
||||
expanded = _ix.moe_expand_input(hidden_states, dst_src, src_dst, top_k)
|
||||
|
||||
# Step 4: group_gemm(w13) — batched GEMM across all experts
|
||||
gate_up = _ix.moe_group_gemm(expanded, w13, expert_sizes)
|
||||
|
||||
# Step 5: activation — SiLU(gate) * up
|
||||
act = act_fn(gate_up)
|
||||
|
||||
# Step 6: group_gemm(w2) — down projection
|
||||
down = _ix.moe_group_gemm(act, w2, expert_sizes)
|
||||
|
||||
# Step 7: combine_result — gather back and weighted sum
|
||||
output = _ix.moe_combine_result(
|
||||
down.view(T, top_k, H), topk_weights)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _tier1_decode_single_token(
|
||||
hidden_states: torch.Tensor, # [1, H]
|
||||
expert_ids: torch.Tensor, # [K]
|
||||
weights: torch.Tensor, # [K]
|
||||
w13: torch.Tensor, # [E, 2*I, H]
|
||||
w2: torch.Tensor, # [E, H, I]
|
||||
act_fn: Callable,
|
||||
) -> torch.Tensor:
|
||||
"""Tier 1: Single-token decode via prebuilt corex_moe_direct_routed.so.
|
||||
|
||||
Only works for T=1 decode. The .so implements fused expert indexing +
|
||||
GEMM + reduction in a single kernel launch.
|
||||
"""
|
||||
gate_up = _corex.moe_direct_routed.w13(hidden_states, w13, expert_ids)
|
||||
act = act_fn(gate_up)
|
||||
return _corex.moe_direct_routed.w2_reduce(act, w2, expert_ids, weights)
|
||||
|
||||
|
||||
def _tier2_pytorch_loop(
|
||||
hidden_states: torch.Tensor, # [T, H]
|
||||
router_logits: torch.Tensor, # [T, E]
|
||||
w13: torch.Tensor, # [E, 2*I, H]
|
||||
w2: torch.Tensor, # [E, H, I]
|
||||
top_k: int,
|
||||
act_fn: Callable,
|
||||
) -> torch.Tensor:
|
||||
"""Tier 2: Pure PyTorch per-expert loop (always works, slowest)."""
|
||||
T, H = hidden_states.shape
|
||||
|
||||
# Softmax → topk
|
||||
topk_logits, topk_ids = torch.topk(router_logits.float(), top_k, dim=-1)
|
||||
topk_weights = torch.softmax(topk_logits, dim=-1).to(hidden_states.dtype)
|
||||
|
||||
if T == 1:
|
||||
# Fast single-token path: batched GEMM
|
||||
eids = topk_ids[0]
|
||||
ws = topk_weights[0]
|
||||
w13_sel = w13[eids]
|
||||
w2_sel = w2[eids]
|
||||
gate_up = F.linear(hidden_states, w13_sel.reshape(-1, H))
|
||||
gate_up = gate_up.view(top_k, -1)
|
||||
act = act_fn(gate_up)
|
||||
expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1)
|
||||
return (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
|
||||
hidden_states.dtype)
|
||||
else:
|
||||
# General prefill path: sorted per-expert loop
|
||||
out = torch.zeros_like(hidden_states)
|
||||
flat_eids = topk_ids.reshape(-1)
|
||||
order = torch.argsort(flat_eids, stable=True)
|
||||
sorted_tok_ids = torch.arange(
|
||||
T, device=topk_ids.device).repeat_interleave(top_k)[order]
|
||||
sorted_weights = topk_weights.reshape(-1)[order]
|
||||
expert_counts = torch.bincount(
|
||||
flat_eids, minlength=w13.shape[0]).tolist()
|
||||
|
||||
start = 0
|
||||
for eid, count in enumerate(expert_counts):
|
||||
if count == 0:
|
||||
continue
|
||||
end = start + count
|
||||
tok_ids = sorted_tok_ids[start:end]
|
||||
tokens = hidden_states[tok_ids]
|
||||
gate_up = F.linear(tokens, w13[eid])
|
||||
act = act_fn(gate_up)
|
||||
expert_out = F.linear(act, w2[eid])
|
||||
weights_e = sorted_weights[start:end].unsqueeze(-1)
|
||||
out.index_add_(0, tok_ids, (expert_out * weights_e).to(out.dtype))
|
||||
start = end
|
||||
return out
|
||||
|
||||
|
||||
def fused_moe_forward(
|
||||
hidden_states: torch.Tensor, # [T, H]
|
||||
router_logits: torch.Tensor, # [T, E]
|
||||
w13: torch.Tensor, # [E, 2*I, H]
|
||||
w2: torch.Tensor, # [E, H, I]
|
||||
top_k: int = 8,
|
||||
num_experts: int = 256,
|
||||
act_fn: Optional[Callable] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Dispatch MoE through Tier 0 → 1 → 2.
|
||||
|
||||
Returns partial output (pre all-reduce), same contract as vllm FusedMoE.
|
||||
"""
|
||||
_lazy_init()
|
||||
|
||||
if act_fn is None:
|
||||
def _default_act(x):
|
||||
gate, up = x.chunk(2, dim=-1)
|
||||
return F.silu(gate) * up
|
||||
act_fn = _default_act
|
||||
|
||||
T = hidden_states.shape[0]
|
||||
|
||||
# Tier 0: full ixformer pipeline (all sizes)
|
||||
if _ix is not None and _ix._bridge is not None:
|
||||
try:
|
||||
return _tier0_fused_moe(hidden_states, router_logits, w13, w2,
|
||||
top_k, num_experts, act_fn)
|
||||
except Exception as e:
|
||||
logger.warning("Tier0 MoE failed (%s), falling to Tier1/2", e)
|
||||
|
||||
# Tier 1: corex direct routed (decode T=1 only)
|
||||
if (T == 1 and _corex is not None
|
||||
and _corex.moe_direct_routed is not None
|
||||
and hidden_states.dtype == torch.float16
|
||||
and w13.dtype == torch.float16
|
||||
and w2.dtype == torch.float16
|
||||
and hidden_states.is_contiguous()
|
||||
and w13.is_contiguous()
|
||||
and w2.is_contiguous()):
|
||||
try:
|
||||
topk_logits, topk_ids = torch.topk(
|
||||
router_logits.float(), top_k, dim=-1)
|
||||
topk_weights = torch.softmax(topk_logits, dim=-1).to(
|
||||
hidden_states.dtype)
|
||||
return _tier1_decode_single_token(
|
||||
hidden_states, topk_ids[0], topk_weights[0],
|
||||
w13, w2, act_fn)
|
||||
except Exception as e:
|
||||
logger.warning("Tier1 MoE failed (%s), falling to Tier2", e)
|
||||
|
||||
# Tier 2: PyTorch fallback
|
||||
return _tier2_pytorch_loop(hidden_states, router_logits, w13, w2,
|
||||
top_k, act_fn)
|
||||
@@ -110,9 +110,10 @@ echo "TRANSFORMERS_ROOT=${TRANSFORMERS_ROOT}"
|
||||
}
|
||||
|
||||
VLLM_OVERRIDE_ROOT="./vendor_overrides/vllm"
|
||||
_HAS_OVERRIDES=true
|
||||
[[ -d "$VLLM_OVERRIDE_ROOT" ]] || {
|
||||
printf 'vLLM override directory missing: %s\n' "$VLLM_OVERRIDE_ROOT" >&2
|
||||
exit 2
|
||||
printf '[WARN] vLLM override directory missing: %s — skipping override installs\n' "$VLLM_OVERRIDE_ROOT" >&2
|
||||
_HAS_OVERRIDES=false
|
||||
}
|
||||
|
||||
# --- Mirror path: base image may have TWO vllm installs ---
|
||||
@@ -142,6 +143,7 @@ deploy_both() {
|
||||
[[ -n "$VLLM2" ]] && cp "$src" "${VLLM2}/${rel}" 2>/dev/null || true
|
||||
}
|
||||
|
||||
if $_HAS_OVERRIDES; then
|
||||
build_stage "installing authoritative vLLM core block overrides"
|
||||
install_patch_file \
|
||||
"${VLLM_OVERRIDE_ROOT}/core/evictor_v2.py" \
|
||||
@@ -170,6 +172,9 @@ install_patch_file \
|
||||
install_patch_file \
|
||||
"${VLLM_OVERRIDE_ROOT}/model_executor/layers/sampler.py" \
|
||||
"${VLLM_ROOT}/model_executor/layers/sampler.py"
|
||||
else
|
||||
build_stage "skipping vLLM core block overrides (vendor_overrides not found)"
|
||||
fi
|
||||
|
||||
build_stage "installing hash-pinned CoreX 3.2.3 extensions"
|
||||
bash ./install_prebuilt_corex.sh "${VLLM_ROOT}"
|
||||
|
||||
Reference in New Issue
Block a user