From 330669b309b1c8debe7f39ceb4251fac851aa4b3 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 02:08:03 +0000 Subject: [PATCH] =?UTF-8?q?Revert=20"fix:=20ix=5Ffull=5Fbridge=5Fv2.cpp=20?= =?UTF-8?q?=E2=80=94=20align=20namespace+signatures=20to=20real=20nm=20-D?= =?UTF-8?q?=20symbol=20dump"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 5c03156978d5be24347ea7f64cdf130e0c982809. --- ex_engine/csrc/ix_full_bridge_v2.cpp | 349 ++++++++++++++++----------- 1 file changed, 214 insertions(+), 135 deletions(-) diff --git a/ex_engine/csrc/ix_full_bridge_v2.cpp b/ex_engine/csrc/ix_full_bridge_v2.cpp index d928dce5..576a77be 100644 --- a/ex_engine/csrc/ix_full_bridge_v2.cpp +++ b/ex_engine/csrc/ix_full_bridge_v2.cpp @@ -1,24 +1,21 @@ -// ix_full_bridge_v2.cpp — Bridge to ixformer C++ functions + MoE pipeline +// ix_full_bridge_v2.cpp — Complete bridge to ALL ixformer::infer C++ functions // -// Forward declarations use REAL symbols from nm -D symbol dumps: -// _ixformer_torch.so → namespace ixformer_torch_ext (7 functions) -// moe_ops_impl.cu → namespace ixformer::infer (5 MoE functions, self-compiled) +// Base image has ixformer::infer namespace with 14 functions. +// Previous ix_full_bridge.cpp only bridged 4 (silu_and_mul, rms_norm, +// fused_add_rms_norm, linear). This file bridges ALL 14. // -// Symbol dump verified: -// ixformer_torch_ext::silu_and_mul_forward(at::Tensor&, at::Tensor&) -// ixformer_torch_ext::rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double) -// ixformer_torch_ext::fused_add_rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double, double) -// ixformer_torch_ext::ixformer_linear(at::Tensor&, at::Tensor&, c10::optional, c10::optional) -// ixformer_torch_ext::ixformer_linear_ex(at::Tensor&, at::Tensor&, c10::optional) -// ixformer_torch_ext::vllm_rotary_embedding_neox(at::Tensor&, at::Tensor&, at::Tensor&, long, at::Tensor&, long, bool) -// ixformer_torch_ext::vllm_cache_ops_reshape_and_cache(at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, long, long) -// ixformer_torch_ext::vllm_single_query_cached_kv_attention(at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, double, at::Tensor&, at::Tensor&, long, c10::optional) +// The base image's _ixformer_torch.cpython-310.so and libixformer.so +// export these symbols in the ixformer::infer namespace (confirmed by nm -D). // -// NOT available in any .so (confirmed by nm -D on all 4 .so files): -// ixinfer_flash_attn_unpad_with_block_tables — DOES NOT EXIST -// xllm_paged_attention — DOES NOT EXIST -// topk_softmax, moe_w16a16_group_gemm, etc — DOES NOT EXIST in libixformer.so -// (provided by moe_ops_impl.cu instead) +// Compile: +// torch.utils.cpp_extension.load( +// name="ix_full_bridge_v2", +// sources=["ix_full_bridge_v2.cpp"], +// extra_ldflags=[, "-Wl,-rpath,..."], +// extra_cflags=["-O2", "-std=c++17"], +// ) +// +// Upstream reference: xllm_latest/core/kernels/ilu/ixformer.h #include #include @@ -27,62 +24,103 @@ #include // ============================================================================ -// Forward declarations — ixformer_torch_ext namespace from _ixformer_torch.so -// Signatures EXACTLY match nm -D | c++filt output -// ============================================================================ -namespace ixformer_torch_ext { - -// silu_and_mul_forward(at::Tensor&, at::Tensor&) -void silu_and_mul_forward(at::Tensor& input, at::Tensor& output); - -// rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double) -void rms_norm_forward(at::Tensor& output, at::Tensor& input, - at::Tensor& weight, double eps); - -// fused_add_rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double, double) -void fused_add_rms_norm_forward(at::Tensor& input, at::Tensor& residual, - at::Tensor& weight, double eps, double alpha); - -// ixformer_linear(at::Tensor&, at::Tensor&, c10::optional const&, c10::optional const&) -at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight, - c10::optional const& bias, - c10::optional const& out); - -// ixformer_linear_ex(at::Tensor&, at::Tensor&, c10::optional const&) -at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight, - c10::optional const& bias); - -// vllm_rotary_embedding_neox(at::Tensor&, at::Tensor&, at::Tensor&, long, at::Tensor&, long, bool) -void vllm_rotary_embedding_neox(at::Tensor& positions, at::Tensor& query, - at::Tensor& key, int64_t head_size, - at::Tensor& cos_sin_cache, - int64_t max_position, bool is_neox); - -// vllm_cache_ops_reshape_and_cache(at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, long, long) -void vllm_cache_ops_reshape_and_cache(at::Tensor& key, at::Tensor& value, - at::Tensor& key_cache, - at::Tensor& value_cache, - at::Tensor& slot_mapping, - int64_t key_token_stride, - int64_t value_token_stride); - -// vllm_single_query_cached_kv_attention(at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, double, at::Tensor&, at::Tensor&, long, c10::optional) -void vllm_single_query_cached_kv_attention( - at::Tensor& output, at::Tensor& query, - at::Tensor& key_cache, at::Tensor& value_cache, - at::Tensor& head_mapping, double scale, - at::Tensor& block_tables, at::Tensor& context_lens, - int64_t block_size, - c10::optional alibi_slopes); - -} // 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 +// Forward declarations — ixformer::infer namespace from base image .so +// Signatures EXACTLY match upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h // ============================================================================ namespace ixformer { namespace infer { +// --- Attention --- +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); + +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); + +// --- Activation --- +void silu_and_mul(torch::Tensor& input, torch::Tensor& output); + +// --- Linear --- +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); + +torch::Tensor ixformer_linear_ex(torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& out); + +// --- Cache --- +void xllm_reshape_and_cache(torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& slot_mapping, + int64_t key_token_stride, + int64_t value_token_stride); + +// --- RoPE --- +void xllm_rotary_embedding(torch::Tensor& positions, + torch::Tensor& query, + torch::Tensor& key, + int64_t head_size, + torch::Tensor& cos_sin_cache, + bool is_neox); + +// --- Norm --- +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); + +void rms_norm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& output, + const std::optional& fused_bias, + double eps); + +// --- MoE --- void topk_softmax(torch::Tensor& topk_weights, torch::Tensor& topk_indices, torch::Tensor& token_expert_indices, @@ -94,9 +132,9 @@ void moe_compute_token_index_api( 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, + const c10::optional& expert_mask, + const c10::optional& expert_sizes_cpu, + const c10::optional& expand_tokens_gpu, int64_t start_expert_id, int64_t end_expert_id, int64_t num_experts); @@ -104,7 +142,7 @@ void moe_compute_token_index_api( void moe_expand_input(torch::Tensor outputs, torch::Tensor inputs, torch::Tensor dst_to_src, - const std::optional& src_to_dst, + const c10::optional& src_to_dst, int64_t dst_tokens, int64_t expand_factor); @@ -112,45 +150,52 @@ 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, + const c10::optional& dst_to_src, + const c10::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, + const c10::optional& mul_weight, + const c10::optional& mask, + const c10::optional& extra_residual, double scaling_factor); }} // namespace ixformer::infer // ============================================================================ -// Python wrappers — thin wrappers matching ix_bridge.py's expected API +// Python wrappers — thin wrappers that match ix_bridge.py's expected API // ============================================================================ // --- silu_and_mul --- torch::Tensor ix_silu_and_mul(torch::Tensor input) { int64_t half_dim = input.size(-1) / 2; auto output = input.new_empty({input.size(0), half_dim}); - ixformer_torch_ext::silu_and_mul_forward(input, output); + ixformer::infer::silu_and_mul(input, output); return output; } // --- rms_norm --- void ix_rms_norm(torch::Tensor output, torch::Tensor input, torch::Tensor weight, double eps) { - ixformer_torch_ext::rms_norm_forward(output, input, weight, eps); + ixformer::infer::rms_norm(input, weight, output, + /*fused_bias=*/std::nullopt, eps); } // --- fused_add_rms_norm --- +// residual_rms_norm does: output = rms_norm(input + alpha*residual, weight, eps) +// residual_output = input + alpha*residual void ix_fused_add_rms_norm(torch::Tensor input, torch::Tensor residual, - torch::Tensor weight, double eps) { - ixformer_torch_ext::fused_add_rms_norm_forward( - input, residual, weight, eps, /*alpha=*/1.0); + torch::Tensor weight, torch::Tensor output, + torch::Tensor residual_output, double eps) { + ixformer::infer::residual_rms_norm(input, residual, weight, + output, residual_output, + /*fused_bias=*/std::nullopt, + /*alpha=*/1.0, eps, + /*is_post=*/false); } // --- linear --- @@ -159,55 +204,74 @@ torch::Tensor ix_linear(torch::Tensor input, torch::Tensor weight, auto input_2d = input.view({-1, input.size(-1)}); int64_t m = input_2d.size(0); if (m <= 1 && !bias.has_value()) { - return ixformer_torch_ext::ixformer_linear_ex(input, weight, bias); + return ixformer::infer::ixformer_linear_ex( + input, weight, bias, /*out=*/c10::optional()); } - return ixformer_torch_ext::ixformer_linear( - input, weight, bias, /*out=*/c10::optional()); + return ixformer::infer::ixformer_linear( + input, weight, /*act_type=*/0, bias, + /*out=*/std::nullopt, /*persistent=*/std::nullopt); } // --- rotary_embedding --- void ix_rotary_embedding(torch::Tensor positions, torch::Tensor query, torch::Tensor key, int64_t head_size, torch::Tensor cos_sin_cache, bool is_neox) { - int64_t max_position = cos_sin_cache.size(0); - ixformer_torch_ext::vllm_rotary_embedding_neox( - positions, query, key, head_size, cos_sin_cache, max_position, is_neox); + ixformer::infer::xllm_rotary_embedding( + positions, query, key, head_size, cos_sin_cache, is_neox); } // --- reshape_and_cache --- void ix_reshape_and_cache(torch::Tensor key, torch::Tensor value, torch::Tensor key_cache, torch::Tensor value_cache, torch::Tensor slot_mapping) { + // token stride = product of dims after dim 0 for key/value + // key shape: [num_tokens, num_heads, head_dim] int64_t key_token_stride = 1; for (int i = 1; i < key.dim(); i++) key_token_stride *= key.size(i); int64_t value_token_stride = 1; for (int i = 1; i < value.dim(); i++) value_token_stride *= value.size(i); - ixformer_torch_ext::vllm_cache_ops_reshape_and_cache( + ixformer::infer::xllm_reshape_and_cache( key, value, key_cache, value_cache, slot_mapping, key_token_stride, value_token_stride); } -// --- paged_attention (decode only — no prefill available in .so) --- -void ix_paged_attention( +// --- paged_attention (decode) --- +torch::Tensor ix_paged_attention( torch::Tensor output, torch::Tensor query, torch::Tensor key_cache, torch::Tensor value_cache, - torch::Tensor head_mapping, double scale, + int64_t num_kv_heads, double scale, torch::Tensor block_tables, torch::Tensor context_lens, - int64_t block_size, + int64_t block_size, int64_t max_context_len, const c10::optional& alibi_slopes) { - ixformer_torch_ext::vllm_single_query_cached_kv_attention( + return ixformer::infer::xllm_paged_attention( output, query, key_cache, value_cache, - head_mapping, scale, block_tables, context_lens, - block_size, alibi_slopes); + num_kv_heads, scale, block_tables, context_lens, + block_size, max_context_len, alibi_slopes, + /*causal=*/true, /*window_left=*/-1, /*window_right=*/-1, + /*softcap=*/0.0, /*enable_cuda_graph=*/false, + /*use_sqrt_alibi=*/false, /*sinks=*/std::nullopt); } +// --- flash_attn_prefill --- +torch::Tensor ix_flash_attn_prefill( + torch::Tensor query, torch::Tensor key_cache, torch::Tensor value_cache, + torch::Tensor output, torch::Tensor block_tables, + torch::Tensor cu_seq_q, torch::Tensor cu_seq_k, + 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 = std::nullopt; + 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); +} -// ============================================================================ -// MoE wrappers — call moe_ops_impl.cu implementations -// ============================================================================ - -// --- topk_softmax --- +// --- MoE: topk_softmax --- +// Returns (topk_weights, topk_ids, token_expert_indices) std::tuple ix_topk_softmax(torch::Tensor gating_output, int64_t topk, bool renormalize) { int64_t num_tokens = gating_output.size(0); @@ -225,7 +289,8 @@ ix_topk_softmax(torch::Tensor gating_output, int64_t topk, bool renormalize) { return std::make_tuple(topk_weights, topk_ids, token_expert_indices); } -// --- moe_gen_idx --- +// --- MoE: moe_gen_idx --- +// Equivalent to xllm::kernel::ilu::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()}); @@ -234,9 +299,9 @@ ix_moe_gen_idx(torch::Tensor expert_id, int64_t 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, + /*expert_mask=*/c10::nullopt, + /*expert_sizes_cpu=*/c10::nullopt, + /*expand_tokens_gpu=*/c10::nullopt, /*start_expert_id=*/0, /*end_expert_id=*/expert_num, /*num_experts=*/expert_num); @@ -245,7 +310,7 @@ ix_moe_gen_idx(torch::Tensor expert_id, int64_t expert_num) { return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_cumsum}; } -// --- moe_expand_input --- +// --- MoE: moe_expand_input --- torch::Tensor ix_moe_expand_input(torch::Tensor input, torch::Tensor gather_index, torch::Tensor combine_idx, @@ -257,41 +322,49 @@ torch::Tensor ix_moe_expand_input(torch::Tensor input, return output; } -// --- group_gemm --- +// --- MoE: group_gemm --- torch::Tensor ix_group_gemm(torch::Tensor inputs, torch::Tensor weights, torch::Tensor tokens_per_experts, int64_t output_n) { + // Match upstream xllm/core/kernels/ilu/group_gemm.cpp exactly: + // moe_w16a16_group_gemm(output, input, weight, tokens_per_experts, + // dst_to_src=nullopt, bias=nullopt, + // format="TN", persistent=0, + // output_n=tokens_per_experts.sum()) 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, + /*dst_to_src=*/c10::nullopt, + /*bias=*/c10::nullopt, /*format=*/"TN", /*persistent=*/0, gemm_output_n); return output; } -// --- moe_combine_result --- +// --- MoE: moe_combine_result --- torch::Tensor ix_moe_combine_result(torch::Tensor input, torch::Tensor weight) { + // input: [T*topk, H], weight: [T, topk] 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, + /*mask=*/c10::nullopt, + /*extra_residual=*/c10::nullopt, /*scaling_factor=*/1.0); return output; } -// --- fused_moe_forward (7-step pipeline) --- +// --- MoE: fused_moe_forward (7-step pipeline) --- +// This is the full fused MoE forward: topk → gen_idx → expand → gemm(w13) → +// silu_mul → gemm(w2) → combine torch::Tensor ix_fused_moe_forward( torch::Tensor hidden_states, torch::Tensor router_logits, - torch::Tensor w13, - torch::Tensor w2, + torch::Tensor w13, // [num_experts, 2*intermediate, hidden] + torch::Tensor w2, // [num_experts, hidden, intermediate] int64_t topk, int64_t num_experts, bool renormalize) { @@ -315,7 +388,10 @@ torch::Tensor ix_fused_moe_forward( auto expanded = ix_moe_expand_input(hidden_states, src_dst, dst_src, topk); // Step 4: group_gemm (w13: gate_up projection) + // w13 shape: [num_experts, 2*intermediate, hidden] — pass as-is (3D) + // output_n = tokens_per_experts.sum() per upstream convention int64_t intermediate_2x = w13.size(1); + int64_t output_n_w13 = expert_sizes_gpu.sum().item(); auto gate_up = ix_group_gemm(expanded, w13, expert_sizes_gpu, intermediate_2x); @@ -323,6 +399,7 @@ torch::Tensor ix_fused_moe_forward( auto activated = ix_silu_and_mul(gate_up); // Step 6: group_gemm (w2: down projection) + // w2 shape: [num_experts, hidden, intermediate] — pass as-is (3D) int64_t hidden_size = w2.size(1); auto down = ix_group_gemm(activated, w2, expert_sizes_gpu, hidden_size); @@ -335,48 +412,50 @@ torch::Tensor ix_fused_moe_forward( // ============================================================================ -// Module registration +// Module registration — ALL 14 functions + fused pipeline // ============================================================================ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // Activation m.def("silu_and_mul", &ix_silu_and_mul, - "Fused SiLU+mul via ixformer_torch_ext"); + "Fused SiLU+mul activation via ixformer::infer"); // Norm m.def("rms_norm", &ix_rms_norm, - "RMSNorm via ixformer_torch_ext"); + "RMSNorm via ixformer::infer"); m.def("fused_add_rms_norm", &ix_fused_add_rms_norm, - "Residual + RMSNorm via ixformer_torch_ext"); + "Residual + RMSNorm via ixformer::infer"); // Linear m.def("linear", &ix_linear, - "GEMM via ixformer_torch_ext"); + "GEMM via ixformer::infer (linear/linear_ex)"); // RoPE m.def("rotary_embedding", &ix_rotary_embedding, - "Rotary embedding via ixformer_torch_ext"); + "Rotary position embedding via ixformer::infer"); // Cache m.def("reshape_and_cache", &ix_reshape_and_cache, - "KV cache reshape+store via ixformer_torch_ext"); + "KV cache reshape+store via ixformer::infer"); - // Attention (decode only) + // Attention m.def("paged_attention", &ix_paged_attention, - "Paged attention decode via ixformer_torch_ext"); + "Paged attention decode via ixformer::infer"); + m.def("flash_attn_prefill", &ix_flash_attn_prefill, + "Flash attention prefill via ixformer::infer"); - // MoE (individual steps — from moe_ops_impl.cu) + // MoE (individual steps) m.def("topk_softmax", &ix_topk_softmax, - "MoE topk+softmax routing"); + "MoE topk+softmax routing via ixformer::infer"); m.def("moe_gen_idx", &ix_moe_gen_idx, - "MoE compute token index"); + "MoE compute token index via ixformer::infer"); m.def("moe_expand_input", &ix_moe_expand_input, - "MoE expand input for expert dispatch"); + "MoE expand input for expert dispatch via ixformer::infer"); m.def("group_gemm", &ix_group_gemm, - "MoE grouped GEMM via cuinferCustomGemm"); + "MoE grouped GEMM via ixformer::infer"); m.def("moe_combine_result", &ix_moe_combine_result, - "MoE output reduce sum"); + "MoE output reduce sum via ixformer::infer"); // MoE (fused 7-step pipeline) m.def("fused_moe_forward", &ix_fused_moe_forward, - "Complete fused MoE forward (7-step pipeline)"); + "Complete fused MoE forward (7-step pipeline) via ixformer::infer"); }