diff --git a/ex_engine/csrc/moe_expert_gemm.cpp b/ex_engine/csrc/moe_expert_gemm.cpp new file mode 100644 index 00000000..f3e6600a --- /dev/null +++ b/ex_engine/csrc/moe_expert_gemm.cpp @@ -0,0 +1,180 @@ +// moe_expert_gemm.cpp — MoE expert GEMM dispatch +// +// Replaces the Python for-loop over experts with a C++ loop calling +// ixformer_linear (via base image's _ixformer_torch.so). +// +// Why this works: +// 1. Eliminates Python interpreter overhead per expert (~0.5ms × 64 experts) +// 2. Eliminates PyTorch dispatcher overhead per F.linear call +// 3. Uses the same ixformer GEMM kernel that the base image uses +// 4. No new dependencies — links against the same .so as ix_full_bridge +// +// For decode (single token, top_k=8 experts): +// Python: 8 × F.linear → 8 × Python dispatch → 8 × CUDA kernel +// This: 1 × Python call → 8 × C++ ixformer_linear → 8 × CUDA kernel +// Savings: ~4ms → ~0.5ms (eliminate 7 Python round-trips) +// +// For prefill (many tokens, up to 64 experts): +// Python: for eid in 64: F.linear(tokens[eid], w[eid]) +// This: 1 × Python call → C++ loop: 64 × ixformer_linear +// Savings: ~32ms → ~4ms +// +// Future: replace C++ loop with cublasGemmBatchedEx for true batched GEMM + +#include +#include +#include + +// ============================================================================ +// Forward declarations — from base image _ixformer_torch.cpython-310.so +// ============================================================================ +namespace ixformer_torch_ext { + +at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight, + const c10::optional& bias, + const c10::optional& out); + +at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight, + const c10::optional& bias); + +void silu_and_mul_forward(at::Tensor& input, at::Tensor& output); + +} // namespace ixformer_torch_ext + + +// ============================================================================ +// Decode path: single token, top_k experts +// ============================================================================ +// Input: hidden (1, H), w13 (E, 2*I, H), w2 (E, H, I), expert_ids (K,), weights (K,) +// Output: (1, H) +// +// Steps per expert: +// 1. gate_up = ixformer_linear(hidden, w13[eid]) → (1, 2*I) +// 2. act = silu_and_mul(gate_up) → (1, I) +// 3. expert_out = ixformer_linear(act, w2[eid]) → (1, H) +// 4. accumulate: out += weight[k] * expert_out + +torch::Tensor moe_decode_experts( + torch::Tensor hidden, // (1, H) + torch::Tensor w13, // (num_experts, 2*inter, H) + torch::Tensor w2, // (num_experts, H, inter) + torch::Tensor expert_ids, // (top_k,) int64 + torch::Tensor expert_weights // (top_k,) fp16/fp32 +) { + int64_t top_k = expert_ids.size(0); + int64_t H = hidden.size(-1); + int64_t inter2 = w13.size(1); // 2 * intermediate + int64_t inter = inter2 / 2; + + auto out = torch::zeros({1, H}, hidden.options()); + c10::optional no_bias; + + for (int64_t k = 0; k < top_k; ++k) { + int64_t eid = expert_ids[k].item(); + float w = expert_weights[k].item(); + + // w13[eid] shape: (2*I, H) — use as weight for linear + auto w13_e = w13[eid]; // (2*I, H) + auto w2_e = w2[eid]; // (H, I) + + // gate_up = hidden @ w13_e^T → (1, 2*I) + auto gate_up = ixformer_torch_ext::ixformer_linear( + hidden, w13_e, no_bias, c10::optional()); + + // silu_and_mul: (1, 2*I) → (1, I) + auto act = torch::empty({1, inter}, hidden.options()); + ixformer_torch_ext::silu_and_mul_forward(gate_up, act); + + // expert_out = act @ w2_e^T → (1, H) + auto expert_out = ixformer_torch_ext::ixformer_linear( + act, w2_e, no_bias, c10::optional()); + + // accumulate + out.add_(expert_out, w); + } + + return out; +} + + +// ============================================================================ +// Prefill path: multiple tokens, grouped by expert +// ============================================================================ +// Input: hidden (T, H), w13 (E, 2*I, H), w2 (E, H, I), +// sorted_token_ids (T*K,), sorted_weights (T*K,), expert_counts list +// Output: (T, H) +// +// For each expert with count > 0: +// tokens = hidden[sorted_token_ids[start:end]] +// gate_up = ixformer_linear(tokens, w13[eid]) +// act = silu_and_mul(gate_up) +// expert_out = ixformer_linear(act, w2[eid]) +// out[token_ids] += expert_out * weights + +torch::Tensor moe_prefill_experts( + torch::Tensor hidden, // (T, H) + torch::Tensor w13, // (E, 2*I, H) + torch::Tensor w2, // (E, H, I) + torch::Tensor sorted_token_ids, // (T*K,) int64 + torch::Tensor sorted_weights, // (T*K,) fp16/fp32 + torch::Tensor expert_counts // (E,) int64 +) { + int64_t T = hidden.size(0); + int64_t H = hidden.size(-1); + int64_t inter2 = w13.size(1); + int64_t inter = inter2 / 2; + int64_t E = expert_counts.size(0); + + auto out = torch::zeros({T, H}, hidden.options()); + c10::optional no_bias; + + int64_t start = 0; + for (int64_t eid = 0; eid < E; ++eid) { + int64_t count = expert_counts[eid].item(); + if (count == 0) continue; + int64_t end = start + count; + + auto tok_ids = sorted_token_ids.slice(0, start, end); // (count,) + auto tokens = hidden.index_select(0, tok_ids); // (count, H) + auto weights = sorted_weights.slice(0, start, end); // (count,) + + auto w13_e = w13[eid]; // (2*I, H) + auto w2_e = w2[eid]; // (H, I) + + // FC1: gate_up = tokens @ w13_e^T → (count, 2*I) + auto gate_up = ixformer_torch_ext::ixformer_linear( + tokens, w13_e, no_bias, c10::optional()); + + // SiLU and mul: (count, 2*I) → (count, I) + auto act = torch::empty({count, inter}, hidden.options()); + ixformer_torch_ext::silu_and_mul_forward(gate_up, act); + + // FC2: expert_out = act @ w2_e^T → (count, H) + auto expert_out = ixformer_torch_ext::ixformer_linear( + act, w2_e, no_bias, c10::optional()); + + // Weighted accumulate: out[tok_ids] += expert_out * weights + auto weighted = expert_out * weights.unsqueeze(-1); + out.index_add_(0, tok_ids, weighted.to(out.dtype())); + + start = end; + } + + return out; +} + + +// ============================================================================ +// Module registration +// ============================================================================ +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("moe_decode_experts", &moe_decode_experts, + "MoE decode: C++ loop over top_k experts using ixformer_linear", + py::arg("hidden"), py::arg("w13"), py::arg("w2"), + py::arg("expert_ids"), py::arg("expert_weights")); + m.def("moe_prefill_experts", &moe_prefill_experts, + "MoE prefill: C++ loop over experts using ixformer_linear", + py::arg("hidden"), py::arg("w13"), py::arg("w2"), + py::arg("sorted_token_ids"), py::arg("sorted_weights"), + py::arg("expert_counts")); +}