feat: Cu10 TensorOp batched HGEMM via Iluvatar CUTLASS framework
From corex-samples batched_gemm.cu, changed: float → half_t, OpClassSimt → OpClassTensorOp, Sm61 → Cu10 Uses __ivcorex_matrix_mad_f32x4_f16x4 via mma_cu10.h Default config: TB<128,128,32> Warp<32,32,32> Inst<16,16,16> Standalone test: correctness + perf for MoE decode (8 × 1x4096@4096x11008)
This commit is contained in:
Claude
25
ex_engine/xllm_kernels/build_test_cutlass_batched.sh
Executable file
25
ex_engine/xllm_kernels/build_test_cutlass_batched.sh
Executable file
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# build_test_cutlass_batched.sh — Compile and test Cu10 TensorOp batched GEMM
|
||||
set -eo pipefail
|
||||
|
||||
SAMPLES="/usr/local/corex-samples-3.2.3_x86_64/samples/cutlass"
|
||||
SRC="ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu"
|
||||
|
||||
echo "=== Compile Cu10 TensorOp batched HGEMM ==="
|
||||
/usr/local/corex/bin/clang++ \
|
||||
--cuda-gpu-arch=ivcore10 --cuda-path=/usr/local/corex \
|
||||
-I"${SAMPLES}/include" \
|
||||
-I/usr/local/corex/include \
|
||||
-L/usr/local/corex/lib64 -lcudart -lcutlass \
|
||||
-DBUILD_STANDALONE_TEST \
|
||||
-O2 -std=c++17 \
|
||||
"$SRC" -o /tmp/test_cutlass_batched 2>&1
|
||||
|
||||
if [ -f /tmp/test_cutlass_batched ]; then
|
||||
echo "Compile: SUCCESS"
|
||||
echo ""
|
||||
echo "=== Run ==="
|
||||
/tmp/test_cutlass_batched
|
||||
else
|
||||
echo "Compile: FAILED"
|
||||
fi
|
||||
233
ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu
Normal file
233
ex_engine/xllm_kernels/cuda/moe_cutlass_batched.cu
Normal file
@@ -0,0 +1,233 @@
|
||||
// moe_cutlass_batched.cu — FP16 Cu10 TensorOp batched GEMM for MoE on BI-V100
|
||||
//
|
||||
// Adapted from corex-samples cutlass/examples/05_batched_gemm/batched_gemm.cu
|
||||
// Changes from original:
|
||||
// 1. float → cutlass::half_t (FP16 data)
|
||||
// 2. arch::OpClassSimt → arch::OpClassTensorOp (use TCU)
|
||||
// 3. arch::Sm61 → arch::Cu10 (BI-V100 arch)
|
||||
// 4. ElementAccumulator = float (FP32 accumulation)
|
||||
// 5. Row-major layout (PyTorch convention) instead of column-major
|
||||
//
|
||||
// Default Cu10 FP16 TensorOp config from default_gemm_configuration.h:
|
||||
// ThreadblockShape = GemmShape<128, 128, 32>
|
||||
// WarpShape = GemmShape<32, 32, 32>
|
||||
// InstructionShape = GemmShape<16, 16, 16>
|
||||
// kStages = 2
|
||||
//
|
||||
// This uses __ivcorex_matrix_mad_f32x4_f16x4 under the hood (via mma_cu10.h).
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/gemm/device/gemm_batched.h"
|
||||
|
||||
// FP16 batched GEMM using Cu10 TensorOp
|
||||
// C[i] = alpha * A[i] @ B[i] + beta * C[i]
|
||||
// All matrices row-major, FP16 in/out, FP32 accumulation.
|
||||
cudaError_t cutlass_batched_hgemm_tensorop(
|
||||
int m, int n, int k,
|
||||
float alpha,
|
||||
cutlass::half_t const *A, int lda, long long int batch_stride_A,
|
||||
cutlass::half_t const *B, int ldb, long long int batch_stride_B,
|
||||
cutlass::half_t *C, int ldc, long long int batch_stride_C,
|
||||
float beta,
|
||||
int batch_count)
|
||||
{
|
||||
using Gemm = cutlass::gemm::device::GemmBatched<
|
||||
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
|
||||
cutlass::arch::OpClassTensorOp, // OperatorClass — use TCU
|
||||
cutlass::arch::Cu10 // ArchTag — BI-V100
|
||||
// Remaining params use defaults from DefaultGemmConfiguration:
|
||||
// ThreadblockShape = <128, 128, 32>
|
||||
// WarpShape = <32, 32, 32>
|
||||
// InstructionShape = <16, 16, 16>
|
||||
// Stages = 2
|
||||
>;
|
||||
|
||||
Gemm gemm_op;
|
||||
|
||||
cutlass::Status status = gemm_op({
|
||||
{m, n, k},
|
||||
{A, lda},
|
||||
batch_stride_A,
|
||||
{B, ldb},
|
||||
batch_stride_B,
|
||||
{C, ldc},
|
||||
batch_stride_C,
|
||||
{C, ldc},
|
||||
batch_stride_C,
|
||||
{alpha, beta},
|
||||
batch_count
|
||||
});
|
||||
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
return cudaErrorUnknown;
|
||||
}
|
||||
|
||||
return cudaSuccess;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Standalone test
|
||||
// ============================================================================
|
||||
#ifdef BUILD_STANDALONE_TEST
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
|
||||
int main() {
|
||||
// Test: 8 batches of (1, 256) @ (256, 128) — simulates decode MoE
|
||||
int m = 1, n = 128, k = 256;
|
||||
int batch_count = 8;
|
||||
float alpha = 1.0f, beta = 0.0f;
|
||||
|
||||
int lda = k; // row-major: (m, k), stride = k
|
||||
int ldb = n; // row-major: (k, n), stride = n
|
||||
int ldc = n; // row-major: (m, n), stride = n
|
||||
|
||||
long long int stride_A = (long long)m * k;
|
||||
long long int stride_B = (long long)k * n;
|
||||
long long int stride_C = (long long)m * n;
|
||||
|
||||
size_t size_A = batch_count * stride_A * sizeof(cutlass::half_t);
|
||||
size_t size_B = batch_count * stride_B * sizeof(cutlass::half_t);
|
||||
size_t size_C = batch_count * stride_C * sizeof(cutlass::half_t);
|
||||
|
||||
// Allocate host
|
||||
std::vector<cutlass::half_t> h_A(batch_count * stride_A);
|
||||
std::vector<cutlass::half_t> h_B(batch_count * stride_B);
|
||||
std::vector<cutlass::half_t> h_C(batch_count * stride_C, cutlass::half_t(0.0f));
|
||||
|
||||
// Fill with small values
|
||||
for (auto &v : h_A) v = cutlass::half_t(0.01f * (rand() % 100 - 50));
|
||||
for (auto &v : h_B) v = cutlass::half_t(0.01f * (rand() % 100 - 50));
|
||||
|
||||
// Allocate device
|
||||
cutlass::half_t *d_A, *d_B, *d_C;
|
||||
cudaMalloc(&d_A, size_A);
|
||||
cudaMalloc(&d_B, size_B);
|
||||
cudaMalloc(&d_C, size_C);
|
||||
|
||||
cudaMemcpy(d_A, h_A.data(), size_A, cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(d_B, h_B.data(), size_B, cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(d_C, h_C.data(), size_C, cudaMemcpyHostToDevice);
|
||||
|
||||
// Run CUTLASS batched GEMM
|
||||
cudaError_t result = cutlass_batched_hgemm_tensorop(
|
||||
m, n, k, alpha,
|
||||
d_A, lda, stride_A,
|
||||
d_B, ldb, stride_B,
|
||||
d_C, ldc, stride_C,
|
||||
beta, batch_count);
|
||||
|
||||
cudaDeviceSynchronize();
|
||||
|
||||
if (result != cudaSuccess) {
|
||||
printf("CUTLASS batched GEMM FAILED: %s\n", cudaGetErrorString(result));
|
||||
cudaError_t last = cudaGetLastError();
|
||||
if (last != cudaSuccess)
|
||||
printf("Last CUDA error: %s\n", cudaGetErrorString(last));
|
||||
cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Copy back
|
||||
cudaMemcpy(h_C.data(), d_C, size_C, cudaMemcpyDeviceToHost);
|
||||
|
||||
// Verify against CPU reference
|
||||
bool pass = true;
|
||||
for (int b = 0; b < batch_count; b++) {
|
||||
for (int i = 0; i < m; i++) {
|
||||
for (int j = 0; j < n; j++) {
|
||||
float ref = 0.0f;
|
||||
for (int p = 0; p < k; p++) {
|
||||
float a_val = float(h_A[b * stride_A + i * k + p]);
|
||||
float b_val = float(h_B[b * stride_B + p * n + j]);
|
||||
ref += a_val * b_val;
|
||||
}
|
||||
float got = float(h_C[b * stride_C + i * n + j]);
|
||||
if (fabs(ref - got) > 1.0f) {
|
||||
printf("MISMATCH batch=%d [%d,%d]: ref=%.4f got=%.4f\n",
|
||||
b, i, j, ref, got);
|
||||
pass = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (pass) {
|
||||
printf("CUTLASS Cu10 TensorOp batched HGEMM: PASSED (%d batches of %dx%d@%dx%d)\n",
|
||||
batch_count, m, k, k, n);
|
||||
}
|
||||
|
||||
// Benchmark
|
||||
cudaEvent_t t0, t1;
|
||||
cudaEventCreate(&t0);
|
||||
cudaEventCreate(&t1);
|
||||
|
||||
// Warmup
|
||||
for (int i = 0; i < 5; i++)
|
||||
cutlass_batched_hgemm_tensorop(m, n, k, alpha,
|
||||
d_A, lda, stride_A, d_B, ldb, stride_B,
|
||||
d_C, ldc, stride_C, beta, batch_count);
|
||||
cudaDeviceSynchronize();
|
||||
|
||||
cudaEventRecord(t0);
|
||||
for (int i = 0; i < 100; i++)
|
||||
cutlass_batched_hgemm_tensorop(m, n, k, alpha,
|
||||
d_A, lda, stride_A, d_B, ldb, stride_B,
|
||||
d_C, ldc, stride_C, beta, batch_count);
|
||||
cudaEventRecord(t1);
|
||||
cudaEventSynchronize(t1);
|
||||
|
||||
float ms;
|
||||
cudaEventElapsedTime(&ms, t0, t1);
|
||||
printf("Perf: %.3f ms/iter (8 batches of 1x256 @ 256x128)\n", ms / 100.0f);
|
||||
|
||||
// Also test MoE-sized: 8 batches of (1, 4096) @ (4096, 11008)
|
||||
int m2 = 1, n2 = 11008, k2 = 4096;
|
||||
long long stride_A2 = (long long)m2 * k2;
|
||||
long long stride_B2 = (long long)k2 * n2;
|
||||
long long stride_C2 = (long long)m2 * n2;
|
||||
|
||||
cutlass::half_t *d_A2, *d_B2, *d_C2;
|
||||
cudaMalloc(&d_A2, batch_count * stride_A2 * sizeof(cutlass::half_t));
|
||||
cudaMalloc(&d_B2, batch_count * stride_B2 * sizeof(cutlass::half_t));
|
||||
cudaMalloc(&d_C2, batch_count * stride_C2 * sizeof(cutlass::half_t));
|
||||
|
||||
for (int i = 0; i < 5; i++)
|
||||
cutlass_batched_hgemm_tensorop(m2, n2, k2, alpha,
|
||||
d_A2, k2, stride_A2, d_B2, n2, stride_B2,
|
||||
d_C2, n2, stride_C2, beta, batch_count);
|
||||
cudaDeviceSynchronize();
|
||||
|
||||
cudaEventRecord(t0);
|
||||
for (int i = 0; i < 20; i++)
|
||||
cutlass_batched_hgemm_tensorop(m2, n2, k2, alpha,
|
||||
d_A2, k2, stride_A2, d_B2, n2, stride_B2,
|
||||
d_C2, n2, stride_C2, beta, batch_count);
|
||||
cudaEventRecord(t1);
|
||||
cudaEventSynchronize(t1);
|
||||
cudaEventElapsedTime(&ms, t0, t1);
|
||||
printf("Perf: %.3f ms/iter (8 batches of 1x4096 @ 4096x11008 — MoE decode)\n", ms / 20.0f);
|
||||
|
||||
cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
|
||||
cudaFree(d_A2); cudaFree(d_B2); cudaFree(d_C2);
|
||||
cudaEventDestroy(t0);
|
||||
cudaEventDestroy(t1);
|
||||
|
||||
return pass ? 0 : -1;
|
||||
}
|
||||
|
||||
#endif // BUILD_STANDALONE_TEST
|
||||
Reference in New Issue
Block a user