349 lines
13 KiB
Python
349 lines
13 KiB
Python
"""
|
|
ix_ops.py — Drop-in operator replacements via ix_full_bridge.so
|
|
|
|
Architecture (CCCL dispatch pattern):
|
|
CCCL: compute_capability → policy_selector → tuned_kernel
|
|
EX: base_image_so → ix_full_bridge → ixformer::infer
|
|
|
|
This module provides torch.nn.Module-compatible replacements for:
|
|
1. RMSNorm → residual_rms_norm / rms_norm (fused kernel)
|
|
2. SiluAndMul → silu_and_mul (fused activation)
|
|
3. RotaryEmbedding → xllm_rotary_embedding (fused RoPE)
|
|
4. reshape_and_cache → xllm_reshape_and_cache (fused KV write)
|
|
5. paged_attention → xllm_paged_attention (fused decode attn)
|
|
6. flash_attn_prefill → ixinfer_flash_attn_unpad (fused prefill attn)
|
|
7. linear → ixformer_linear / linear_ex (GEMM)
|
|
|
|
Loading: tries prebuilt ix_full_bridge.so first, then JIT-compiles
|
|
ix_full_bridge_v2.cpp as fallback.
|
|
|
|
Source mapping:
|
|
upstream_ref/xllm_latest/core/kernels/ilu/*.cpp → this file (Python side)
|
|
ex_engine/csrc/ix_full_bridge_v2.cpp → .so (C++ side)
|
|
ixformer::infer namespace (base image) → actual CUDA kernels
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import logging
|
|
import importlib
|
|
import importlib.util
|
|
import glob
|
|
import torch
|
|
from typing import Optional, Tuple, List
|
|
|
|
logger = logging.getLogger("ex_engine.ix_ops")
|
|
|
|
# =========================================================================
|
|
# Bridge loader
|
|
# =========================================================================
|
|
_bridge = None
|
|
_loaded = False
|
|
_available = False
|
|
|
|
|
|
def _try_prebuilt():
|
|
"""Load prebuilt ix_full_bridge.so."""
|
|
search = [
|
|
# Deployed by patch_ops.sh into vllm package
|
|
"/usr/local/corex/lib/python3/dist-packages/vllm/ix_full_bridge.so",
|
|
]
|
|
# Also check vllm package dir
|
|
try:
|
|
import vllm
|
|
vd = os.path.dirname(vllm.__file__)
|
|
search.insert(0, os.path.join(vd, "ix_full_bridge.so"))
|
|
except ImportError:
|
|
pass
|
|
# Check prebuilt dir
|
|
here = os.path.dirname(os.path.abspath(__file__))
|
|
search.append(os.path.join(here, "..", "..", "qwen3_6_scripts", "prebuilt",
|
|
"corex-3.2.3-ivcore10", "ix_full_bridge.so"))
|
|
|
|
for path in search:
|
|
path = os.path.normpath(path)
|
|
if not os.path.isfile(path):
|
|
continue
|
|
try:
|
|
spec = importlib.util.spec_from_file_location("ix_full_bridge", path)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
fns = [x for x in dir(mod) if not x.startswith("_")]
|
|
logger.info("ix_ops: loaded prebuilt %s: %s", path, fns)
|
|
return mod
|
|
except Exception as e:
|
|
logger.debug("ix_ops: prebuilt %s failed: %s", path, e)
|
|
return None
|
|
|
|
|
|
def _try_jit():
|
|
"""JIT compile ix_full_bridge_v2.cpp."""
|
|
here = os.path.dirname(os.path.abspath(__file__))
|
|
cpp_candidates = [
|
|
os.path.join(here, "..", "csrc", "ix_full_bridge_v2.cpp"),
|
|
os.path.join(here, "..", "csrc", "ix_full_bridge.cpp"),
|
|
"/workspace/ex_engine/csrc/ix_full_bridge_v2.cpp",
|
|
"/workspace/qwen3_6_scripts/ix_full_bridge_v2.cpp",
|
|
]
|
|
cpp_file = None
|
|
for c in cpp_candidates:
|
|
c = os.path.normpath(c)
|
|
if os.path.isfile(c):
|
|
cpp_file = c
|
|
break
|
|
if cpp_file is None:
|
|
return None
|
|
|
|
extra_ldflags = []
|
|
# Link ixformer .so libraries
|
|
try:
|
|
import ixformer
|
|
ixf_dir = os.path.dirname(ixformer.__file__)
|
|
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
|
|
extra_ldflags.append(so)
|
|
extra_ldflags.append(f"-Wl,-rpath,{ixf_dir}")
|
|
except ImportError:
|
|
pass
|
|
# Also link corex libraries
|
|
corex_lib = "/usr/local/corex/lib64"
|
|
if os.path.isdir(corex_lib):
|
|
for lib in ["libixattn.so", "libixformer.so", "libcublas.so"]:
|
|
p = os.path.join(corex_lib, lib)
|
|
if os.path.isfile(p):
|
|
extra_ldflags.append(p)
|
|
extra_ldflags.append(f"-Wl,-rpath,{corex_lib}")
|
|
|
|
try:
|
|
from torch.utils.cpp_extension import load
|
|
logger.info("ix_ops: JIT compiling %s", cpp_file)
|
|
mod = load(
|
|
name="ix_full_bridge_v2",
|
|
sources=[cpp_file],
|
|
extra_cflags=["-O2", "-std=c++17"],
|
|
extra_ldflags=extra_ldflags,
|
|
verbose=False,
|
|
)
|
|
fns = [x for x in dir(mod) if not x.startswith("_")]
|
|
logger.info("ix_ops: JIT compiled: %s", fns)
|
|
return mod
|
|
except Exception as e:
|
|
logger.warning("ix_ops: JIT compile failed: %s", e)
|
|
return None
|
|
|
|
|
|
def _ensure_loaded():
|
|
global _bridge, _loaded, _available
|
|
if _loaded:
|
|
return _available
|
|
_loaded = True
|
|
_bridge = _try_prebuilt()
|
|
if _bridge is None:
|
|
_bridge = _try_jit()
|
|
_available = _bridge is not None
|
|
if _available:
|
|
logger.info("ix_ops: bridge available with %d functions",
|
|
len([x for x in dir(_bridge) if not x.startswith("_")]))
|
|
else:
|
|
logger.warning("ix_ops: bridge NOT available, all ops will be no-op")
|
|
return _available
|
|
|
|
|
|
def is_available() -> bool:
|
|
return _ensure_loaded()
|
|
|
|
|
|
def get_bridge():
|
|
if not _ensure_loaded():
|
|
raise RuntimeError("ix_ops bridge not available")
|
|
return _bridge
|
|
|
|
|
|
# =========================================================================
|
|
# Feature probes — check what the loaded bridge supports
|
|
# =========================================================================
|
|
def has_silu_and_mul() -> bool:
|
|
return is_available() and hasattr(_bridge, "silu_and_mul")
|
|
|
|
def has_rms_norm() -> bool:
|
|
return is_available() and hasattr(_bridge, "rms_norm")
|
|
|
|
def has_fused_add_rms_norm() -> bool:
|
|
return is_available() and hasattr(_bridge, "fused_add_rms_norm")
|
|
|
|
def has_rotary_embedding() -> bool:
|
|
return is_available() and hasattr(_bridge, "rotary_embedding")
|
|
|
|
def has_reshape_and_cache() -> bool:
|
|
return is_available() and hasattr(_bridge, "reshape_and_cache")
|
|
|
|
def has_paged_attention() -> bool:
|
|
return is_available() and hasattr(_bridge, "paged_attention")
|
|
|
|
def has_flash_attn_prefill() -> bool:
|
|
return is_available() and hasattr(_bridge, "flash_attn_prefill")
|
|
|
|
def has_linear() -> bool:
|
|
return is_available() and hasattr(_bridge, "linear")
|
|
|
|
def has_topk_softmax() -> bool:
|
|
return is_available() and hasattr(_bridge, "topk_softmax")
|
|
|
|
def has_fused_moe_forward() -> bool:
|
|
return is_available() and hasattr(_bridge, "fused_moe_forward")
|
|
|
|
|
|
# =========================================================================
|
|
# Op wrappers — match xllm upstream signatures
|
|
# Source: upstream_ref/xllm_latest/core/kernels/ilu/*.cpp
|
|
# =========================================================================
|
|
|
|
def silu_and_mul(input: torch.Tensor) -> torch.Tensor:
|
|
"""Fused SiLU activation + element-wise multiply.
|
|
|
|
Source: xllm/core/kernels/ilu/activation.cpp → infer::silu_and_mul
|
|
input: (T, 2*I) → output: (T, I)
|
|
"""
|
|
return _bridge.silu_and_mul(input)
|
|
|
|
|
|
def rms_norm(output: torch.Tensor, input: torch.Tensor,
|
|
weight: torch.Tensor, eps: float = 1e-6) -> None:
|
|
"""RMSNorm: output = rms_norm(input, weight, eps).
|
|
|
|
Source: xllm/core/kernels/ilu/norm.cpp → infer::rms_norm
|
|
"""
|
|
_bridge.rms_norm(output, input, weight, eps)
|
|
|
|
|
|
def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
|
|
weight: torch.Tensor, output: torch.Tensor,
|
|
residual_output: torch.Tensor,
|
|
eps: float = 1e-6) -> None:
|
|
"""Fused residual addition + RMSNorm.
|
|
|
|
Source: xllm/core/kernels/ilu/norm.cpp → infer::residual_rms_norm
|
|
The C++ function is in-place: modifies input → rms_norm(input+residual)*weight,
|
|
and residual → input+residual. We copy results to output/residual_output.
|
|
"""
|
|
# C++ signature: fused_add_rms_norm_forward(input, residual, weight, eps, alpha)
|
|
# It modifies input and residual in-place.
|
|
inp_clone = input.clone()
|
|
res_clone = residual.clone()
|
|
_bridge.fused_add_rms_norm(inp_clone, res_clone, weight, eps)
|
|
output.copy_(inp_clone)
|
|
residual_output.copy_(res_clone)
|
|
|
|
|
|
def rotary_embedding(positions: torch.Tensor, query: torch.Tensor,
|
|
key: torch.Tensor, head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
is_neox: bool = True) -> None:
|
|
"""Fused rotary position embedding (in-place on query and key).
|
|
|
|
Source: xllm/core/kernels/ilu/rope.cpp → infer::xllm_rotary_embedding
|
|
"""
|
|
_bridge.rotary_embedding(positions, query, key, head_size,
|
|
cos_sin_cache, is_neox)
|
|
|
|
|
|
def reshape_and_cache(key: torch.Tensor, value: torch.Tensor,
|
|
key_cache: torch.Tensor, value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor) -> None:
|
|
"""Write KV to paged cache.
|
|
|
|
Source: xllm/core/kernels/ilu/attention.cpp → infer::xllm_reshape_and_cache
|
|
"""
|
|
_bridge.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)
|
|
|
|
|
|
def paged_attention(output: torch.Tensor, query: torch.Tensor,
|
|
key_cache: torch.Tensor, value_cache: torch.Tensor,
|
|
num_kv_heads: int, scale: float,
|
|
block_tables: torch.Tensor, seq_lens: torch.Tensor,
|
|
block_size: int, max_context_len: int,
|
|
alibi_slopes: Optional[torch.Tensor] = None
|
|
) -> torch.Tensor:
|
|
"""Paged attention decode.
|
|
|
|
Source: xllm/core/kernels/ilu/attention.cpp → infer::xllm_paged_attention
|
|
"""
|
|
return _bridge.paged_attention(
|
|
output, query, key_cache, value_cache,
|
|
num_kv_heads, scale, block_tables, seq_lens,
|
|
block_size, max_context_len, alibi_slopes)
|
|
|
|
|
|
def flash_attn_prefill(query: torch.Tensor, key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor, output: torch.Tensor,
|
|
block_tables: torch.Tensor,
|
|
cu_seq_q: torch.Tensor, cu_seq_k: torch.Tensor,
|
|
max_query_len: int, max_seq_len: int,
|
|
scale: float, is_causal: bool = True,
|
|
window_left: int = -1,
|
|
window_right: int = -1) -> torch.Tensor:
|
|
"""Flash attention prefill with paged KV cache.
|
|
|
|
Source: xllm/core/kernels/ilu/attention.cpp →
|
|
infer::ixinfer_flash_attn_unpad_with_block_tables
|
|
"""
|
|
return _bridge.flash_attn_prefill(
|
|
query, key_cache, value_cache, output, block_tables,
|
|
cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
|
|
scale, is_causal, window_left, window_right)
|
|
|
|
|
|
def linear(input: torch.Tensor, weight: torch.Tensor,
|
|
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
|
|
"""GEMM via ixformer (auto-selects linear vs linear_ex).
|
|
|
|
Source: xllm/core/kernels/ilu/matmul.cpp → infer::ixformer_linear[_ex]
|
|
"""
|
|
return _bridge.linear(input, weight, bias)
|
|
|
|
|
|
# =========================================================================
|
|
# MoE ops — full 7-step pipeline
|
|
# Source: xllm/core/layers/ilu/fused_moe.cpp
|
|
# =========================================================================
|
|
def topk_softmax(gating_output: torch.Tensor, topk: int,
|
|
renormalize: bool = True):
|
|
"""Fused topk + softmax routing."""
|
|
return _bridge.topk_softmax(gating_output, topk, renormalize)
|
|
|
|
|
|
def moe_gen_idx(expert_id: torch.Tensor, expert_num: int):
|
|
"""Build expert permutation maps."""
|
|
return _bridge.moe_gen_idx(expert_id, expert_num)
|
|
|
|
|
|
def moe_expand_input(input: torch.Tensor, gather_index: torch.Tensor,
|
|
combine_idx: torch.Tensor, topk: int):
|
|
"""Expand input tokens by expert assignment."""
|
|
return _bridge.moe_expand_input(input, gather_index, combine_idx, topk)
|
|
|
|
|
|
def group_gemm(inputs: torch.Tensor, weights: torch.Tensor,
|
|
token_count: torch.Tensor, output_n: int):
|
|
"""Batched expert GEMM."""
|
|
return _bridge.group_gemm(inputs, weights, token_count, output_n)
|
|
|
|
|
|
def moe_combine_result(input: torch.Tensor, weight: torch.Tensor):
|
|
"""Weighted scatter-back of expert outputs."""
|
|
return _bridge.moe_combine_result(input, weight)
|
|
|
|
|
|
def fused_moe_forward(hidden_states: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
w13: torch.Tensor, w2: torch.Tensor,
|
|
topk: int, num_experts: int,
|
|
renormalize: bool = True) -> torch.Tensor:
|
|
"""Full fused MoE forward (7-step pipeline).
|
|
|
|
Source: xllm/core/layers/ilu/fused_moe.cpp → FusedMoEImpl::forward_experts
|
|
Pipeline: topk → gen_idx → expand → gemm1(w13) → silu → gemm2(w2) → combine
|
|
"""
|
|
return _bridge.fused_moe_forward(
|
|
hidden_states, router_logits, w13, w2,
|
|
topk, num_experts, renormalize)
|