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

feat: xllm_ops NO-FALLBACK kernel loader + 6 missing .so build targets + hot-path patcher

Browse Source 
Build infrastructure:
- build_xllm_kernels.sh: add 5 missing build targets (norm, rope, activation, cache, moe)
  Previously only built xllm_fused_qknorm_rope.so, now builds all 6 .so files

Kernel loader (xllm_ops.py):
- NO-FALLBACK architecture matching xllm/core/kernels/ilu/ dispatch chain
- Loads: xllm_norm.so, xllm_rope.so, xllm_activation.so, xllm_cache.so,
         xllm_moe.so, ix_full_bridge.so, xllm_fused_qknorm_rope.so
- check_all(strict=True) verifies ALL .so at startup

Hot-path patcher (patch_vllm_hot_path.py):
- Monkey-patches vllm._custom_ops to route through xllm .so
- Critical fix: topk_softmax patch prevents comp 168 cascade
- Patches: topk_softmax, rms_norm, silu_and_mul, rotary_embedding, reshape_and_cache

Source mapping: xllm/core/kernels/ilu/*.cpp -> our xllm_*.so files
This commit is contained in:
Claude
2026-08-15 14:13:13 +00:00
parent e873e5f27b
commit 52e2ef31a8
3 changed files with 495 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
ex_engine/build_xllm_kernels.sh
View File

@@ -98,6 +98,56 @@ echo "============================================================"
build_so "xllm_fused_qknorm_rope" \
"ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp"
# 2. xllm_norm — RMSNorm + Fused Add RMSNorm
# Source: upstream xllm norm.cu
# Hot path: called 2× per decoder layer = 72× per forward pass
echo ""
echo "============================================================"
echo " 2. xllm_norm.so"
echo "============================================================"
build_so "xllm_norm" \
"ex_engine/xllm_kernels/cuda/norm.cu ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp"
# 3. xllm_rope — Rotary Position Embedding
# Source: upstream xllm rope.cu
# Hot path: called 1× per attention layer = 36× per forward pass
echo ""
echo "============================================================"
echo " 3. xllm_rope.so"
echo "============================================================"
build_so "xllm_rope" \
"ex_engine/xllm_kernels/cuda/rope.cu ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp"
# 4. xllm_activation — SiLU-and-Mul fused activation
# Source: upstream xllm activation.cu
# Hot path: called 1× per MLP = 36× per forward pass
echo ""
echo "============================================================"
echo " 4. xllm_activation.so"
echo "============================================================"
build_so "xllm_activation" \
"ex_engine/xllm_kernels/cuda/activation.cu ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp"
# 5. xllm_cache — Reshape + block copy for KV cache
# Source: upstream xllm reshape_paged_cache.cu + block_copy.cu
# Hot path: called every prefill + decode step
echo ""
echo "============================================================"
echo " 5. xllm_cache.so"
echo "============================================================"
build_so "xllm_cache" \
"ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu ex_engine/xllm_kernels/cuda/block_copy.cu ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp"
# 6. xllm_moe — MoE topk + index + combine + fused pipeline
# Source: upstream xllm moe_fused_topk.cu + moe_compute_index.cu + moe_combine.cu + fused_moe.cpp
# THE critical .so: replaces Python for-loop over 64 experts
echo ""
echo "============================================================"
echo " 6. xllm_moe.so"
echo "============================================================"
build_so "xllm_moe" \
"ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu ex_engine/xllm_kernels/cuda/moe/moe_combine.cu ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp"
echo ""
echo "============================================================"
echo " Build complete. Output:"