dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Revert "fix: ix_full_bridge_v2.cpp — align namespace+signatures to real nm -D symbol dump"

Browse Source 
This reverts commit 5c03156978.
This commit is contained in:
Claude
2026-08-17 02:08:03 +00:00
parent 5c03156978
commit 330669b309
1 changed files with 214 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

349
ex_engine/csrc/ix_full_bridge_v2.cpp
View File

@@ -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<at::Tensor>, c10::optional<at::Tensor>)
// ixformer_torch_ext::ixformer_linear_ex(at::Tensor&, at::Tensor&, c10::optional<at::Tensor>)
// 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<at::Tensor>)
// 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=[<all ixformer .so files>, "-Wl,-rpath,..."],
// extra_cflags=["-O2", "-std=c++17"],
// )
//
// Upstream reference: xllm_latest/core/kernels/ilu/ixformer.h
#include <torch/extension.h>
#include <optional>
@@ -27,62 +24,103 @@
#include <vector>
// ============================================================================
// 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<at::Tensor> const&, c10::optional<at::Tensor> const&)
at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight,
c10::optional<at::Tensor> const& bias,
c10::optional<at::Tensor> const& out);
// ixformer_linear_ex(at::Tensor&, at::Tensor&, c10::optional<at::Tensor> const&)
at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight,
c10::optional<at::Tensor> 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<at::Tensor>)
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<at::Tensor> 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<torch::Tensor>& alibi_slopes,
const std::optional<torch::Tensor>& sinks,
std::optional<torch::Tensor>& 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<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);
// --- 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<torch::Tensor>& bias,
const std::optional<torch::Tensor>& out,
const std::optional<bool> persistent);
torch::Tensor ixformer_linear_ex(torch::Tensor& input,
torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& 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<torch::Tensor>& fused_bias,
double alpha,
double eps,
bool is_post);
void rms_norm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& output,
const std::optional<torch::Tensor>& 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<torch::Tensor>& expert_mask,
const std::optional<torch::Tensor>& expert_sizes_cpu,
const std::optional<torch::Tensor>& expand_tokens_gpu,
const c10::optional<torch::Tensor>& expert_mask,
const c10::optional<torch::Tensor>& expert_sizes_cpu,
const c10::optional<torch::Tensor>& 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<torch::Tensor>& src_to_dst,
const c10::optional<torch::Tensor>& 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<torch::Tensor>& dst_to_src,
const std::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& dst_to_src,
const c10::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 std::optional<torch::Tensor>& mul_weight,
const std::optional<torch::Tensor>& mask,
const std::optional<torch::Tensor>& extra_residual,
const c10::optional<torch::Tensor>& mul_weight,
const c10::optional<torch::Tensor>& mask,
const c10::optional<torch::Tensor>& 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<torch::Tensor>());
}
return ixformer_torch_ext::ixformer_linear(
input, weight, bias, /*out=*/c10::optional<at::Tensor>());
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<torch::Tensor>& 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<torch::Tensor> 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<torch::Tensor, torch::Tensor, torch::Tensor>
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<torch::Tensor>
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<int64_t>();
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<int64_t>();
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");
}