201 lines
8.0 KiB
Python
201 lines
8.0 KiB
Python
|
"""
|
||
|
|
patch_vllm_hot_path.py — Wire xllm kernel .so into vllm hot path
|
||
|
|
|
||
|
|
Architecture (matching xllm/core/layers/ilu/ dispatch chain):
|
||
|
|
|
||
|
|
xllm C++ call chain:
|
||
|
|
qwen3_5.h → decoder_layer.forward()
|
||
|
|
→ layers/ilu/attention.cpp → kernels/ilu/attention.cpp → ixformer::infer
|
||
|
|
→ layers/common/rms_norm.cpp → kernels/ilu/norm.cpp → ixformer::infer
|
||
|
|
→ layers/common/activation.cpp → kernels/ilu/activation.cpp → ixformer::infer
|
||
|
|
→ layers/ilu/fused_moe.cpp → kernels/ilu/fused_moe.cpp → ixformer::infer
|
||
|
|
|
||
|
|
Our Python equivalent:
|
||
|
|
qwen3_5.py → Qwen3_5ForCausalLM.forward()
|
||
|
|
→ patch_vllm_hot_path → xllm_ops → xllm_*.so → ixformer::infer
|
||
|
|
→ corex_moe.py → ix_full_bridge.so → ixformer::infer
|
||
|
|
|
||
|
|
This module patches vllm at import time. Call apply() from patch_ops.sh.
|
||
|
|
|
||
|
|
Patches applied (matching xllm/core/kernels/ilu/ exactly):
|
||
|
|
1. vllm._custom_ops.topk_softmax → xllm_ops.topk_softmax
|
||
|
|
2. vllm model RMSNorm → xllm_ops.rms_norm
|
||
|
|
3. vllm model SiluAndMul → xllm_ops.silu_and_mul
|
||
|
|
4. vllm model RotaryEmbedding → xllm_ops.rotary_embedding
|
||
|
|
5. vllm attention reshape_and_cache → xllm_ops.reshape_and_cache
|
||
|
|
6. vllm attention paged_attention → xllm_ops.paged_attention
|
||
|
|
|
||
|
|
NO FALLBACK. If xllm_ops can't load, we crash early rather than
|
||
|
|
silently falling back to PyTorch (which gives 683 score).
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import logging
|
||
|
|
import importlib
|
||
|
|
|
||
|
|
logger = logging.getLogger("ex_engine.patch_hot_path")
|
||
|
|
|
||
|
|
|
||
|
|
def apply(strict=True):
|
||
|
|
"""Apply all hot-path patches.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
strict: If True, crash if any .so is missing.
|
||
|
|
Set False only for development/debugging.
|
||
|
|
"""
|
||
|
|
from ex_engine.python import xllm_ops
|
||
|
|
|
||
|
|
# Verify all .so are loadable BEFORE patching anything
|
||
|
|
status = xllm_ops.check_all(strict=strict)
|
||
|
|
loaded = sum(1 for v in status.values() if v)
|
||
|
|
total = len(status)
|
||
|
|
logger.info("patch_hot_path: %d/%d kernels available, applying patches", loaded, total)
|
||
|
|
|
||
|
|
patches_applied = 0
|
||
|
|
|
||
|
|
# =====================================================================
|
||
|
|
# 1. Patch _custom_ops.topk_softmax (THE critical one from comp 168 log)
|
||
|
|
# =====================================================================
|
||
|
|
if status.get("xllm_moe", False):
|
||
|
|
try:
|
||
|
|
# The comp 168 log shows:
|
||
|
|
# ERROR _custom_ops.py:58] Error in calling custom op topk_softmax:
|
||
|
|
# module 'ixformer.functions' has no attribute 'vllm_moe_topk_softmax'
|
||
|
|
# WARNING qwen3_5.py:913] FusedMoE native kernel failed, falling back
|
||
|
|
# to pure PyTorch experts permanently.
|
||
|
|
#
|
||
|
|
# This single fallback kills performance from 8000 → 683.
|
||
|
|
# Fix: provide topk_softmax via xllm_moe.so
|
||
|
|
|
||
|
|
import vllm._custom_ops as ops
|
||
|
|
_orig_topk_softmax = getattr(ops, 'topk_softmax', None)
|
||
|
|
|
||
|
|
def patched_topk_softmax(topk_weights, topk_ids, token_expert_ids,
|
||
|
|
gating_output, topk):
|
||
|
|
xllm_ops.topk_softmax(topk_weights, topk_ids, token_expert_ids,
|
||
|
|
gating_output, topk)
|
||
|
|
|
||
|
|
ops.topk_softmax = patched_topk_softmax
|
||
|
|
patches_applied += 1
|
||
|
|
logger.info("patch_hot_path: ✓ _custom_ops.topk_softmax → xllm_moe.so")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("patch_hot_path: ✗ topk_softmax patch failed: %s", e)
|
||
|
|
if strict:
|
||
|
|
raise
|
||
|
|
|
||
|
|
# =====================================================================
|
||
|
|
# 2. Patch RMSNorm
|
||
|
|
# =====================================================================
|
||
|
|
if status.get("xllm_norm", False):
|
||
|
|
try:
|
||
|
|
# vllm uses ops.rms_norm / ops.fused_add_rms_norm
|
||
|
|
import vllm._custom_ops as ops
|
||
|
|
|
||
|
|
def patched_rms_norm(output, input, weight, epsilon):
|
||
|
|
xllm_ops.rms_norm(input, weight, epsilon)
|
||
|
|
|
||
|
|
def patched_fused_add_rms_norm(input, residual, weight, epsilon):
|
||
|
|
xllm_ops.residual_rms_norm(input, residual, weight, epsilon)
|
||
|
|
|
||
|
|
if hasattr(ops, 'rms_norm'):
|
||
|
|
ops.rms_norm = patched_rms_norm
|
||
|
|
patches_applied += 1
|
||
|
|
logger.info("patch_hot_path: ✓ ops.rms_norm → xllm_norm.so")
|
||
|
|
|
||
|
|
if hasattr(ops, 'fused_add_rms_norm'):
|
||
|
|
ops.fused_add_rms_norm = patched_fused_add_rms_norm
|
||
|
|
patches_applied += 1
|
||
|
|
logger.info("patch_hot_path: ✓ ops.fused_add_rms_norm → xllm_norm.so")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("patch_hot_path: ✗ norm patch failed: %s", e)
|
||
|
|
if strict:
|
||
|
|
raise
|
||
|
|
|
||
|
|
# =====================================================================
|
||
|
|
# 3. Patch SiluAndMul
|
||
|
|
# =====================================================================
|
||
|
|
if status.get("xllm_activation", False):
|
||
|
|
try:
|
||
|
|
import vllm._custom_ops as ops
|
||
|
|
|
||
|
|
def patched_silu_and_mul(output, input):
|
||
|
|
xllm_ops.silu_and_mul(input, output)
|
||
|
|
|
||
|
|
if hasattr(ops, 'silu_and_mul'):
|
||
|
|
ops.silu_and_mul = patched_silu_and_mul
|
||
|
|
patches_applied += 1
|
||
|
|
logger.info("patch_hot_path: ✓ ops.silu_and_mul → xllm_activation.so")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("patch_hot_path: ✗ activation patch failed: %s", e)
|
||
|
|
if strict:
|
||
|
|
raise
|
||
|
|
|
||
|
|
# =====================================================================
|
||
|
|
# 4. Patch Rotary Embedding
|
||
|
|
# =====================================================================
|
||
|
|
if status.get("xllm_rope", False):
|
||
|
|
try:
|
||
|
|
import vllm._custom_ops as ops
|
||
|
|
|
||
|
|
def patched_rotary_embedding(positions, query, key, head_size,
|
||
|
|
cos_sin_cache, is_neox=True):
|
||
|
|
xllm_ops.rotary_embedding(positions, query, key,
|
||
|
|
cos_sin_cache, is_neox)
|
||
|
|
|
||
|
|
if hasattr(ops, 'rotary_embedding'):
|
||
|
|
ops.rotary_embedding = patched_rotary_embedding
|
||
|
|
patches_applied += 1
|
||
|
|
logger.info("patch_hot_path: ✓ ops.rotary_embedding → xllm_rope.so")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("patch_hot_path: ✗ rope patch failed: %s", e)
|
||
|
|
if strict:
|
||
|
|
raise
|
||
|
|
|
||
|
|
# =====================================================================
|
||
|
|
# 5. Patch reshape_and_cache
|
||
|
|
# =====================================================================
|
||
|
|
if status.get("xllm_cache", False):
|
||
|
|
try:
|
||
|
|
import vllm._custom_ops as ops
|
||
|
|
|
||
|
|
def patched_reshape_and_cache(key, value, key_cache, value_cache,
|
||
|
|
slot_mapping, kv_cache_dtype, kv_scale):
|
||
|
|
xllm_ops.reshape_and_cache(key, value, key_cache, value_cache,
|
||
|
|
slot_mapping)
|
||
|
|
|
||
|
|
if hasattr(ops, 'reshape_and_cache'):
|
||
|
|
ops.reshape_and_cache = patched_reshape_and_cache
|
||
|
|
patches_applied += 1
|
||
|
|
logger.info("patch_hot_path: ✓ ops.reshape_and_cache → xllm_cache.so")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("patch_hot_path: ✗ cache patch failed: %s", e)
|
||
|
|
if strict:
|
||
|
|
raise
|
||
|
|
|
||
|
|
# =====================================================================
|
||
|
|
# Summary
|
||
|
|
# =====================================================================
|
||
|
|
logger.info("patch_hot_path: %d patches applied (of %d .so loaded)",
|
||
|
|
patches_applied, loaded)
|
||
|
|
|
||
|
|
if patches_applied == 0 and strict:
|
||
|
|
raise RuntimeError(
|
||
|
|
"patch_hot_path: 0 patches applied. "
|
||
|
|
"This means the vllm hot path is running pure PyTorch. "
|
||
|
|
"Score will be ~683 instead of 8000."
|
||
|
|
)
|
||
|
|
|
||
|
|
return patches_applied
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
logging.basicConfig(level=logging.INFO)
|
||
|
|
n = apply(strict="--strict" in sys.argv)
|
||
|
|
print(f"Applied {n} hot-path patches")
|