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

221
ex_engine/csrc/ix_attn_bridge.cpp Normal file
View File

@@ -0,0 +1,221 @@
// ix_attn_bridge.cpp — Bridge to ixformer::infer attention + linear functions
//
// Exposes functions from ixformer.h that are NOT available via ixformer.functions:
// 1. ixinfer_flash_attn_unpad_with_block_tables — fused prefill attention
// 2. xllm_paged_attention — fused paged decode attention
// 3. ixformer_linear — fused linear (matmul + optional activation)
// 4. ixformer_linear_ex — simple fused linear
// 5. residual_rms_norm — fused residual + RMS norm (NOT in ixformer_torch_ext)
//
// Source: xllm/xllm/core/kernels/ilu/ixformer.h
// Usage: xllm/xllm/core/kernels/ilu/attention.cpp
// xllm/xllm/core/layers/ilu/attention.cpp
#include <torch/extension.h>
#include <optional>
namespace ixformer {
namespace infer {
// Prefill: flash attention with block tables (variable-length batched)
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);
// Decode: paged attention (single-step cached KV)
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);
// Fused linear: matmul + optional activation
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);
// Simple linear
torch::Tensor ixformer_linear_ex(
torch::Tensor& input,
torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& out);
// Fused residual + RMS norm (not in ixformer_torch_ext, only in ixformer::infer)
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);
} // namespace infer
} // namespace ixformer
// ============================================================================
// Python-facing wrappers
// Port from: xllm/xllm/core/kernels/ilu/attention.cpp
// ============================================================================
// Prefill attention via flash_attn_unpad_with_block_tables
torch::Tensor ix_prefill_attention(
torch::Tensor query, // (total_q_tokens, num_heads, head_dim)
torch::Tensor key_cache, // (num_blocks, num_heads, block_size, head_dim)
torch::Tensor value_cache, // (num_blocks, num_heads, block_size, head_dim)
torch::Tensor output, // (total_q_tokens, num_heads, head_dim)
torch::Tensor block_tables, // (batch, max_blocks)
torch::Tensor cu_seq_q, // (batch+1,)
torch::Tensor cu_seq_k, // (batch+1,)
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;
return ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables(
query, key_cache, value_cache, 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=*/std::nullopt,
/*sinks=*/std::nullopt,
lse);
}
// Decode attention via xllm_paged_attention
torch::Tensor ix_decode_attention(
torch::Tensor output, // (num_seqs, num_heads, head_dim)
torch::Tensor query, // (num_seqs, num_heads, head_dim)
torch::Tensor key_cache, // (num_blocks, num_kv_heads, block_size, head_dim)
torch::Tensor value_cache, // (num_blocks, num_kv_heads, block_size, head_dim)
int64_t num_kv_heads,
double scale,
torch::Tensor block_tables, // (num_seqs, max_blocks)
torch::Tensor seq_lens, // (num_seqs,)
int64_t block_size,
int64_t max_context_len) {
return 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=*/std::nullopt,
/*causal=*/true,
/*window_left=*/-1,
/*window_right=*/-1,
/*softcap=*/0.0,
/*enable_cuda_graph=*/false,
/*use_sqrt_alibi=*/false,
/*sinks=*/std::nullopt);
}
// Fused linear (matmul + optional activation)
// act_type: 0=none, 1=silu, 2=gelu, 3=gelu_tanh
torch::Tensor ix_linear(
torch::Tensor input,
torch::Tensor weight,
int64_t act_type) {
return ixformer::infer::ixformer_linear(
input, weight, act_type,
/*bias=*/std::nullopt,
/*out=*/std::nullopt,
/*persistent=*/std::nullopt);
}
// Fused residual + RMS norm
// Port from: xllm/xllm/core/kernels/ilu/norm.cpp residual_layer_norm()
std::tuple<torch::Tensor, torch::Tensor> ix_residual_rms_norm(
torch::Tensor input,
torch::Tensor residual,
torch::Tensor weight,
double eps) {
auto output = torch::zeros_like(input);
auto residual_output = torch::zeros_like(input);
ixformer::infer::residual_rms_norm(
input, residual, weight, output, residual_output,
/*fused_bias=*/std::nullopt,
/*alpha=*/1.0,
eps,
/*is_post=*/false);
return std::make_tuple(output, residual_output);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("prefill_attention", &ix_prefill_attention,
"Fused prefill attention via ixformer flash_attn_unpad_with_block_tables",
py::arg("query"), py::arg("key_cache"), py::arg("value_cache"),
py::arg("output"), py::arg("block_tables"),
py::arg("cu_seq_q"), py::arg("cu_seq_k"),
py::arg("max_query_len"), py::arg("max_seq_len"),
py::arg("scale"),
py::arg("is_causal") = true,
py::arg("window_left") = -1,
py::arg("window_right") = -1);
m.def("decode_attention", &ix_decode_attention,
"Paged decode attention via ixformer xllm_paged_attention",
py::arg("output"), py::arg("query"),
py::arg("key_cache"), py::arg("value_cache"),
py::arg("num_kv_heads"), py::arg("scale"),
py::arg("block_tables"), py::arg("seq_lens"),
py::arg("block_size"), py::arg("max_context_len"));
m.def("linear", &ix_linear,
"Fused linear via ixformer (matmul + optional activation)",
py::arg("input"), py::arg("weight"), py::arg("act_type") = 0);
m.def("residual_rms_norm", &ix_residual_rms_norm,
"Fused residual + RMS norm via ixformer",
py::arg("input"), py::arg("residual"),
py::arg("weight"), py::arg("eps") = 1e-6);
}

2
ex_engine/csrc/ix_moe_bridge.cpp
View File

@@ -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;

71
qwen3_6_scripts/build_ix_attn_bridge.sh Normal file
View File

@@ -0,0 +1,71 @@
#!/bin/bash
# build_ix_attn_bridge.sh — Build ix_attn_bridge.so on real BI-V100
#
# Compiles ix_attn_bridge.cpp → prebuilt .so for Docker deployment.
# Functions: prefill_attention, decode_attention, linear, residual_rms_norm
#
# Run on real machine: bash qwen3_6_scripts/build_ix_attn_bridge.sh
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
CPP_SOURCE="${SCRIPT_DIR}/ix_attn_bridge.cpp"
PREBUILT_DIR="${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10"
if [ ! -f "$CPP_SOURCE" ]; then
echo "ERROR: ix_attn_bridge.cpp not found at $CPP_SOURCE"
exit 1
fi
echo "=== Building ix_attn_bridge.so ==="
echo "Source: $CPP_SOURCE"
python3 -c "
import os, sys, glob, shutil
from torch.utils.cpp_extension import load
cpp_source = '$CPP_SOURCE'
extra_ldflags = []
try:
import ixformer
ixf_dir = os.path.dirname(ixformer.__file__)
for so in glob.glob(os.path.join(ixf_dir, '*.so')):
extra_ldflags.append(so)
extra_ldflags.append(f'-Wl,-rpath,{ixf_dir}')
except ImportError:
pass
corex_lib = '/usr/local/corex/lib64'
if os.path.isdir(corex_lib):
for lib in ['libixattn.so', 'libixformer.so', 'libcublas.so']:
p = os.path.join(corex_lib, lib)
if os.path.isfile(p):
extra_ldflags.append(p)
extra_ldflags.append(f'-Wl,-rpath,{corex_lib}')
print(f'Linking: {extra_ldflags}')
mod = load(
name='ix_attn_bridge',
sources=[cpp_source],
extra_cflags=['-O2', '-std=c++17'],
extra_ldflags=extra_ldflags,
verbose=True,
)
import torch.utils.cpp_extension as ext
build_dir = ext._get_build_directory('ix_attn_bridge', verbose=False)
for f in glob.glob(os.path.join(build_dir, '*.so')):
dst = '$PREBUILT_DIR/ix_attn_bridge.so'
os.makedirs(os.path.dirname(dst), exist_ok=True)
shutil.copy2(f, dst)
sz = os.path.getsize(dst)
print(f'✓ ix_attn_bridge.so ({sz} bytes) → {dst}')
break
fns = [x for x in dir(mod) if not x.startswith('_')]
print(f'Functions: {fns}')
print('=== Build SUCCESS ===')
"

84
qwen3_6_scripts/build_ix_moe_bridge.sh Normal file
View File

@@ -0,0 +1,84 @@
#!/bin/bash
# build_ix_moe_bridge.sh — Build ix_moe_bridge.so on real BI-V100
#
# This compiles ex_engine/csrc/ix_moe_bridge.cpp into a prebuilt .so
# that can be deployed without JIT compilation in Docker.
#
# Run on real machine: bash qwen3_6_scripts/build_ix_moe_bridge.sh
# Output: qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ix_moe_bridge.so
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
CPP_SOURCE="${PROJECT_DIR}/ex_engine/csrc/ix_moe_bridge.cpp"
PREBUILT_DIR="${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10"
if [ ! -f "$CPP_SOURCE" ]; then
# Also try the local copy
CPP_SOURCE="${SCRIPT_DIR}/ix_moe_bridge.cpp"
fi
if [ ! -f "$CPP_SOURCE" ]; then
echo "ERROR: ix_moe_bridge.cpp not found"
exit 1
fi
echo "=== Building ix_moe_bridge.so ==="
echo "Source: $CPP_SOURCE"
echo "Output: $PREBUILT_DIR/ix_moe_bridge.so"
python3 -c "
import os, sys, glob
from torch.utils.cpp_extension import load
cpp_source = '$CPP_SOURCE'
extra_ldflags = []
# Find ixformer .so to link against
try:
import ixformer
ixf_dir = os.path.dirname(ixformer.__file__)
for so in glob.glob(os.path.join(ixf_dir, '*.so')):
extra_ldflags.append(so)
extra_ldflags.append(f'-Wl,-rpath,{ixf_dir}')
except ImportError:
pass
corex_lib = '/usr/local/corex/lib64'
if os.path.isdir(corex_lib):
for lib in ['libixattn.so', 'libixformer.so', 'libcublas.so']:
p = os.path.join(corex_lib, lib)
if os.path.isfile(p):
extra_ldflags.append(p)
extra_ldflags.append(f'-Wl,-rpath,{corex_lib}')
print(f'Linking: {extra_ldflags}')
mod = load(
name='ix_moe_bridge',
sources=[cpp_source],
extra_cflags=['-O2', '-std=c++17'],
extra_ldflags=extra_ldflags,
verbose=True,
)
# Find the compiled .so and copy to prebuilt
import torch.utils.cpp_extension as ext
build_dir = ext._get_build_directory('ix_moe_bridge', verbose=False)
print(f'Build dir: {build_dir}')
import shutil
for f in glob.glob(os.path.join(build_dir, '*.so')):
dst = '$PREBUILT_DIR/ix_moe_bridge.so'
os.makedirs(os.path.dirname(dst), exist_ok=True)
shutil.copy2(f, dst)
sz = os.path.getsize(dst)
print(f'✓ ix_moe_bridge.so ({sz} bytes) → {dst}')
break
# Verify
fns = [x for x in dir(mod) if not x.startswith('_')]
print(f'Functions: {fns}')
print('=== Build SUCCESS ===')
"

221
qwen3_6_scripts/ix_attn_bridge.cpp Normal file
View File

@@ -0,0 +1,221 @@
// ix_attn_bridge.cpp — Bridge to ixformer::infer attention + linear functions
//
// Exposes functions from ixformer.h that are NOT available via ixformer.functions:
// 1. ixinfer_flash_attn_unpad_with_block_tables — fused prefill attention
// 2. xllm_paged_attention — fused paged decode attention
// 3. ixformer_linear — fused linear (matmul + optional activation)
// 4. ixformer_linear_ex — simple fused linear
// 5. residual_rms_norm — fused residual + RMS norm (NOT in ixformer_torch_ext)
//
// Source: xllm/xllm/core/kernels/ilu/ixformer.h
// Usage: xllm/xllm/core/kernels/ilu/attention.cpp
// xllm/xllm/core/layers/ilu/attention.cpp
#include <torch/extension.h>
#include <optional>
namespace ixformer {
namespace infer {
// Prefill: flash attention with block tables (variable-length batched)
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);
// Decode: paged attention (single-step cached KV)
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);
// Fused linear: matmul + optional activation
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);
// Simple linear
torch::Tensor ixformer_linear_ex(
torch::Tensor& input,
torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& out);
// Fused residual + RMS norm (not in ixformer_torch_ext, only in ixformer::infer)
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);
} // namespace infer
} // namespace ixformer
// ============================================================================
// Python-facing wrappers
// Port from: xllm/xllm/core/kernels/ilu/attention.cpp
// ============================================================================
// Prefill attention via flash_attn_unpad_with_block_tables
torch::Tensor ix_prefill_attention(
torch::Tensor query, // (total_q_tokens, num_heads, head_dim)
torch::Tensor key_cache, // (num_blocks, num_heads, block_size, head_dim)
torch::Tensor value_cache, // (num_blocks, num_heads, block_size, head_dim)
torch::Tensor output, // (total_q_tokens, num_heads, head_dim)
torch::Tensor block_tables, // (batch, max_blocks)
torch::Tensor cu_seq_q, // (batch+1,)
torch::Tensor cu_seq_k, // (batch+1,)
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;
return ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables(
query, key_cache, value_cache, 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=*/std::nullopt,
/*sinks=*/std::nullopt,
lse);
}
// Decode attention via xllm_paged_attention
torch::Tensor ix_decode_attention(
torch::Tensor output, // (num_seqs, num_heads, head_dim)
torch::Tensor query, // (num_seqs, num_heads, head_dim)
torch::Tensor key_cache, // (num_blocks, num_kv_heads, block_size, head_dim)
torch::Tensor value_cache, // (num_blocks, num_kv_heads, block_size, head_dim)
int64_t num_kv_heads,
double scale,
torch::Tensor block_tables, // (num_seqs, max_blocks)
torch::Tensor seq_lens, // (num_seqs,)
int64_t block_size,
int64_t max_context_len) {
return 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=*/std::nullopt,
/*causal=*/true,
/*window_left=*/-1,
/*window_right=*/-1,
/*softcap=*/0.0,
/*enable_cuda_graph=*/false,
/*use_sqrt_alibi=*/false,
/*sinks=*/std::nullopt);
}
// Fused linear (matmul + optional activation)
// act_type: 0=none, 1=silu, 2=gelu, 3=gelu_tanh
torch::Tensor ix_linear(
torch::Tensor input,
torch::Tensor weight,
int64_t act_type) {
return ixformer::infer::ixformer_linear(
input, weight, act_type,
/*bias=*/std::nullopt,
/*out=*/std::nullopt,
/*persistent=*/std::nullopt);
}
// Fused residual + RMS norm
// Port from: xllm/xllm/core/kernels/ilu/norm.cpp residual_layer_norm()
std::tuple<torch::Tensor, torch::Tensor> ix_residual_rms_norm(
torch::Tensor input,
torch::Tensor residual,
torch::Tensor weight,
double eps) {
auto output = torch::zeros_like(input);
auto residual_output = torch::zeros_like(input);
ixformer::infer::residual_rms_norm(
input, residual, weight, output, residual_output,
/*fused_bias=*/std::nullopt,
/*alpha=*/1.0,
eps,
/*is_post=*/false);
return std::make_tuple(output, residual_output);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("prefill_attention", &ix_prefill_attention,
"Fused prefill attention via ixformer flash_attn_unpad_with_block_tables",
py::arg("query"), py::arg("key_cache"), py::arg("value_cache"),
py::arg("output"), py::arg("block_tables"),
py::arg("cu_seq_q"), py::arg("cu_seq_k"),
py::arg("max_query_len"), py::arg("max_seq_len"),
py::arg("scale"),
py::arg("is_causal") = true,
py::arg("window_left") = -1,
py::arg("window_right") = -1);
m.def("decode_attention", &ix_decode_attention,
"Paged decode attention via ixformer xllm_paged_attention",
py::arg("output"), py::arg("query"),
py::arg("key_cache"), py::arg("value_cache"),
py::arg("num_kv_heads"), py::arg("scale"),
py::arg("block_tables"), py::arg("seq_lens"),
py::arg("block_size"), py::arg("max_context_len"));
m.def("linear", &ix_linear,
"Fused linear via ixformer (matmul + optional activation)",
py::arg("input"), py::arg("weight"), py::arg("act_type") = 0);
m.def("residual_rms_norm", &ix_residual_rms_norm,
"Fused residual + RMS norm via ixformer",
py::arg("input"), py::arg("residual"),
py::arg("weight"), py::arg("eps") = 1e-6);
}

196
qwen3_6_scripts/ix_fused_moe.py Normal file
View File

@@ -0,0 +1,196 @@
"""
ix_fused_moe.py — Fused MoE pipeline via ixformer C++ API
Replaces the entire Python expert loop in qwen3_5.py with xllm's 7-step
fused pipeline:
topk_softmax → gen_idx → expand → group_gemm → silu → group_gemm → combine
Source: upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
Bridge: ex_engine/csrc/ix_moe_bridge.cpp → ixformer::infer namespace
Loading strategy:
1. Try prebuilt ix_moe_bridge.so from known locations
2. Try JIT compile from .cpp source
3. Return unavailable (caller falls back to Python loop)
"""
import os
import logging
import importlib
import torch
logger = logging.getLogger("ix_fused_moe")
_bridge = None
_loaded = False
def _try_load_prebuilt():
"""Load prebuilt ix_moe_bridge.so without JIT compilation."""
search_paths = [
# Deployed by patch_ops.sh into vllm package
os.path.join(os.path.dirname(__file__), "ix_moe_bridge.so"),
# Prebuilt directory
os.path.join(os.path.dirname(__file__), "prebuilt",
"corex-3.2.3-ivcore10", "ix_moe_bridge.so"),
# Workspace deployment
"/workspace/qwen3_6_scripts/ix_moe_bridge.so",
]
# Also check the vllm package directory
try:
import vllm
vllm_dir = os.path.dirname(vllm.__file__)
search_paths.append(os.path.join(vllm_dir, "ix_moe_bridge.so"))
except ImportError:
pass
for path in search_paths:
if os.path.isfile(path):
try:
spec = importlib.util.spec_from_file_location(
"ix_moe_bridge", path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
fns = [x for x in dir(mod) if not x.startswith("_")]
logger.info("ix_moe_bridge loaded from %s: %s", path, fns)
return mod
except Exception as e:
logger.debug("Failed to load %s: %s", path, e)
return None
def _try_jit_compile():
"""JIT compile ix_moe_bridge.cpp using torch.utils.cpp_extension."""
import glob
cpp_paths = [
os.path.join(os.path.dirname(__file__), "..", "ex_engine",
"csrc", "ix_moe_bridge.cpp"),
os.path.join(os.path.dirname(__file__), "ix_moe_bridge.cpp"),
"/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
"/workspace/qwen3_6_scripts/ix_moe_bridge.cpp",
]
cpp_file = None
for p in cpp_paths:
p = os.path.normpath(p)
if os.path.isfile(p):
cpp_file = p
break
if cpp_file is None:
logger.debug("ix_moe_bridge.cpp not found in search paths")
return None
# Build link flags to find ixformer symbols
extra_ldflags = []
try:
import ixformer
ixf_dir = os.path.dirname(ixformer.__file__)
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
extra_ldflags.append(so)
extra_ldflags.append(f"-Wl,-rpath,{ixf_dir}")
except ImportError:
pass
corex_lib = "/usr/local/corex/lib64"
if os.path.isdir(corex_lib):
for lib in ["libixattn.so", "libixformer.so", "libcublas.so"]:
p = os.path.join(corex_lib, lib)
if os.path.isfile(p):
extra_ldflags.append(p)
extra_ldflags.append(f"-Wl,-rpath,{corex_lib}")
try:
from torch.utils.cpp_extension import load
logger.info("JIT compiling ix_moe_bridge from %s", cpp_file)
mod = load(
name="ix_moe_bridge",
sources=[cpp_file],
extra_cflags=["-O2", "-std=c++17"],
extra_ldflags=extra_ldflags,
verbose=False,
)
fns = [x for x in dir(mod) if not x.startswith("_")]
logger.info("ix_moe_bridge JIT compiled: %s", fns)
return mod
except Exception as e:
logger.warning("JIT compile failed: %s", e)
return None
def _ensure_loaded():
global _bridge, _loaded
if _loaded:
return _bridge is not None
_loaded = True
_bridge = _try_load_prebuilt()
if _bridge is not None:
return True
_bridge = _try_jit_compile()
if _bridge is not None:
return True
logger.info("ix_moe_bridge unavailable — MoE will use Python loop fallback")
return False
def is_available():
"""Check if the fused MoE bridge is available."""
return _ensure_loaded()
# =========================================================================
# Public API — matches xllm's 7-step pipeline
# =========================================================================
def fused_moe_forward(
hidden_states: torch.Tensor, # (T, H)
router_logits: torch.Tensor, # (T, E)
w13: torch.Tensor, # (E, 2*I, H)
w2: torch.Tensor, # (E, H, I)
topk: int,
num_experts: int,
renormalize: bool = True,
) -> torch.Tensor:
"""Full fused MoE forward — replaces _pure_pytorch_experts().
Pipeline (matching xllm/core/layers/ilu/fused_moe.cpp):
1. topk_softmax — router_logits → (weights, expert_ids)
2. moe_gen_idx — expert_ids → permutation maps
3. moe_expand_input — gather tokens by expert
4. group_gemm 1 — w13 projection (gate+up)
5. silu_and_mul — fused activation
6. group_gemm 2 — w2 projection (down)
7. combine_result — weighted scatter back
"""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.fused_moe_forward(
hidden_states, router_logits, w13, w2,
topk, num_experts, renormalize)
def topk_softmax(gating_output, topk, renormalize=True):
"""Fused topk + softmax routing."""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.topk_softmax(gating_output, topk, renormalize)
def moe_gen_idx(expert_id, expert_num):
"""Build expert permutation maps."""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.moe_gen_idx(expert_id, expert_num)
def group_gemm(inputs, weights, token_count, output_n):
"""Batched expert GEMM via ixformer."""
if _bridge is None:
raise RuntimeError("ix_fused_moe not loaded")
return _bridge.group_gemm(inputs, weights, token_count, output_n)

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;

12
qwen3_6_scripts/patch_ops.sh
View File

@@ -185,8 +185,20 @@ build_stage "installing vLLM Qwen3.6 model implementation"
# --- vllm model: Qwen3.6-35B-A3B (Qwen3_5 MoE arch) -------------------------
cp ./mamba_cache.py "${VLLM_ROOT}/model_executor/models/"
cp ./qwen3_5.py "${VLLM_ROOT}/model_executor/models/qwen3_5.py"
cp ./ix_fused_moe.py "${VLLM_ROOT}/model_executor/models/ix_fused_moe.py" || true
python3 ./patch_vllm_qwen3_5.py
# --- Deploy prebuilt .so into vllm package for import -----------------------
PREBUILT_DIR="./prebuilt/corex-3.2.3-ivcore10"
if [ -d "$PREBUILT_DIR" ]; then
for so_file in "$PREBUILT_DIR"/*.so; do
base=$(basename "$so_file" .so)
# Deploy corex_*.so as vllm submodules (import from vllm import corex_xxx)
cp "$so_file" "${VLLM_ROOT}/${base}.so" 2>/dev/null || true
echo "[patch_ops] deployed ${base}.so → ${VLLM_ROOT}/"
done
fi
# --- sequence.py: fix completion_tokens inflation under chunked prefill ------
# Bug: get_output_token_ids_to_return(delta=True) with num_new_tokens=0
# returns _cached_all_token_ids[-0:] == [0:] (the ENTIRE prompt+output list).

39
qwen3_6_scripts/qwen3_5.py
View File

@@ -204,6 +204,26 @@ _USE_COREX_MOE_INDEX_COMBINE = (
and env_bool("BI100_MOE_COREX_INDEX_COMBINE", True))
_USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True)
# ix_fused_moe: full 7-step fused MoE pipeline via ixformer C++ API
# Source: xllm/core/layers/ilu/fused_moe.cpp → ix_moe_bridge.so
try:
from vllm.model_executor.models import ix_fused_moe as _ix_fused_moe
_HAS_IX_FUSED_MOE = _ix_fused_moe.is_available()
except ImportError:
try:
import ix_fused_moe as _ix_fused_moe
_HAS_IX_FUSED_MOE = _ix_fused_moe.is_available()
except ImportError:
_ix_fused_moe = None
_HAS_IX_FUSED_MOE = False
_USE_IX_FUSED_MOE = (
_HAS_IX_FUSED_MOE
and env_bool("BI100_MOE_IX_FUSED", True))
if _USE_IX_FUSED_MOE:
logger.info("ix_fused_moe ENABLED — full 7-step fused MoE pipeline")
else:
logger.info("ix_fused_moe unavailable — using point-optimized Python MoE")
# ---------------------------------------------------------------------------
# Qwen3.6 vision tower and vLLM 0.6 multimodal input integration
@@ -1620,13 +1640,30 @@ class Qwen3_5MoeSparseBlock(nn.Module):
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
) -> torch.Tensor:
"""Pure-PyTorch MoE (ixformer has no MoE kernels on BI-V100).
"""MoE expert dispatch — fused C++ pipeline when available.
w13_weight: (num_experts, 2*inter_per_partition, hidden) [TP-sharded]
w2_weight: (num_experts, hidden, inter_per_partition) [TP-sharded]
Output is partial (pre-all-reduce), same contract as FusedMoE
with reduce_results=False.
"""
# ---------------------------------------------------------------
# Tier 0: Full fused MoE via ix_moe_bridge (xllm 7-step pipeline)
# topk → gen_idx → expand → group_gemm → silu → group_gemm → combine
# Source: xllm/core/layers/ilu/fused_moe.cpp
# ---------------------------------------------------------------
if _USE_IX_FUSED_MOE:
w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I)
return _ix_fused_moe.fused_moe_forward(
hidden_states, router_logits,
w13, w2,
self.top_k, w13.shape[0],
True) # renormalize
# ---------------------------------------------------------------
# Tier 1: Point-optimized Python loop (individual corex .so)
# ---------------------------------------------------------------
# Fused topk+softmax: single CUB kernel vs 2 PyTorch ops.
# Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh
if _USE_COREX_MOE_TOPK_SOFTMAX:

34
upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h
View File

@@ -19,6 +19,7 @@ limitations under the License.
#include <algorithm>
#include <memory>
#include <optional>
#include <string>
#include <vector>
@@ -38,6 +39,8 @@ limitations under the License.
#include "core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h"
#elif defined(USE_MLU)
#include "core/layers/mlu/qwen3_5/qwen3_5_hybrid_decoder_layer_base.h"
#elif defined(USE_MUSA)
#include "core/layers/musa/qwen3_next_hybrid_decoder_layer_base.h"
#endif
namespace xllm {
@@ -105,12 +108,30 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
}
}
layer::AttentionMetadataBuildOptions metadata_build_options;
#if defined(USE_NPU)
// Native NPU GDN consumes the canonical host mask directly. Avoid
// materializing the unused device bool tensor inside ACL graph capture.
metadata_build_options.materialize_linear_state_validity =
!input_params.enable_graph;
#endif
#if defined(USE_MUSA)
layer::AttentionMetadata attn_metadata =
layer::AttentionMetadataBuilder::build(input_params,
model_args_.enable_mla(),
/*attn_mask=*/std::nullopt,
/*device=*/device_,
metadata_build_options);
attn_metadata.fa3_metadata.share_fa3_scheduler_metadata = true;
#else
layer::AttentionMetadata attn_metadata =
layer::AttentionMetadataBuilder::build(
input_params,
model_args_.enable_mla(),
build_attention_mask(input_params),
/*device=*/device_);
/*device=*/device_,
metadata_build_options);
#endif
const int32_t num_tokens = static_cast<int32_t>(tokens.size(0));
const auto& batch_forward_type = input_params.meta.batch_forward_type;
const bool is_prefill_side = batch_forward_type.no_decode();
@@ -137,6 +158,17 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
std::optional<torch::Tensor> residual = std::nullopt;
for (size_t i = 0; i < layers_.size(); i++) {
#if defined(USE_MUSA)
if (attn_metadata.plan_info != nullptr) {
attn_metadata.plan_info->layer_id = static_cast<int32_t>(i);
}
if (attn_metadata.shared_plan_info != nullptr) {
attn_metadata.shared_plan_info->layer_id = static_cast<int32_t>(i);
}
if (attn_metadata.unshared_plan_info != nullptr) {
attn_metadata.unshared_plan_info->layer_id = static_cast<int32_t>(i);
}
#endif
auto& layer = layers_[i];
h = layer->forward(h,
residual,