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

ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp

@@ -0,0 +1,34 @@
+// xllm_moe_bind.cpp — pybind11 for MoE CUDA kernels
+#include <torch/extension.h>
+#include <optional>
+#include <tuple>
+
+namespace xllm::kernel::cuda {
+std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
+    torch::Tensor& gating_output, int64_t topk, bool renormalize,
+    const std::optional<torch::Tensor>& correction_bias,
+    const std::string& scoring_func);
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> 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"));
+}

ex_engine/xllm_kernels/cuda/moe/moe_combine.cu

@@ -28,7 +28,7 @@ limitations under the License.
 #include <c10/cuda/CUDAGuard.h>
 
 #include "device_utils.cuh"
-#include "device_utils.cuh"
+#include <torch/extension.h>
 
 namespace xllm::kernel::cuda {
 

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 <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
 
 #include <cub/block/block_scan.cuh>
 
@@ -115,7 +116,7 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> moe_compute_index(
   auto stream = at::cuda::getCurrentCUDAStream();
   int64_t N = expert_id.numel();
   int32_t E = static_cast<int32_t>(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();
 

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 <torch/extension.h>
 #endif
 #include "moe_topk_sigmoid_kernels.cuh"
 #include "moe_topk_softmax_kernels.cuh"
@@ -49,8 +50,7 @@ std::tuple<torch::Tensor, torch::Tensor> 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);

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"

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