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

feat(EX): ix_full_bridge — all 14 ixformer::infer functions bridged

Browse Source 
Upstream source: xllm/core/kernels/ilu/ixformer.h (Apache 2.0)
Wrapper patterns: xllm/core/kernels/ilu/{attention,norm,rope,activation,fused_moe,group_gemm}.cpp

Complete bridge (ix_full_bridge.cpp, 331 lines):
  MoE:       topk_softmax, gen_idx, expand, group_gemm, silu_mul, combine, fused_forward
  Attention: paged_attention (decode), flash_attn_prefill (prefill)
  Norm:      rms_norm, fused_add_rms_norm
  RoPE:      rotary_embedding
  Cache:     reshape_and_cache
  Linear:    ixformer_linear

ix_bridge.py: tries ix_full_bridge first, falls back to ix_moe_bridge
patch_ops.sh: deploys both .cpp files to all JIT search paths
Copied ixformer.h + utils.h headers for reference
This commit is contained in:
EX Engine
2026-08-10 04:01:35 +00:00
parent 5efb0fcc35
commit f955dd127e
5 changed files with 664 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
ex_engine/csrc/ilu/ixformer.h Normal file
View File

@@ -0,0 +1,147 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <torch/all.h>
#include "ATen/Tensor.h"
#include "utils.h"
namespace ixformer::infer {
torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& out,
torch::Tensor& block_tables,
torch::Tensor& cu_seq_q,
torch::Tensor& cu_seq_k,
int64_t max_seq_q,
int64_t max_seq_k,
bool is_causal,
int64_t window_left,
int64_t window_right,
double scale,
double softcap,
bool sqrt_alibi,
const std::optional<torch::Tensor>& alibi_slopes,
const std::optional<torch::Tensor>& sinks,
std::optional<torch::Tensor>& lse);
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
torch::Tensor xllm_paged_attention(
torch::Tensor& out,
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
int64_t num_kv_heads,
double scale,
torch::Tensor& block_tables,
torch::Tensor& context_lens,
int64_t block_size,
int64_t max_context_len,
const std::optional<torch::Tensor>& alibi_slopes,
bool causal,
int32_t window_left,
int32_t window_right,
double softcap,
bool enable_cuda_graph,
bool use_sqrt_alibi,
const std::optional<torch::Tensor>& sinks);
torch::Tensor ixformer_linear(torch::Tensor& input,
torch::Tensor& weight,
int64_t act_type,
const std::optional<torch::Tensor>& bias,
const std::optional<torch::Tensor>& out,
const std::optional<bool> persistent);
torch::Tensor ixformer_linear_ex(torch::Tensor& input,
torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& out);
void xllm_reshape_and_cache(torch::Tensor& key,
torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
int64_t key_token_stride,
int64_t value_token_stride);
void xllm_rotary_embedding(torch::Tensor& positions,
torch::Tensor& query,
torch::Tensor& key,
int64_t head_size,
torch::Tensor& cos_sin_cache,
bool is_neox);
void residual_rms_norm(torch::Tensor& input,
torch::Tensor& residual,
torch::Tensor& weight,
torch::Tensor& output,
torch::Tensor& residual_output,
const std::optional<torch::Tensor>& fused_bias,
double alpha,
double eps,
bool is_post);
void rms_norm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& output,
const std::optional<torch::Tensor>& fused_bias,
double eps);
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);
} // namespace ixformer::infer

63
ex_engine/csrc/ilu/utils.h Normal file
View File

@@ -0,0 +1,63 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
namespace xllm::kernel::ilu {
#undef check_tensor_contiguous
#define check_tensor_contiguous(x, type) \
TORCH_CHECK(x.scalar_type() == type); \
TORCH_CHECK(x.is_cuda()); \
TORCH_CHECK(x.is_contiguous());
#undef check_tensor_half_bf_float
#define check_tensor_half_bf_float(x) \
TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || \
x.scalar_type() == at::ScalarType::Float || \
x.scalar_type() == at::ScalarType::BFloat16); \
TORCH_CHECK(x.is_cuda());
// from torchCheckMsgImpl
inline const char* ixformer_check_msg_impl(const char* msg) { return msg; }
// // If there is just 1 user-provided C-string argument, use it.
#define IXFORMER_CHECK_MSG(cond, type, ...) \
(ixformer_check_msg_impl( \
"Expected " #cond \
" to be true, but got false. " \
"(Could this error message be improved? If so, " \
"please report an enhancement request to ixformer.)", \
##__VA_ARGS__))
#define IXFORMER_CHECK(cond, ...) \
{ \
if (!(cond)) { \
std::cerr << __FILE__ << " (" << __LINE__ << ")" \
<< "-" << __FUNCTION__ << " : " \
<< IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \
throw std::runtime_error("IXFORMER_CHECK ERROR"); \
} \
}
#undef CUINFER_CHECK
#define CUINFER_CHECK(func) \
do { \
cuinferStatus_t status = (func); \
if (status != CUINFER_STATUS_SUCCESS) { \
std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \
<< ": " << cuinferGetErrorString(status) << std::endl; \
throw std::runtime_error("CUINFER_CHECK ERROR"); \
} \
} while (0)
} // namespace xllm::kernel::ilu

331
ex_engine/csrc/ix_full_bridge.cpp Normal file
View File

@@ -0,0 +1,331 @@
// ix_full_bridge.cpp — Complete ixformer::infer bridge for BI-V100
//
// Exposes ALL 14 ixformer C++ functions to Python via pybind11.
// Header source: upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h
//
// This replaces the partial ix_moe_bridge.cpp with the full set:
// MoE pipeline: topk_softmax, moe_compute_token_index_api, moe_expand_input,
// moe_w16a16_group_gemm, silu_and_mul, moe_output_reduce_sum
// Attention: ixinfer_flash_attn_unpad_with_block_tables, xllm_paged_attention
// Norm: rms_norm, residual_rms_norm
// RoPE: xllm_rotary_embedding
// Linear: ixformer_linear, ixformer_linear_ex
// Cache: xllm_reshape_and_cache
#include <torch/extension.h>
#include <tuple>
#include <vector>
// ============================================================================
// Forward-declare ixformer::infer namespace — matches ixformer.h exactly
// We forward-declare instead of #include to avoid build-time dependency
// on internal headers (ixinfer.h etc) that may not be on include path.
// The symbols resolve at link time against the base image's libixattn.so etc.
// ============================================================================
namespace ixformer {
namespace infer {
// --- Attention ---
torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
torch::Tensor& query, torch::Tensor& key_cache, torch::Tensor& value_cache,
torch::Tensor& out, torch::Tensor& block_tables,
torch::Tensor& cu_seq_q, torch::Tensor& cu_seq_k,
int64_t max_seq_q, int64_t max_seq_k, bool is_causal,
int64_t window_left, int64_t window_right,
double scale, double softcap, bool sqrt_alibi,
const std::optional<torch::Tensor>& alibi_slopes,
const std::optional<torch::Tensor>& sinks,
std::optional<torch::Tensor>& lse);
torch::Tensor xllm_paged_attention(
torch::Tensor& out, torch::Tensor& query,
torch::Tensor& key_cache, torch::Tensor& value_cache,
int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& context_lens,
int64_t block_size, int64_t max_context_len,
const std::optional<torch::Tensor>& alibi_slopes,
bool causal, int32_t window_left, int32_t window_right,
double softcap, bool enable_cuda_graph, bool use_sqrt_alibi,
const std::optional<torch::Tensor>& sinks);
// --- Norm ---
void residual_rms_norm(
torch::Tensor& input, torch::Tensor& residual, torch::Tensor& weight,
torch::Tensor& output, torch::Tensor& residual_output,
const std::optional<torch::Tensor>& fused_bias,
double alpha, double eps, bool is_post);
void rms_norm(
torch::Tensor& input, torch::Tensor& weight, torch::Tensor& output,
const std::optional<torch::Tensor>& fused_bias, double eps);
// --- Activation ---
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
// --- RoPE ---
void xllm_rotary_embedding(
torch::Tensor& positions, torch::Tensor& query, torch::Tensor& key,
int64_t head_size, torch::Tensor& cos_sin_cache, bool is_neox);
// --- KV Cache ---
void xllm_reshape_and_cache(
torch::Tensor& key, torch::Tensor& value,
torch::Tensor& key_cache, torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
int64_t key_token_stride, int64_t value_token_stride);
// --- Linear ---
torch::Tensor ixformer_linear(
torch::Tensor& input, torch::Tensor& weight, int64_t act_type,
const std::optional<torch::Tensor>& bias,
const std::optional<torch::Tensor>& out,
const std::optional<bool> persistent);
torch::Tensor ixformer_linear_ex(
torch::Tensor& input, torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& out);
// --- MoE ---
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);
} // namespace infer
} // namespace ixformer
// ============================================================================
// Python wrappers — thin wrappers matching upstream xllm ILU kernel layer
// Source: upstream_ref/xllm/xllm/core/kernels/ilu/*.cpp
// ============================================================================
// --- MoE: topk_softmax (from ilu/fused_moe.cpp moe_active_topk) ---
std::tuple<torch::Tensor, torch::Tensor> ix_topk_softmax(
torch::Tensor gating_output, 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, false);
if (renormalize) {
topk_weights = topk_weights / topk_weights.sum(-1, /*keepdim=*/true);
}
return std::make_tuple(topk_weights, topk_indices);
}
// --- MoE: gen_idx (from ilu/fused_moe.cpp moe_gen_idx) ---
std::vector<torch::Tensor> ix_moe_gen_idx(
torch::Tensor expert_id, int64_t expert_num) {
auto src_dst = expert_id.new_empty({expert_id.numel()});
auto dst_src = torch::empty_like(src_dst);
auto expert_sizes_gpu = expert_id.new_empty({expert_num});
ixformer::infer::moe_compute_token_index_api(
expert_id, src_dst, dst_src, expert_sizes_gpu,
c10::nullopt, c10::nullopt, c10::nullopt, 0, expert_num, expert_num);
auto expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
}
// --- MoE: expand_input ---
torch::Tensor ix_moe_expand_input(
torch::Tensor input, torch::Tensor gather_index,
torch::Tensor combine_idx, int64_t topk) {
int64_t dst_tokens = input.size(0) * topk;
auto output = input.new_empty({dst_tokens, input.size(1)});
ixformer::infer::moe_expand_input(
output, input, combine_idx, gather_index, dst_tokens, topk);
return output;
}
// --- MoE: group_gemm ---
torch::Tensor ix_group_gemm(
torch::Tensor inputs, torch::Tensor weights,
torch::Tensor token_count, int64_t output_n) {
int64_t total_tokens = inputs.size(0);
auto output = inputs.new_empty({total_tokens, output_n});
ixformer::infer::moe_w16a16_group_gemm(
output, inputs, weights, token_count,
c10::nullopt, c10::nullopt, "NT", 0, output_n);
return output;
}
// --- MoE: silu_and_mul ---
torch::Tensor ix_silu_and_mul(torch::Tensor input) {
int64_t half_dim = input.size(-1) / 2;
auto output = input.new_empty({input.size(0), half_dim});
ixformer::infer::silu_and_mul(input, output);
return output;
}
// --- MoE: combine_result ---
torch::Tensor ix_moe_combine_result(torch::Tensor input, torch::Tensor weight) {
input = input.view({-1, weight.size(1), input.size(1)});
auto output = input.new_empty({input.size(0), input.size(2)});
ixformer::infer::moe_output_reduce_sum(
output, input, weight, c10::nullopt, c10::nullopt, 1.0);
return output;
}
// --- MoE: full fused forward (from ilu/layers/fused_moe.cpp) ---
torch::Tensor ix_fused_moe_forward(
torch::Tensor hidden_states, torch::Tensor router_logits,
torch::Tensor w13, torch::Tensor w2,
int64_t topk, int64_t num_experts, bool renormalize) {
auto [topk_weights, topk_ids] = ix_topk_softmax(router_logits, topk, renormalize);
auto idx = ix_moe_gen_idx(topk_ids.view({-1}), num_experts);
auto expanded = ix_moe_expand_input(hidden_states, idx[0], idx[1], topk);
int64_t gate_up_dim = w13.size(1);
auto gemm1_out = ix_group_gemm(expanded, w13, idx[2], gate_up_dim);
auto act_out = ix_silu_and_mul(gemm1_out);
int64_t hidden_dim = w2.size(1);
auto gemm2_out = ix_group_gemm(act_out, w2, idx[2], hidden_dim);
return ix_moe_combine_result(gemm2_out, topk_weights);
}
// --- Attention: paged decode (from ilu/attention.cpp batch_decode) ---
void ix_paged_attention(
torch::Tensor output, torch::Tensor query,
torch::Tensor key_cache, torch::Tensor value_cache,
int64_t num_kv_heads, double scale,
torch::Tensor block_tables, torch::Tensor seq_lens,
int64_t block_size, int64_t max_context_len,
const std::optional<torch::Tensor>& alibi_slopes) {
if (query.dim() == 4) {
query = query.view({query.size(0)*query.size(1), query.size(2), query.size(3)}).contiguous();
}
if (output.dim() == 4) {
output = output.view({output.size(0)*output.size(1), output.size(2), output.size(3)}).contiguous();
}
ixformer::infer::xllm_paged_attention(
output, query, key_cache, value_cache, num_kv_heads, scale,
block_tables, seq_lens, block_size, max_context_len,
alibi_slopes, /*causal=*/true, /*window_left=*/-1, /*window_right=*/-1,
/*softcap=*/0.0, /*enable_cuda_graph=*/false, /*use_sqrt_alibi=*/false,
/*sinks=*/c10::nullopt);
}
// --- Attention: prefill flash (from ilu/attention.cpp batch_prefill) ---
void ix_flash_attn_prefill(
torch::Tensor query, torch::Tensor key, torch::Tensor value,
torch::Tensor output, torch::Tensor block_tables,
torch::Tensor cu_seq_q, torch::Tensor cu_seq_k,
int64_t max_query_len, int64_t max_seq_len,
double scale, bool is_causal,
int64_t window_left, int64_t window_right) {
std::optional<torch::Tensor> lse = c10::nullopt;
ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables(
query, key, value, output, block_tables,
cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
is_causal, window_left, window_right,
scale, /*softcap=*/0.0, /*sqrt_alibi=*/false,
/*alibi_slopes=*/c10::nullopt, /*sinks=*/c10::nullopt, lse);
}
// --- Norm: rms_norm (from ilu/norm.cpp) ---
void ix_rms_norm(
torch::Tensor output, torch::Tensor input,
torch::Tensor weight, double eps) {
ixformer::infer::rms_norm(input, weight, output, c10::nullopt, eps);
}
// --- Norm: fused residual + rms_norm (from ilu/norm.cpp) ---
void ix_fused_add_rms_norm(
torch::Tensor input, torch::Tensor residual,
torch::Tensor weight, torch::Tensor output,
torch::Tensor residual_output, double eps) {
ixformer::infer::residual_rms_norm(
input, residual, weight, output, residual_output,
c10::nullopt, 1.0, eps, false);
}
// --- RoPE (from ilu/rope.cpp) ---
void ix_rotary_embedding(
torch::Tensor positions, torch::Tensor query, torch::Tensor key,
int64_t head_size, torch::Tensor cos_sin_cache, bool is_neox) {
ixformer::infer::xllm_rotary_embedding(
positions, query, key, head_size, cos_sin_cache, is_neox);
}
// --- KV Cache reshape (from ilu/attention.cpp reshape_paged_cache) ---
void ix_reshape_and_cache(
torch::Tensor key, torch::Tensor value,
torch::Tensor key_cache, torch::Tensor value_cache,
torch::Tensor slot_mapping) {
slot_mapping = slot_mapping.to(torch::kLong);
int64_t key_stride = key.stride(0);
int64_t val_stride = value.stride(0);
ixformer::infer::xllm_reshape_and_cache(
key, value, key_cache, value_cache, slot_mapping,
key_stride, val_stride);
}
// --- Linear ---
torch::Tensor ix_linear(
torch::Tensor input, torch::Tensor weight,
const std::optional<torch::Tensor>& bias) {
return ixformer::infer::ixformer_linear(
input, weight, /*act_type=*/0, bias, c10::nullopt, c10::nullopt);
}
// ============================================================================
// Module registration
// ============================================================================
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
// MoE
m.def("topk_softmax", &ix_topk_softmax, "Fused topk+softmax",
py::arg("gating_output"), py::arg("topk"), py::arg("renormalize")=true);
m.def("moe_gen_idx", &ix_moe_gen_idx);
m.def("moe_expand_input", &ix_moe_expand_input);
m.def("group_gemm", &ix_group_gemm);
m.def("silu_and_mul", &ix_silu_and_mul);
m.def("moe_combine_result", &ix_moe_combine_result);
m.def("fused_moe_forward", &ix_fused_moe_forward,
py::arg("hidden_states"), py::arg("router_logits"),
py::arg("w13"), py::arg("w2"),
py::arg("topk"), py::arg("num_experts"), py::arg("renormalize")=true);
// Attention
m.def("paged_attention", &ix_paged_attention);
m.def("flash_attn_prefill", &ix_flash_attn_prefill);
// Norm
m.def("rms_norm", &ix_rms_norm);
m.def("fused_add_rms_norm", &ix_fused_add_rms_norm);
// RoPE
m.def("rotary_embedding", &ix_rotary_embedding);
// Cache
m.def("reshape_and_cache", &ix_reshape_and_cache);
// Linear
m.def("linear", &ix_linear);
}

250
ex_engine/python/ix_bridge.py
View File

@@ -1,15 +1,17 @@
"""
ix_bridge.py — Full ixformer MoE pipeline bridge.
ix_bridge.py — Full ixformer bridge loader.
Loads ix_moe_bridge.so via JIT and exposes both individual ops and the full
fused MoE forward pass that replaces the Python for-loop in qwen3_5.py.
Loads ix_full_bridge.so (all 14 ixformer::infer functions) or falls back
to ix_moe_bridge.so (MoE-only 6 functions).
Pipeline (mirrors xllm/core/layers/ilu/fused_moe.cpp):
topk_softmax → moe_gen_idx → moe_expand_input → group_gemm(w13)
→ silu_and_mul → group_gemm(w2) → moe_combine_result
All 6 ixformer::infer C++ functions are called through ix_moe_bridge.cpp
which forward-declares them and links against the base image SDK.
Functions exposed:
MoE: topk_softmax, moe_gen_idx, moe_expand_input, group_gemm,
silu_and_mul, moe_combine_result, fused_moe_forward
Attention: paged_attention, flash_attn_prefill
Norm: rms_norm, fused_add_rms_norm
RoPE: rotary_embedding
Cache: reshape_and_cache
Linear: linear
"""
import os
@@ -19,23 +21,22 @@ from typing import Tuple, Optional, List
logger = logging.getLogger("ex_engine.ix_bridge")
_ix_bridge = None
_ix_bridge_loaded = False # True after attempt, even if failed
_ix_bridge_available = False
_bridge = None
_loaded = False
_available = False
# All .cpp sources to try, in priority order
_CPP_NAMES = ["ix_full_bridge.cpp", "ix_moe_bridge.cpp"]
def _find_cpp_source():
"""Find ix_moe_bridge.cpp in multiple locations."""
candidates = []
# 1. Relative to this file: ex_engine/csrc/
def _find_cpp(name):
here = os.path.dirname(os.path.abspath(__file__))
candidates.append(os.path.join(here, "..", "csrc", "ix_moe_bridge.cpp"))
# 2. Deployed path inside vllm model dir
candidates.append(os.path.join(here, "ix_moe_bridge.cpp"))
# 3. /workspace paths
candidates.append("/workspace/ex_engine/csrc/ix_moe_bridge.cpp")
candidates.append("/workspace/qwen3_6_scripts/ix_moe_bridge.cpp")
candidates = [
os.path.join(here, "..", "csrc", name),
os.path.join(here, name),
os.path.join("/workspace/ex_engine/csrc", name),
os.path.join("/workspace/qwen3_6_scripts", name),
]
for c in candidates:
p = os.path.normpath(c)
if os.path.exists(p):
@@ -44,136 +45,117 @@ def _find_cpp_source():
def _load_bridge():
"""JIT-compile and load ix_moe_bridge.so — called once."""
global _ix_bridge, _ix_bridge_loaded, _ix_bridge_available
if _ix_bridge_loaded:
return _ix_bridge_available
_ix_bridge_loaded = True
global _bridge, _loaded, _available
if _loaded:
return _available
_loaded = True
cpp_file = _find_cpp_source()
if cpp_file is None:
logger.warning("ix_moe_bridge.cpp not found in any search path")
return False
from torch.utils.cpp_extension import load
try:
from torch.utils.cpp_extension import load
logger.info("JIT-compiling ix_moe_bridge.cpp from %s ...", cpp_file)
_ix_bridge = load(
name="ix_moe_bridge",
sources=[cpp_file],
extra_cflags=["-O2", "-std=c++17"],
verbose=False,
)
_ix_bridge_available = True
fns = [x for x in dir(_ix_bridge) if not x.startswith("_")]
logger.info("ix_moe_bridge loaded: %s", fns)
return True
except Exception as e:
logger.warning("ix_moe_bridge JIT compile failed: %s", e)
return False
for cpp_name in _CPP_NAMES:
cpp_path = _find_cpp(cpp_name)
if cpp_path is None:
continue
mod_name = cpp_name.replace(".cpp", "").replace(".", "_")
try:
logger.info("JIT-compiling %s from %s ...", cpp_name, cpp_path)
_bridge = load(
name=mod_name,
sources=[cpp_path],
extra_cflags=["-O2", "-std=c++17"],
verbose=False,
)
_available = True
fns = [x for x in dir(_bridge) if not x.startswith("_")]
logger.info("ix_bridge loaded (%s): %s", cpp_name, fns)
return True
except Exception as e:
logger.warning("JIT compile %s failed: %s — trying next", cpp_name, e)
logger.warning("All ix_bridge sources failed to compile")
return False
def is_available() -> bool:
"""Check if bridge is available (lazy-load on first call)."""
if not _ix_bridge_loaded:
if not _loaded:
_load_bridge()
return _ix_bridge_available
return _available
def _get():
if not is_available():
raise RuntimeError("ix_bridge not available")
return _bridge
# =========================================================================
# Individual ops (thin wrappers with type safety)
# MoE
# =========================================================================
def topk_softmax(gating_output, topk, renormalize=True):
return _get().topk_softmax(gating_output, topk, renormalize)
def topk_softmax(
gating_output: torch.Tensor,
topk: int,
renormalize: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Fused topk+softmax via ixformer::infer::topk_softmax.
Returns: (topk_weights [T, K] fp32, topk_ids [T, K] int32)
"""
if not is_available():
raise RuntimeError("ix_moe_bridge not available — JIT compile failed")
return _ix_bridge.topk_softmax(gating_output, topk, renormalize)
def moe_gen_idx(expert_id, expert_num):
return _get().moe_gen_idx(expert_id, expert_num)
def moe_expand_input(input, gather_index, combine_idx, topk):
return _get().moe_expand_input(input, gather_index, combine_idx, topk)
def moe_gen_idx(
expert_id: torch.Tensor,
expert_num: int,
) -> List[torch.Tensor]:
"""
Build expert permutation maps.
Returns: [src_dst, dst_src, expert_sizes, cumsum]
"""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.moe_gen_idx(expert_id, expert_num)
def group_gemm(inputs, weights, token_count, output_n):
return _get().group_gemm(inputs, weights, token_count, output_n)
def silu_and_mul(input):
return _get().silu_and_mul(input)
def moe_expand_input(
input: torch.Tensor,
gather_index: torch.Tensor,
combine_idx: torch.Tensor,
topk: int,
) -> torch.Tensor:
"""Gather tokens by expert assignment."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.moe_expand_input(input, gather_index, combine_idx, topk)
def group_gemm(
inputs: torch.Tensor,
weights: torch.Tensor,
token_count: torch.Tensor,
output_n: int,
) -> torch.Tensor:
"""Batched expert GEMM via ixformer."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.group_gemm(inputs, weights, token_count, output_n)
def silu_and_mul(input: torch.Tensor) -> torch.Tensor:
"""Fused SiLU gate activation."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.silu_and_mul(input)
def moe_combine_result(
input: torch.Tensor,
weight: torch.Tensor,
) -> torch.Tensor:
"""Weighted reduce for MoE output."""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.moe_combine_result(input, weight)
def moe_combine_result(input, weight):
return _get().moe_combine_result(input, weight)
def fused_moe_forward(hidden_states, router_logits, w13, w2,
topk, num_experts, renormalize=True):
return _get().fused_moe_forward(
hidden_states, router_logits, w13, w2, topk, num_experts, renormalize)
# =========================================================================
# Full fused MoE forward — replaces _pure_pytorch_experts() entirely
# Attention
# =========================================================================
def paged_attention(output, query, key_cache, value_cache,
num_kv_heads, scale, block_tables, seq_lens,
block_size, max_context_len, alibi_slopes=None):
return _get().paged_attention(
output, query, key_cache, value_cache,
num_kv_heads, scale, block_tables, seq_lens,
block_size, max_context_len, alibi_slopes)
def fused_moe_forward(
hidden_states: torch.Tensor, # (T, H)
router_logits: torch.Tensor, # (T, E)
w13: torch.Tensor, # (E, 2*I, H) gate_up
w2: torch.Tensor, # (E, H, I) down
topk: int,
num_experts: int,
renormalize: bool = True,
) -> torch.Tensor:
"""
Full fused MoE forward via ixformer C++ pipeline.
def flash_attn_prefill(query, key, value, output, block_tables,
cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
scale, is_causal=True, window_left=-1, window_right=-1):
return _get().flash_attn_prefill(
query, key, value, output, block_tables,
cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
scale, is_causal, window_left, window_right)
Pipeline: topk → gen_idx → expand → gemm1(w13) → silu → gemm2(w2) → combine
# =========================================================================
# Norm
# =========================================================================
def rms_norm(output, input, weight, eps=1e-6):
return _get().rms_norm(output, input, weight, eps)
Returns: (T, H) — partial output, needs all-reduce after.
"""
if not is_available():
raise RuntimeError("ix_moe_bridge not available")
return _ix_bridge.fused_moe_forward(
hidden_states, router_logits, w13, w2, topk, num_experts, renormalize
)
def fused_add_rms_norm(input, residual, weight, output, residual_output, eps=1e-6):
return _get().fused_add_rms_norm(input, residual, weight, output, residual_output, eps)
# =========================================================================
# RoPE
# =========================================================================
def rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox=True):
return _get().rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox)
# =========================================================================
# Cache
# =========================================================================
def reshape_and_cache(key, value, key_cache, value_cache, slot_mapping):
return _get().reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)
# =========================================================================
# Linear
# =========================================================================
def linear(input, weight, bias=None):
return _get().linear(input, weight, bias)

12
qwen3_6_scripts/patch_ops.sh
View File

@@ -197,11 +197,13 @@ if [ -d "$EX_ENGINE_SRC/python" ]; then
EX_DST="$VLLM/model_executor/models/ex_engine"
mkdir -p "$EX_DST/python" "$EX_DST/csrc"
cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST/python/" 2>/dev/null || true
# ix_moe_bridge.cpp for JIT compile — deploy to ALL search paths
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
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "/workspace/ex_engine/csrc/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "/workspace/qwen3_6_scripts/" 2>/dev/null || true
# ix_full_bridge.cpp + ix_moe_bridge.cpp for JIT compile — deploy to ALL search paths
for _BRIDGE in ix_full_bridge.cpp ix_moe_bridge.cpp; do
cp "$EX_ENGINE_SRC/csrc/$_BRIDGE" "$EX_DST/csrc/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/$_BRIDGE" "$EX_DST/python/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/$_BRIDGE" "/workspace/ex_engine/csrc/" 2>/dev/null || true
cp "$EX_ENGINE_SRC/csrc/$_BRIDGE" "/workspace/qwen3_6_scripts/" 2>/dev/null || true
done
touch "$EX_DST/__init__.py"
touch "$EX_DST/python/__init__.py"
# Copy built .so files