344 lines
14 KiB
Python
344 lines
14 KiB
Python
"""ix_unified.py — Unified Python interface to all ixformer::infer APIs.
|
|
|
|
Dispatch hierarchy (CCCL policy_selector pattern):
|
|
Tier 0: ix_unified_bridge.so (C++ direct call to ixformer::infer)
|
|
Tier 1: ixformer.functions.* (base image Python bindings, partial)
|
|
Tier 2: PyTorch fallback (always works, slowest)
|
|
|
|
Usage:
|
|
from ex_engine.python.ix_unified import ix
|
|
out = ix.silu_and_mul(input)
|
|
ix.rms_norm(output, input, weight, eps)
|
|
weights, indices = ix.moe_topk_softmax(gating, topk, renorm)
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import importlib
|
|
import importlib.util
|
|
import torch
|
|
import logging
|
|
|
|
logger = logging.getLogger("ix_unified")
|
|
|
|
_bridge = None
|
|
|
|
|
|
def _load_bridge():
|
|
"""Load ix_unified_bridge.so from known locations."""
|
|
global _bridge
|
|
if _bridge is not None:
|
|
return _bridge
|
|
|
|
# Pre-load ixformer .so symbols into GLOBAL symbol table.
|
|
# ix_unified_bridge.so has undefined ixformer::infer::* symbols that get
|
|
# resolved at runtime. Python default import uses RTLD_LOCAL, so we must
|
|
# force RTLD_GLOBAL on the ixformer .so files BEFORE loading our bridge.
|
|
try:
|
|
import ctypes
|
|
|
|
# Phase 0: Load torch core libs first — ixformer depends on libc10.so etc.
|
|
try:
|
|
import torch as _torch
|
|
_torch_lib = os.path.join(os.path.dirname(_torch.__file__), "lib")
|
|
for _name in ["libc10.so", "libtorch_cpu.so", "libtorch.so",
|
|
"libc10_cuda.so", "libtorch_cuda.so", "libtorch_python.so"]:
|
|
_p = os.path.join(_torch_lib, _name)
|
|
if os.path.isfile(_p):
|
|
try:
|
|
ctypes.CDLL(_p, mode=ctypes.RTLD_GLOBAL)
|
|
except Exception:
|
|
pass
|
|
except ImportError:
|
|
pass
|
|
|
|
# Phase 1: libixformer.so (CUDA kernels)
|
|
# Phase 2: _ixformer_torch.so (torch extension with ixformer_torch_ext::*)
|
|
# ONLY these two — do NOT recursively load unknown .so (causes segfault)
|
|
_ixf_base = "/usr/local/corex/lib64/python3/dist-packages/ixformer"
|
|
if os.path.isdir(_ixf_base):
|
|
for _name in ["libixformer.so",
|
|
"_ixformer_torch.cpython-310-x86_64-linux-gnu.so"]:
|
|
_p = os.path.join(_ixf_base, _name)
|
|
if os.path.isfile(_p):
|
|
try:
|
|
ctypes.CDLL(_p, mode=ctypes.RTLD_GLOBAL)
|
|
logger.info("Preloaded: %s", _name)
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
search_paths = []
|
|
|
|
# 1. Same directory as this file
|
|
here = os.path.dirname(os.path.abspath(__file__))
|
|
search_paths.append(os.path.join(here, "..", "build"))
|
|
search_paths.append(here)
|
|
|
|
# 2. Workspace build dirs (Docker / real machine)
|
|
search_paths.append("/workspace/ex_engine/build")
|
|
search_paths.append("/home/dylan/project_6/ex_engine/build")
|
|
|
|
# 2. vllm install root (where prebuilt .so are deployed)
|
|
for p in sys.path:
|
|
if "vllm" in p or "dist-packages" in p:
|
|
search_paths.append(p)
|
|
|
|
# 3. Explicit env var
|
|
env_path = os.getenv("IX_BRIDGE_PATH")
|
|
if env_path:
|
|
search_paths.insert(0, env_path)
|
|
|
|
for search_dir in search_paths:
|
|
for name in ["ix_unified_bridge.so",
|
|
"ix_unified_bridge.cpython-310-x86_64-linux-gnu.so"]:
|
|
so_path = os.path.join(search_dir, name)
|
|
if os.path.isfile(so_path):
|
|
try:
|
|
spec = importlib.util.spec_from_file_location(
|
|
"ix_unified_bridge", so_path)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
_bridge = mod
|
|
logger.info("ix_unified_bridge loaded from %s", so_path)
|
|
return _bridge
|
|
except (ImportError, OSError, SystemError) as e:
|
|
logger.warning("Bridge load failed (expected if ixformer "
|
|
"namespace mismatch): %s: %s",
|
|
os.path.basename(so_path), e)
|
|
continue
|
|
except Exception as e:
|
|
logger.warning("Bridge load unexpected error: %s", e)
|
|
continue
|
|
|
|
logger.info("ix_unified_bridge.so not found, using fallback dispatch")
|
|
return None
|
|
|
|
|
|
def _try_ixformer_functions():
|
|
"""Try importing ixformer.functions from base image."""
|
|
try:
|
|
import ixformer.functions as ixf
|
|
return ixf
|
|
except (ImportError, AttributeError):
|
|
return None
|
|
|
|
|
|
# ============================================================================
|
|
# Dispatch class
|
|
# ============================================================================
|
|
|
|
class IXDispatch:
|
|
"""Three-tier dispatch for all ixformer ops."""
|
|
|
|
def __init__(self):
|
|
self._bridge = _load_bridge()
|
|
self._ixf = _try_ixformer_functions()
|
|
tier = ("Tier0:bridge" if self._bridge else
|
|
"Tier1:ixformer" if self._ixf else "Tier2:pytorch")
|
|
logger.info("IXDispatch initialized: %s", tier)
|
|
|
|
# --- Activation -----------------------------------------------------------
|
|
def silu_and_mul(self, input: torch.Tensor) -> torch.Tensor:
|
|
if self._bridge:
|
|
return self._bridge.silu_and_mul(input)
|
|
if self._ixf and hasattr(self._ixf, 'silu_and_mul'):
|
|
d = input.size(-1) // 2
|
|
out = input.new_empty([input.size(0), d])
|
|
self._ixf.silu_and_mul(input, out)
|
|
return out
|
|
# PyTorch fallback
|
|
d = input.size(-1) // 2
|
|
x, gate = input[..., :d], input[..., d:]
|
|
return x * torch.sigmoid(gate)
|
|
|
|
# --- Norm -----------------------------------------------------------------
|
|
def rms_norm(self, output: torch.Tensor, input: torch.Tensor,
|
|
weight: torch.Tensor, eps: float):
|
|
if self._bridge:
|
|
self._bridge.rms_norm(output, input, weight, eps)
|
|
return
|
|
if self._ixf and hasattr(self._ixf, 'rms_norm'):
|
|
self._ixf.rms_norm(input, weight, output, eps)
|
|
return
|
|
# PyTorch fallback
|
|
variance = input.float().pow(2).mean(-1, keepdim=True)
|
|
normed = input * torch.rsqrt(variance + eps)
|
|
output.copy_(normed * weight)
|
|
|
|
def fused_add_rms_norm(self, input: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
weight: torch.Tensor, eps: float):
|
|
if self._bridge:
|
|
self._bridge.fused_add_rms_norm(input, residual, weight, eps)
|
|
return
|
|
if self._ixf and hasattr(self._ixf, 'fused_add_rms_norm'):
|
|
self._ixf.fused_add_rms_norm(input, residual, weight, eps, 1.0)
|
|
return
|
|
# PyTorch fallback
|
|
hidden = input + residual
|
|
residual.copy_(hidden)
|
|
variance = hidden.float().pow(2).mean(-1, keepdim=True)
|
|
normed = hidden * torch.rsqrt(variance + eps)
|
|
input.copy_(normed * weight)
|
|
|
|
# --- Linear ---------------------------------------------------------------
|
|
def linear(self, input: torch.Tensor, weight: torch.Tensor,
|
|
bias=None) -> torch.Tensor:
|
|
if self._bridge:
|
|
return self._bridge.linear(input, weight, bias)
|
|
# PyTorch fallback
|
|
out = torch.nn.functional.linear(input, weight, bias)
|
|
return out
|
|
|
|
# --- RoPE -----------------------------------------------------------------
|
|
def rotary_embedding(self, positions, query, key, head_size,
|
|
cos_sin_cache, is_neox=True):
|
|
if self._bridge:
|
|
self._bridge.rotary_embedding(positions, query, key, head_size,
|
|
cos_sin_cache, is_neox)
|
|
return
|
|
if self._ixf and hasattr(self._ixf, 'vllm_rotary_embedding_neox'):
|
|
self._ixf.vllm_rotary_embedding_neox(
|
|
positions, query, key, head_size, cos_sin_cache, is_neox)
|
|
return
|
|
# No PyTorch fallback — this is handled by vllm's own rope
|
|
|
|
# --- KV Cache -------------------------------------------------------------
|
|
def reshape_and_cache(self, key, value, key_cache, value_cache,
|
|
slot_mapping):
|
|
if self._bridge:
|
|
self._bridge.reshape_and_cache(key, value, key_cache, value_cache,
|
|
slot_mapping)
|
|
return
|
|
if self._ixf and hasattr(self._ixf, 'vllm_cache_ops_reshape_and_cache'):
|
|
self._ixf.vllm_cache_ops_reshape_and_cache(
|
|
key, value, key_cache, value_cache, slot_mapping)
|
|
return
|
|
# PyTorch fallback — slot-by-slot copy
|
|
for i, slot in enumerate(slot_mapping):
|
|
if slot < 0:
|
|
continue
|
|
block_idx = slot // key_cache.size(2)
|
|
block_off = slot % key_cache.size(2)
|
|
key_cache[block_idx, :, block_off, :] = key[i]
|
|
value_cache[block_idx, :, block_off, :] = value[i]
|
|
|
|
# --- Attention: prefill ---------------------------------------------------
|
|
def flash_attn_prefill(self, query, key_cache, value_cache, output,
|
|
block_tables, cu_seq_q, cu_seq_k,
|
|
max_seq_q, max_seq_k, is_causal, scale):
|
|
if self._bridge:
|
|
return self._bridge.flash_attn_prefill(
|
|
query, key_cache, value_cache, output, block_tables,
|
|
cu_seq_q, cu_seq_k, max_seq_q, max_seq_k, is_causal, scale)
|
|
if self._ixf and hasattr(self._ixf, 'ixinfer_flash_attn_unpad'):
|
|
return self._ixf.ixinfer_flash_attn_unpad(
|
|
query, key_cache, value_cache, output, block_tables,
|
|
cu_seq_q, cu_seq_k, max_seq_q, max_seq_k,
|
|
is_causal, -1, -1, scale, 0.0, False, None, None, None)
|
|
raise RuntimeError("flash_attn_prefill: no backend available")
|
|
|
|
# --- Attention: decode (paged) -------------------------------------------
|
|
def paged_attention(self, output, query, key_cache, value_cache,
|
|
num_kv_heads, scale, block_tables, context_lens,
|
|
block_size, max_context_len):
|
|
if self._bridge:
|
|
return self._bridge.paged_attention(
|
|
output, query, key_cache, value_cache,
|
|
num_kv_heads, scale, block_tables, context_lens,
|
|
block_size, max_context_len)
|
|
if self._ixf and hasattr(self._ixf,
|
|
'vllm_single_query_cached_kv_attention_v2'):
|
|
return self._ixf.vllm_single_query_cached_kv_attention_v2(
|
|
output, query, key_cache, value_cache,
|
|
num_kv_heads, scale, block_tables, context_lens,
|
|
block_size, max_context_len, None)
|
|
raise RuntimeError("paged_attention: no backend available")
|
|
|
|
# --- MoE: topk_softmax ---------------------------------------------------
|
|
def moe_topk_softmax(self, gating_output: torch.Tensor,
|
|
topk: int, renormalize: bool = True):
|
|
if self._bridge:
|
|
return self._bridge.moe_topk_softmax(
|
|
gating_output, topk, renormalize)
|
|
# PyTorch fallback
|
|
scores = torch.softmax(gating_output.float(), dim=-1)
|
|
topk_weights, topk_indices = torch.topk(scores, k=topk, dim=-1)
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1,
|
|
keepdim=True)
|
|
return topk_weights, topk_indices.to(torch.int32)
|
|
|
|
# --- MoE: gen_idx ---------------------------------------------------------
|
|
def moe_gen_idx(self, expert_ids: torch.Tensor, num_experts: int):
|
|
if self._bridge:
|
|
return self._bridge.moe_gen_idx(expert_ids, num_experts)
|
|
# PyTorch fallback: compute scatter/gather indices
|
|
flat = expert_ids.view(-1)
|
|
n = flat.numel()
|
|
src_dst = torch.empty(n, dtype=flat.dtype, device=flat.device)
|
|
dst_src = torch.empty(n, dtype=flat.dtype, device=flat.device)
|
|
expert_sizes = torch.zeros(num_experts, dtype=flat.dtype,
|
|
device=flat.device)
|
|
# Simple counting sort
|
|
for i in range(n):
|
|
expert_sizes[flat[i].item()] += 1
|
|
cumsum = expert_sizes.cumsum(-1)
|
|
offsets = torch.zeros_like(expert_sizes)
|
|
offsets[1:] = cumsum[:-1]
|
|
counts = torch.zeros_like(expert_sizes)
|
|
for i in range(n):
|
|
e = flat[i].item()
|
|
pos = (offsets[e] + counts[e]).item()
|
|
src_dst[i] = pos
|
|
dst_src[pos] = i
|
|
counts[e] += 1
|
|
return [src_dst, dst_src, expert_sizes, cumsum]
|
|
|
|
# --- MoE: expand_input ----------------------------------------------------
|
|
def moe_expand_input(self, input: torch.Tensor,
|
|
gather_index: torch.Tensor,
|
|
combine_idx: torch.Tensor, topk: int):
|
|
if self._bridge:
|
|
return self._bridge.moe_expand_input(
|
|
input, gather_index, combine_idx, topk)
|
|
# PyTorch fallback
|
|
return input.index_select(0, combine_idx.view(-1).long())
|
|
|
|
# --- MoE: group_gemm -----------------------------------------------------
|
|
def moe_group_gemm(self, input: torch.Tensor, weight: torch.Tensor,
|
|
tokens_per_experts: torch.Tensor):
|
|
if self._bridge:
|
|
return self._bridge.moe_group_gemm(
|
|
input, weight, tokens_per_experts)
|
|
# PyTorch fallback: sequential per-expert GEMM
|
|
outputs = []
|
|
offset = 0
|
|
for e in range(tokens_per_experts.size(0)):
|
|
count = tokens_per_experts[e].item()
|
|
if count == 0:
|
|
continue
|
|
inp_e = input[offset:offset + count]
|
|
w_e = weight[e] # [out_features, in_features]
|
|
outputs.append(inp_e @ w_e.t())
|
|
offset += count
|
|
if outputs:
|
|
return torch.cat(outputs, dim=0)
|
|
return input.new_empty(0, weight.size(-2))
|
|
|
|
# --- MoE: combine_result -------------------------------------------------
|
|
def moe_combine_result(self, expert_output: torch.Tensor,
|
|
weights: torch.Tensor):
|
|
if self._bridge:
|
|
return self._bridge.moe_combine_result(expert_output, weights)
|
|
# PyTorch fallback: weighted sum
|
|
# expert_output: [n_tokens, topk, hidden]
|
|
# weights: [n_tokens, topk]
|
|
return (expert_output * weights.unsqueeze(-1)).sum(dim=1)
|
|
|
|
|
|
# Singleton
|
|
ix = IXDispatch()
|