d21b2505bb
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
96 lines
3.8 KiB
C++
96 lines
3.8 KiB
C++
// 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);
|
|
}
|