fix: wire MoE topk via ixformer C++ bridge + disable broken flash_qla GDN
Two call chain breaks fixed: 1. MoE routing (2304 calls/token): BEFORE: torch.softmax + torch.topk (3 Python GPU ops, no ixformer) AFTER: ix_bridge.py → ix_moe_bridge.cpp → ixformer::infer::topk_softmax() Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp line 46 The C++ API exists in base image SDK (ixformer.h declares it), only the Python binding (ixformer.functions) was missing. 2. GDN prefill (4 layers, 99.98% NaN): BEFORE: flash_qla SM70 kernel → abs mean=inf → nan_to_num → zeros AFTER: skip flash_qla, use _pytorch_forward directly Source: upstream_ref/xllm qwen3_gated_delta_net_base.cpp uses identical PyTorch chunked logic (no flash_qla). Sub168 (working build) never deployed flash_qla either. Files: - ex_engine/csrc/ix_moe_bridge.cpp: torch C++ extension calling ixformer C++ API - ex_engine/python/ix_bridge.py: JIT-compile loader with PyTorch fallback - qwen3_5.py: import ix_bridge for MoE, disable flash_qla for GDN - patch_ops.sh: deploy ix_bridge .cpp + .py into vllm model dir
This commit is contained in:
EX Engine
95
ex_engine/csrc/ix_moe_bridge.cpp
Normal file
95
ex_engine/csrc/ix_moe_bridge.cpp
Normal file
@@ -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 <torch/extension.h>
|
||||
|
||||
// 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<torch::Tensor>& expert_mask,
|
||||
const c10::optional<torch::Tensor>& expert_sizes_cpu,
|
||||
const c10::optional<torch::Tensor>& 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<torch::Tensor>& 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<torch::Tensor>& dst_to_src,
|
||||
const c10::optional<torch::Tensor>& 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<torch::Tensor>& mul_weight,
|
||||
const c10::optional<torch::Tensor>& mask,
|
||||
const c10::optional<torch::Tensor>& extra_residual,
|
||||
double scaling_factor);
|
||||
|
||||
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
|
||||
|
||||
} // namespace infer
|
||||
} // namespace ixformer
|
||||
|
||||
// Python-callable wrappers
|
||||
std::tuple<torch::Tensor, torch::Tensor> 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);
|
||||
}
|
||||
77
ex_engine/python/ix_bridge.py
Normal file
77
ex_engine/python/ix_bridge.py
Normal file
@@ -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)
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user