ViperEkura
5244f1a8fc
- per-tensor scales applied inside cublasLt via A_SCALE/B_SCALE - delayed scaling: weight amax history ring, refresh every 16 steps - quantize kernels emit atomic amax, device-side scale updates - fp8_autocast context toggles aten::linear dispatch like torch.autocast - fallback to bf16 when M/N not 16-aligned (fp8 gemm constraint) - x/g scales delayed one step, reuse free atomic amax (no abs/max reduce)
85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
"""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_backward, fp8_linear_forward
|
|
from astrai.extension.fp8_state import fp8_autocast, fp8_state
|
|
|
|
|
|
def fp8_linear_enable(enabled: bool = True) -> None:
|
|
"""Toggle fp8 dispatch for aten::linear on this thread."""
|
|
fp8_state().enabled = enabled
|
|
|
|
|
|
def fp8_linear_enabled() -> bool:
|
|
return fp8_state().enabled
|
|
|
|
|
|
def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool:
|
|
"""cuBLASLt fp8 requires M % 16 == 0 and N % 16 == 0 (K is padded); else fall back."""
|
|
m = x.numel() // x.size(-1)
|
|
return m % 16 == 0 and w.size(0) % 16 == 0
|
|
|
|
|
|
def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
|
|
if (
|
|
fp8_linear_enabled()
|
|
and x.dtype == torch.bfloat16
|
|
and w.dtype == torch.bfloat16
|
|
and _fp8_supported(x, w)
|
|
):
|
|
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_tensor, grad_output, weight, output_mask):
|
|
# VariableType wraps aten::linear; its backward runs aten::linear_backward
|
|
# with schema (self, grad_output, weight, mask). When fp8 is enabled the
|
|
# fused CUDA backward runs in one call (scale-corrected); otherwise the
|
|
# plain bf16/fp32 math, dtype aligned to the leaf weight:
|
|
# grad_input = g @ W, grad_weight = g^T @ X, grad_bias = sum(g, dim=0)
|
|
if (
|
|
fp8_linear_enabled()
|
|
and weight.dtype == torch.bfloat16
|
|
and _fp8_supported(grad_output, weight)
|
|
):
|
|
return fp8_linear_backward(grad_output, input_tensor, weight, list(output_mask))
|
|
compute_dtype = weight.dtype
|
|
grad = grad_output.to(compute_dtype)
|
|
grad_2d = grad.reshape(-1, weight.size(0))
|
|
input_2d = input_tensor.reshape(-1, input_tensor.size(-1)).to(compute_dtype)
|
|
grad_input = (
|
|
torch.mm(grad_2d, weight)
|
|
if output_mask[0]
|
|
else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
|
|
)
|
|
grad_weight = (
|
|
torch.mm(grad_2d.t(), input_2d)
|
|
if output_mask[1]
|
|
else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
|
|
)
|
|
grad_bias = (
|
|
grad.sum(dim=0)
|
|
if output_mask[2]
|
|
else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
|
|
)
|
|
return grad_input.reshape_as(input_tensor), grad_weight, grad_bias
|
|
|
|
|
|
_lib = Library("aten", "IMPL", "CUDA")
|
|
_lib.impl("linear", _linear_cuda_impl)
|
|
_lib.impl("linear_backward", _linear_backward_cuda_impl)
|