ViperEkura
15862d4b56
- fp8_linear_forward: cast + cublasLt GEMM + transpose + bias in one call - fp8_linear_backward: scale-free, dtype derived from input tensor - drops per-op Python dispatch (was ~6-8 launches per linear) and amax syncs - 1024x1024 linear: 6.8x slow -> 0.67x (36.7us vs 24.8us bf16) - small-model e2e still 1.71x slow; 15bt estimate ~0.78x (linear-heavy)
87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
"""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), fused in one CUDA call.
|
|
|
|
x: [..., K] bf16 (any leading dims), w: [N,K] bf16 (in_dim=K).
|
|
The kernel pipeline (scale cast -> cublasLt fp8 GEMM -> unscale + bias ->
|
|
transpose -> bf16) runs inside a single extension call, so Python-side
|
|
dispatch overhead is paid once per linear instead of per operator.
|
|
"""
|
|
if bias is None:
|
|
bias = torch.empty(0, device=x.device, dtype=x.dtype)
|
|
return get_module("fp8_mm").fp8_linear_forward(x, w, bias)
|
|
|
|
|
|
def fp8_linear_backward(g, x, w, masks):
|
|
"""Fused linear backward (dX/dW/dB in one CUDA call, scale-corrected)."""
|
|
return get_module("fp8_mm").fp8_linear_backward(g, x, w, masks)
|
|
|
|
|
|
def fp8_available() -> bool:
|
|
return is_available("fp8_mm")
|