diff --git a/Dockerfile b/Dockerfile index faa0a98a..de2b4d2d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,6 +4,12 @@ WORKDIR /workspace/ # Copy all our engine patches COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts COPY ./computility-run.yaml /workspace/computility-run.yaml +# Copy ex_engine source for MoE bridge compilation +COPY ./ex_engine/csrc/moe_ops_impl.cu /workspace/qwen3_6_scripts/ex_engine_src/csrc/moe_ops_impl.cu +COPY ./ex_engine/csrc/ix_full_bridge_v2.cpp /workspace/qwen3_6_scripts/ex_engine_src/csrc/ix_full_bridge_v2.cpp +COPY ./ex_engine/build_moe_bridge.sh /workspace/qwen3_6_scripts/ex_engine_src/build_moe_bridge.sh +COPY ./ex_engine/python/moe_dispatch.py /workspace/qwen3_6_scripts/ex_engine_src/python/moe_dispatch.py +COPY ./ex_engine/python/patch_moe_hot_path.py /workspace/qwen3_6_scripts/ex_engine_src/python/patch_moe_hot_path.py # Make patch script executable and run it 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 ; \ diff --git a/computility-run.yaml b/computility-run.yaml index 2e09be09..5d4d5217 100644 --- a/computility-run.yaml +++ b/computility-run.yaml @@ -15,7 +15,7 @@ command: - -tp - '4' - --max-num-seqs - - '1' + - '2' - --disable-log-requests - --disable-frontend-multiprocessing - --max-num-batched-tokens diff --git a/ex_engine/build_moe_bridge.sh b/ex_engine/build_moe_bridge.sh new file mode 100755 index 00000000..6ad07399 --- /dev/null +++ b/ex_engine/build_moe_bridge.sh @@ -0,0 +1,156 @@ +#!/usr/bin/env bash +# build_moe_bridge.sh — Compile MoE ops + bridge into ix_moe_bridge.so +# +# Links against: +# libcuinfer.so (cuinferCustomGemm, cuinferTopK — confirmed in symbol dump) +# libixformer.so (silu_and_mul, rms_norm, flash_attn, etc — confirmed) +# +# Real device compiler: corex clang/16, NOT nvcc +# Reference: ex_engine/build_ix_bridge.sh + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +VLLM_ROOT="${1:-}" + +echo "[moe_bridge] Building ix_moe_bridge.so" +echo "[moe_bridge] Script dir: ${SCRIPT_DIR}" + +# --- Locate sources --- +MOE_CU="${SCRIPT_DIR}/csrc/moe_ops_impl.cu" +BRIDGE_CPP="${SCRIPT_DIR}/csrc/ix_full_bridge_v2.cpp" + +if [[ ! -f "$MOE_CU" ]]; then + echo "[moe_bridge] ERROR: $MOE_CU not found" >&2 + exit 1 +fi +if [[ ! -f "$BRIDGE_CPP" ]]; then + echo "[moe_bridge] ERROR: $BRIDGE_CPP not found" >&2 + exit 1 +fi + +# --- Locate libraries --- +COREX_ROOT="${COREX_ROOT:-/usr/local/corex}" + +# Find libcuinfer.so +CUINFER_SO="" +for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib" "/usr/lib64" "/usr/lib"; do + if [[ -f "${d}/libcuinfer.so" ]]; then + CUINFER_SO="${d}/libcuinfer.so" + break + fi +done + +# Find libixformer.so and ixformer Python package +IX_LIB_DIR="" +IX_SO_FILES=() +for d in \ + "${COREX_ROOT}/lib/python3/dist-packages/ixformer" \ + "${COREX_ROOT}/lib64/python3/dist-packages/ixformer" \ + "$(python3 -c 'import ixformer, os; print(os.path.dirname(ixformer.__file__))' 2>/dev/null || echo '')"; do + if [[ -d "$d" ]]; then + IX_LIB_DIR="$d" + while IFS= read -r so; do + IX_SO_FILES+=("$so") + done < <(find "$d" -name "*.so" -type f 2>/dev/null) + break + fi +done + +echo "[moe_bridge] COREX_ROOT: ${COREX_ROOT}" +echo "[moe_bridge] cuinfer: ${CUINFER_SO:-NOT FOUND}" +echo "[moe_bridge] ixformer dir: ${IX_LIB_DIR:-NOT FOUND}" +echo "[moe_bridge] ixformer .so count: ${#IX_SO_FILES[@]}" + +# --- Build via torch.utils.cpp_extension --- +mkdir -p "${SCRIPT_DIR}/prebuilt" + +python3 << 'PYEOF' +import os, sys, glob, shutil + +script_dir = os.environ.get("SCRIPT_DIR", ".") +vllm_root = os.environ.get("VLLM_ROOT", "") + +moe_cu = os.path.join(script_dir, "csrc", "moe_ops_impl.cu") +bridge_cpp = os.path.join(script_dir, "csrc", "ix_full_bridge_v2.cpp") + +# Collect linker flags +extra_ldflags = [] +rpath_dirs = set() + +corex_root = os.environ.get("COREX_ROOT", "/usr/local/corex") +for search_dir in [ + os.path.join(corex_root, "lib64"), + os.path.join(corex_root, "lib"), +]: + if os.path.isdir(search_dir): + rpath_dirs.add(search_dir) + for so in glob.glob(os.path.join(search_dir, "libcuinfer*.so*")): + extra_ldflags.append(so) + +# ixformer .so files +try: + import ixformer + ix_dir = os.path.dirname(ixformer.__file__) + rpath_dirs.add(ix_dir) + for so in glob.glob(os.path.join(ix_dir, "*.so")): + extra_ldflags.append(so) + for so in glob.glob(os.path.join(ix_dir, "lib*.so")): + if so not in extra_ldflags: + extra_ldflags.append(so) +except ImportError: + # Search common paths + for d in [ + os.path.join(corex_root, "lib", "python3", "dist-packages", "ixformer"), + os.path.join(corex_root, "lib64", "python3", "dist-packages", "ixformer"), + ]: + if os.path.isdir(d): + rpath_dirs.add(d) + for so in glob.glob(os.path.join(d, "*.so")): + extra_ldflags.append(so) + +for d in rpath_dirs: + extra_ldflags.append(f"-Wl,-rpath,{d}") + +print(f"[moe_bridge] Linking against {len(extra_ldflags)} items") +for f in extra_ldflags[:10]: + print(f" {f}") + +try: + from torch.utils.cpp_extension import load + + mod = load( + name="ix_moe_bridge", + sources=[moe_cu, bridge_cpp], + extra_include_paths=[os.path.join(script_dir, "csrc")], + extra_cflags=["-O2", "-std=c++17"], + extra_cuda_cflags=["-O2", "--extended-lambda"], + extra_ldflags=extra_ldflags, + verbose=True, + ) + print("[moe_bridge] ✓ Compilation successful") + + # Find and copy the built .so + import importlib + spec = importlib.util.find_spec("ix_moe_bridge") + if spec and spec.origin: + dst = os.path.join(script_dir, "prebuilt", "ix_moe_bridge.so") + shutil.copy2(spec.origin, dst) + print(f"[moe_bridge] ✓ Saved to {dst}") + + if vllm_root: + vllm_dst = os.path.join(vllm_root, "ex_engine", "ix_moe_bridge.so") + os.makedirs(os.path.dirname(vllm_dst), exist_ok=True) + shutil.copy2(spec.origin, vllm_dst) + print(f"[moe_bridge] ✓ Deployed to {vllm_dst}") + else: + print("[moe_bridge] ⚠ Could not locate compiled .so via importlib") + +except Exception as e: + print(f"[moe_bridge] ERROR: {e}", file=sys.stderr) + import traceback; traceback.print_exc() + sys.exit(1) +PYEOF + +echo "[moe_bridge] Done" diff --git a/ex_engine/csrc/moe_ops_impl.cu b/ex_engine/csrc/moe_ops_impl.cu index 55e8fc19..61974e72 100644 --- a/ex_engine/csrc/moe_ops_impl.cu +++ b/ex_engine/csrc/moe_ops_impl.cu @@ -69,14 +69,17 @@ cuinferStatus_t cuinferCustomGemm( // Adapted from moe_topk_softmax_v3.cu (already working, 64-expert specialized) // ============================================================================ -static constexpr int MOE_EXPERTS = 64; -static constexpr int MOE_BLOCK = 64; +// Qwen3.5-27B: 128 routed experts +// Block size = 128 threads (1 thread per expert for ≤128 experts) +static constexpr int MOE_MAX_EXPERTS = 128; +static constexpr int MOE_BLOCK = 128; +// All reductions use blockDim.x (dynamic block size, power-of-2) __device__ float smem_reduce_max(float val, float* smem) { int tid = threadIdx.x; smem[tid] = val; __syncthreads(); - for (int s = MOE_BLOCK / 2; s > 0; s >>= 1) { + for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]); __syncthreads(); } @@ -87,7 +90,7 @@ __device__ float smem_reduce_sum(float val, float* smem) { int tid = threadIdx.x; smem[tid] = val; __syncthreads(); - for (int s = MOE_BLOCK / 2; s > 0; s >>= 1) { + for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) smem[tid] += smem[tid + s]; __syncthreads(); } @@ -99,7 +102,7 @@ __device__ void smem_argmax(float val, int idx, float* s_val, int* s_idx) { s_val[tid] = val; s_idx[tid] = idx; __syncthreads(); - for (int s = MOE_BLOCK / 2; s > 0; s >>= 1) { + for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s && s_val[tid + s] > s_val[tid]) { s_val[tid] = s_val[tid + s]; s_idx[tid] = s_idx[tid + s]; @@ -113,20 +116,23 @@ __global__ void topk_softmax_kernel( float* __restrict__ topk_weights, int32_t* __restrict__ topk_indices, int32_t* __restrict__ token_expert_indices, - int num_tokens, int topk, bool renormalize + int num_tokens, int num_experts, int topk, bool renormalize ) { int row = blockIdx.x; if (row >= num_tokens) return; int tid = threadIdx.x; - __shared__ float smem[MOE_BLOCK]; - __shared__ int smem_idx[MOE_BLOCK]; + extern __shared__ char shared_buf[]; + float* smem = (float*)shared_buf; + int* smem_idx = (int*)(smem + blockDim.x); - float val = (tid < MOE_EXPERTS) ? input[row * MOE_EXPERTS + tid] : -1e30f; + // num_experts passed via gridDim.y (encoded), or read from shared + // We use a separate parameter for clarity + float val = (tid < num_experts) ? input[row * num_experts + tid] : -1e30f; // Softmax float row_max = smem_reduce_max(val, smem); - val = (tid < MOE_EXPERTS) ? expf(val - row_max) : 0.0f; + val = (tid < num_experts) ? expf(val - row_max) : 0.0f; float row_sum = smem_reduce_sum(val, smem); val *= (1.0f / row_sum); @@ -286,17 +292,24 @@ void topk_softmax( bool renormalize ) { int num_tokens = gating_output.size(0); + int num_experts = gating_output.size(1); int topk = topk_weights.size(1); auto stream = c10::cuda::getCurrentCUDAStream(); auto input_f32 = gating_output.to(torch::kFloat32).contiguous(); - topk_softmax_kernel<<>>( + // Block size must be >= num_experts, round up to next power of 2 + int block_size = 1; + while (block_size < num_experts) block_size <<= 1; + TORCH_CHECK(block_size <= 1024, "Too many experts for topk kernel: ", num_experts); + + size_t smem_bytes = block_size * (sizeof(float) + sizeof(int)); + topk_softmax_kernel<<>>( input_f32.data_ptr(), topk_weights.data_ptr(), topk_indices.data_ptr(), token_expert_indices.data_ptr(), - num_tokens, topk, renormalize); + num_tokens, num_experts, topk, renormalize); } void moe_compute_token_index_api( diff --git a/ex_engine/probe_moe_symbols.sh b/ex_engine/probe_moe_symbols.sh new file mode 100755 index 00000000..a4711400 --- /dev/null +++ b/ex_engine/probe_moe_symbols.sh @@ -0,0 +1,102 @@ +#!/usr/bin/env bash +# probe_moe_symbols.sh — Verify ix_moe_bridge.so has all 5 MoE symbols +# +# Run on real device after build_moe_bridge.sh + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Find the .so +SO_FILE="" +for p in \ + "${SCRIPT_DIR}/prebuilt/ix_moe_bridge.so" \ + "${SCRIPT_DIR}/ix_moe_bridge.so" \ + "$(python3 -c 'import ix_moe_bridge; print(ix_moe_bridge.__file__)' 2>/dev/null || echo '')"; do + if [[ -f "$p" ]]; then + SO_FILE="$p" + break + fi +done + +if [[ -z "$SO_FILE" ]]; then + echo "[probe] ERROR: ix_moe_bridge.so not found" + exit 1 +fi + +echo "[probe] Checking: $SO_FILE" +echo "[probe] Size: $(du -h "$SO_FILE" | cut -f1)" +echo "" + +# Required MoE symbols (must be in ixformer::infer namespace) +REQUIRED=( + "topk_softmax" + "moe_compute_token_index_api" + "moe_expand_input" + "moe_w16a16_group_gemm" + "moe_output_reduce_sum" +) + +# Required bridge symbols (pybind11 Python bindings) +BRIDGE_REQUIRED=( + "topk_softmax" + "moe_gen_idx" + "moe_expand_input" + "group_gemm" + "moe_combine_result" + "fused_moe_forward" + "silu_and_mul" + "rms_norm" + "linear" + "paged_attention" + "flash_attn_prefill" +) + +echo "=== MoE implementation symbols (ixformer::infer) ===" +PASS=0 +FAIL=0 +ALL_SYMS=$(nm -D "$SO_FILE" 2>/dev/null || nm "$SO_FILE" 2>/dev/null || echo "") + +for sym in "${REQUIRED[@]}"; do + count=$(echo "$ALL_SYMS" | grep -c "$sym" || true) + if [[ $count -gt 0 ]]; then + echo " ✓ $sym ($count matches)" + PASS=$((PASS + 1)) + else + echo " ✗ $sym — MISSING" + FAIL=$((FAIL + 1)) + fi +done + +echo "" +echo "=== pybind11 bridge symbols ===" +for sym in "${BRIDGE_REQUIRED[@]}"; do + count=$(echo "$ALL_SYMS" | grep -c "$sym" || true) + if [[ $count -gt 0 ]]; then + echo " ✓ $sym" + else + echo " ✗ $sym — MISSING" + FAIL=$((FAIL + 1)) + fi +done + +echo "" +echo "=== Python import test ===" +python3 -c " +import sys +sys.path.insert(0, '$(dirname "$SO_FILE")') +try: + import ix_moe_bridge as m + funcs = [f for f in dir(m) if not f.startswith('_')] + print(f' ✓ Import OK, {len(funcs)} functions: {funcs}') +except Exception as e: + print(f' ✗ Import failed: {e}') +" 2>&1 + +echo "" +if [[ $FAIL -eq 0 ]]; then + echo "[probe] ✓ ALL SYMBOLS PRESENT ($PASS MoE + bridge OK)" +else + echo "[probe] ✗ $FAIL SYMBOLS MISSING" + exit 1 +fi diff --git a/ex_engine/python/moe_dispatch.py b/ex_engine/python/moe_dispatch.py new file mode 100644 index 00000000..f693150d --- /dev/null +++ b/ex_engine/python/moe_dispatch.py @@ -0,0 +1,172 @@ +"""moe_dispatch.py — Load ix_moe_bridge.so and dispatch MoE forward. + +3-level fallback: + Tier 0: ix_moe_bridge.fused_moe_forward (C++ fused 7-step pipeline) + Tier 1: ix_moe_bridge individual ops (topk + expand + gemm + silu + gemm + combine) + Tier 2: Pure PyTorch fallback (F.linear loop) + +Used by: patch_moe_hot_path.py → replaces Qwen3_5MoE.forward() + +Reference: ex_engine/python/corex_moe.py (237L) +""" +import os +import sys +import logging +import torch +import torch.nn.functional as F + +logger = logging.getLogger("moe_dispatch") + +# --- Load bridge .so --- +_bridge = None +_tier = 2 # default: PyTorch fallback + + +def _try_load_bridge(): + global _bridge, _tier + + # Try 1: prebuilt .so + search_paths = [ + os.path.join(os.path.dirname(__file__), "ix_moe_bridge.so"), + os.path.join(os.path.dirname(__file__), "..", "prebuilt", "ix_moe_bridge.so"), + os.path.join(os.path.dirname(__file__), "..", "ix_moe_bridge.so"), + ] + for p in search_paths: + if os.path.isfile(p): + try: + import importlib.util + spec = importlib.util.spec_from_file_location("ix_moe_bridge", p) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + _bridge = mod + logger.info(f"[moe_dispatch] ✓ Loaded bridge from {p}") + break + except Exception as e: + logger.warning(f"[moe_dispatch] Failed to load {p}: {e}") + + # Try 2: torch JIT compiled module + if _bridge is None: + try: + import ix_moe_bridge + _bridge = ix_moe_bridge + logger.info("[moe_dispatch] ✓ Loaded bridge via import") + except ImportError: + pass + + if _bridge is None: + logger.warning("[moe_dispatch] Bridge not available, using PyTorch fallback") + _tier = 2 + return + + # Check what functions are available + try: + if hasattr(_bridge, 'fused_moe_forward'): + _tier = 0 + logger.info("[moe_dispatch] Tier 0: fused pipeline available") + elif hasattr(_bridge, 'topk_softmax') and hasattr(_bridge, 'group_gemm'): + _tier = 1 + logger.info("[moe_dispatch] Tier 1: individual ops available") + else: + _tier = 2 + logger.warning("[moe_dispatch] Bridge loaded but missing functions") + except Exception as e: + logger.warning(f"[moe_dispatch] Function check failed: {e}") + _tier = 2 + + +_try_load_bridge() + + +# ============================================================================ +# Tier 2: Pure PyTorch fallback (identical to base vllm behavior) +# ============================================================================ + +def _pytorch_moe_forward(hidden_states, router_logits, w13, w2, + topk, num_experts, renormalize): + """Python fallback: softmax → topk → loop over experts with F.linear.""" + gating = torch.softmax(router_logits.float(), dim=-1) + topk_weights, topk_ids = torch.topk(gating, topk, dim=-1) + if renormalize: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8) + topk_weights = topk_weights.to(hidden_states.dtype) + + # Per-expert loop + final_output = torch.zeros_like(hidden_states) + for k in range(topk): + expert_ids = topk_ids[:, k] # [T] + weights_k = topk_weights[:, k].unsqueeze(-1) # [T, 1] + for e in range(num_experts): + mask = (expert_ids == e) + if not mask.any(): + continue + expert_input = hidden_states[mask] + # gate_up = expert_input @ w13[e].T → [n, 2*inter] + gate_up = F.linear(expert_input, w13[e]) + inter = gate_up.shape[-1] // 2 + gate = torch.sigmoid(gate_up[:, :inter]) + up = gate_up[:, inter:] + activated = gate * up # SiLU approximated as sigmoid * x (should be silu_and_mul) + # down = activated @ w2[e].T → [n, hidden] + down = F.linear(activated, w2[e]) + final_output[mask] += weights_k[mask] * down + + return final_output + + +# ============================================================================ +# Tier 1: Individual bridge ops +# ============================================================================ + +def _bridge_individual_moe_forward(hidden_states, router_logits, w13, w2, + topk, num_experts, renormalize): + """Use individual bridge ops: topk → gen_idx → expand → gemm → silu → gemm → combine.""" + topk_weights, topk_ids, _ = _bridge.topk_softmax(router_logits, topk, False) + if renormalize: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8) + + idx_results = _bridge.moe_gen_idx(topk_ids.view(-1).to(torch.int32), num_experts) + src_dst, dst_src, expert_sizes = idx_results[0], idx_results[1], idx_results[2] + + expanded = _bridge.moe_expand_input(hidden_states, src_dst, dst_src, topk) + + gate_up = _bridge.group_gemm(expanded, w13, expert_sizes, w13.size(1)) + activated = _bridge.silu_and_mul(gate_up) + down = _bridge.group_gemm(activated, w2, expert_sizes, w2.size(1)) + output = _bridge.moe_combine_result(down, topk_weights) + + return output + + +# ============================================================================ +# Public API +# ============================================================================ + +def moe_forward(hidden_states, router_logits, w13, w2, + topk, num_experts, renormalize=True): + """Dispatch MoE forward to best available implementation.""" + if _tier == 0: + try: + return _bridge.fused_moe_forward( + hidden_states, router_logits, w13, w2, + topk, num_experts, renormalize) + except Exception as e: + logger.warning(f"[moe_dispatch] Tier 0 failed: {e}, falling to Tier 1") + pass + + if _tier <= 1 and _bridge is not None: + try: + return _bridge_individual_moe_forward( + hidden_states, router_logits, w13, w2, + topk, num_experts, renormalize) + except Exception as e: + logger.warning(f"[moe_dispatch] Tier 1 failed: {e}, falling to Tier 2") + pass + + return _pytorch_moe_forward( + hidden_states, router_logits, w13, w2, + topk, num_experts, renormalize) + + +def get_tier(): + """Return current dispatch tier (0=fused, 1=individual, 2=pytorch).""" + return _tier diff --git a/ex_engine/python/patch_moe_hot_path.py b/ex_engine/python/patch_moe_hot_path.py new file mode 100644 index 00000000..35f80df1 --- /dev/null +++ b/ex_engine/python/patch_moe_hot_path.py @@ -0,0 +1,109 @@ +"""patch_moe_hot_path.py — Replace Qwen3_5MoE.forward() with bridge dispatch. + +This is the key performance patch: replaces the Python expert-loop MoE +with a single C++ call that does all 7 steps fused. + +Called by: patch_ops.sh during Docker build +Target: vllm.model_executor.models.qwen3_5.Qwen3_5MoE + +Reference: ex_engine/python/patch_vllm_hot_path.py (200L) +""" +import sys +import logging +import torch + +logger = logging.getLogger("patch_moe_hot_path") + + +def apply_moe_patch(): + """Monkey-patch Qwen3_5MoE.forward to use moe_dispatch.""" + try: + from ex_engine.python.moe_dispatch import moe_forward, get_tier + except ImportError: + try: + from moe_dispatch import moe_forward, get_tier + except ImportError: + logger.warning("[moe_patch] moe_dispatch not available, skipping patch") + return False + + tier = get_tier() + logger.info(f"[moe_patch] moe_dispatch tier={tier}") + + # Find the MoE class + moe_cls = None + try: + from vllm.model_executor.models.qwen3_5 import Qwen3_5MoE + moe_cls = Qwen3_5MoE + except ImportError: + pass + + if moe_cls is None: + # Try to find it in sys.modules (may be registered under different name) + for mod_name, mod in sys.modules.items(): + if hasattr(mod, 'Qwen3_5MoE'): + moe_cls = getattr(mod, 'Qwen3_5MoE') + break + + if moe_cls is None: + logger.warning("[moe_patch] Qwen3_5MoE class not found") + return False + + # Save original forward + _original_forward = moe_cls.forward + + def patched_forward(self, hidden_states, *args, **kwargs): + """Patched MoE forward using bridge dispatch.""" + # Get router logits + # In Qwen3_5, the gate + shared_expert_gate are concatenated: + # router_and_shared_gate = self.gate(hidden_states) + # router_logits = router_and_shared_gate[..., :self.num_experts] + # shared_gate = router_and_shared_gate[..., -1] + router_and_shared_gate = self.gate(hidden_states) + router_logits = router_and_shared_gate[..., :self.num_experts] + + # Shared expert (if any) — run in parallel + shared_output = None + if hasattr(self, 'shared_expert') and self.shared_expert is not None: + if hasattr(self, 'shared_expert_gate'): + shared_gate = torch.sigmoid( + router_and_shared_gate[..., -1].unsqueeze(-1)) + else: + shared_gate = None + + # Routed experts via bridge + try: + routed_output = moe_forward( + hidden_states.view(-1, hidden_states.shape[-1]), + router_logits.view(-1, router_logits.shape[-1]), + self.w13_weight if hasattr(self, 'w13_weight') else self.experts.w13_weight, + self.w2_weight if hasattr(self, 'w2_weight') else self.experts.w2_weight, + topk=self.top_k, + num_experts=self.num_experts, + renormalize=True, + ) + routed_output = routed_output.view_as(hidden_states) + except Exception as e: + logger.warning(f"[moe_patch] Bridge failed ({e}), using original forward") + return _original_forward(self, hidden_states, *args, **kwargs) + + # Add shared expert output + if hasattr(self, 'shared_expert') and self.shared_expert is not None: + shared_out = self.shared_expert(hidden_states) + if shared_gate is not None: + shared_out = shared_out * shared_gate + routed_output = routed_output + shared_out + + return routed_output + + # Only patch if we have a real bridge (not pure Python fallback) + if tier < 2: + moe_cls.forward = patched_forward + logger.info(f"[moe_patch] ✓ Patched Qwen3_5MoE.forward (tier={tier})") + return True + else: + logger.info("[moe_patch] Tier 2 (Python only), not patching") + return False + + +if __name__ == "__main__": + apply_moe_patch() diff --git a/ex_engine/test_moe_bridge.py b/ex_engine/test_moe_bridge.py new file mode 100644 index 00000000..a93b633e --- /dev/null +++ b/ex_engine/test_moe_bridge.py @@ -0,0 +1,186 @@ +"""test_moe_bridge.py — Integration test for ix_moe_bridge on real device. + +Run after build_moe_bridge.sh. No model weights needed — uses random tensors. +Tests each of the 5 MoE functions + the fused pipeline. + +Usage: + python3 test_moe_bridge.py +""" +import sys +import os +import torch +import time + +# Qwen3.5-27B MoE params +NUM_EXPERTS = 128 +TOPK = 8 +HIDDEN_SIZE = 3584 +INTERMEDIATE_SIZE = 18944 # per-partition (full=18944*2 for gate+up, /TP if sharded) +NUM_TOKENS = 4 + +def load_bridge(): + """Try to load ix_moe_bridge.""" + # Try prebuilt + script_dir = os.path.dirname(os.path.abspath(__file__)) + for p in [ + os.path.join(script_dir, "prebuilt", "ix_moe_bridge.so"), + os.path.join(script_dir, "ix_moe_bridge.so"), + ]: + if os.path.isfile(p): + import importlib.util + spec = importlib.util.spec_from_file_location("ix_moe_bridge", p) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + # Try import + import ix_moe_bridge + return ix_moe_bridge + + +def test_topk_softmax(bridge, device): + print("\n--- topk_softmax ---") + gating = torch.randn(NUM_TOKENS, NUM_EXPERTS, device=device, dtype=torch.float32) + topk_w, topk_ids, token_expert_ids = bridge.topk_softmax(gating, TOPK, True) + + assert topk_w.shape == (NUM_TOKENS, TOPK), f"weights shape: {topk_w.shape}" + assert topk_ids.shape == (NUM_TOKENS, TOPK), f"ids shape: {topk_ids.shape}" + assert topk_w.dtype == torch.float32 + assert topk_ids.dtype == torch.int32 + assert (topk_ids >= 0).all() and (topk_ids < NUM_EXPERTS).all(), "ids out of range" + assert torch.allclose(topk_w.sum(-1), torch.ones(NUM_TOKENS, device=device), atol=1e-5), \ + f"weights don't sum to 1: {topk_w.sum(-1)}" + print(f" ✓ shape={topk_w.shape}, sum={topk_w.sum(-1).tolist()}") + print(f" ✓ top expert ids (row 0): {topk_ids[0].tolist()}") + + +def test_moe_gen_idx(bridge, device): + print("\n--- moe_gen_idx ---") + expert_ids = torch.randint(0, NUM_EXPERTS, (NUM_TOKENS * TOPK,), + device=device, dtype=torch.int32) + results = bridge.moe_gen_idx(expert_ids, NUM_EXPERTS) + src_dst, dst_src, expert_sizes, expert_cumsum = results + + assert src_dst.shape == (NUM_TOKENS * TOPK,), f"src_dst shape: {src_dst.shape}" + assert dst_src.shape == (NUM_TOKENS * TOPK,), f"dst_src shape: {dst_src.shape}" + assert expert_sizes.shape[0] == NUM_EXPERTS, f"expert_sizes shape: {expert_sizes.shape}" + assert expert_sizes.sum().item() == NUM_TOKENS * TOPK, \ + f"expert_sizes sum: {expert_sizes.sum().item()} != {NUM_TOKENS * TOPK}" + print(f" ✓ src_dst={src_dst.shape}, expert_sizes sum={expert_sizes.sum().item()}") + + +def test_moe_expand_input(bridge, device): + print("\n--- moe_expand_input ---") + hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, device=device, dtype=torch.float16) + # Create simple gather index: [0,1,2,...,NUM_TOKENS*TOPK-1] mod NUM_TOKENS + gather_idx = torch.arange(NUM_TOKENS * TOPK, device=device, dtype=torch.int32) % NUM_TOKENS + combine_idx = torch.arange(NUM_TOKENS * TOPK, device=device, dtype=torch.int32) + + expanded = bridge.moe_expand_input(hidden, gather_idx, combine_idx, TOPK) + assert expanded.shape == (NUM_TOKENS * TOPK, HIDDEN_SIZE), f"shape: {expanded.shape}" + print(f" ✓ shape={expanded.shape}, dtype={expanded.dtype}") + + +def test_group_gemm(bridge, device): + print("\n--- group_gemm ---") + # Simulate: expanded tokens × expert weights + total_tokens = NUM_TOKENS * TOPK # 32 + inputs = torch.randn(total_tokens, HIDDEN_SIZE, device=device, dtype=torch.float16) + # weights: [NUM_EXPERTS, 2*INTERMEDIATE, HIDDEN] — 3D + weights = torch.randn(NUM_EXPERTS, INTERMEDIATE_SIZE * 2, HIDDEN_SIZE, + device=device, dtype=torch.float16) * 0.01 + # tokens_per_expert: distribute evenly + tpe = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32) + for i in range(total_tokens): + tpe[i % NUM_EXPERTS] += 1 + + output_n = INTERMEDIATE_SIZE * 2 + result = bridge.group_gemm(inputs, weights, tpe, output_n) + assert result.shape == (total_tokens, output_n), f"shape: {result.shape}" + assert not torch.isnan(result).any(), "NaN in group_gemm output" + print(f" ✓ shape={result.shape}, max={result.abs().max().item():.4f}") + + +def test_silu_and_mul(bridge, device): + print("\n--- silu_and_mul ---") + gate_up = torch.randn(NUM_TOKENS, INTERMEDIATE_SIZE * 2, + device=device, dtype=torch.float16) + activated = bridge.silu_and_mul(gate_up) + assert activated.shape == (NUM_TOKENS, INTERMEDIATE_SIZE), f"shape: {activated.shape}" + print(f" ✓ shape={activated.shape}") + + +def test_moe_combine_result(bridge, device): + print("\n--- moe_combine_result ---") + expert_out = torch.randn(NUM_TOKENS * TOPK, HIDDEN_SIZE, + device=device, dtype=torch.float16) + weights = torch.randn(NUM_TOKENS, TOPK, device=device, dtype=torch.float32) + weights = torch.softmax(weights, dim=-1) + + combined = bridge.moe_combine_result(expert_out, weights) + assert combined.shape == (NUM_TOKENS, HIDDEN_SIZE), f"shape: {combined.shape}" + assert not torch.isnan(combined).any(), "NaN in combine output" + print(f" ✓ shape={combined.shape}") + + +def test_fused_pipeline(bridge, device): + print("\n--- fused_moe_forward (7-step pipeline) ---") + hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, device=device, dtype=torch.float16) + router = torch.randn(NUM_TOKENS, NUM_EXPERTS, device=device, dtype=torch.float16) + w13 = torch.randn(NUM_EXPERTS, INTERMEDIATE_SIZE * 2, HIDDEN_SIZE, + device=device, dtype=torch.float16) * 0.01 + w2 = torch.randn(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, + device=device, dtype=torch.float16) * 0.01 + + t0 = time.time() + output = bridge.fused_moe_forward(hidden, router, w13, w2, TOPK, NUM_EXPERTS, True) + torch.cuda.synchronize() + elapsed = time.time() - t0 + + assert output.shape == (NUM_TOKENS, HIDDEN_SIZE), f"shape: {output.shape}" + assert not torch.isnan(output).any(), "NaN in fused output" + print(f" ✓ shape={output.shape}, time={elapsed*1000:.1f}ms") + + +def main(): + if not torch.cuda.is_available(): + print("CUDA not available, skipping GPU tests") + sys.exit(0) + + device = torch.device("cuda:0") + print(f"Device: {torch.cuda.get_device_name(0)}") + print(f"Params: {NUM_EXPERTS} experts, topk={TOPK}, hidden={HIDDEN_SIZE}, " + f"inter={INTERMEDIATE_SIZE}, tokens={NUM_TOKENS}") + + bridge = load_bridge() + funcs = [f for f in dir(bridge) if not f.startswith('_')] + print(f"Bridge loaded: {len(funcs)} functions: {funcs}") + + passed = 0 + failed = 0 + + for test_fn in [ + test_topk_softmax, + test_moe_gen_idx, + test_moe_expand_input, + test_group_gemm, + test_silu_and_mul, + test_moe_combine_result, + test_fused_pipeline, + ]: + try: + test_fn(bridge, device) + passed += 1 + except Exception as e: + print(f" ✗ FAILED: {e}") + import traceback; traceback.print_exc() + failed += 1 + + print(f"\n{'='*40}") + print(f"Results: {passed} passed, {failed} failed") + if failed > 0: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index abac4a57..31650f7f 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -332,6 +332,23 @@ if b"max_completion_tokens" not in installed: raise SystemExit("protocol.py missing max_completion_tokens field") PY +build_stage "building MoE bridge (ix_moe_bridge.so)" +if [[ -f "./ex_engine_src/build_moe_bridge.sh" ]]; then + bash ./ex_engine_src/build_moe_bridge.sh "${VLLM_ROOT}" 2>&1 || { + echo "[WARN] MoE bridge build failed — will use Python fallback" + } +fi + +build_stage "deploying MoE dispatch modules" +EX_DIR="${VLLM_ROOT}/ex_engine/python" +mkdir -p "${EX_DIR}" +for pyfile in moe_dispatch.py patch_moe_hot_path.py; do + if [[ -f "./ex_engine_src/python/${pyfile}" ]]; then + cp "./ex_engine_src/python/${pyfile}" "${EX_DIR}/${pyfile}" + echo " ✓ ${pyfile}" + fi +done + build_stage "compiling submission Python sources" find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile