test: sweep 6 kernel 10 configs + cublas baseline — find best params for warp64
This commit is contained in:
Claude
289
qwen3_6_scripts/probe_k10_configs.sh
Executable file
289
qwen3_6_scripts/probe_k10_configs.sh
Executable file
@@ -0,0 +1,289 @@
|
||||
#!/bin/bash
|
||||
# probe_k10_configs.sh — Test multiple kernel 10 configs on BI-V100
|
||||
set -eo pipefail
|
||||
|
||||
cat > /tmp/probe_k10_configs.cu << 'CUDA'
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cstdio>
|
||||
|
||||
#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
|
||||
const int WARPSIZE = 64;
|
||||
|
||||
// Same kernel code as before
|
||||
namespace wt {
|
||||
template <const int BM, const int BN, const int BK, const int rowStrideA,
|
||||
const int rowStrideB>
|
||||
__device__ void loadFromGmem(int N, int K, const __half *A, const __half *B,
|
||||
__half *As, __half *Bs, int innerRowA, int innerColA,
|
||||
int innerRowB, int innerColB) {
|
||||
for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
|
||||
__half a0 = A[(innerRowA + offset) * K + innerColA * 4 + 0];
|
||||
__half a1 = A[(innerRowA + offset) * K + innerColA * 4 + 1];
|
||||
__half a2 = A[(innerRowA + offset) * K + innerColA * 4 + 2];
|
||||
__half a3 = A[(innerRowA + offset) * K + innerColA * 4 + 3];
|
||||
As[(innerColA * 4 + 0) * BM + innerRowA + offset] = a0;
|
||||
As[(innerColA * 4 + 1) * BM + innerRowA + offset] = a1;
|
||||
As[(innerColA * 4 + 2) * BM + innerRowA + offset] = a2;
|
||||
As[(innerColA * 4 + 3) * BM + innerRowA + offset] = a3;
|
||||
}
|
||||
for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
|
||||
Bs[(innerRowB + offset) * BN + innerColB * 4 + 0] =
|
||||
B[(innerRowB + offset) * N + innerColB * 4 + 0];
|
||||
Bs[(innerRowB + offset) * BN + innerColB * 4 + 1] =
|
||||
B[(innerRowB + offset) * N + innerColB * 4 + 1];
|
||||
Bs[(innerRowB + offset) * BN + innerColB * 4 + 2] =
|
||||
B[(innerRowB + offset) * N + innerColB * 4 + 2];
|
||||
Bs[(innerRowB + offset) * BN + innerColB * 4 + 3] =
|
||||
B[(innerRowB + offset) * N + innerColB * 4 + 3];
|
||||
}
|
||||
}
|
||||
|
||||
template <const int BM, const int BN, const int BK, const int WM, const int WN,
|
||||
const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
|
||||
const int TM, const int TN>
|
||||
__device__ void
|
||||
processFromSmem(float *regM, float *regN, float *threadResults, const __half *As,
|
||||
const __half *Bs, const uint warpRow, const uint warpCol,
|
||||
const uint threadRowInWarp, const uint threadColInWarp) {
|
||||
for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
|
||||
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
|
||||
for (uint i = 0; i < TM; ++i) {
|
||||
regM[wSubRowIdx * TM + i] = __half2float(
|
||||
As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
|
||||
threadRowInWarp * TM + i]);
|
||||
}
|
||||
}
|
||||
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
|
||||
for (uint i = 0; i < TN; ++i) {
|
||||
regN[wSubColIdx * TN + i] = __half2float(
|
||||
Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
|
||||
threadColInWarp * TN + i]);
|
||||
}
|
||||
}
|
||||
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
|
||||
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
|
||||
for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
|
||||
for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
|
||||
threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
|
||||
(wSubColIdx * TN) + resIdxN] +=
|
||||
regM[wSubRowIdx * TM + resIdxM] *
|
||||
regN[wSubColIdx * TN + resIdxN];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace wt
|
||||
|
||||
template <const int BM, const int BN, const int BK, const int WM, const int WN,
|
||||
const int WNITER, const int TM, const int TN, const int NUM_THREADS>
|
||||
__global__ void __launch_bounds__(NUM_THREADS)
|
||||
hgemmWarptiling(int M, int N, int K, float alpha, const __half *A,
|
||||
const __half *B, float beta, __half *C) {
|
||||
const uint cRow = blockIdx.y;
|
||||
const uint cCol = blockIdx.x;
|
||||
const uint warpIdx = threadIdx.x / WARPSIZE;
|
||||
const uint warpCol = warpIdx % (BN / WN);
|
||||
const uint warpRow = warpIdx / (BN / WN);
|
||||
|
||||
constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
|
||||
constexpr uint WSUBM = WM / WMITER;
|
||||
constexpr uint WSUBN = WN / WNITER;
|
||||
|
||||
const uint threadIdxInWarp = threadIdx.x % WARPSIZE;
|
||||
const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN);
|
||||
const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN);
|
||||
|
||||
__shared__ __half As[BM * BK];
|
||||
__shared__ __half Bs[BK * BN];
|
||||
|
||||
A += cRow * BM * K;
|
||||
B += cCol * BN;
|
||||
C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
|
||||
|
||||
const uint innerRowA = threadIdx.x / (BK / 4);
|
||||
const uint innerColA = threadIdx.x % (BK / 4);
|
||||
constexpr uint rowStrideA = (NUM_THREADS * 4) / BK;
|
||||
const uint innerRowB = threadIdx.x / (BN / 4);
|
||||
const uint innerColB = threadIdx.x % (BN / 4);
|
||||
constexpr uint rowStrideB = NUM_THREADS / (BN / 4);
|
||||
|
||||
float threadResults[WMITER * TM * WNITER * TN] = {0.0f};
|
||||
float regM[WMITER * TM] = {0.0f};
|
||||
float regN[WNITER * TN] = {0.0f};
|
||||
|
||||
for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
|
||||
wt::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
|
||||
N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB);
|
||||
__syncthreads();
|
||||
wt::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
|
||||
TN>(regM, regN, threadResults, As, Bs, warpRow, warpCol,
|
||||
threadRowInWarp, threadColInWarp);
|
||||
A += BK;
|
||||
B += BK * N;
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
|
||||
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
|
||||
__half *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
|
||||
for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
|
||||
for (uint resIdxN = 0; resIdxN < TN; resIdxN += 1) {
|
||||
uint idx = (threadRowInWarp * TM + resIdxM) * N +
|
||||
threadColInWarp * TN + resIdxN;
|
||||
float c_old = __half2float(C_interim[idx]);
|
||||
const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
|
||||
wSubColIdx * TN + resIdxN;
|
||||
C_interim[idx] = __float2half(alpha * threadResults[i] + beta * c_old);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int BM, int BN, int BK, int WM, int WN, int WNITER, int TM, int TN, int NT>
|
||||
float bench(int M, int N, int K, const __half *A, const __half *B, __half *C) {
|
||||
dim3 grid(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
|
||||
dim3 block(NT);
|
||||
|
||||
// warmup
|
||||
for (int i = 0; i < 3; i++)
|
||||
hgemmWarptiling<BM, BN, BK, WM, WN, WNITER, TM, TN, NT>
|
||||
<<<grid, block>>>(M, N, K, 1.0f, A, B, 0.0f, C);
|
||||
cudaDeviceSynchronize();
|
||||
|
||||
cudaEvent_t t0, t1;
|
||||
cudaEventCreate(&t0);
|
||||
cudaEventCreate(&t1);
|
||||
cudaEventRecord(t0);
|
||||
for (int i = 0; i < 10; i++)
|
||||
hgemmWarptiling<BM, BN, BK, WM, WN, WNITER, TM, TN, NT>
|
||||
<<<grid, block>>>(M, N, K, 1.0f, A, B, 0.0f, C);
|
||||
cudaEventRecord(t1);
|
||||
cudaEventSynchronize(t1);
|
||||
float ms;
|
||||
cudaEventElapsedTime(&ms, t0, t1);
|
||||
cudaEventDestroy(t0);
|
||||
cudaEventDestroy(t1);
|
||||
|
||||
cudaError_t err = cudaGetLastError();
|
||||
if (err != cudaSuccess) {
|
||||
printf(" CUDA error: %s\n", cudaGetErrorString(err));
|
||||
return -1.0f;
|
||||
}
|
||||
return ms / 10.0f;
|
||||
}
|
||||
|
||||
float bench_cublas(int M, int N, int K, const __half *A, const __half *B, __half *C) {
|
||||
cublasHandle_t handle;
|
||||
cublasCreate(&handle);
|
||||
cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH);
|
||||
|
||||
__half alpha_h = __float2half(1.0f), beta_h = __float2half(0.0f);
|
||||
|
||||
for (int i = 0; i < 3; i++)
|
||||
cublasHgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
|
||||
N, M, K, &alpha_h, B, N, A, K, &beta_h, C, N);
|
||||
cudaDeviceSynchronize();
|
||||
|
||||
cudaEvent_t t0, t1;
|
||||
cudaEventCreate(&t0);
|
||||
cudaEventCreate(&t1);
|
||||
cudaEventRecord(t0);
|
||||
for (int i = 0; i < 10; i++)
|
||||
cublasHgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
|
||||
N, M, K, &alpha_h, B, N, A, K, &beta_h, C, N);
|
||||
cudaEventRecord(t1);
|
||||
cudaEventSynchronize(t1);
|
||||
float ms;
|
||||
cudaEventElapsedTime(&ms, t0, t1);
|
||||
cudaEventDestroy(t0);
|
||||
cudaEventDestroy(t1);
|
||||
cublasDestroy(handle);
|
||||
return ms / 10.0f;
|
||||
}
|
||||
|
||||
int main() {
|
||||
int M = 256, N = 256, K = 256;
|
||||
__half *dA, *dB, *dC;
|
||||
cudaMalloc(&dA, M*K*sizeof(__half));
|
||||
cudaMalloc(&dB, K*N*sizeof(__half));
|
||||
cudaMalloc(&dC, M*N*sizeof(__half));
|
||||
cudaMemset(dA, 0, M*K*sizeof(__half));
|
||||
cudaMemset(dB, 0, K*N*sizeof(__half));
|
||||
|
||||
float ms_cublas = bench_cublas(M, N, K, dA, dB, dC);
|
||||
printf("cublas baseline 256x256: %.3f ms\n\n", ms_cublas);
|
||||
|
||||
// Config A: current (broken)
|
||||
printf("Config A: BM128 BN128 BK16 WM64 WN128 WNITER4 TM4 TN4 NT128\n");
|
||||
float msA = bench<128,128,16, 64,128, 4, 4,4, 128>(M,N,K,dA,dB,dC);
|
||||
printf(" %.3f ms (%.1fx cublas)\n\n", msA, msA/ms_cublas);
|
||||
|
||||
// Config B: fewer WNITER, bigger TM
|
||||
printf("Config B: BM128 BN128 BK16 WM64 WN64 WNITER2 TM8 TN4 NT128\n");
|
||||
float msB = bench<128,128,16, 64,64, 2, 8,4, 128>(M,N,K,dA,dB,dC);
|
||||
printf(" %.3f ms (%.1fx cublas)\n\n", msB, msB/ms_cublas);
|
||||
|
||||
// Config C: 256 threads (4 warps of 64)
|
||||
printf("Config C: BM128 BN128 BK16 WM64 WN64 WNITER2 TM4 TN4 NT256\n");
|
||||
float msC = bench<128,128,16, 64,64, 2, 4,4, 256>(M,N,K,dA,dB,dC);
|
||||
printf(" %.3f ms (%.1fx cublas)\n\n", msC, msC/ms_cublas);
|
||||
|
||||
// Config D: smaller block, more blocks for 16 SMs
|
||||
printf("Config D: BM64 BN64 BK16 WM64 WN64 WNITER4 TM4 TN4 NT64\n");
|
||||
float msD = bench<64,64,16, 64,64, 4, 4,4, 64>(M,N,K,dA,dB,dC);
|
||||
printf(" %.3f ms (%.1fx cublas)\n\n", msD, msD/ms_cublas);
|
||||
|
||||
// Config E: WMITER=1 by design
|
||||
printf("Config E: BM128 BN64 BK16 WM64 WN64 WNITER1 TM4 TN4 NT128\n");
|
||||
float msE = bench<128,64,16, 64,64, 1, 4,4, 128>(M,N,K,dA,dB,dC);
|
||||
printf(" %.3f ms (%.1fx cublas)\n\n", msE, msE/ms_cublas);
|
||||
|
||||
// Config F: bigger BK=32
|
||||
printf("Config F: BM128 BN128 BK32 WM64 WN128 WNITER4 TM4 TN4 NT256\n");
|
||||
float msF = bench<128,128,32, 64,128, 4, 4,4, 256>(M,N,K,dA,dB,dC);
|
||||
printf(" %.3f ms (%.1fx cublas)\n\n", msF, msF/ms_cublas);
|
||||
|
||||
cudaFree(dA); cudaFree(dB); cudaFree(dC);
|
||||
|
||||
// Big matrix
|
||||
M = 256; N = 11008; K = 4096;
|
||||
cudaMalloc(&dA, (long long)M*K*sizeof(__half));
|
||||
cudaMalloc(&dB, (long long)K*N*sizeof(__half));
|
||||
cudaMalloc(&dC, (long long)M*N*sizeof(__half));
|
||||
cudaMemset(dA, 0, (long long)M*K*sizeof(__half));
|
||||
cudaMemset(dB, 0, (long long)K*N*sizeof(__half));
|
||||
|
||||
printf("=== Big matrix 256x4096 @ 4096x11008 ===\n");
|
||||
ms_cublas = bench_cublas(M, N, K, dA, dB, dC);
|
||||
printf("cublas: %.3f ms\n", ms_cublas);
|
||||
|
||||
msA = bench<128,128,16, 64,128, 4, 4,4, 128>(M,N,K,dA,dB,dC);
|
||||
printf("Config A: %.3f ms (%.1fx)\n", msA, msA/ms_cublas);
|
||||
msB = bench<128,128,16, 64,64, 2, 8,4, 128>(M,N,K,dA,dB,dC);
|
||||
printf("Config B: %.3f ms (%.1fx)\n", msB, msB/ms_cublas);
|
||||
msC = bench<128,128,16, 64,64, 2, 4,4, 256>(M,N,K,dA,dB,dC);
|
||||
printf("Config C: %.3f ms (%.1fx)\n", msC, msC/ms_cublas);
|
||||
msF = bench<128,128,32, 64,128, 4, 4,4, 256>(M,N,K,dA,dB,dC);
|
||||
printf("Config F: %.3f ms (%.1fx)\n", msF, msF/ms_cublas);
|
||||
|
||||
cudaFree(dA); cudaFree(dB); cudaFree(dC);
|
||||
return 0;
|
||||
}
|
||||
CUDA
|
||||
|
||||
echo "=== Compiling ==="
|
||||
/usr/local/corex/bin/clang++ --cuda-gpu-arch=ivcore10 --cuda-path=/usr/local/corex \
|
||||
-I/usr/local/corex/include -L/usr/local/corex/lib64 -lcudart -lcublas \
|
||||
-O2 /tmp/probe_k10_configs.cu -o /tmp/probe_k10_configs 2>&1
|
||||
|
||||
if [ -f /tmp/probe_k10_configs ]; then
|
||||
echo "Compile: SUCCESS"
|
||||
/tmp/probe_k10_configs
|
||||
else
|
||||
echo "Compile: FAILED"
|
||||
fi
|
||||
Reference in New Issue
Block a user