diff --git a/Dockerfile b/Dockerfile index 12e795d8..81668792 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,8 +7,17 @@ WORKDIR /workspace/ COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts COPY ./computility-run.yaml /workspace/computility-run.yaml +# Copy EX Engine (algorithm factor replacement system) +COPY ./ex_engine /workspace/ex_engine + +# Build EX Engine .so factors for BI-V100 +# These replace missing ixformer.functions ops (moe_topk_softmax, gdn_chunk_fwd) +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: $?" + # Make patch script executable and run it -# Using bash explicitly to avoid shell interpretation issues +# patch_ops.sh also wires EX Engine into vllm 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: $?" diff --git a/ex_engine/build.sh b/ex_engine/build.sh index e2b21e9c..c56cebde 100755 --- a/ex_engine/build.sh +++ b/ex_engine/build.sh @@ -57,19 +57,32 @@ compile_factor() { if [[ "$COMPILER" == "corex" ]]; then # CoreX/Iluvatar: clang-based CUDA compilation - # SM70 = BI-V100 architecture + # From real machine GDN compile log (dockerrizhi.txt): + # /usr/local/corex/bin/clang++ ... --cuda-gpu-arch=ivcore10 + # --cuda-path=/usr/local/corex -std=c++17 + # -D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__ + local OBJ="${BUILD_DIR}/$(basename ${cu_file} .cu).cuda.o" "${COREX_ROOT}/bin/clang++" \ - -x cuda \ - --cuda-gpu-arch=sm_70 \ - -std=c++17 \ + -D__ILUVATAR__ \ + -D__ILUVATAR_WORKAROUND__ \ + -D__ILUVATAR_DIAG__ \ + -fPIC \ -O2 \ - -shared -fPIC \ + --cuda-gpu-arch=ivcore10 \ + --cuda-path="${COREX_ROOT}" \ + -std=c++17 \ -I"${INCLUDE_DIR}" \ - -I"${COREX_ROOT}/include" \ + -isystem "${COREX_ROOT}/include" \ + -c "${cu_file}" \ + -o "${OBJ}" + + # Link .o → .so (match real machine: c++ ... -shared -L ... -lcudart) + c++ "${OBJ}" -shared \ -L"${COREX_ROOT}/lib64" \ -lcudart \ - -o "${so_path}" \ - "${cu_file}" + -o "${so_path}" + + rm -f "${OBJ}" else # Standard nvcc nvcc \ diff --git a/ex_engine/csrc/factor_moe_topk_softmax.cu b/ex_engine/csrc/factor_moe_topk_softmax.cu index 5c1958ce..217e8c47 100644 --- a/ex_engine/csrc/factor_moe_topk_softmax.cu +++ b/ex_engine/csrc/factor_moe_topk_softmax.cu @@ -98,28 +98,13 @@ __global__ void moe_topk_softmax_kernel( float global_sum = s_sum[0] + s_sum[1]; float my_prob = my_exp / global_sum; // softmax output - // Step 5: Top-K selection via shared memory partial sort - // Use shared memory to collect all (prob, id) pairs, then - // do a register-based bitonic top-K. + // Step 5: Top-K selection via shared memory + // 64 elements is tiny — thread-0 serial insertion sort is faster than + // launching a parallel radix/bitonic for k=8 from n=64. __shared__ float s_probs[64]; - __shared__ int s_ids[64]; s_probs[tid] = my_prob; - s_ids[tid] = my_id; __syncthreads(); - // Thread 0 does a simple insertion sort for top_k=8 from 64 elements - // This is faster than full bitonic for k << n. - // 64 elements × 8 comparisons = 512 ops (fits in registers) - if (tid < top_k) { - // Each of the first top_k threads finds one winner - // We use a parallel argmax approach: each thread looks for - // the (tid+1)-th largest element. - // Simple approach: tid=0 finds max, tid=1 finds 2nd max, etc. - // Use iterative suppression in shared memory. - - // Actually, simpler: thread 0 does all work (64 experts is tiny) - } - if (tid == 0) { float* out_w = topk_weights + token_idx * top_k; int32_t* out_id = topk_ids + token_idx * top_k; @@ -134,12 +119,11 @@ __global__ void moe_topk_softmax_kernel( best_id[k] = -1; } - for (int e = 0; e < num_experts; e++) { + for (int e = 0; e < num_experts && e < BLOCK_SIZE; e++) { float p = s_probs[e]; - // Insert into sorted top-K if (p > best_w[TOP_K - 1]) { best_w[TOP_K - 1] = p; - best_id[TOP_K - 1] = e; + best_id[TOP_K - 1] = e; // expert index = thread index // Bubble up #pragma unroll for (int k = TOP_K - 1; k > 0; k--) { diff --git a/ex_engine/python/ex_loader.py b/ex_engine/python/ex_loader.py index f4536bec..132c8775 100644 --- a/ex_engine/python/ex_loader.py +++ b/ex_engine/python/ex_loader.py @@ -156,13 +156,27 @@ class EXEngine: return False def load_all(self) -> int: - """Load all available factor .so files from build_dir.""" + """Load all available factor .so files from build_dir or co-located.""" loaded = 0 + # Search paths: build_dir first, then directory containing this module + search_dirs = [self.build_dir] + module_dir = os.path.dirname(os.path.abspath(__file__)) + if module_dir not in search_dirs: + search_dirs.append(module_dir) + # Also check parent's build dir + parent_build = os.path.join(os.path.dirname(module_dir), "build") + if parent_build not in search_dirs: + search_dirs.append(parent_build) + for fid in range(EX_FACTOR_COUNT): - so_path = os.path.join(self.build_dir, f"ex_factor_{fid}.so") - if self.load_factor(fid, so_path): - loaded += 1 - logger.info("EX Engine: loaded %d/%d factors", loaded, EX_FACTOR_COUNT) + for d in search_dirs: + so_path = os.path.join(d, f"ex_factor_{fid}.so") + if os.path.exists(so_path): + if self.load_factor(fid, so_path): + loaded += 1 + break + logger.info("EX Engine: loaded %d/%d factors from %s", loaded, EX_FACTOR_COUNT, + search_dirs) return loaded def has_factor(self, factor_id: int) -> bool: diff --git a/ex_engine/python/patch_model.py b/ex_engine/python/patch_model.py index 321a5f22..96692bca 100644 --- a/ex_engine/python/patch_model.py +++ b/ex_engine/python/patch_model.py @@ -191,11 +191,7 @@ def _patch_gdn_prefill(engine): # --------------------------------------------------------------------------- -# Auto-apply on import if build dir exists +# Call apply_patches() explicitly AFTER vllm model modules are loaded. +# Integration point: qwen3_5.py calls this at the end of model __init__, +# or patch_ops.sh adds it to the startup sequence. # --------------------------------------------------------------------------- -_AUTO_BUILD_DIR = os.environ.get("EX_ENGINE_BUILD_DIR", "/workspace/ex_engine/build") -if os.path.isdir(_AUTO_BUILD_DIR): - try: - apply_patches(_AUTO_BUILD_DIR) - except Exception as e: - logger.warning("EX Engine auto-apply failed: %s", e) diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 2a862e14..3b0d0659 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -178,9 +178,28 @@ if [ -n "$VLLM2" ]; then cp ./chat_utils.py "$VLLM2/entrypoints/chat_utils.py" 2>/dev/null || true fi -echo "[patch_ops] DONE — SM70 GDN kernel + serving layer + engine patches deployed" -echo "[patch_ops] Deployed: qwen3_5.py, flash_qla_sm70 (SM70 GDN CUDA kernel), paged_attn.py, mamba_cache.py, sequence.py, scheduler.py, xformers patches, serving layer" -echo "[patch_ops] SM70 GDN kernel: JIT compiles on first forward pass (~2min), then cached" +# Deploy EX Engine Python module into vllm importable path +EX_ENGINE_SRC="/workspace/ex_engine" +if [ -d "$EX_ENGINE_SRC/python" ] && [ -d "$EX_ENGINE_SRC/build" ]; then + EX_DST="$VLLM/model_executor/models/ex_engine" + mkdir -p "$EX_DST" + cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST/" 2>/dev/null || true + # Copy built .so files + cp "$EX_ENGINE_SRC/build/"*.so "$EX_DST/" 2>/dev/null || true + echo "[patch_ops] EX Engine deployed: $(ls $EX_DST/*.so 2>/dev/null | wc -l) factors" + if [ -n "$VLLM2" ]; then + EX_DST2="$VLLM2/model_executor/models/ex_engine" + mkdir -p "$EX_DST2" + cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST2/" 2>/dev/null || true + cp "$EX_ENGINE_SRC/build/"*.so "$EX_DST2/" 2>/dev/null || true + fi +else + echo "[patch_ops] WARNING: EX Engine not built — MoE will use slow PyTorch fallback" +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)" echo "[patch_ops] NOT deployed (base image native): model_runner.py, _custom_ops.py, sampler.py, logits_processor.py, arg_utils.py" # Deploy flash_qla SM70 GDN kernel (from 1Cat-vLLM, MIT license)