70c898ac8b
真机验证发现的问题: 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)
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"));
|
|
}
|