diff --git a/ex_engine/csrc/ix_attn_bridge.cpp b/ex_engine/csrc/ix_attn_bridge.cpp new file mode 100644 index 00000000..0e99ad0f --- /dev/null +++ b/ex_engine/csrc/ix_attn_bridge.cpp @@ -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 +#include + +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& alibi_slopes, + const std::optional& sinks, + std::optional& 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& 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& sinks); + +// Fused linear: matmul + optional activation +torch::Tensor ixformer_linear( + torch::Tensor& input, + torch::Tensor& weight, + int64_t act_type, + const std::optional& bias, + const std::optional& out, + const std::optional persistent); + +// Simple linear +torch::Tensor ixformer_linear_ex( + torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& 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& 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 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 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); +} diff --git a/ex_engine/csrc/ix_moe_bridge.cpp b/ex_engine/csrc/ix_moe_bridge.cpp index 6e294984..d56d6880 100644 --- a/ex_engine/csrc/ix_moe_bridge.cpp +++ b/ex_engine/csrc/ix_moe_bridge.cpp @@ -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; diff --git a/qwen3_6_scripts/build_ix_attn_bridge.sh b/qwen3_6_scripts/build_ix_attn_bridge.sh new file mode 100644 index 00000000..66c23041 --- /dev/null +++ b/qwen3_6_scripts/build_ix_attn_bridge.sh @@ -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 ===') +" diff --git a/qwen3_6_scripts/build_ix_moe_bridge.sh b/qwen3_6_scripts/build_ix_moe_bridge.sh new file mode 100644 index 00000000..ff5fa361 --- /dev/null +++ b/qwen3_6_scripts/build_ix_moe_bridge.sh @@ -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 ===') +" diff --git a/qwen3_6_scripts/ix_attn_bridge.cpp b/qwen3_6_scripts/ix_attn_bridge.cpp new file mode 100644 index 00000000..0e99ad0f --- /dev/null +++ b/qwen3_6_scripts/ix_attn_bridge.cpp @@ -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 +#include + +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& alibi_slopes, + const std::optional& sinks, + std::optional& 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& 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& sinks); + +// Fused linear: matmul + optional activation +torch::Tensor ixformer_linear( + torch::Tensor& input, + torch::Tensor& weight, + int64_t act_type, + const std::optional& bias, + const std::optional& out, + const std::optional persistent); + +// Simple linear +torch::Tensor ixformer_linear_ex( + torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& 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& 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 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 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); +} diff --git a/qwen3_6_scripts/ix_fused_moe.py b/qwen3_6_scripts/ix_fused_moe.py new file mode 100644 index 00000000..135ef329 --- /dev/null +++ b/qwen3_6_scripts/ix_fused_moe.py @@ -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) diff --git a/qwen3_6_scripts/ix_moe_bridge.cpp b/qwen3_6_scripts/ix_moe_bridge.cpp index 3c228580..d56d6880 100644 --- a/qwen3_6_scripts/ix_moe_bridge.cpp +++ b/qwen3_6_scripts/ix_moe_bridge.cpp @@ -17,7 +17,7 @@ #include #include -static const c10::optional kNoneTensor = {}; +static const std::optional 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& expert_mask, - const c10::optional& expert_sizes_cpu, - const c10::optional& expand_tokens_gpu, + const std::optional& expert_mask, + const std::optional& expert_sizes_cpu, + const std::optional& 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& src_to_dst, + const std::optional& 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& dst_to_src, - const c10::optional& bias, + const std::optional& dst_to_src, + const std::optional& 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& mul_weight, - const c10::optional& mask, - const c10::optional& extra_residual, + const std::optional& mul_weight, + const std::optional& mask, + const std::optional& 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; diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index b4aa0fbb..44737c26 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -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). diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 8103948e..03bff5cc 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -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: diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h b/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h index 83e42e59..50b4bc8e 100644 --- a/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h +++ b/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include #include @@ -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(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 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(i); + } + if (attn_metadata.shared_plan_info != nullptr) { + attn_metadata.shared_plan_info->layer_id = static_cast(i); + } + if (attn_metadata.unshared_plan_info != nullptr) { + attn_metadata.unshared_plan_info->layer_id = static_cast(i); + } +#endif auto& layer = layers_[i]; h = layer->forward(h, residual,