[init] baseline7 from project_6

ex_engine/csrc/moe_v055/moe_pybind.cpp

@@ -0,0 +1,42 @@
+/*
+ * moe_pybind.cpp — pybind11 entry for vllm MoE CUDA kernels
+ * 
+ * Compiled via torch.utils.cpp_extension.load() on BI-V100 (CoreX)
+ * Exposes:
+ *   - topk_softmax(topk_weights, topk_indices, token_expert_indices, gating_output)
+ *   - moe_align_block_size(topk_ids, num_experts, block_size, sorted_token_ids, experts_ids, num_tokens_post_pad)
+ *
+ * Source: vllm v0.5.5 csrc/moe/ (torch::Tensor API, pre-libtorch_stable)
+ */
+
+#include <torch/extension.h>
+
+// Forward declarations matching vllm v0.5.5 signatures
+void topk_softmax(torch::Tensor& topk_weights,
+                  torch::Tensor& topk_indices,
+                  torch::Tensor& token_expert_indices,
+                  torch::Tensor& gating_output);
+
+void moe_align_block_size(torch::Tensor topk_ids,
+                           int64_t num_experts,
+                           int64_t block_size,
+                           torch::Tensor sorted_token_ids,
+                           torch::Tensor experts_ids,
+                           torch::Tensor num_tokens_post_pad);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("topk_softmax", &topk_softmax,
+          "MoE topk softmax (vllm v0.5.5 CUDA kernel)",
+          py::arg("topk_weights"),
+          py::arg("topk_indices"),
+          py::arg("token_expert_indices"),
+          py::arg("gating_output"));
+    m.def("moe_align_block_size", &moe_align_block_size,
+          "MoE align block size (vllm v0.5.5 CUDA kernel)",
+          py::arg("topk_ids"),
+          py::arg("num_experts"),
+          py::arg("block_size"),
+          py::arg("sorted_token_ids"),
+          py::arg("experts_ids"),
+          py::arg("num_tokens_post_pad"));
+}