dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat: ix_ops integration layer — wire ix_full_bridge.so into vllm hot path

Browse Source 
Architecture (CCCL dispatch pattern):
  base_image ixformer::infer → ix_full_bridge.so → ix_ops.py → vllm patches

New files:
  ex_engine/python/ix_ops.py          — Python API for all 14 ixformer::infer ops
  ex_engine/python/patch_vllm_ops.py  — monkey-patch vllm GemmaRMSNorm, SiluAndMul
  ex_engine/deploy_ix_bridge.sh       — build-time deployment script

Modified:
  qwen3_6_scripts/patch_ops.sh        — integrated ix_bridge deployment + startup hook

Call chain: DecoderLayer.forward → GemmaRMSNorm → ix_ops.fused_add_rms_norm
            → ixformer::infer::residual_rms_norm (fused C++ kernel)
This commit is contained in:
dylan
2026-08-15 06:15:17 +00:00
parent 109d29fa60
commit e8f0948fe1
4 changed files with 757 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

343
ex_engine/python/ix_ops.py Normal file
View File

@@ -0,0 +1,343 @@
"""
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
output = rms_norm(input + residual, weight, eps)
residual_output = input + residual
"""
_bridge.fused_add_rms_norm(input, residual, weight, output,
residual_output, eps)
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)

201
ex_engine/python/patch_vllm_ops.py Normal file
View File

@@ -0,0 +1,201 @@
"""
patch_vllm_ops.py — Wire ix_full_bridge C++ kernels into vllm's hot path.
Architecture (CCCL policy_selector pattern):
Base image provides fused C++ kernels in ixformer::infer namespace.
ix_full_bridge.so wraps these with pybind11.
This module monkey-patches vllm's Python operators to call the bridge
instead of PyTorch fallback code.
Problem statement (683 → 8000 gap):
vllm's _custom_ops.py fails to load on BI-V100 (no vllm C++ extensions).
Without patches, EVERY norm/activation/rope/cache/attention call goes
through pure PyTorch — multiple kernel launches per op instead of 1.
Sub168 (competitor): all ops fused via xllm C++ engine → 11.9 TPS
Sub655 (us without patches): Python fallback → 2.6 TPS
Solution:
Patch vllm's operator dispatch points so they call our bridge .so,
which links against the SAME ixformer .so files in the base image.
Patched modules and their vllm paths:
1. vllm.model_executor.layers.layernorm.GemmaRMSNorm
→ ix_ops.rms_norm / ix_ops.fused_add_rms_norm
2. vllm.model_executor.layers.activation.SiluAndMul
→ ix_ops.silu_and_mul
3. vllm._custom_ops (ops fallback registry)
→ ix_ops for all registered ops
Source mapping:
upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp → rms_norm patch
upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp → silu_and_mul patch
upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp → rotary_embedding patch
upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp → cache/attention patch
"""
import os
import sys
import logging
import torch
from typing import Optional, Tuple
logger = logging.getLogger("ex_engine.patch_vllm_ops")
_patched = False
def apply_all_patches() -> int:
"""Apply all available patches. Returns count of patches applied."""
global _patched
if _patched:
return 0
_patched = True
from ex_engine.python import ix_ops
if not ix_ops.is_available():
logger.warning("ix_ops bridge not available — no patches applied")
return 0
n = 0
n += _patch_layernorm()
n += _patch_silu_and_mul()
n += _patch_custom_ops()
logger.info("patch_vllm_ops: %d patches applied", n)
return n
# =========================================================================
# Patch 1: GemmaRMSNorm → fused C++ kernel
# =========================================================================
def _patch_layernorm() -> int:
"""Replace GemmaRMSNorm.forward with ix_ops.rms_norm."""
from ex_engine.python import ix_ops
if not ix_ops.has_rms_norm():
logger.debug("ix_ops missing rms_norm, skip layernorm patch")
return 0
try:
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
except ImportError:
logger.debug("Cannot import GemmaRMSNorm, skip")
return 0
_orig_forward = GemmaRMSNorm.forward
def _patched_forward(self, x, residual=None):
if residual is not None:
# fused_add_rms_norm: norm(x + residual) → (normed, new_residual)
if ix_ops.has_fused_add_rms_norm():
out = torch.empty_like(x)
residual_out = torch.empty_like(x)
ix_ops.fused_add_rms_norm(
x, residual, self.weight, out, residual_out,
self.variance_epsilon)
return out, residual_out
else:
# Two-step fallback using just rms_norm
new_residual = x + residual
out = torch.empty_like(x)
ix_ops.rms_norm(out, new_residual, self.weight,
self.variance_epsilon)
return out, new_residual
else:
out = torch.empty_like(x)
ix_ops.rms_norm(out, x, self.weight, self.variance_epsilon)
return out
GemmaRMSNorm.forward = _patched_forward
logger.info("PATCHED: GemmaRMSNorm.forward → ix_ops.rms_norm")
return 1
# =========================================================================
# Patch 2: SiluAndMul → fused C++ kernel
# =========================================================================
def _patch_silu_and_mul() -> int:
"""Replace SiluAndMul.forward with ix_ops.silu_and_mul."""
from ex_engine.python import ix_ops
if not ix_ops.has_silu_and_mul():
logger.debug("ix_ops missing silu_and_mul, skip activation patch")
return 0
try:
from vllm.model_executor.layers.activation import SiluAndMul
except ImportError:
logger.debug("Cannot import SiluAndMul, skip")
return 0
def _patched_forward(self, x):
return ix_ops.silu_and_mul(x)
SiluAndMul.forward = _patched_forward
logger.info("PATCHED: SiluAndMul.forward → ix_ops.silu_and_mul")
return 1
# =========================================================================
# Patch 3: _custom_ops fallback registry
# =========================================================================
def _patch_custom_ops() -> int:
"""Patch vllm's _custom_ops to use ix_ops for registered ops."""
from ex_engine.python import ix_ops
count = 0
try:
import vllm._custom_ops as ops
except ImportError:
logger.debug("Cannot import vllm._custom_ops, skip")
return 0
# Patch silu_and_mul
if ix_ops.has_silu_and_mul() and hasattr(ops, 'silu_and_mul'):
def _silu_and_mul(out, x):
result = ix_ops.silu_and_mul(x)
out.copy_(result)
ops.silu_and_mul = _silu_and_mul
count += 1
logger.info("PATCHED: _custom_ops.silu_and_mul → ix_ops")
# Patch rms_norm
if ix_ops.has_rms_norm() and hasattr(ops, 'rms_norm'):
def _rms_norm(out, input, weight, eps):
ix_ops.rms_norm(out, input, weight, eps)
ops.rms_norm = _rms_norm
count += 1
logger.info("PATCHED: _custom_ops.rms_norm → ix_ops")
# Patch fused_add_rms_norm
if ix_ops.has_fused_add_rms_norm() and hasattr(ops, 'fused_add_rms_norm'):
def _fused_add_rms_norm(input, residual, weight, eps):
out = torch.empty_like(input)
residual_out = torch.empty_like(input)
ix_ops.fused_add_rms_norm(input, residual, weight,
out, residual_out, eps)
input.copy_(out)
residual.copy_(residual_out)
ops.fused_add_rms_norm = _fused_add_rms_norm
count += 1
logger.info("PATCHED: _custom_ops.fused_add_rms_norm → ix_ops")
# Patch rotary_embedding
if ix_ops.has_rotary_embedding() and hasattr(ops, 'rotary_embedding'):
def _rotary_embedding(positions, query, key, head_size,
cos_sin_cache, is_neox):
ix_ops.rotary_embedding(positions, query, key, head_size,
cos_sin_cache, is_neox)
ops.rotary_embedding = _rotary_embedding
count += 1
logger.info("PATCHED: _custom_ops.rotary_embedding → ix_ops")
return count
# =========================================================================
# Auto-apply on import if requested
# =========================================================================
if os.environ.get("IX_OPS_AUTO_PATCH", "0") == "1":
try:
apply_all_patches()
except Exception as e:
logger.warning("ix_ops auto-patch failed: %s", e)