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

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,