diff --git a/ex_engine/python/xllm_ops.py b/ex_engine/python/xllm_ops.py index d8aa3d68..15217662 100644 --- a/ex_engine/python/xllm_ops.py +++ b/ex_engine/python/xllm_ops.py @@ -8,6 +8,13 @@ Function name mapping (verified via `nm -D` + `strings` on real BI-V100): xllm_moe.so: moe_compute_index (NOT moe_compute_token_index) ix_moe_bridge.so: ix_paged_attention, ix_linear (NOT in ix_full_bridge.so) +C++ argument order verified against *_bind.cpp pybind11 source: + xllm_norm_bind.cpp: rms_norm(output, input, weight, eps) + xllm_activation_bind.cpp: silu_and_mul(out, input) + xllm_cache_bind.cpp: reshape_paged_cache(slot_ids, keys, values, kc, vc) + ix_full_bridge_v2.cpp: ix_paged_attention(out, q, kc, vc, head_mapping, scale, ...) + xllm_moe_bind.cpp: moe_fused_topk(gating, topk) → returns (w, ids) + NO FALLBACK: If a .so fails to load, we raise immediately. """ @@ -15,6 +22,8 @@ import os import sys import importlib.util import logging + +import torch from typing import Optional, Dict, Any logger = logging.getLogger("ex_engine.xllm_ops") @@ -92,49 +101,78 @@ def _get(name: str) -> Any: # ========================================================================= -# Public API — matches xllm/core/kernels/ilu/ function signatures +# Public API — C++ signatures verified against *_bind.cpp pybind source # ========================================================================= -# --- Norm (xllm_norm.so: rms_norm, fused_add_rms_norm) --- +# --- Norm (xllm_norm.so) --- +# C++ rms_norm(output, input, weight, eps) — output FIRST def rms_norm(input, weight, epsilon): - """RMSNorm. Maps to ixformer::infer::rms_norm.""" - return _get("xllm_norm").rms_norm(input, weight, epsilon) + """RMSNorm. C++ takes (output, input, weight, eps).""" + output = torch.empty_like(input) + _get("xllm_norm").rms_norm(output, input, weight, epsilon) + return output +# C++ fused_add_rms_norm(input&, residual&, weight&, epsilon) — in-place def residual_rms_norm(input, residual, weight, epsilon): - """Fused residual + RMSNorm. .so export: fused_add_rms_norm.""" - return _get("xllm_norm").fused_add_rms_norm(input, residual, weight, epsilon) + """Fused residual + RMSNorm. Modifies input and residual in-place.""" + _get("xllm_norm").fused_add_rms_norm(input, residual, weight, epsilon) + return input, residual -# --- RoPE (xllm/core/kernels/ilu/rope.cpp) --- +# --- RoPE (xllm_rope.so) --- +# C++ rotary_embedding(positions, query, key, cos_sin_cache, is_neox) def rotary_embedding(positions, query, key, cos_sin_cache, is_neox=True): - """Fused rotary embedding. Maps to ixformer::infer::xllm_rotary_embedding.""" + """Fused rotary embedding. Signature matches C++ directly.""" return _get("xllm_rope").rotary_embedding(positions, query, key, cos_sin_cache, is_neox) -# --- Activation (xllm/core/kernels/ilu/activation.cpp) --- +# --- Activation (xllm_activation.so) --- +# C++ silu_and_mul(out, input) — out FIRST def silu_and_mul(input, output=None): - """Fused SiLU activation. Maps to ixformer::infer::silu_and_mul.""" - return _get("xllm_activation").silu_and_mul(input, output) + """Fused SiLU activation. C++ takes (out, input).""" + if output is None: + d = input.shape[-1] // 2 + output = torch.empty(*input.shape[:-1], d, dtype=input.dtype, + device=input.device) + _get("xllm_activation").silu_and_mul(output, input) + return output +# C++ gelu_and_mul(out, input) — out FIRST def gelu_and_mul(input, output=None): - """Fused GeLU activation.""" - return _get("xllm_activation").gelu_and_mul(input, output) + """Fused GeLU activation. C++ takes (out, input).""" + if output is None: + d = input.shape[-1] // 2 + output = torch.empty(*input.shape[:-1], d, dtype=input.dtype, + device=input.device) + _get("xllm_activation").gelu_and_mul(output, input) + return output -# --- Cache (xllm_cache.so: reshape_paged_cache) --- +# --- Cache (xllm_cache.so) --- +# C++ reshape_paged_cache(slot_ids, keys, values, key_cache, value_cache) +# — slot_ids FIRST (not last!) def reshape_and_cache(key, value, key_cache, value_cache, slot_mapping): - """Write KV to paged cache. .so export: reshape_paged_cache.""" + """Write KV to paged cache. C++ takes slot_ids as FIRST arg.""" return _get("xllm_cache").reshape_paged_cache(slot_mapping, key, value, key_cache, value_cache) -# --- Attention (ix_moe_bridge.so: ix_paged_attention) --- +# --- Attention (ix_moe_bridge.so) --- +# C++ ix_paged_attention(output, query, key_cache, value_cache, +# head_mapping, scale, block_tables, context_lens, +# block_size, max_context_len, num_kv_heads, +# alibi_slopes) def paged_attention(out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, context_lens, block_size, max_context_len, alibi_slopes=None): - """Paged attention decode. .so export: ix_paged_attention in ix_moe_bridge.so.""" + """Paged attention decode. C++ needs head_mapping tensor at position 5.""" bridge = _get("ix_moe_bridge") + num_q_heads = query.shape[1] + head_mapping = torch.arange(num_q_heads, dtype=torch.int32, + device=query.device) + if num_kv_heads != num_q_heads: + head_mapping = head_mapping // (num_q_heads // num_kv_heads) return bridge.paged_attention( out, query, key_cache, value_cache, - scale, block_tables, context_lens, - block_size, max_context_len, max_context_len, alibi_slopes + head_mapping, scale, block_tables, context_lens, + block_size, max_context_len, num_kv_heads, alibi_slopes ) def flash_attn_prefill(query, key_cache, value_cache, out, @@ -147,20 +185,29 @@ def flash_attn_prefill(query, key_cache, value_cache, out, block_tables, cu_seq_q, max_seq_q, scale ) -# --- MoE (xllm_moe.so: moe_fused_topk, moe_compute_index) --- +# --- MoE (xllm_moe.so) --- +# C++ moe_fused_topk(gating_output, topk, renormalize=true, +# correction_bias=None, scoring_func="softmax") +# → returns (topk_weights, topk_ids) (C++ allocates internally) def topk_softmax(topk_weights, topk_ids, token_expert_ids, gating_output, topk): - """MoE topk + softmax. .so export: moe_fused_topk.""" - return _get("xllm_moe").moe_fused_topk( - topk_weights, topk_ids, token_expert_ids, gating_output, topk - ) + """MoE topk+softmax. C++ returns new tensors; we copy into pre-allocated.""" + weights, ids = _get("xllm_moe").moe_fused_topk(gating_output, topk) + topk_weights.copy_(weights) + topk_ids.copy_(ids) + return topk_weights, topk_ids, token_expert_ids +# C++ moe_compute_index(expert_id, num_experts) +# → returns (sorted_token_ids, expert_ids, num_tokens_post_padded) def moe_compute_token_index(sorted_token_ids, expert_ids, num_tokens_post_padded, token_expert_ids, num_experts, block_size): - """MoE token routing. .so export: moe_compute_index.""" - return _get("xllm_moe").moe_compute_index( - sorted_token_ids, expert_ids, num_tokens_post_padded, - token_expert_ids, num_experts, block_size + """MoE token routing. C++ takes only (expert_id, num_experts).""" + s_ids, e_ids, n_post = _get("xllm_moe").moe_compute_index( + token_expert_ids, num_experts ) + sorted_token_ids.copy_(s_ids[:sorted_token_ids.numel()].reshape_as(sorted_token_ids)) + expert_ids.copy_(e_ids[:expert_ids.numel()].reshape_as(expert_ids)) + num_tokens_post_padded.copy_(n_post[:num_tokens_post_padded.numel()].reshape_as(num_tokens_post_padded)) + return sorted_token_ids, expert_ids, num_tokens_post_padded # --- Linear (ix_moe_bridge.so: ix_linear) --- def ixformer_linear(input, weight, act_type=0, bias=None, out=None): @@ -232,4 +279,4 @@ def check_all(strict=True): total = len(results) logger.info("xllm_ops: %d/%d .so loaded", loaded, total) - return results + return results \ No newline at end of file diff --git a/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py b/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py index 1829846c..43b72474 100644 --- a/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py +++ b/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py @@ -94,9 +94,13 @@ def apply(strict=True): import vllm._custom_ops as ops def patched_rms_norm(output, input, weight, epsilon): - xllm_ops.rms_norm(input, weight, epsilon) + # vllm: ops.rms_norm(output, input, weight, eps) + # xllm_norm.so: rms_norm(output, input, weight, eps) — same + xllm_ops._get("xllm_norm").rms_norm(output, input, weight, epsilon) def patched_fused_add_rms_norm(input, residual, weight, epsilon): + # vllm: ops.fused_add_rms_norm(input, residual, weight, eps) + # xllm_norm.so: same signature, in-place xllm_ops.residual_rms_norm(input, residual, weight, epsilon) if hasattr(ops, 'rms_norm'): @@ -122,7 +126,9 @@ def apply(strict=True): import vllm._custom_ops as ops def patched_silu_and_mul(output, input): - xllm_ops.silu_and_mul(input, output) + # vllm: ops.silu_and_mul(output, input) + # xllm_activation.so: silu_and_mul(out, input) — same order + xllm_ops._get("xllm_activation").silu_and_mul(output, input) if hasattr(ops, 'silu_and_mul'): ops.silu_and_mul = patched_silu_and_mul @@ -198,4 +204,4 @@ def apply(strict=True): if __name__ == "__main__": logging.basicConfig(level=logging.INFO) n = apply(strict="--strict" in sys.argv) - print(f"Applied {n} hot-path patches") + print(f"Applied {n} hot-path patches") \ No newline at end of file diff --git a/qwen3_6_scripts/ex_engine/python/xllm_ops.py b/qwen3_6_scripts/ex_engine/python/xllm_ops.py index d8aa3d68..15217662 100644 --- a/qwen3_6_scripts/ex_engine/python/xllm_ops.py +++ b/qwen3_6_scripts/ex_engine/python/xllm_ops.py @@ -8,6 +8,13 @@ Function name mapping (verified via `nm -D` + `strings` on real BI-V100): xllm_moe.so: moe_compute_index (NOT moe_compute_token_index) ix_moe_bridge.so: ix_paged_attention, ix_linear (NOT in ix_full_bridge.so) +C++ argument order verified against *_bind.cpp pybind11 source: + xllm_norm_bind.cpp: rms_norm(output, input, weight, eps) + xllm_activation_bind.cpp: silu_and_mul(out, input) + xllm_cache_bind.cpp: reshape_paged_cache(slot_ids, keys, values, kc, vc) + ix_full_bridge_v2.cpp: ix_paged_attention(out, q, kc, vc, head_mapping, scale, ...) + xllm_moe_bind.cpp: moe_fused_topk(gating, topk) → returns (w, ids) + NO FALLBACK: If a .so fails to load, we raise immediately. """ @@ -15,6 +22,8 @@ import os import sys import importlib.util import logging + +import torch from typing import Optional, Dict, Any logger = logging.getLogger("ex_engine.xllm_ops") @@ -92,49 +101,78 @@ def _get(name: str) -> Any: # ========================================================================= -# Public API — matches xllm/core/kernels/ilu/ function signatures +# Public API — C++ signatures verified against *_bind.cpp pybind source # ========================================================================= -# --- Norm (xllm_norm.so: rms_norm, fused_add_rms_norm) --- +# --- Norm (xllm_norm.so) --- +# C++ rms_norm(output, input, weight, eps) — output FIRST def rms_norm(input, weight, epsilon): - """RMSNorm. Maps to ixformer::infer::rms_norm.""" - return _get("xllm_norm").rms_norm(input, weight, epsilon) + """RMSNorm. C++ takes (output, input, weight, eps).""" + output = torch.empty_like(input) + _get("xllm_norm").rms_norm(output, input, weight, epsilon) + return output +# C++ fused_add_rms_norm(input&, residual&, weight&, epsilon) — in-place def residual_rms_norm(input, residual, weight, epsilon): - """Fused residual + RMSNorm. .so export: fused_add_rms_norm.""" - return _get("xllm_norm").fused_add_rms_norm(input, residual, weight, epsilon) + """Fused residual + RMSNorm. Modifies input and residual in-place.""" + _get("xllm_norm").fused_add_rms_norm(input, residual, weight, epsilon) + return input, residual -# --- RoPE (xllm/core/kernels/ilu/rope.cpp) --- +# --- RoPE (xllm_rope.so) --- +# C++ rotary_embedding(positions, query, key, cos_sin_cache, is_neox) def rotary_embedding(positions, query, key, cos_sin_cache, is_neox=True): - """Fused rotary embedding. Maps to ixformer::infer::xllm_rotary_embedding.""" + """Fused rotary embedding. Signature matches C++ directly.""" return _get("xllm_rope").rotary_embedding(positions, query, key, cos_sin_cache, is_neox) -# --- Activation (xllm/core/kernels/ilu/activation.cpp) --- +# --- Activation (xllm_activation.so) --- +# C++ silu_and_mul(out, input) — out FIRST def silu_and_mul(input, output=None): - """Fused SiLU activation. Maps to ixformer::infer::silu_and_mul.""" - return _get("xllm_activation").silu_and_mul(input, output) + """Fused SiLU activation. C++ takes (out, input).""" + if output is None: + d = input.shape[-1] // 2 + output = torch.empty(*input.shape[:-1], d, dtype=input.dtype, + device=input.device) + _get("xllm_activation").silu_and_mul(output, input) + return output +# C++ gelu_and_mul(out, input) — out FIRST def gelu_and_mul(input, output=None): - """Fused GeLU activation.""" - return _get("xllm_activation").gelu_and_mul(input, output) + """Fused GeLU activation. C++ takes (out, input).""" + if output is None: + d = input.shape[-1] // 2 + output = torch.empty(*input.shape[:-1], d, dtype=input.dtype, + device=input.device) + _get("xllm_activation").gelu_and_mul(output, input) + return output -# --- Cache (xllm_cache.so: reshape_paged_cache) --- +# --- Cache (xllm_cache.so) --- +# C++ reshape_paged_cache(slot_ids, keys, values, key_cache, value_cache) +# — slot_ids FIRST (not last!) def reshape_and_cache(key, value, key_cache, value_cache, slot_mapping): - """Write KV to paged cache. .so export: reshape_paged_cache.""" + """Write KV to paged cache. C++ takes slot_ids as FIRST arg.""" return _get("xllm_cache").reshape_paged_cache(slot_mapping, key, value, key_cache, value_cache) -# --- Attention (ix_moe_bridge.so: ix_paged_attention) --- +# --- Attention (ix_moe_bridge.so) --- +# C++ ix_paged_attention(output, query, key_cache, value_cache, +# head_mapping, scale, block_tables, context_lens, +# block_size, max_context_len, num_kv_heads, +# alibi_slopes) def paged_attention(out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, context_lens, block_size, max_context_len, alibi_slopes=None): - """Paged attention decode. .so export: ix_paged_attention in ix_moe_bridge.so.""" + """Paged attention decode. C++ needs head_mapping tensor at position 5.""" bridge = _get("ix_moe_bridge") + num_q_heads = query.shape[1] + head_mapping = torch.arange(num_q_heads, dtype=torch.int32, + device=query.device) + if num_kv_heads != num_q_heads: + head_mapping = head_mapping // (num_q_heads // num_kv_heads) return bridge.paged_attention( out, query, key_cache, value_cache, - scale, block_tables, context_lens, - block_size, max_context_len, max_context_len, alibi_slopes + head_mapping, scale, block_tables, context_lens, + block_size, max_context_len, num_kv_heads, alibi_slopes ) def flash_attn_prefill(query, key_cache, value_cache, out, @@ -147,20 +185,29 @@ def flash_attn_prefill(query, key_cache, value_cache, out, block_tables, cu_seq_q, max_seq_q, scale ) -# --- MoE (xllm_moe.so: moe_fused_topk, moe_compute_index) --- +# --- MoE (xllm_moe.so) --- +# C++ moe_fused_topk(gating_output, topk, renormalize=true, +# correction_bias=None, scoring_func="softmax") +# → returns (topk_weights, topk_ids) (C++ allocates internally) def topk_softmax(topk_weights, topk_ids, token_expert_ids, gating_output, topk): - """MoE topk + softmax. .so export: moe_fused_topk.""" - return _get("xllm_moe").moe_fused_topk( - topk_weights, topk_ids, token_expert_ids, gating_output, topk - ) + """MoE topk+softmax. C++ returns new tensors; we copy into pre-allocated.""" + weights, ids = _get("xllm_moe").moe_fused_topk(gating_output, topk) + topk_weights.copy_(weights) + topk_ids.copy_(ids) + return topk_weights, topk_ids, token_expert_ids +# C++ moe_compute_index(expert_id, num_experts) +# → returns (sorted_token_ids, expert_ids, num_tokens_post_padded) def moe_compute_token_index(sorted_token_ids, expert_ids, num_tokens_post_padded, token_expert_ids, num_experts, block_size): - """MoE token routing. .so export: moe_compute_index.""" - return _get("xllm_moe").moe_compute_index( - sorted_token_ids, expert_ids, num_tokens_post_padded, - token_expert_ids, num_experts, block_size + """MoE token routing. C++ takes only (expert_id, num_experts).""" + s_ids, e_ids, n_post = _get("xllm_moe").moe_compute_index( + token_expert_ids, num_experts ) + sorted_token_ids.copy_(s_ids[:sorted_token_ids.numel()].reshape_as(sorted_token_ids)) + expert_ids.copy_(e_ids[:expert_ids.numel()].reshape_as(expert_ids)) + num_tokens_post_padded.copy_(n_post[:num_tokens_post_padded.numel()].reshape_as(num_tokens_post_padded)) + return sorted_token_ids, expert_ids, num_tokens_post_padded # --- Linear (ix_moe_bridge.so: ix_linear) --- def ixformer_linear(input, weight, act_type=0, bias=None, out=None): @@ -232,4 +279,4 @@ def check_all(strict=True): total = len(results) logger.info("xllm_ops: %d/%d .so loaded", loaded, total) - return results + return results \ No newline at end of file