project6-dev
b806b15688
feat(moe): corex_moe_topk_softmax — CUB-based fused topk+softmax kernel
Upstream source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh (855 lines)
Adapted from vllm v0.7.3 / TensorRT-LLM v0.7.1.
Files added from upstream:
- moe_topk_softmax_kernels.cuh — CUB-based fused topk+softmax kernel
- moe_topk_sigmoid_kernels.cuh — sigmoid variant (dependency)
- device_utils.cuh — CUB warp utilities
- corex_moe_topk_softmax.cu — pybind11 wrapper (.so entry point)
- build_corex_moe_topk_softmax.sh — corex clang++ build script
Files modified:
- qwen3_5.py — import corex_moe_topk_softmax, dispatch in _pure_pytorch_experts
- patch_ops.sh — compile .so during docker build
Dispatch: if BI100_MOE_COREX_TOPK_SOFTMAX=1 and .so loaded, use fused kernel.
Fallback: torch.topk + torch.softmax (existing behavior, no breakage).
Replaces 2 PyTorch ops with 1 fused CUB kernel for MoE routing.
2026-08-12 06:15:15 +00:00
..
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 06:15:15 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 06:15:15 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 06:15:15 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 06:15:15 +00:00
2026-08-12 06:15:15 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 06:15:15 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 06:15:15 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 03:31:05 +00:00
2026-08-12 04:22:34 +00:00
2026-08-12 03:31:05 +00:00