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

perf: output fp8 gemm in bf16 instead of fp32

Browse Source 
- cublasLt C layout and buffer switched to CUDA_R_16BF, halving output bandwidth
- downstream ops (RMSNorm etc.) keep matching bf16 dtype, fused kernels stay
- numeric error unchanged (0.19% vs fp32 ref on quantized inputs)
This commit is contained in:
ViperEkura
2026-08-14 01:08:26 +08:00
parent c6a82a5029
commit f9efb705b8
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8_mm.cu
+2 -2
View File
@@ -36,7 +36,7 @@ static void ensure_cublas_lt() {
CUBLAS_STATUS_SUCCESS); CUBLAS_STATUS_SUCCESS);
TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_b, CUDA_R_8F_E4M3, 1, 1, 1) == TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_b, CUDA_R_8F_E4M3, 1, 1, 1) ==
CUBLAS_STATUS_SUCCESS); CUBLAS_STATUS_SUCCESS);
TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_c, CUDA_R_32F, 1, 1, 1) == TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_c, CUDA_R_16BF, 1, 1, 1) ==
CUBLAS_STATUS_SUCCESS); CUBLAS_STATUS_SUCCESS);
TORCH_CHECK(cublasLtMatmulPreferenceCreate(&g_pref) == CUBLAS_STATUS_SUCCESS); TORCH_CHECK(cublasLtMatmulPreferenceCreate(&g_pref) == CUBLAS_STATUS_SUCCESS);
size_t ws = 16 * 1024 * 1024; size_t ws = 16 * 1024 * 1024;
@@ -71,7 +71,7 @@ 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); 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"); TORCH_CHECK(b_c.size(1) == k, "inner dim mismatch");
auto buf = torch::empty({n, m}, a_c.options().dtype(torch::kFloat32)); auto buf = torch::empty({n, m}, a_c.options().dtype(torch::kBFloat16));
ensure_cublas_lt(); 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_a, k, m, k); // A col-major [K,M] (a row-major, op=T)