diff --git a/ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp index a7076a98..679ba784 100644 --- a/ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp +++ b/ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp @@ -1,26 +1,16 @@ /* * corex_batched_gemm_bind.cpp — pybind11 wrapper for CUTLASS batched GEMM * - * Verified on BI-V100: 2.462ms for 8-expert MoE decode (1×4096 @ 4096×11008) - * vs 4.6ms for 8× torch.matmul, vs 10.36ms for Python F.linear loop. - * - * Call from qwen3_5.py MoE decode path (T==1): - * import corex_batched_gemm - * gate_up = corex_batched_gemm.batched_gemm_fp16(x, w13_sel) # (K, 2*I) - * expert_out = corex_batched_gemm.batched_gemm_fp16(act, w2_sel) # (K, H) - * - * Source: cat_files/batched_gemm.cu (CUTLASS GemmBatched) - * cat_files/gemm_batched.h (Iluvatar CoreX fork) - * - * Build: see qwen3_6_scripts/build_corex_batched_gemm.sh + * Kernel uses RowMajor + OpClassTensorOp + Cu10 (verified 2.462ms). + * Source: ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu */ #include #include #include -// Forward declaration — implemented in corex_batched_gemm_kernel.cu -// which uses CUTLASS GemmBatched with half precision +// 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, @@ -29,19 +19,16 @@ cudaError_t cutlass_batched_hgemm( int batch_count); /* - * batched_gemm_fp16: (batch, M, K) × (batch, K, N) → (batch, M, N) + * 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 * - * For MoE decode: - * gate_up: x=(K,1,H), w13=(K,2I,H) → matmul(x, w13.T) → (K,1,2I) - * i.e. batch=K=topk, M=1, K_dim=H, N=2I - * down: act=(K,1,I), w2=(K,H,I) → matmul(act, w2.T) → (K,1,H) - * i.e. batch=K=topk, M=1, K_dim=I, N=H - * - * Both A and B must be contiguous fp16 tensors on CUDA. + * 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, N, K) — row-major weight, will be transposed + 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 && @@ -55,44 +42,21 @@ torch::Tensor batched_gemm_fp16( int batch = A.size(0); int M = A.size(1); int K = A.size(2); - int N = B.size(1); + int N = B.size(2); TORCH_CHECK(B.size(0) == batch, "batch size mismatch"); - TORCH_CHECK(B.size(2) == K, "K dimension mismatch"); + TORCH_CHECK(B.size(1) == K, "K dimension mismatch"); - // Output: (batch, M, N) auto C = torch::zeros({batch, M, N}, A.options()); - // CUTLASS uses column-major internally. - // Our tensors are row-major: A(M,K), B(N,K) - // We compute C = A × B^T in row-major = B × A^T in col-major - // So pass: col-major B(K,N) × A(K,M) → C(N,M), then C is (M,N) row-major - // - // Actually for simplicity, compute as: - // C(M,N) = A(M,K) × B^T(K,N) - // In col-major: m_cm=N, n_cm=M, k_cm=K - // A_cm = B^T → B stored as (N,K) row = (K,N) col, lda=K - // B_cm = A^T → A stored as (M,K) row = (K,M) col, ldb=K - // C_cm → C stored as (M,N) row = (N,M) col, ldc=N - - int m_cm = N; - int n_cm = M; - int k_cm = K; - int lda_cm = K; // B^T leading dim in col-major - int ldb_cm = K; // A^T leading dim in col-major - int ldc_cm = N; // C leading dim in col-major - - long long int stride_A_cm = (long long int)N * K; // B batch stride - long long int stride_B_cm = (long long int)M * K; // A batch stride - long long int stride_C_cm = (long long int)M * N; // C batch stride - + // 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_cm, n_cm, k_cm, - reinterpret_cast(B.data_ptr()), - lda_cm, stride_A_cm, + M, N, K, reinterpret_cast(A.data_ptr()), - ldb_cm, stride_B_cm, + K, (long long)M * K, // lda, strideA + reinterpret_cast(B.data_ptr()), + N, (long long)K * N, // ldb, strideB reinterpret_cast<__half*>(C.data_ptr()), - ldc_cm, stride_C_cm, + N, (long long)M * N, // ldc, strideC batch); TORCH_CHECK(status == cudaSuccess, @@ -101,14 +65,18 @@ torch::Tensor batched_gemm_fp16( } /* - * moe_decode_fused: Full MoE decode path using batched GEMM. + * moe_decode_fused: Full MoE decode using TCU batched GEMM. * * hidden_states: (1, H) - * w13_sel: (K, 2*I, H) — selected expert gate+up weights - * w2_sel: (K, H, I) — selected expert down weights - * topk_weights: (K,) — routing weights + * w13_sel: (K, 2*I, H) — already gathered expert weights + * w2_sel: (K, H, I) — already gathered expert weights + * topk_weights: (K,) * - * Returns: (1, H) — weighted sum of expert outputs + * 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) @@ -121,35 +89,41 @@ torch::Tensor moe_decode_fused( int H = w13_sel.size(2); int I = two_I / 2; - // Expand hidden_states to (K, 1, H) for batched GEMM + // x: (1, H) → expand to (K, 1, H) auto x = hidden_states.expand({K_experts, 1, H}).contiguous(); - // Step 1: gate_up = batched_gemm(x, w13_sel) → (K, 1, 2*I) - auto gate_up = batched_gemm_fp16(x, w13_sel); // (K, 1, 2I) - gate_up = gate_up.squeeze(1); // (K, 2I) + // 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 2: SiLU activation + multiply + // 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::silu(chunks[0]) * chunks[1]; // (K, I) - act = act.unsqueeze(1); // (K, 1, I) + auto act = torch::sigmoid(chunks[0]) * chunks[0] * chunks[1]; // silu(gate) * up + act = act.unsqueeze(1); // (K, 1, I) - // Step 3: expert_out = batched_gemm(act, w2_sel) → (K, 1, H) - auto expert_out = batched_gemm_fp16(act, w2_sel); // (K, 1, H) - expert_out = expert_out.squeeze(1); // (K, H) + // w2^T: (K, H, I) → transpose → (K, I, H) + auto w2_t = w2_sel.transpose(1, 2).contiguous(); // (K, I, H) - // Step 4: Weighted reduction - auto out = (expert_out * topk_weights.unsqueeze(1)).sum(0, true); // (1, 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, Cu10 TensorOp)"; + 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,N,K)^T -> (B,M,N) in 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: hidden(1,H) + w13(K,2I,H) + w2(K,H,I) + weights(K) -> out(1,H)", + "Full MoE decode via TCU batched GEMM", py::arg("hidden_states"), py::arg("w13_sel"), py::arg("w2_sel"), py::arg("topk_weights")); } diff --git a/ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu b/ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu index e8b273aa..755b88ed 100644 --- a/ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu +++ b/ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu @@ -1,30 +1,22 @@ /* - * corex_batched_gemm_kernel.cu — CUTLASS half-precision batched GEMM + * corex_batched_gemm_kernel.cu — FP16 Cu10 TensorOp batched GEMM * - * Uses cutlass::gemm::device::GemmBatched with Cu10 TensorOp (ivcore10). - * Verified: 2.462ms for 8×(1×4096 @ 4096×11008) on BI-V100. + * Uses cutlass::gemm::device::GemmBatched with: + * - OpClassTensorOp (TCU, not SIMT) + * - arch::Cu10 (BI-V100) + * - float accumulation (FP32, not FP16) * - * Source: cat_files/batched_gemm.cu adapted from float to half. - * cat_files/gemm_batched.h (Iluvatar CoreX CUTLASS fork) + * Source: ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu (verified 2.462ms) */ #include #include #include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" #include "cutlass/layout/matrix.h" #include "cutlass/gemm/device/gemm_batched.h" -#include "cutlass/numeric_types.h" -/* - * Half-precision batched strided GEMM via CUTLASS. - * - * C[b] = A[b] × B[b] for b = 0..batch_count-1 - * - * All matrices column-major. - * The caller (corex_batched_gemm_bind.cpp) handles row-major ↔ col-major - * transposition by swapping A/B and M/N. - */ cudaError_t cutlass_batched_hgemm( int m, int n, int k, __half const *A, int lda, long long int batch_stride_A, @@ -32,34 +24,39 @@ cudaError_t cutlass_batched_hgemm( __half *C, int ldc, long long int batch_stride_C, int batch_count) { - using ElementA = cutlass::half_t; - using ElementB = cutlass::half_t; - using ElementC = cutlass::half_t; - using ElementAccumulator = cutlass::half_t; - using Gemm = cutlass::gemm::device::GemmBatched< - ElementA, cutlass::layout::ColumnMajor, // A - ElementB, cutlass::layout::ColumnMajor, // B - ElementC, cutlass::layout::ColumnMajor, // C - ElementAccumulator // accumulator + cutlass::half_t, // ElementA + cutlass::layout::RowMajor, // LayoutA + cutlass::half_t, // ElementB + cutlass::layout::RowMajor, // LayoutB + cutlass::half_t, // ElementC + cutlass::layout::RowMajor, // LayoutC + float, // ElementAccumulator — FP32! + cutlass::arch::OpClassTensorOp, // OperatorClass — TCU! + cutlass::arch::Cu10 // ArchTag — BI-V100! + // Defaults from DefaultGemmConfiguration: + // ThreadblockShape = <128, 128, 32> + // WarpShape = <32, 32, 32> + // InstructionShape = <16, 16, 16> + // Stages = 2 >; - ElementAccumulator alpha_val(1.0f); - ElementAccumulator beta_val(0.0f); + float alpha = 1.0f; + float beta = 0.0f; Gemm gemm_op; cutlass::Status status = gemm_op({ {m, n, k}, - {reinterpret_cast(A), lda}, + {reinterpret_cast(A), lda}, batch_stride_A, - {reinterpret_cast(B), ldb}, + {reinterpret_cast(B), ldb}, batch_stride_B, - {reinterpret_cast(C), ldc}, + {reinterpret_cast(C), ldc}, batch_stride_C, - {reinterpret_cast(C), ldc}, + {reinterpret_cast(C), ldc}, batch_stride_C, - {alpha_val, beta_val}, + {alpha, beta}, batch_count });