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
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)
|
||||
Reference in New Issue
Block a user