diff --git a/csrc/kernels/fp8_mm.cu b/csrc/kernels/fp8_mm.cu index 5ce5829..8e4a670 100644 --- a/csrc/kernels/fp8_mm.cu +++ b/csrc/kernels/fp8_mm.cu @@ -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(); } - 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(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 shape(x.sizes().begin(), x.sizes().end() - 1); @@ -279,8 +244,6 @@ std::tuple 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,