diff --git a/Dockerfile b/Dockerfile index 832acb2f..19f1d298 100644 --- a/Dockerfile +++ b/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: $?" diff --git a/Dockerfile.broken_head b/Dockerfile.broken_head new file mode 100644 index 00000000..832acb2f --- /dev/null +++ b/Dockerfile.broken_head @@ -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: $?" diff --git a/computility-run.yaml b/computility-run.yaml index 41351f46..491b3d11 100644 --- a/computility-run.yaml +++ b/computility-run.yaml @@ -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' diff --git a/ex_engine/python/corex_so_loader.py b/ex_engine/python/corex_so_loader.py new file mode 100644 index 00000000..7bcc78ae --- /dev/null +++ b/ex_engine/python/corex_so_loader.py @@ -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() diff --git a/ex_engine/python/moe_fused_dispatch.py b/ex_engine/python/moe_fused_dispatch.py new file mode 100644 index 00000000..7030a5b0 --- /dev/null +++ b/ex_engine/python/moe_fused_dispatch.py @@ -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) diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index b0d016a0..310d7dca 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -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}"