EngineX
|
7839982707
|
feat(EX): wire xllm CUB topk_softmax kernel into MoE routing
Upstream: xllm/kernels/cuda/moe/moe_topk_softmax_kernels.cuh (Apache 2.0)
Adapted: CHECK→TORCH_CHECK, include path fix, cuda/functional guard, pybind11
Call chain now:
qwen3_5.py:_pure_pytorch_experts()
→ _ex_moe_topk_softmax (fused CUB kernel, 1 launch)
→ fallback: torch.softmax + torch.topk (3 launches)
Files:
ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh — xllm kernel (adapted)
ex_engine/csrc/moe/device_utils.cuh — xllm device utils
ex_engine/csrc/moe/moe_topk_softmax_ext.cu — pybind11 wrapper
ex_engine/python/moe_topk.py — JIT loader (same pattern as flash_qla_sm70)
qwen3_5.py — import + use in _pure_pytorch_experts()
patch_ops.sh — deploy kernel sources for JIT
|
2026-08-10 03:10:58 +00:00 |
|