43 lines
1.6 KiB
C++
43 lines
1.6 KiB
C++
|
/*
|
||
|
|
* 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"));
|
||
|
|
}
|