[init] baseline7 from project_6

ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp

@@ -0,0 +1,129 @@
+/*
+ * corex_batched_gemm_bind.cpp — pybind11 wrapper for CUTLASS batched GEMM
+ *
+ * Kernel uses RowMajor + OpClassTensorOp + Cu10 (verified 2.462ms).
+ * Source: ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu
+ */
+
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+// Implemented in corex_batched_gemm_kernel.cu
+// RowMajor, FP16 data, FP32 accumulation, TCU, Cu10
+cudaError_t cutlass_batched_hgemm(
+    int m, int n, int k,
+    __half const *A, int lda, long long int batch_stride_A,
+    __half const *B, int ldb, long long int batch_stride_B,
+    __half *C, int ldc, long long int batch_stride_C,
+    int batch_count);
+
+/*
+ * batched_gemm_fp16: C[i] = A[i] @ B[i]
+ *   A: (batch, M, K) row-major
+ *   B: (batch, K, N) row-major
+ *   C: (batch, M, N) row-major
+ *
+ * Both A and B must be contiguous fp16 CUDA tensors.
+ */
+torch::Tensor batched_gemm_fp16(
+    torch::Tensor A,   // (batch, M, K)
+    torch::Tensor B)   // (batch, K, N)
+{
+    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA tensors");
+    TORCH_CHECK(A.scalar_type() == torch::kFloat16 &&
+                B.scalar_type() == torch::kFloat16,
+                "inputs must be float16");
+    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
+                "inputs must be contiguous");
+    TORCH_CHECK(A.dim() == 3 && B.dim() == 3,
+                "inputs must be 3D (batch, rows, cols)");
+
+    int batch = A.size(0);
+    int M = A.size(1);
+    int K = A.size(2);
+    int N = B.size(2);
+    TORCH_CHECK(B.size(0) == batch, "batch size mismatch");
+    TORCH_CHECK(B.size(1) == K, "K dimension mismatch");
+
+    auto C = torch::zeros({batch, M, N}, A.options());
+
+    // RowMajor: A is (M,K) with lda=K, B is (K,N) with ldb=N, C is (M,N) with ldc=N
+    auto status = cutlass_batched_hgemm(
+        M, N, K,
+        reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
+        K, (long long)M * K,    // lda, strideA
+        reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
+        N, (long long)K * N,    // ldb, strideB
+        reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
+        N, (long long)M * N,    // ldc, strideC
+        batch);
+
+    TORCH_CHECK(status == cudaSuccess,
+                "CUTLASS batched HGEMM failed: ", cudaGetErrorString(status));
+    return C;
+}
+
+/*
+ * moe_decode_fused: Full MoE decode using TCU batched GEMM.
+ *
+ * hidden_states: (1, H)
+ * w13_sel: (K, 2*I, H)  — already gathered expert weights
+ * w2_sel:  (K, H, I)    — already gathered expert weights
+ * topk_weights: (K,)
+ *
+ * Pipeline:
+ *   1. gate_up = x @ w13^T  via batched GEMM  (K, 1, 2I)
+ *   2. act = silu(gate) * up
+ *   3. down = act @ w2^T    via batched GEMM  (K, 1, H)
+ *   4. out = weighted sum
+ */
+torch::Tensor moe_decode_fused(
+    torch::Tensor hidden_states,   // (1, H)
+    torch::Tensor w13_sel,         // (K, 2*I, H)
+    torch::Tensor w2_sel,          // (K, H, I)
+    torch::Tensor topk_weights)    // (K,)
+{
+    int K_experts = w13_sel.size(0);
+    int two_I = w13_sel.size(1);
+    int H = w13_sel.size(2);
+    int I = two_I / 2;
+
+    // x: (1, H) → expand to (K, 1, H)
+    auto x = hidden_states.expand({K_experts, 1, H}).contiguous();
+
+    // w13^T: (K, 2I, H) → transpose last two dims → (K, H, 2I)
+    auto w13_t = w13_sel.transpose(1, 2).contiguous();  // (K, H, 2I)
+
+    // Step 1: gate_up = x @ w13^T → (K, 1, 2I)
+    auto gate_up = batched_gemm_fp16(x, w13_t);
+    gate_up = gate_up.squeeze(1);  // (K, 2I)
+
+    // Step 2: silu activation
+    auto chunks = gate_up.chunk(2, /*dim=*/1);
+    auto act = torch::sigmoid(chunks[0]) * chunks[0] * chunks[1];  // silu(gate) * up
+    act = act.unsqueeze(1);  // (K, 1, I)
+
+    // w2^T: (K, H, I) → transpose → (K, I, H)
+    auto w2_t = w2_sel.transpose(1, 2).contiguous();  // (K, I, H)
+
+    // Step 3: down = act @ w2^T → (K, 1, H)
+    auto down = batched_gemm_fp16(act, w2_t);
+    down = down.squeeze(1);  // (K, H)
+
+    // Step 4: weighted sum
+    auto out = (down * topk_weights.unsqueeze(1)).sum(0, true);
+    return out.to(hidden_states.dtype());
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "CUTLASS batched GEMM for MoE decode (BI-V100 TCU, Cu10 TensorOp)";
+    m.def("batched_gemm_fp16", &batched_gemm_fp16,
+          "Batched GEMM: (B,M,K) x (B,K,N) -> (B,M,N) in fp16 via TCU",
+          py::arg("A"), py::arg("B"));
+    m.def("moe_decode_fused", &moe_decode_fused,
+          "Full MoE decode via TCU batched GEMM",
+          py::arg("hidden_states"), py::arg("w13_sel"),
+          py::arg("w2_sel"), py::arg("topk_weights"));
+}

ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp

@@ -0,0 +1,135 @@
+// hgemm_bind.cpp — pybind11 bindings for hgemm_blocktiling.cu
+//
+// Exports:
+//   hgemm(A, B, M, N, K) → C
+//   moe_expert_gemm(input, weights, expert_counts) → output
+
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAStream.h>
+#include <vector>
+
+// Forward declarations from hgemm_blocktiling.cu
+void launch_hgemm_blocktiling(
+    int M, int N, int K,
+    const __half* alpha, const __half* A, int lda,
+    const __half* B, int ldb,
+    const __half* beta, __half* C, int ldc,
+    cudaStream_t stream);
+
+void launch_moe_expert_hgemm(
+    int num_experts,
+    const int* expert_counts,
+    const int* expert_offsets,
+    int N, int K,
+    const __half* input,
+    const __half* weights,
+    __half* output,
+    cudaStream_t stream);
+
+
+// ============================================================================
+// Python-facing wrappers
+// ============================================================================
+
+// Simple GEMM: C = A @ B
+// A: (M, K) fp16, B: (K, N) fp16 → C: (M, N) fp16
+torch::Tensor hgemm(torch::Tensor A, torch::Tensor B) {
+    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
+    TORCH_CHECK(A.scalar_type() == torch::kHalf, "A must be fp16");
+    TORCH_CHECK(B.scalar_type() == torch::kHalf, "B must be fp16");
+    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D");
+    TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions must match");
+
+    int M = A.size(0);
+    int K = A.size(1);
+    int N = B.size(1);
+
+    auto C = torch::zeros({M, N}, A.options());
+
+    __half alpha = __float2half(1.0f);
+    __half beta  = __float2half(0.0f);
+
+    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
+
+    launch_hgemm_blocktiling(
+        M, N, K, &alpha,
+        reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
+        A.size(1),
+        reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
+        B.size(1),
+        &beta,
+        reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
+        C.size(1),
+        stream);
+
+    return C;
+}
+
+
+// MoE expert GEMM: for each expert e, compute
+//   output[offset_e : offset_e + count_e] = input[offset_e : offset_e + count_e] @ weights[e].T
+//
+// input:          (total_tokens, K) fp16
+// weights:        (num_experts, N, K) fp16 — weight layout matches vllm w13/w2 convention
+// expert_counts:  (num_experts,) int32 — number of tokens per expert
+//
+// Returns: output (total_tokens, N) fp16
+torch::Tensor moe_expert_gemm(
+    torch::Tensor input,
+    torch::Tensor weights,
+    torch::Tensor expert_counts
+) {
+    TORCH_CHECK(input.is_cuda() && weights.is_cuda(), "Inputs must be CUDA");
+    TORCH_CHECK(input.scalar_type() == torch::kHalf, "input must be fp16");
+    TORCH_CHECK(weights.scalar_type() == torch::kHalf, "weights must be fp16");
+    TORCH_CHECK(expert_counts.scalar_type() == torch::kInt32 ||
+                expert_counts.scalar_type() == torch::kInt64,
+                "expert_counts must be int32 or int64");
+
+    int total_tokens = input.size(0);
+    int K = input.size(1);
+    int num_experts = weights.size(0);
+    int N = weights.size(1);  // output dim
+
+    TORCH_CHECK(weights.size(2) == K, "weights K dim must match input");
+
+    auto output = torch::zeros({total_tokens, N}, input.options());
+
+    // Convert expert_counts to host int array
+    auto counts_cpu = expert_counts.to(torch::kCPU).to(torch::kInt32).contiguous();
+    std::vector<int> counts(num_experts);
+    std::vector<int> offsets(num_experts);
+    int cumsum = 0;
+    for (int i = 0; i < num_experts; i++) {
+        counts[i] = counts_cpu.data_ptr<int32_t>()[i];
+        offsets[i] = cumsum;
+        cumsum += counts[i];
+    }
+
+    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
+
+    launch_moe_expert_hgemm(
+        num_experts,
+        counts.data(),
+        offsets.data(),
+        N, K,
+        reinterpret_cast<const __half*>(input.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(weights.data_ptr<at::Half>()),
+        reinterpret_cast<__half*>(output.data_ptr<at::Half>()),
+        stream);
+
+    return output;
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("hgemm", &hgemm,
+          "FP16 GEMM: C = A @ B (adapted from siboehm kernel 6 for BI-V100)",
+          py::arg("A"), py::arg("B"));
+    m.def("moe_expert_gemm", &moe_expert_gemm,
+          "MoE expert GEMM: per-expert matmul with variable token counts",
+          py::arg("input"), py::arg("weights"), py::arg("expert_counts"));
+}

ex_engine/xllm_kernels/cuda/bindings/hgemm_warp_bind.cpp

@@ -0,0 +1,36 @@
+// hgemm_warp_bind.cpp — pybind11 for hgemm_warptiling (kernel 10, warp64)
+
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAStream.h>
+
+void launch_hgemm_warptiling(
+    int M, int N, int K, float alpha,
+    const __half* A, const __half* B,
+    float beta, __half* C, cudaStream_t stream);
+
+torch::Tensor hgemm_warp(torch::Tensor A, torch::Tensor B) {
+    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
+    TORCH_CHECK(A.scalar_type() == torch::kHalf, "A must be fp16");
+    TORCH_CHECK(B.scalar_type() == torch::kHalf, "B must be fp16");
+    TORCH_CHECK(A.size(1) == B.size(0), "Inner dims must match");
+
+    int M = A.size(0), K = A.size(1), N = B.size(1);
+    auto C = torch::zeros({M, N}, A.options());
+
+    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
+    launch_hgemm_warptiling(M, N, K, 1.0f,
+        reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
+        0.0f,
+        reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
+        stream);
+    return C;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("hgemm_warp", &hgemm_warp,
+          "FP16 GEMM warp-tiling (siboehm K10, WARPSIZE=64 for BI-V100)");
+}

ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp

@@ -0,0 +1,18 @@
+// xllm_activation_bind.cpp
+#include <torch/extension.h>
+
+namespace xllm::kernel::cuda {
+void act_and_mul(torch::Tensor out, torch::Tensor input,
+                 const std::string& act_mode);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("silu_and_mul", [](torch::Tensor out, torch::Tensor input) {
+        xllm::kernel::cuda::act_and_mul(out, input, "silu");
+    }, "SiLU and Mul", py::arg("out"), py::arg("input"));
+    m.def("gelu_and_mul", [](torch::Tensor out, torch::Tensor input) {
+        xllm::kernel::cuda::act_and_mul(out, input, "gelu");
+    }, "GELU and Mul", py::arg("out"), py::arg("input"));
+    m.def("act_and_mul", &xllm::kernel::cuda::act_and_mul,
+          "Activation and Mul", py::arg("out"), py::arg("input"), py::arg("act_mode"));
+}

ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp

@@ -0,0 +1,19 @@
+// xllm_cache_bind.cpp
+#include <torch/extension.h>
+
+namespace xllm::kernel::cuda {
+void reshape_paged_cache(torch::Tensor slot_ids, torch::Tensor keys,
+                         torch::Tensor values, torch::Tensor key_cache,
+                         torch::Tensor value_cache);
+void block_copy(torch::Tensor key_cache_ptrs, torch::Tensor value_cache_ptrs,
+                torch::Tensor src_block_indices, torch::Tensor dst_block_indices,
+                torch::Tensor cum_sum, int64_t numel_per_block,
+                torch::ScalarType cache_dtype);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("reshape_paged_cache", &xllm::kernel::cuda::reshape_paged_cache,
+          "Reshape Paged KV Cache");
+    m.def("block_copy", &xllm::kernel::cuda::block_copy,
+          "Block Copy for KV Cache");
+}

ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp

@@ -0,0 +1,38 @@
+// xllm_fused_qknorm_rope_bind.cpp — pybind11 for fused QK-Norm + RoPE kernel
+// Source: upstream_ref/xllm/xllm/core/kernels/cuda/fused_qknorm_rope.cu
+// Saves 4 kernel launches per layer (separate q_norm, k_norm, q_rope, k_rope)
+// Qwen3.5 has 32 full-attention layers → saves 128 kernel launches per forward
+
+#include <torch/extension.h>
+
+namespace xllm::kernel::cuda {
+void fused_qk_norm_rope(
+    torch::Tensor& qkv,
+    int64_t num_heads_q,
+    int64_t num_heads_k,
+    int64_t num_heads_v,
+    int64_t head_dim,
+    double eps,
+    const torch::Tensor& q_weight,
+    const torch::Tensor& k_weight,
+    const torch::Tensor& cos_sin_cache,
+    bool interleaved,
+    const torch::Tensor& position_ids);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fused_qk_norm_rope",
+          &xllm::kernel::cuda::fused_qk_norm_rope,
+          "Fused QK-Norm + RoPE (xllm CUDA kernel)",
+          py::arg("qkv"),
+          py::arg("num_heads_q"),
+          py::arg("num_heads_k"),
+          py::arg("num_heads_v"),
+          py::arg("head_dim"),
+          py::arg("eps") = 1e-6,
+          py::arg("q_weight"),
+          py::arg("k_weight"),
+          py::arg("cos_sin_cache"),
+          py::arg("interleaved") = false,
+          py::arg("position_ids"));
+}

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/bindings/xllm_norm_bind.cpp

@@ -0,0 +1,24 @@
+// xllm_norm_bind.cpp — pybind11 entry point for xllm norm kernels
+// Compiled together with norm.cu to produce xllm_norm.so
+//
+// Exports: rms_norm, fused_add_rms_norm
+
+#include <torch/extension.h>
+
+namespace xllm::kernel::cuda {
+void rms_norm(torch::Tensor output, torch::Tensor input,
+              torch::Tensor weight, double eps);
+void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
+                        torch::Tensor& weight, double epsilon);
+}  // namespace xllm::kernel::cuda
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("rms_norm", &xllm::kernel::cuda::rms_norm,
+          "RMS Norm (xllm CUDA kernel)",
+          py::arg("output"), py::arg("input"),
+          py::arg("weight"), py::arg("eps") = 1e-6);
+    m.def("fused_add_rms_norm", &xllm::kernel::cuda::fused_add_rms_norm,
+          "Fused Add + RMS Norm (xllm CUDA kernel)",
+          py::arg("input"), py::arg("residual"),
+          py::arg("weight"), py::arg("epsilon") = 1e-6);
+}

ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp

@@ -0,0 +1,17 @@
+// xllm_rope_bind.cpp
+#include <torch/extension.h>
+#include <optional>
+
+namespace xllm::kernel::cuda {
+void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
+                      std::optional<torch::Tensor> key,
+                      torch::Tensor& cos_sin_cache, bool is_neox);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("rotary_embedding", &xllm::kernel::cuda::rotary_embedding,
+          "Rotary Position Embedding (xllm CUDA kernel)",
+          py::arg("positions"), py::arg("query"),
+          py::arg("key"), py::arg("cos_sin_cache"),
+          py::arg("is_neox") = true);
+}