fix: corex_batched_gemm use TCU OpClassTensorOp + Cu10 + float accum
Previous version used default SIMT path (25ms). Fixed version matches moe_cutlass_batched.cu (2.462ms): - ElementAccumulator = float (was half_t) - OpClassTensorOp (was default OpClassSimt) - arch::Cu10 (was default Sm61) - RowMajor layout (was ColumnMajor) - torch::sigmoid(x)*x instead of torch::silu (not in corex torch) Also fixed bind.cpp: removed col-major transposition logic, kernel now RowMajor so A(M,K) @ B(K,N) = C(M,N) directly.
This commit is contained in:
dylan
@@ -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 <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
// 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<const __half*>(B.data_ptr<at::Half>()),
|
||||
lda_cm, stride_A_cm,
|
||||
M, N, K,
|
||||
reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
|
||||
ldb_cm, stride_B_cm,
|
||||
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>()),
|
||||
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"));
|
||||
}
|
||||
|
||||
@@ -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 <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#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<OpClassTensorOp, Cu10, half, half, half, float>:
|
||||
// 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<ElementA const *>(A), lda},
|
||||
{reinterpret_cast<cutlass::half_t const *>(A), lda},
|
||||
batch_stride_A,
|
||||
{reinterpret_cast<ElementB const *>(B), ldb},
|
||||
{reinterpret_cast<cutlass::half_t const *>(B), ldb},
|
||||
batch_stride_B,
|
||||
{reinterpret_cast<ElementC const *>(C), ldc},
|
||||
{reinterpret_cast<cutlass::half_t const *>(C), ldc},
|
||||
batch_stride_C,
|
||||
{reinterpret_cast<ElementC *>(C), ldc},
|
||||
{reinterpret_cast<cutlass::half_t *>(C), ldc},
|
||||
batch_stride_C,
|
||||
{alpha_val, beta_val},
|
||||
{alpha, beta},
|
||||
batch_count
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user