ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: eliminate gemm output transpose via A/B swap

Browse Source 
- pass w as param A (op=T) and x as param B (op=N) so the col-major [N,M] output storage is row-major C[M,N] directly, zero copy
- transpose_bias_cast kernel becomes a plain bias+write kernel
- fp8 e2e now beats bf16: 1.09x at M=4096, 1.06x at M=8192 (was 0.88x)
This commit is contained in:
ViperEkura
2026-08-14 01:42:09 +08:00
parent 15862d4b56
commit a711d9f478
1 changed files with 34 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8_mm.cu
+34 -71
View File
@@ -46,6 +46,9 @@ static void ensure_cublas_lt() {
CUBLAS_STATUS_SUCCESS);
}
static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n,
cublasLtMatmulAlgo_t* algo);
static void set_layout(cublasLtMatrixLayout_t layout, int64_t rows, int64_t cols,
int64_t ld) {
TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_ROWS,
@@ -72,72 +75,30 @@ torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b) {
int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0);
TORCH_CHECK(b_c.size(1) == k, "inner dim mismatch");
auto buf = torch::empty({n, m}, a_c.options().dtype(torch::kBFloat16));
// A/B swapped so the col-major [N,M] output storage IS row-major C[M,N]:
// param A = b (op=T -> [N,K]), param B = a (op=N -> [K,M]), output no copy.
auto buf = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16));
ensure_cublas_lt();
set_layout(g_layout_a, k, m, k); // A col-major [K,M] (a row-major, op=T)
set_layout(g_layout_b, k, n, k); // B col-major [K,N] (wT row-major, op=N)
set_layout(g_layout_c, m, n, m); // C col-major [M,N]
set_layout(g_layout_a, k, n, k); // A col-major [K,N] (b row-major, op=T)
set_layout(g_layout_b, k, m, k); // B col-major [K,M] (a row-major, op=N)
set_layout(g_layout_c, n, m, n); // C col-major [N,M], ld=N
float alpha = 1.0f, beta = 0.0f;
cublasLtMatmulHeuristicResult_t heur;
int returned = 0;
cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(
g_handle, g_desc, g_layout_a, g_layout_b, g_layout_c, g_layout_c, g_pref, 1,
&heur, &returned);
cublasLtMatmulAlgo_t algo;
cublasStatus_t st = get_algo_cached(m, k, n, &algo);
TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
"cublasLtMatmulAlgoGetHeuristic failed: ", cublasLtGetStatusName(st));
if (heur.workspaceSize > g_ws_size) {
if (g_workspace) {
cudaFree(g_workspace);
}
TORCH_CHECK(cudaMalloc(&g_workspace, heur.workspaceSize) == cudaSuccess);
g_ws_size = heur.workspaceSize;
}
st = cublasLtMatmul(
g_handle, g_desc, &alpha, a_c.data_ptr(), g_layout_a, b_c.data_ptr(),
g_handle, g_desc, &alpha, b_c.data_ptr(), g_layout_a, a_c.data_ptr(),
g_layout_b, &beta, buf.data_ptr(), g_layout_c, buf.data_ptr(), g_layout_c,
&heur.algo, g_workspace, g_ws_size, stream.stream());
&algo, g_workspace, g_ws_size, stream.stream());
TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
"cublasLtMatmul failed: ", cublasLtGetStatusName(st));
return buf.transpose(0, 1).contiguous();
return buf;
}
// Debug variant: return the raw col-major buffer WITHOUT the transpose copy,
// so the cost of the transpose can be measured in isolation.
torch::Tensor fp8_mm_view(torch::Tensor a, torch::Tensor b) {
TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn, "a must be float8_e4m3fn");
TORCH_CHECK(b.scalar_type() == torch::kFloat8_e4m3fn, "b must be float8_e4m3fn");
const at::cuda::OptionalCUDAGuard guard(a.device());
auto stream = at::cuda::getCurrentCUDAStream();
auto a_c = a.contiguous();
auto b_c = b.contiguous();
int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0);
TORCH_CHECK(b_c.size(1) == k, "inner dim mismatch");
auto buf = torch::empty({n, m}, a_c.options().dtype(torch::kBFloat16));
ensure_cublas_lt();
set_layout(g_layout_a, k, m, k);
set_layout(g_layout_b, k, n, k);
set_layout(g_layout_c, m, n, m);
float alpha = 1.0f, beta = 0.0f;
cublasLtMatmulHeuristicResult_t heur;
int returned = 0;
cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(
g_handle, g_desc, g_layout_a, g_layout_b, g_layout_c, g_layout_c, g_pref, 1,
&heur, &returned);
TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
"cublasLtMatmulAlgoGetHeuristic failed: ", cublasLtGetStatusName(st));
st = cublasLtMatmul(g_handle, g_desc, &alpha, a_c.data_ptr(), g_layout_a,
b_c.data_ptr(), g_layout_b, &beta, buf.data_ptr(), g_layout_c,
buf.data_ptr(), g_layout_c, &heur.algo, g_workspace, g_ws_size,
stream.stream());
TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
"cublasLtMatmul failed: ", cublasLtGetStatusName(st));
return buf; // col-major [M,N] storage, no transpose
}
// ---------------------------------------------------------------------------
// Fused FP8 linear forward: one call = scale cast x8/w8 -> cublasLt GEMM
// (bf16 output) -> transpose + unscale + bias -> bf16 [..., N].
@@ -151,16 +112,14 @@ __global__ void cast_bf16_to_fp8_kernel(
dst[i] = __nv_fp8_e4m3(__bfloat162float(src[i]));
}
__global__ void transpose_bias_cast_kernel(
__global__ void bias_cast_kernel(
const __nv_bfloat16* __restrict__ src, __nv_bfloat16* __restrict__ dst,
const float* __restrict__ bias, int64_t m, int64_t n) {
// src is col-major [M,N] (= row-major C^T[N,M]); write row-major C[M,N].
const float* __restrict__ bias, int64_t total, int64_t n) {
// Same layout both sides (row-major [M,N]); bias added per column.
int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
int64_t total = m * n;
if (idx >= total) return;
int64_t i = idx / n, j = idx % n;
float v = __bfloat162float(src[j * m + i]);
if (bias) v += bias[j];
float v = __bfloat162float(src[idx]);
if (bias) v += bias[idx % n];
dst[idx] = __float2bfloat16(v);
}
@@ -178,7 +137,13 @@ static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n,
cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(
g_handle, g_desc, g_layout_a, g_layout_b, g_layout_c, g_layout_c, g_pref, 1,
&heur, &returned);
if (st != CUBLAS_STATUS_SUCCESS || returned == 0) return st;
if (st != CUBLAS_STATUS_SUCCESS || returned == 0)
return CUBLAS_STATUS_NOT_SUPPORTED;
if (heur.workspaceSize > g_ws_size) {
if (g_workspace) cudaFree(g_workspace);
TORCH_CHECK(cudaMalloc(&g_workspace, heur.workspaceSize) == cudaSuccess);
g_ws_size = heur.workspaceSize;
}
g_last_algo = heur.algo;
g_last_m = m; g_last_k = k; g_last_n = n;
*algo = heur.algo;
@@ -200,9 +165,9 @@ torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w,
auto out = torch::empty({m, n}, x_c.options());
ensure_cublas_lt();
set_layout(g_layout_a, k, m, k);
set_layout(g_layout_b, k, n, k);
set_layout(g_layout_c, m, n, m); // bf16 col-major [M,N] output
set_layout(g_layout_a, k, n, k); // param A = w8 (op=T -> [N,K])
set_layout(g_layout_b, k, m, k); // param B = x8 (op=N -> [K,M])
set_layout(g_layout_c, n, m, n); // col-major [N,M] == row-major C[M,N]
auto x8 = torch::empty({m, k}, x_c.options().dtype(torch::kFloat8_e4m3fn));
auto w8 = torch::empty({n, k}, w_c.options().dtype(torch::kFloat8_e4m3fn));
@@ -215,14 +180,14 @@ torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w,
reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), n * k);
C10_CUDA_CHECK(cudaGetLastError());
auto buf = torch::empty({n, m}, out.options()); // col-major [M,N] = C^T
auto buf = torch::empty({m, n}, out.options()); // row-major C[M,N] direct
float alpha = 1.0f, beta = 0.0f;
cublasLtMatmulAlgo_t algo;
cublasStatus_t st = get_algo_cached(m, k, n, &algo);
TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
"cublasLtMatmulAlgoGetHeuristic failed: ", cublasLtGetStatusName(st));
st = cublasLtMatmul(g_handle, g_desc, &alpha, x8.data_ptr(), g_layout_a,
w8.data_ptr(), g_layout_b, &beta, buf.data_ptr(), g_layout_c,
st = cublasLtMatmul(g_handle, g_desc, &alpha, w8.data_ptr(), g_layout_a,
x8.data_ptr(), g_layout_b, &beta, buf.data_ptr(), g_layout_c,
buf.data_ptr(), g_layout_c, &algo, g_workspace, g_ws_size,
stream.stream());
TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
@@ -234,9 +199,9 @@ torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w,
bias_f = bias.to(torch::kFloat32).contiguous();
bias_ptr = bias_f.data_ptr<float>();
}
transpose_bias_cast_kernel<<<(unsigned)((m * n + block - 1) / block), block, 0, stream>>>(
bias_cast_kernel<<<(unsigned)((m * n + block - 1) / block), block, 0, stream>>>(
reinterpret_cast<const __nv_bfloat16*>(buf.data_ptr()),
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), bias_ptr, m, n);
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), bias_ptr, m * n, n);
C10_CUDA_CHECK(cudaGetLastError());
std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1);
@@ -279,8 +244,6 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward(
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fp8_mm_view", &fp8_mm_view, py::arg("a"), py::arg("b"),
"FP8 GEMM returning raw col-major buffer (debug)");
m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"),
"FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)");
m.def("fp8_linear_forward", &fp8_linear_forward,