From a50adefdfca2200568b97b7bedb46bae2ae2acb8 Mon Sep 17 00:00:00 2001 From: claude Date: Fri, 14 Aug 2026 11:23:49 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20xllm=20MoE=20CUDA=20kernels=20=E2=80=94?= =?UTF-8?q?=20fused=5Ftopk=20+=20compute=5Findex=20+=20combine?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../cuda/bindings/xllm_moe_bind.cpp | 34 +++++++++++++ .../xllm_kernels/cuda/moe/moe_combine.cu | 2 +- .../cuda/moe/moe_compute_index.cu | 3 +- .../xllm_kernels/cuda/moe/moe_fused_topk.cu | 4 +- qwen3_6_scripts/build_xllm_kernels.sh | 3 ++ qwen3_6_scripts/verify_ast_chain.py | 51 ++++++++++++++++++- 6 files changed, 91 insertions(+), 6 deletions(-) create mode 100644 ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp diff --git a/ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp new file mode 100644 index 00000000..58053d1a --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp @@ -0,0 +1,34 @@ +// xllm_moe_bind.cpp — pybind11 for MoE CUDA kernels +#include +#include +#include + +namespace xllm::kernel::cuda { +std::tuple moe_fused_topk( + torch::Tensor& gating_output, int64_t topk, bool renormalize, + const std::optional& correction_bias, + const std::string& scoring_func); + +std::tuple moe_compute_index( + const torch::Tensor& expert_id, int64_t num_experts); + +torch::Tensor moe_combine_result( + const torch::Tensor& gemm2, const torch::Tensor& reduce_weight, + int64_t N, int32_t topk); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("moe_fused_topk", &xllm::kernel::cuda::moe_fused_topk, + "MoE fused topk (softmax or sigmoid routing)", + py::arg("gating_output"), py::arg("topk"), + py::arg("renormalize") = true, + py::arg("correction_bias") = py::none(), + py::arg("scoring_func") = "softmax"); + m.def("moe_compute_index", &xllm::kernel::cuda::moe_compute_index, + "MoE compute permutation index (histogram + prefix_sum + place)", + py::arg("expert_id"), py::arg("num_experts")); + m.def("moe_combine_result", &xllm::kernel::cuda::moe_combine_result, + "MoE combine (reorder + weighted sum)", + py::arg("gemm2"), py::arg("reduce_weight"), + py::arg("N"), py::arg("topk")); +} diff --git a/ex_engine/xllm_kernels/cuda/moe/moe_combine.cu b/ex_engine/xllm_kernels/cuda/moe/moe_combine.cu index 4be4078c..8fb585b3 100755 --- a/ex_engine/xllm_kernels/cuda/moe/moe_combine.cu +++ b/ex_engine/xllm_kernels/cuda/moe/moe_combine.cu @@ -28,7 +28,7 @@ limitations under the License. #include #include "device_utils.cuh" -#include "device_utils.cuh" +#include namespace xllm::kernel::cuda { diff --git a/ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu b/ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu index b7eae3bf..e4bfeb0b 100644 --- a/ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu +++ b/ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu @@ -24,6 +24,7 @@ limitations under the License. // expert_offsets = exclusive prefix sum of counts (scratch, reused) #include +#include #include @@ -115,7 +116,7 @@ std::tuple moe_compute_index( auto stream = at::cuda::getCurrentCUDAStream(); int64_t N = expert_id.numel(); int32_t E = static_cast(num_experts); - CHECK_LE(E, kMoeIndexBlock) << "num_experts cannot exceed " << kMoeIndexBlock; + TORCH_CHECK(E <= kMoeIndexBlock, "num_experts cannot exceed ", kMoeIndexBlock); auto expert_id_i32 = expert_id.to(torch::kInt32).contiguous(); auto opt_i32 = expert_id_i32.options(); diff --git a/ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu b/ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu index 186be86e..21808f1d 100644 --- a/ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu +++ b/ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu @@ -16,6 +16,7 @@ limitations under the License. #include "kernels/dcu/dcu_ops_api.h" #else #include "device_utils.cuh" +#include #endif #include "moe_topk_sigmoid_kernels.cuh" #include "moe_topk_softmax_kernels.cuh" @@ -49,8 +50,7 @@ std::tuple moe_fused_topk( topk_sigmoid( topk_weights, topk_ids, gating_output, renormalize, correction_bias); } else { - LOG(FATAL) << "Unsupported scoring function for moe topk: " << scoring_func - << "only softmax and sigmoid are supported"; + TORCH_CHECK(false, "Unsupported scoring function: ", scoring_func); } return std::make_tuple(topk_weights, topk_ids); diff --git a/qwen3_6_scripts/build_xllm_kernels.sh b/qwen3_6_scripts/build_xllm_kernels.sh index 4537b3c8..ce8f7653 100644 --- a/qwen3_6_scripts/build_xllm_kernels.sh +++ b/qwen3_6_scripts/build_xllm_kernels.sh @@ -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" diff --git a/qwen3_6_scripts/verify_ast_chain.py b/qwen3_6_scripts/verify_ast_chain.py index 9efa7784..026a6e6e 100644 --- a/qwen3_6_scripts/verify_ast_chain.py +++ b/qwen3_6_scripts/verify_ast_chain.py @@ -179,7 +179,51 @@ def test_cache(): report("cache.block_copy", "PASS", "loaded OK (complex setup needed for full test)") # ========================================================================= -# 5. Compare against ixformer (base image) if available +# 5. xllm_moe — moe_fused_topk, moe_compute_index, moe_combine_result +# ========================================================================= +def test_moe(): + mod = load_so("xllm_moe") + if mod is None: + report("xllm_moe", "SKIP", "not found") + return + + num_tokens = 8 + num_experts = 64 + topk = 8 + H = 256 + + # --- moe_fused_topk --- + gating = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda") + weights, ids = mod.moe_fused_topk(gating, topk, True, None, "softmax") + assert weights.shape == (num_tokens, topk), f"weights shape {weights.shape}" + assert ids.shape == (num_tokens, topk), f"ids shape {ids.shape}" + w_sum_err = (weights.sum(-1) - 1.0).abs().max().item() + report("moe.fused_topk", "PASS" if w_sum_err < 0.01 else "FAIL", + f"shape=({num_tokens},{topk}) weight_sum_err={w_sum_err:.6f}") + + # --- moe_compute_index --- + expert_ids = ids.reshape(-1) # (num_tokens * topk,) + src_dst, dst_src, expert_sizes = mod.moe_compute_index(expert_ids, num_experts) + total = expert_sizes.sum().item() + report("moe.compute_index", "PASS" if total == num_tokens * topk else "FAIL", + f"total={total} expected={num_tokens * topk}") + + # --- moe_combine_result --- + gemm2 = torch.randn(num_tokens * topk, H, dtype=torch.float16, device="cuda") + rw = weights # (num_tokens, topk) + out = mod.moe_combine_result(gemm2, rw, num_tokens, topk) + assert out.shape == (num_tokens, H), f"output shape {out.shape}" + # Reference: manual weighted sum + ref = torch.zeros(num_tokens, H, dtype=torch.float32, device="cuda") + for i in range(num_tokens): + for k in range(topk): + ref[i] += rw[i, k] * gemm2[i * topk + k].float() + err = (out.float() - ref).abs().max().item() + report("moe.combine_result", "PASS" if err < 0.1 else "FAIL", + f"max_err={err:.6f}") + +# ========================================================================= +# 6. Compare against ixformer (base image) if available # ========================================================================= def test_vs_ixformer(): """Compare our xllm .so output against ixformer's implementation.""" @@ -244,7 +288,10 @@ if __name__ == "__main__": print("[4/5] xllm_cache") test_cache() - print("[5/5] vs ixformer (base image)") + print("[5/6] xllm_moe") + test_moe() + + print("[6/6] vs ixformer (base image)") test_vs_ixformer() elapsed = time.time() - t0