fix: ex_engine.python subpackage + flash_qla_sm70 deploy + vllm v0.5.5 MoE kernels
真机验证发现的问题: 1. qwen3_5.py 做 'from ex_engine.python.ix_bridge' 但包结构是 ex_engine.ix_bridge → 创建 python/ 子目录 + symlinks 2. flash_qla_sm70 只部署到 /workspace 没有到 vllm models 目录 → 显式 cp -r 到 VLLM/model_executor/models/ 3. 从 vllm v0.5.5 搬 MoE CUDA kernels (torch::Tensor API): - topk_softmax_kernels.cu (506行, CUB BlockReduce) - moe_align_block_size_kernels.cu (134行) - moe_pybind.cpp (pybind11 入口) 真机验证结果: ✓ ix_bridge import OK, available=True ✓ topk_softmax (64 experts, top8) OK — CUDA kernel 命中 ✓ ix_full_bridge silu_and_mul OK ✓ qwen3_5.py import OK ✓ ex_engine build 2/2 factors ✓ moe_topk_softmax_v3.so 编译成功 ✓ flash_qla_sm70_gdn_strided.so 编译成功 ✗ 单卡 32GB OOM (正常, 竞赛 4卡 tp=4)
This commit is contained in:
42
ex_engine/csrc/moe_v055/moe_pybind.cpp
Normal file
42
ex_engine/csrc/moe_v055/moe_pybind.cpp
Normal file
@@ -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"));
|
||||
}
|
||||
Reference in New Issue
Block a user