diff --git a/ex_engine/csrc/ix_moe_bridge.cpp b/ex_engine/csrc/ix_moe_bridge.cpp new file mode 100644 index 00000000..834fc10c --- /dev/null +++ b/ex_engine/csrc/ix_moe_bridge.cpp @@ -0,0 +1,95 @@ +// ix_moe_bridge.cpp — Bridge to ixformer C++ topk_softmax +// +// Problem: ixformer Python (ixformer.functions) lacks vllm_moe_topk_softmax +// Solution: Call ixformer::infer::topk_softmax() directly via C++ torch extension +// +// Source: upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h declares: +// void topk_softmax(torch::Tensor&, torch::Tensor&, torch::Tensor&, +// torch::Tensor&, bool); +// Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp shows usage: +// infer::topk_softmax(reduce_weight, topk_indices, token_expert_indices, input_, false); + +#include + +// Forward-declare ixformer C++ API (from ixformer.h in base image SDK) +namespace ixformer { +namespace infer { +void topk_softmax(torch::Tensor& topk_weights, + torch::Tensor& topk_indices, + torch::Tensor& token_expert_indices, + torch::Tensor& gating_output, + bool renormalize); + +void moe_compute_token_index_api( + torch::Tensor& topk_ids, + torch::Tensor& src_dst, + torch::Tensor& dst_src, + torch::Tensor& expert_sizes_gpu, + const c10::optional& expert_mask, + const c10::optional& expert_sizes_cpu, + const c10::optional& expand_tokens_gpu, + int64_t start_expert_id, + int64_t end_expert_id, + int64_t num_experts); + +void moe_expand_input(torch::Tensor outputs, + torch::Tensor inputs, + torch::Tensor dst_to_src, + const c10::optional& src_to_dst, + int64_t dst_tokens, + int64_t expand_factor); + +void moe_w16a16_group_gemm(torch::Tensor output, + torch::Tensor inputs, + torch::Tensor weights, + torch::Tensor tokens_per_experts, + const c10::optional& dst_to_src, + const c10::optional& bias, + std::string format, + int64_t persistent, + int64_t output_n); + +void moe_output_reduce_sum(torch::Tensor outputs, + torch::Tensor inputs, + const c10::optional& mul_weight, + const c10::optional& mask, + const c10::optional& extra_residual, + double scaling_factor); + +void silu_and_mul(torch::Tensor& input, torch::Tensor& output); + +} // namespace infer +} // namespace ixformer + +// Python-callable wrappers +std::tuple ix_moe_topk_softmax( + torch::Tensor gating_output, // (num_tokens, num_experts) float32 + int64_t topk, + bool renormalize) { + auto input = gating_output.to(torch::kFloat32).contiguous(); + int64_t num_tokens = input.size(0); + + auto topk_weights = torch::empty({num_tokens, topk}, + torch::dtype(torch::kFloat32).device(input.device())); + auto topk_indices = torch::empty({num_tokens, topk}, + torch::dtype(torch::kInt32).device(input.device())); + auto token_expert_indices = torch::empty({num_tokens, topk}, + torch::dtype(torch::kInt32).device(input.device())); + + ixformer::infer::topk_softmax( + topk_weights, topk_indices, token_expert_indices, input, renormalize); + + // Renormalize if not done by kernel (match xllm behavior) + if (!renormalize) { + auto row_sum = topk_weights.sum(-1, /*keepdim=*/true); + topk_weights = topk_weights / row_sum; + } + + return std::make_tuple(topk_weights, topk_indices); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("topk_softmax", &ix_moe_topk_softmax, + "Fused topk+softmax via ixformer C++ API (bypasses missing Python binding)", + py::arg("gating_output"), py::arg("topk"), py::arg("renormalize") = true); +} diff --git a/ex_engine/python/ix_bridge.py b/ex_engine/python/ix_bridge.py new file mode 100644 index 00000000..f39c06c2 --- /dev/null +++ b/ex_engine/python/ix_bridge.py @@ -0,0 +1,77 @@ +""" +ix_bridge.py — Load ix_moe_bridge C++ extension at runtime. + +Calls ixformer::infer::topk_softmax() via C++ torch extension, +bypassing the missing Python binding in ixformer.functions. + +Build: JIT-compiled on first import via torch.utils.cpp_extension.load() + (same mechanism as flash_qla_sm70 GDN kernel — proven to work on BI-V100) +""" + +import os +import logging +import torch + +logger = logging.getLogger("ex_engine.ix_bridge") + +_ix_bridge = None +_ix_bridge_available = False + +def _load_bridge(): + """JIT-compile and load ix_moe_bridge.so""" + global _ix_bridge, _ix_bridge_available + if _ix_bridge is not None: + return _ix_bridge_available + + csrc_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "csrc") + cpp_file = os.path.join(csrc_dir, "ix_moe_bridge.cpp") + + if not os.path.exists(cpp_file): + # Try deployed path (inside vllm model dir) + alt_dir = os.path.dirname(os.path.abspath(__file__)) + cpp_file = os.path.join(alt_dir, "ix_moe_bridge.cpp") + + if not os.path.exists(cpp_file): + logger.warning("ix_moe_bridge.cpp not found at %s", cpp_file) + _ix_bridge_available = False + return False + + try: + from torch.utils.cpp_extension import load + logger.info("JIT-compiling ix_moe_bridge.cpp ...") + _ix_bridge = load( + name="ix_moe_bridge", + sources=[cpp_file], + extra_cflags=["-O2"], + verbose=False, + ) + _ix_bridge_available = True + logger.info("ix_moe_bridge loaded successfully: %s", dir(_ix_bridge)) + return True + except Exception as e: + logger.warning("ix_moe_bridge JIT compile failed: %s", e) + _ix_bridge_available = False + return False + +def topk_softmax(gating_output: torch.Tensor, topk: int, renormalize: bool = True): + """ + Fused topk+softmax via ixformer C++ API. + + Args: + gating_output: (num_tokens, num_experts) router logits + topk: number of experts to select + renormalize: whether to renormalize weights + + Returns: + (topk_weights, topk_indices) — both (num_tokens, topk) + """ + if not _ix_bridge_available: + if not _load_bridge(): + # Fallback to pure PyTorch + probs = torch.softmax(gating_output.float(), dim=-1) + topk_w, topk_ids = torch.topk(probs, topk, dim=-1) + if renormalize: + topk_w = topk_w / topk_w.sum(dim=-1, keepdim=True) + return topk_w, topk_ids.to(torch.int32) + + return _ix_bridge.topk_softmax(gating_output, topk, renormalize) diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 3b0d0659..95a27715 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -178,23 +178,33 @@ if [ -n "$VLLM2" ]; then cp ./chat_utils.py "$VLLM2/entrypoints/chat_utils.py" 2>/dev/null || true fi -# Deploy EX Engine Python module into vllm importable path +# Deploy EX Engine Python module + C++ bridge into vllm importable path EX_ENGINE_SRC="/workspace/ex_engine" -if [ -d "$EX_ENGINE_SRC/python" ] && [ -d "$EX_ENGINE_SRC/build" ]; then +if [ -d "$EX_ENGINE_SRC/python" ]; then + # Deploy into vllm's model dir so qwen3_5.py can import it EX_DST="$VLLM/model_executor/models/ex_engine" - mkdir -p "$EX_DST" - cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST/" 2>/dev/null || true - # Copy built .so files - cp "$EX_ENGINE_SRC/build/"*.so "$EX_DST/" 2>/dev/null || true - echo "[patch_ops] EX Engine deployed: $(ls $EX_DST/*.so 2>/dev/null | wc -l) factors" + mkdir -p "$EX_DST/python" + mkdir -p "$EX_DST/csrc" + cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST/python/" 2>/dev/null || true + # ix_moe_bridge.cpp needs to be next to the python module for JIT compile + cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/csrc/" 2>/dev/null || true + cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/python/" 2>/dev/null || true + # Also make ex_engine importable from Python path + touch "$EX_DST/__init__.py" + touch "$EX_DST/python/__init__.py" + # Copy built .so files if they exist + if [ -d "$EX_ENGINE_SRC/build" ]; then + cp "$EX_ENGINE_SRC/build/"*.so "$EX_DST/" 2>/dev/null || true + fi + echo "[patch_ops] EX Engine deployed to $EX_DST" + ls -la "$EX_DST/csrc/" 2>/dev/null || true if [ -n "$VLLM2" ]; then EX_DST2="$VLLM2/model_executor/models/ex_engine" - mkdir -p "$EX_DST2" - cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST2/" 2>/dev/null || true - cp "$EX_ENGINE_SRC/build/"*.so "$EX_DST2/" 2>/dev/null || true + mkdir -p "$EX_DST2/python" "$EX_DST2/csrc" + cp -r "$EX_DST/"* "$EX_DST2/" 2>/dev/null || true fi else - echo "[patch_ops] WARNING: EX Engine not built — MoE will use slow PyTorch fallback" + echo "[patch_ops] WARNING: EX Engine not found — MoE uses slow PyTorch fallback" fi echo "[patch_ops] DONE — EX Engine + SM70 GDN kernel + serving layer deployed" diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index b156ee4c..39abe28f 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -71,8 +71,25 @@ except ImportError: # corex_gdn/corex_moe: these are custom modules that teams package into their # Docker image. If present, they provide fused GDN/MoE kernels. -_corex_gdn_module = None -_corex_moe_module = None +# ix_bridge: C++ bridge to ixformer::infer::topk_softmax (bypasses missing Python binding) +_ix_bridge_module = None +_ix_bridge_available = False +try: + from ex_engine.python.ix_bridge import topk_softmax as _ix_topk_softmax + _ix_bridge_available = True + logger.info("ix_bridge: ixformer C++ topk_softmax available") +except ImportError: + try: + # Try deployed path inside vllm models dir + import importlib, sys + _ex_dir = os.path.join(os.path.dirname(__file__), "ex_engine") + if os.path.isdir(_ex_dir) and _ex_dir not in sys.path: + sys.path.insert(0, os.path.dirname(_ex_dir)) + from ex_engine.python.ix_bridge import topk_softmax as _ix_topk_softmax + _ix_bridge_available = True + logger.info("ix_bridge: ixformer C++ topk_softmax available (deployed path)") + except ImportError: + logger.info("ix_bridge: not available, MoE uses PyTorch topk") _corex_gdn_available = False _corex_moe_available = False @@ -469,17 +486,13 @@ class GatedDeltaNet(nn.Module): "CoreX GDN forward failed (%s), falling back", e) self._use_corex_gdn = False # permanent fallback - # FlashQLA SM70 dispatch: fused CUDA kernel for prefill - # Decode stays PyTorch (SM70 decode kernel needs different state layout) - if _flash_qla_available and attn_metadata.num_prefill_tokens > 0: - try: - return self._flash_qla_prefill( - hidden_states, attn_metadata, conv_state, temporal_state) - except Exception as e: - if self.layer_idx == 0: - logger.warning( - "FlashQLA SM70 prefill failed (%s), falling back to PyTorch", e) - # Don't disable permanently — may work for different shapes + # flash_qla SM70 DISABLED: produces inf on BI-V100 (abs mean=inf from real test) + # xllm uses equivalent PyTorch chunked path (qwen3_gated_delta_net_base.cpp) + # which works correctly in fp32. Keeping PyTorch path only. + # + # if _flash_qla_available and attn_metadata.num_prefill_tokens > 0: + # try: + # return self._flash_qla_prefill(...) return self._pytorch_forward( hidden_states, attn_metadata, conv_state, temporal_state) @@ -1048,12 +1061,18 @@ class Qwen3_5MoeSparseBlock(nn.Module): Output is partial (pre-all-reduce), same contract as FusedMoE with reduce_results=False. """ - # Routing: softmax → topk → renormalise - routing_weights = _ix_softmax(router_logits.float(), dim=-1) - topk_weights, topk_ids = torch.topk( - routing_weights, self.top_k, dim=-1) # (T, top_k) - topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) - topk_weights = topk_weights.to(hidden_states.dtype) + # Routing: fused topk+softmax via ixformer C++ bridge (if available) + # Falls back to PyTorch softmax → topk → renormalize + if _ix_bridge_available: + topk_weights, topk_ids = _ix_topk_softmax( + router_logits, self.top_k, renormalize=True) + topk_weights = topk_weights.to(hidden_states.dtype) + else: + routing_weights = _ix_softmax(router_logits.float(), dim=-1) + topk_weights, topk_ids = torch.topk( + routing_weights, self.top_k, dim=-1) # (T, top_k) + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + topk_weights = topk_weights.to(hidden_states.dtype) w13 = self.experts.w13_weight # (E, 2*I, H) w2 = self.experts.w2_weight # (E, H, I)