feat: xllm MoE CUDA kernels — fused_topk + compute_index + combine

3 MoE kernel files adapted for corex:
  moe_fused_topk.cu: LOG(FATAL)→TORCH_CHECK, +torch/extension.h
  moe_compute_index.cu: CHECK_LE→TORCH_CHECK, uses cub::BlockScan (corex CUB)
  moe_combine.cu: fixed duplicate include, +torch/extension.h

New pybind binding: xllm_moe_bind.cpp
  → moe_fused_topk(gating, topk, renormalize, bias, scoring_func)
  → moe_compute_index(expert_id, num_experts)
  → moe_combine_result(gemm2, weights, N, topk)

AST verification added for all 3 functions
This commit is contained in:
claude
2026-08-14 11:23:49 +00:00
parent 49cd7def89
commit a50adefdfc
6 changed files with 91 additions and 6 deletions

View File

@@ -64,5 +64,8 @@ build_kernel "xllm_rope" \
build_kernel "xllm_cache" \
"${CUDA_DIR}/reshape_paged_cache.cu" "${CUDA_DIR}/block_copy.cu" "${BIND_DIR}/xllm_cache_bind.cpp"
build_kernel "xllm_moe" \
"${CUDA_DIR}/moe/moe_fused_topk.cu" "${CUDA_DIR}/moe/moe_compute_index.cu" "${CUDA_DIR}/moe/moe_combine.cu" "${BIND_DIR}/xllm_moe_bind.cpp"
echo "=== All kernels built ==="
ls -lh "${PREBUILT_DIR}"/xllm_*.so 2>/dev/null || echo "No .so files found"