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

feat: ix_moe_bridge + ix_attn_bridge — dlopen bridges for full ixformer::infer API

Browse Source 
Bridge architecture (from xllm/core/kernels/ilu/ixformer.h):

ix_moe_bridge.so (MoE 7-step fused pipeline):
  - topk_softmax → moe_compute_token_index_api → moe_expand_input
  - moe_w16a16_group_gemm (x2) → silu_and_mul → moe_output_reduce_sum
  - fused_moe_forward(): replaces entire Python expert loop
  - Fix: group_gemm format NT→TN (match xllm trans_b=true)

ix_attn_bridge.so (attention + linear):
  - ixinfer_flash_attn_unpad_with_block_tables (fused prefill)
  - xllm_paged_attention (fused paged decode)
  - ixformer_linear (matmul + activation)
  - residual_rms_norm (fused residual + norm)

Integration:
  - ix_fused_moe.py: Python loader (prebuilt .so → JIT → unavailable)
  - qwen3_5.py: Tier 0 dispatch in _pure_pytorch_experts()
  - patch_ops.sh: deploys ix_fused_moe.py + all prebuilt/*.so

Source: jd-opensource/xllm (fresh clone, all ILU kernels verified SAME)
Sync: upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h (+32 lines)

Build on real machine:
  bash qwen3_6_scripts/build_ix_moe_bridge.sh
  bash qwen3_6_scripts/build_ix_attn_bridge.sh
This commit is contained in:
claude
2026-08-14 07:32:31 +00:00
parent 5e9b7c292a
commit 051b02d3cd
10 changed files with 888 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
qwen3_6_scripts/ix_moe_bridge.cpp
View File

@@ -17,7 +17,7 @@
#include <vector>
#include <optional>
static const c10::optional<torch::Tensor> kNoneTensor = {};
static const std::optional<torch::Tensor> kNoneTensor = {};
// Forward-declare ixformer C++ API (from base image SDK)
namespace ixformer {
@@ -34,9 +34,9 @@ void moe_compute_token_index_api(
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,
const std::optional<torch::Tensor>& expert_mask,
const std::optional<torch::Tensor>& expert_sizes_cpu,
const std::optional<torch::Tensor>& expand_tokens_gpu,
int64_t start_expert_id,
int64_t end_expert_id,
int64_t num_experts);
@@ -44,7 +44,7 @@ void moe_compute_token_index_api(
void moe_expand_input(torch::Tensor outputs,
torch::Tensor inputs,
torch::Tensor dst_to_src,
const c10::optional<torch::Tensor>& src_to_dst,
const std::optional<torch::Tensor>& src_to_dst,
int64_t dst_tokens,
int64_t expand_factor);
@@ -52,17 +52,17 @@ 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,
const std::optional<torch::Tensor>& dst_to_src,
const std::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,
const std::optional<torch::Tensor>& mul_weight,
const std::optional<torch::Tensor>& mask,
const std::optional<torch::Tensor>& extra_residual,
double scaling_factor);
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
@@ -149,7 +149,7 @@ torch::Tensor ix_group_gemm(
output, inputs, weights, token_count,
/*dst_to_src=*/kNoneTensor,
/*bias=*/kNoneTensor,
/*format=*/"NT",
/*format=*/"TN",
/*persistent=*/0,
/*output_n=*/output_n);
return output;