diff --git a/qwen3_6_scripts/ex_engine/csrc/ix_full_bridge_v2.cpp b/qwen3_6_scripts/ex_engine/csrc/ix_full_bridge_v2.cpp index 22b94ead..1bd213f0 100644 --- a/qwen3_6_scripts/ex_engine/csrc/ix_full_bridge_v2.cpp +++ b/qwen3_6_scripts/ex_engine/csrc/ix_full_bridge_v2.cpp @@ -84,54 +84,10 @@ void vllm_single_query_cached_kv_attention( } // namespace ixformer_torch_ext // ============================================================================ -// Forward declarations — ixformer::infer namespace from moe_ops_impl.cu -// These 5 MoE functions are compiled from our own CUDA code, NOT from .so +// MoE functions are NOT available in this POC image's ixformer .so files. +// MoE path uses separate prebuilt .so (corex_moe_*.so, gemm_grouped.so) +// dispatched via ix_fused_moe.py at runtime. // ============================================================================ -namespace ixformer { namespace infer { - -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 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); - -void moe_expand_input(torch::Tensor outputs, - torch::Tensor inputs, - torch::Tensor dst_to_src, - const std::optional& 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 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 std::optional& mul_weight, - const std::optional& mask, - const std::optional& extra_residual, - double scaling_factor); - -}} // namespace ixformer::infer // ============================================================================ @@ -211,138 +167,6 @@ void ix_paged_attention( /*is_neox=*/true, alibi_slopes); } - -// ============================================================================ -// MoE wrappers — call moe_ops_impl.cu implementations -// ============================================================================ - -// --- topk_softmax --- -std::tuple -ix_topk_softmax(torch::Tensor gating_output, int64_t topk, bool renormalize) { - int64_t num_tokens = gating_output.size(0); - auto topk_weights = torch::empty({num_tokens, topk}, - torch::dtype(torch::kFloat32).device(gating_output.device())); - auto topk_ids = torch::empty({num_tokens, topk}, - torch::dtype(torch::kInt32).device(gating_output.device())); - auto token_expert_indices = torch::empty({num_tokens, topk}, - torch::dtype(torch::kInt32).device(gating_output.device())); - - auto gating_f32 = gating_output.to(torch::kFloat32); - ixformer::infer::topk_softmax( - topk_weights, topk_ids, token_expert_indices, gating_f32, renormalize); - - return std::make_tuple(topk_weights, topk_ids, token_expert_indices); -} - -// --- moe_gen_idx --- -std::vector -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, - /*expert_mask=*/std::nullopt, - /*expert_sizes_cpu=*/std::nullopt, - /*expand_tokens_gpu=*/std::nullopt, - /*start_expert_id=*/0, - /*end_expert_id=*/expert_num, - /*num_experts=*/expert_num); - - auto expert_sizes_cumsum = expert_sizes_gpu.cumsum(-1); - return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_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; -} - -// --- group_gemm --- -torch::Tensor ix_group_gemm(torch::Tensor inputs, torch::Tensor weights, - torch::Tensor tokens_per_experts, - int64_t output_n) { - int64_t total_tokens = inputs.size(0); - auto output = inputs.new_empty({total_tokens, output_n}); - int64_t gemm_output_n = tokens_per_experts.sum().item(); - ixformer::infer::moe_w16a16_group_gemm( - output, inputs, weights, tokens_per_experts, - /*dst_to_src=*/std::nullopt, - /*bias=*/std::nullopt, - /*format=*/"TN", - /*persistent=*/0, - gemm_output_n); - return output; -} - -// --- moe_combine_result --- -torch::Tensor ix_moe_combine_result(torch::Tensor input, torch::Tensor weight) { - auto input_3d = input.view({-1, weight.size(1), input.size(1)}); - auto output = input.new_empty({input_3d.size(0), input_3d.size(2)}); - ixformer::infer::moe_output_reduce_sum( - output, input_3d, weight, - /*mask=*/std::nullopt, - /*extra_residual=*/std::nullopt, - /*scaling_factor=*/1.0); - return output; -} - -// --- fused_moe_forward (7-step pipeline) --- -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) { - - // Step 1: topk_softmax - auto [topk_weights, topk_ids, token_expert_indices] = - ix_topk_softmax(router_logits, topk, renormalize); - - if (renormalize) { - auto sum = topk_weights.sum(-1, /*keepdim=*/true); - topk_weights = topk_weights / sum; - } - - // Step 2: moe_gen_idx - auto idx_results = ix_moe_gen_idx(topk_ids.view({-1}), num_experts); - auto& src_dst = idx_results[0]; - auto& dst_src = idx_results[1]; - auto& expert_sizes_gpu = idx_results[2]; - - // Step 3: moe_expand_input - auto expanded = ix_moe_expand_input(hidden_states, src_dst, dst_src, topk); - - // Step 4: group_gemm (w13: gate_up projection) - int64_t intermediate_2x = w13.size(1); - auto gate_up = ix_group_gemm(expanded, w13, - expert_sizes_gpu, intermediate_2x); - - // Step 5: silu_and_mul - auto activated = ix_silu_and_mul(gate_up); - - // Step 6: group_gemm (w2: down projection) - int64_t hidden_size = w2.size(1); - auto down = ix_group_gemm(activated, w2, - expert_sizes_gpu, hidden_size); - - // Step 7: moe_combine_result - auto output = ix_moe_combine_result(down, topk_weights); - - return output; -} - - // ============================================================================ // Module registration // ============================================================================ @@ -373,19 +197,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("paged_attention", &ix_paged_attention, "Paged attention decode via ixformer_torch_ext"); - // MoE (individual steps — from moe_ops_impl.cu) - m.def("topk_softmax", &ix_topk_softmax, - "MoE topk+softmax routing"); - m.def("moe_gen_idx", &ix_moe_gen_idx, - "MoE compute token index"); - m.def("moe_expand_input", &ix_moe_expand_input, - "MoE expand input for expert dispatch"); - m.def("group_gemm", &ix_group_gemm, - "MoE grouped GEMM via cuinferCustomGemm"); - m.def("moe_combine_result", &ix_moe_combine_result, - "MoE output reduce sum"); - - // MoE (fused 7-step pipeline) - m.def("fused_moe_forward", &ix_fused_moe_forward, - "Complete fused MoE forward (7-step pipeline)"); + // MoE ops are NOT in this bridge — they use separate prebuilt .so files: + // corex_moe_topk_softmax.so, corex_moe_index_combine.so, + // gemm_grouped.so, corex_moe_exact_reduce.so, etc. + // Dispatched via ix_fused_moe.py at runtime. } \ No newline at end of file diff --git a/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ix_full_bridge.so b/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ix_full_bridge.so index 5d12c19f..7c2610eb 100755 Binary files a/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ix_full_bridge.so and b/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ix_full_bridge.so differ