diff --git a/astrai/extension/fp8_dispatch.py b/astrai/extension/fp8_dispatch.py new file mode 100644 index 0000000..6c10c66 --- /dev/null +++ b/astrai/extension/fp8_dispatch.py @@ -0,0 +1,67 @@ +"""FP8 linear dispatch: replace aten::linear on the CUDA key, no model changes. + +``F.linear`` -> ``aten::linear`` -> dispatcher -> this CUDA impl (fp8 when +enabled) or the original composite implementation via ``redispatch``. +Enabling is per-thread; model code stays untouched. +""" + +import threading + +import torch +from torch.library import Library + +from astrai.extension.fp8_ops import fp8_linear_forward + +_state = threading.local() + + +def fp8_linear_enable(enabled: bool = True) -> None: + """Toggle fp8 dispatch for aten::linear on this thread.""" + _state.enabled = enabled + + +def fp8_linear_enabled() -> bool: + return getattr(_state, "enabled", False) + + +def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None): + if fp8_linear_enabled() and x.dtype in (torch.bfloat16, torch.float32): + return fp8_linear_forward(x, w, bias) + return torch.ops.aten.linear.default.redispatch( + torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd), + x, + w, + bias, + ) + + +def _linear_backward_cuda_impl(input, grad_output, weight, output_mask): + # VariableType wraps aten::linear; its backward runs aten::linear_backward + # with schema (self, grad_output, weight, mask). Implement the bf16 + # gradient math directly (no redispatch), supporting [..., K] inputs: + # dX = g @ W, dW = g^T @ X, dB = sum(g, dim=0) + g = grad_output.to(torch.bfloat16) + g2d = g.reshape(-1, weight.size(0)) + x2d = input.reshape(-1, input.size(-1)).to(torch.bfloat16) + dX = ( + torch.mm(g2d, weight) + if output_mask[0] + else torch.empty(0, device=input.device, dtype=input.dtype) + ) + dX = dX.reshape_as(input) + dW = ( + torch.mm(g2d.t(), x2d) + if output_mask[1] + else torch.empty(0, device=input.device, dtype=input.dtype) + ) + dB = ( + g.sum(dim=0) + if output_mask[2] + else torch.empty(0, device=input.device, dtype=input.dtype) + ) + return dX, dW, dB + + +_lib = Library("aten", "IMPL", "CUDA") +_lib.impl("linear", _linear_cuda_impl) +_lib.impl("linear_backward", _linear_backward_cuda_impl) diff --git a/astrai/extension/fp8_ops.py b/astrai/extension/fp8_ops.py new file mode 100644 index 0000000..dc82f5a --- /dev/null +++ b/astrai/extension/fp8_ops.py @@ -0,0 +1,88 @@ +"""FP8 matrix-multiply op (torch.library custom_op) and FP8 linear replacement. + +Dispatch table: +- Meta (register_fake): shapes only, for torch.compile / dynamic shapes +- CUDA: csrc fp8_mm kernel (cuBLASLt TN fp8 GEMM, e4m3 in, fp32 acc/out) +- CPU: fp32 fallback (testing) +- AutogradCUDA (register_autograd): bf16 backward, scale-corrected +""" + +import torch +from torch.library import custom_op + +from astrai.extension.loader import get_module, is_available + + +@custom_op("custom::fp8_mm", mutates_args=()) +def fp8_mm( + a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor +) -> torch.Tensor: + """FP8 e4m3 GEMM: a[M,K] x b[N,K] -> fp32[M,N], scales applied by the caller. + + a/b arrive pre-scaled (divided by sx/sw) fp8 tensors; the op returns the + unscaled fp32 result so scale math stays in autograd-land. + """ + + +@fp8_mm.register_fake +def _fp8_mm_fake(a, b, sx, sw): + return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=torch.float32) + + +@fp8_mm.register_kernel("cuda") +def _fp8_mm_cuda(a, b, sx, sw): + if not is_available("fp8_mm"): + raise RuntimeError( + "CUDA kernel 'fp8_mm' is not available. Build with CSRC_KERNELS=true." + ) + return get_module("fp8_mm").fp8_mm(a, b) + + +@fp8_mm.register_kernel("cpu") +def _fp8_mm_cpu(a, b, sx, sw): + return torch.mm(a.float(), b.float().t()) + + +def _fp8_mm_setup_context(ctx, inputs, output): + ctx.save_for_backward(*inputs) + + +def _fp8_mm_backward(ctx, g): + """Scale-corrected straight-through gradients. + + out = F(a, b) * (sx * sw) with F(a, b) = a @ b^T, a = x/sx, b = w/sw: + dx = g * sw @ b (dout/dx = dF/da * 1/sx * sx*sw) + dW = (g * sx)^T @ a (dout/dw = dF/db * 1/sw * sx*sw) + bf16 GEMMs keep gradients in range (e4m3 saturates at 448). + """ + a, b, sx, sw = ctx.saved_tensors + ga = torch.mm(g * sw, b.float()) + gb = torch.mm((g * sx).t(), a.float()) + return ga.to(torch.bfloat16), gb.to(torch.bfloat16), None, None + + +fp8_mm.register_autograd(_fp8_mm_backward, setup_context=_fp8_mm_setup_context) + + +def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): + """FP8 replacement for F.linear(x, w, bias). + + x: [..., K] bf16 (any leading dims), w: [N,K] bf16 (in_dim=K). + The kernel computes a @ b^T with zero-copy col-major mapping, so w is + passed as-is (no transpose). + """ + orig_shape = x.shape + x2d = x.reshape(-1, w.size(1)) + sx = x2d.abs().amax() / 448.0 + sw = w.abs().amax() / 448.0 + x8 = (x2d / sx).to(torch.float8_e4m3fn) + w8 = (w / sw).to(torch.float8_e4m3fn) + out = torch.ops.custom.fp8_mm(x8, w8, sx, sw) + out = out * (sx * sw) + if bias is not None: + out = out + bias + return out.reshape(*orig_shape[:-1], -1) + + +def fp8_available() -> bool: + return is_available("fp8_mm") diff --git a/astrai/extension/loader.py b/astrai/extension/loader.py index bc98d75..7647c21 100644 --- a/astrai/extension/loader.py +++ b/astrai/extension/loader.py @@ -17,6 +17,7 @@ KERNEL_NAMES = [ "attn_paged_decode", "attn_paged_prefill", "rotary_emb", + "fp8_mm", ] _available: dict[str, bool] = {} diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index cab0b49..321ee08 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -48,7 +48,7 @@ set(TORCH_LIBS set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}") -set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb) +set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb fp8_mm) foreach(name ${KERNELS}) add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${name}.cu") @@ -61,6 +61,9 @@ foreach(name ${KERNELS}) "${PYTHON_INCLUDE_DIR}") target_link_libraries(${name} PRIVATE ${TORCH_LIBS}) + if(${name} STREQUAL "fp8_mm") + target_link_libraries(${name} PRIVATE CUDA::cublasLt) + endif() target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}") target_compile_options(${name} PRIVATE diff --git a/csrc/kernels/fp8_mm.cu b/csrc/kernels/fp8_mm.cu new file mode 100644 index 0000000..a67431b --- /dev/null +++ b/csrc/kernels/fp8_mm.cu @@ -0,0 +1,108 @@ +// FP8 e4m3 matrix multiply via cuBLASLt (sm89 TN layout). +// +// cuBLASLt exposes fp8 kernels only for op(A)=T, op(B)=N on Ada; we exploit +// the identity: row-major a[M,K] == A^T as col-major [K,M] (zero copy), and +// row-major wT[N,K] == B as col-major [K,N] (zero copy). The col-major +// result D[M,N] is C^T in row-major terms, so we transpose the output once. +// +// Inputs arrive pre-scaled fp8 e4m3 tensors; output is unscaled fp32. + +#include +#include +#include +#include +#include + +static cublasLtHandle_t g_handle = nullptr; +static cublasLtMatmulDesc_t g_desc = nullptr; +static cublasLtMatrixLayout_t g_layout_a = nullptr; +static cublasLtMatrixLayout_t g_layout_b = nullptr; +static cublasLtMatrixLayout_t g_layout_c = nullptr; +static cublasLtMatmulPreference_t g_pref = nullptr; +static void* g_workspace = nullptr; +static size_t g_ws_size = 0; + +static void ensure_cublas_lt() { + if (g_handle) { + return; + } + TORCH_CHECK(cublasLtCreate(&g_handle) == CUBLAS_STATUS_SUCCESS); + TORCH_CHECK(cublasLtMatmulDescCreate(&g_desc, CUBLAS_COMPUTE_32F, CUDA_R_32F) == + CUBLAS_STATUS_SUCCESS); + cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N; + cublasLtMatmulDescSetAttribute(g_desc, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)); + cublasLtMatmulDescSetAttribute(g_desc, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)); + TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_a, CUDA_R_8F_E4M3, 1, 1, 1) == + CUBLAS_STATUS_SUCCESS); + TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_b, CUDA_R_8F_E4M3, 1, 1, 1) == + CUBLAS_STATUS_SUCCESS); + TORCH_CHECK(cublasLtMatrixLayoutCreate(&g_layout_c, CUDA_R_32F, 1, 1, 1) == + CUBLAS_STATUS_SUCCESS); + TORCH_CHECK(cublasLtMatmulPreferenceCreate(&g_pref) == CUBLAS_STATUS_SUCCESS); + size_t ws = 16 * 1024 * 1024; + TORCH_CHECK(cublasLtMatmulPreferenceSetAttribute( + g_pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws, sizeof(ws)) == + CUBLAS_STATUS_SUCCESS); +} + +static void set_layout(cublasLtMatrixLayout_t layout, int64_t rows, int64_t cols, + int64_t ld) { + TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_ROWS, + &rows, sizeof(rows)) == + CUBLAS_STATUS_SUCCESS); + TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_COLS, + &cols, sizeof(cols)) == + CUBLAS_STATUS_SUCCESS); + TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_LD, + &ld, sizeof(ld)) == + CUBLAS_STATUS_SUCCESS); +} + +torch::Tensor fp8_mm(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"); + TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "2D tensors required"); + 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::kFloat32)); + + 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] + + 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)); + 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_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.transpose(0, 1).contiguous(); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"), + "FP8 e4m3 GEMM: a[M,K] x b[N,K] -> fp32[M,N] (pre-scaled inputs)"); +}