diff --git a/astrai/extension/fp8_dispatch.py b/astrai/extension/fp8_dispatch.py index 45d8ad8..5153f61 100644 --- a/astrai/extension/fp8_dispatch.py +++ b/astrai/extension/fp8_dispatch.py @@ -11,21 +11,31 @@ import torch from torch.library import Library from astrai.extension.fp8_ops import fp8_linear_backward, fp8_linear_forward - -_state = threading.local() +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.""" - _state.enabled = enabled + fp8_state().enabled = enabled def fp8_linear_enabled() -> bool: - return getattr(_state, "enabled", False) + 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: + 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), @@ -41,7 +51,11 @@ def _linear_backward_cuda_impl(input_tensor, grad_output, weight, output_mask): # 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: + 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) diff --git a/astrai/extension/fp8_ops.py b/astrai/extension/fp8_ops.py index d041af4..ace30c5 100644 --- a/astrai/extension/fp8_ops.py +++ b/astrai/extension/fp8_ops.py @@ -65,21 +65,17 @@ 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. + """TE-style scaled fp8 linear forward (delegates to fp8_state).""" + from astrai.extension.fp8_state import fp8_linear_forward as _f - 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) + return _f(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) + """TE-style scaled fp8 linear backward (delegates to fp8_state).""" + from astrai.extension.fp8_state import fp8_linear_backward as _b + + return _b(g, x, w, masks) def fp8_available() -> bool: diff --git a/astrai/extension/fp8_state.py b/astrai/extension/fp8_state.py new file mode 100644 index 0000000..bdd4f4e --- /dev/null +++ b/astrai/extension/fp8_state.py @@ -0,0 +1,204 @@ +"""FP8 training state: per-tensor scales, amax history, delayed scaling. + +TE-style (TransformerEngine) delayed scaling: +- weight tensors carry an ``FP8TensorMeta`` keyed by (data_ptr, shape) with a + fixed scale derived from a 16-step amax history window; +- activations/gradients reuse the quantize kernel's free atomic amax, delayed + one step (scale updated after each call, used by the next call); +- ``fp8_autocast()`` context manager toggles fp8 dispatch (like + ``torch.autocast``) and advances the scale-update counter once per step. + Entering it also ensures the aten::linear CUDA impl is registered, so + ``import astrai.extension.fp8_dispatch`` is not required by callers. +""" + +from contextlib import contextmanager + +import torch + +E4M3_MAX = 448.0 + +# FP8 GEMM layout: D = A_SCALE * B_SCALE * A * B, so the per-tensor scales are +# amax/448 (e4m3) and the quantization divides by scale (multiplies by 1/scale). + + +class FP8TensorMeta: + """Scales + amax state for one weight tensor and its paired activations. + + - weight: delayed scale from a 16-step amax history window (TE style) + - x/g: delayed one step, reuse the quantize kernel's free atomic amax + """ + + __slots__ = ( + "scale", + "scale_inv", + "amax_history", + "idx", + "x_scale", + "x_scale_inv", + "g_scale", + "g_scale_inv", + ) + + def __init__(self, device: torch.device, update_interval: int): + self.scale = torch.ones(1, device=device, dtype=torch.float32) + self.scale_inv = torch.ones(1, device=device, dtype=torch.float32) + self.amax_history = torch.ones( + update_interval, device=device, dtype=torch.float32 + ) + self.idx = 0 + self.x_scale = torch.ones(1, device=device, dtype=torch.float32) + self.x_scale_inv = torch.ones(1, device=device, dtype=torch.float32) + self.g_scale = torch.ones(1, device=device, dtype=torch.float32) + self.g_scale_inv = torch.ones(1, device=device, dtype=torch.float32) + + def record(self, amax: torch.Tensor) -> None: + """Push the latest amax into the ring buffer (device-side copy, no sync).""" + self.amax_history[self.idx] = amax.reshape(()) + self.idx = (self.idx + 1) % self.amax_history.numel() + + def refresh(self) -> None: + """Recompute scale from the amax history window (delayed scaling).""" + amax = self.amax_history.max() + if amax > 0: + self.scale.copy_(amax / E4M3_MAX) + self.scale_inv.copy_(E4M3_MAX / amax) + + +class FP8State: + """Global fp8 training state, TE-style.""" + + def __init__(self, update_interval: int = 16): + self.enabled = False + self.update_interval = update_interval + self.step_count = 0 + self._metas: dict[tuple, FP8TensorMeta] = {} + self._last_device: torch.device | None = None + + def _get_device(self, t: torch.Tensor) -> torch.device: + if self._last_device is None: + self._last_device = t.device + return t.device + + def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta: + key = (w.data_ptr(), w.shape, w.dtype) + meta = self._metas.get(key) + if meta is None: + meta = FP8TensorMeta(self._get_device(w), self.update_interval) + self._metas[key] = meta + return meta + + def step(self) -> None: + """Advance the counter and refresh all weight scales every N steps.""" + self.step_count += 1 + if self.step_count % self.update_interval == 0: + for meta in self._metas.values(): + meta.refresh() + + def reset(self) -> None: + self.enabled = False + self.step_count = 0 + self._metas.clear() + self._last_device = None + + +# Global singleton: autograd backward runs on the engine worker threads, so +# thread-local state would lose the fp8 flag during loss.backward(). The GIL +# protects Python-side mutation; the CUDA kernels take their own mutex. +_state = FP8State() + + +def fp8_state() -> FP8State: + return _state + + +@contextmanager +def fp8_autocast(enabled: bool = True, update_interval: int = 16): + """Autocast-style context: fp8 linear dispatch on this thread. + + Usage:: + + with fp8_autocast(enabled=True): + logits = model(input_ids) # aten::linear -> fp8 path + loss.backward() + + The scale-update counter advances once per ``enter`` (one training step), + refreshing weight scales from their amax history every ``update_interval``. + """ + state = fp8_state() + prev_enabled = state.enabled + prev_interval = state.update_interval + state.enabled = enabled + state.update_interval = update_interval + try: + if enabled: + state.step() + yield + finally: + state.enabled = prev_enabled + state.update_interval = prev_interval + + +def _update_delayed_scale(scale, scale_inv, amax) -> None: + """scale = amax / 448 for the *next* call (device-side, no sync).""" + amax_f = amax.reshape(()).to(torch.float32).clamp_min(1e-12) + scale.copy_(amax_f / E4M3_MAX) + scale_inv.copy_(E4M3_MAX / amax_f) + + +def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): + """TE-style scaled fp8 linear forward (called from the aten::linear impl). + + x uses the delayed scale of its paired weight meta (amax from the previous + forward of this linear); the quantize kernel emits the current amax for the + next step. No extra abs/max reduce. + """ + if bias is None: + bias = torch.empty(0, device=x.device, dtype=x.dtype) + state = fp8_state() + mod = _mod() + meta = state.get_weight_meta(w) + amax_x = torch.empty(1, device=x.device, dtype=torch.float32) + amax_w = torch.empty(1, device=x.device, dtype=torch.float32) + out = mod.fp8_linear_forward_scaled( + x, + w, + bias, + meta.x_scale, + meta.scale, + meta.x_scale_inv, + meta.scale_inv, + amax_x, + amax_w, + ) + meta.record(amax_w) + _update_delayed_scale(meta.x_scale, meta.x_scale_inv, amax_x) + return out + + +def fp8_linear_backward(g, x, w, masks): + """TE-style scaled fp8 linear backward (called from aten::linear_backward).""" + state = fp8_state() + mod = _mod() + meta = state.get_weight_meta(w) + amax_g = torch.empty(1, device=g.device, dtype=torch.float32) + out = mod.fp8_linear_backward_scaled( + g, + x, + w, + list(masks), + meta.g_scale, + meta.scale, + meta.x_scale, + meta.g_scale_inv, + meta.scale_inv, + meta.x_scale_inv, + amax_g, + ) + _update_delayed_scale(meta.g_scale, meta.g_scale_inv, amax_g) + return out + + +def _mod(): + from astrai.extension.loader import get_module + + return get_module("fp8_mm") diff --git a/csrc/kernels/fp8_mm.cu b/csrc/kernels/fp8_mm.cu index 10152fc..b6e248d 100644 --- a/csrc/kernels/fp8_mm.cu +++ b/csrc/kernels/fp8_mm.cu @@ -87,7 +87,11 @@ static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n, cublasLtMatmulAlgo_t* algo); static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out, - int64_t m, int64_t k, int64_t n, cudaStream_t stream); + int64_t m, int64_t k, int64_t n, + const float* a_scale, const float* b_scale, + cudaStream_t stream); + +static const float k_scale_one = 1.0f; static void set_layout(cublasLtMatrixLayout_t layout, int64_t rows, int64_t cols, int64_t ld) { @@ -117,34 +121,70 @@ torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b) { auto buf = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16)); ensure_cublas_lt(); - fp8_gemm_into(a_c, b_c, buf, m, k, n, stream.stream()); + fp8_gemm_into(a_c, b_c, buf, m, k, n, &k_scale_one, &k_scale_one, + stream.stream()); return buf; } // --------------------------------------------------------------------------- -// Fused FP8 linear forward: one call = scale cast x8/w8 -> cublasLt GEMM -// (bf16 output) -> bias in-place -> bf16 [..., N]. +// Quantize: bf16 * scale_inv -> fp8, one atomicMax amax per kernel call. +// amax_ptr must be zeroed before launch; float-bits atomicMax works because +// |v| >= 0 has a monotonic IEEE bit pattern. // --------------------------------------------------------------------------- -__global__ void cast_bf16_to_fp8_kernel( - const __nv_bfloat16* __restrict__ src, __nv_fp8_e4m3* __restrict__ dst, - int64_t n) { - int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; - if (i >= n) return; - dst[i] = __nv_fp8_e4m3(__bfloat162float(src[i])); +template +__device__ __forceinline__ T8 cast_fp8(float v); + +template <> +__device__ __forceinline__ __nv_fp8_e4m3 cast_fp8<__nv_fp8_e4m3>(float v) { + return __nv_fp8_e4m3(v); } -__global__ void transpose_cast_bf16_to_fp8_kernel( - const __nv_bfloat16* __restrict__ src, __nv_fp8_e4m3* __restrict__ dst, - int64_t rows, int64_t cols) { - __shared__ __nv_fp8_e4m3 tile[32][33]; +template <> +__device__ __forceinline__ __nv_fp8_e5m2 cast_fp8<__nv_fp8_e5m2>(float v) { + return __nv_fp8_e5m2(v); +} + +template +__global__ void quantize_kernel(const __nv_bfloat16* __restrict__ src, + const float* __restrict__ scale_inv, + T8* __restrict__ dst, + float* __restrict__ amax_ptr, int64_t n) { + int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; + float amax = 0.f; + if (i < n) { + float v = __bfloat162float(src[i]) * *scale_inv; + dst[i] = cast_fp8(v); + amax = fabsf(v); + } + for (int off = 16; off; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off)); + __shared__ float sm[8]; + if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax; + __syncthreads(); + if (threadIdx.x == 0) { + float m = 0.f; + for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]); + atomicMax(reinterpret_cast(amax_ptr), __float_as_uint(m)); + } +} + +// Same but with a transpose (rows x cols bf16 row-major -> fp8 [cols, rows]). +template +__global__ void transpose_quantize_kernel( + const __nv_bfloat16* __restrict__ src, const float* __restrict__ scale_inv, + T8* __restrict__ dst, float* __restrict__ amax_ptr, int64_t rows, + int64_t cols) { + __shared__ T8 tile[32][33]; int64_t x = blockIdx.x * 32 + threadIdx.x; int64_t y = blockIdx.y * 32 + threadIdx.y; + float amax = 0.f; for (int j = 0; j < 32; j += 8) { if (x < cols && y + j < rows) { - tile[threadIdx.y + j][threadIdx.x] = - __nv_fp8_e4m3(__bfloat162float(src[(y + j) * cols + x])); + float v = __bfloat162float(src[(y + j) * cols + x]) * *scale_inv; + tile[threadIdx.y + j][threadIdx.x] = cast_fp8(v); + amax = fmaxf(amax, fabsf(v)); } } __syncthreads(); @@ -156,6 +196,16 @@ __global__ void transpose_cast_bf16_to_fp8_kernel( dst[(y + j) * rows + x] = tile[threadIdx.x][threadIdx.y + j]; } } + for (int off = 16; off; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off)); + __shared__ float sm[8]; + if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax; + __syncthreads(); + if (threadIdx.x == 0) { + float m = 0.f; + for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]); + atomicMax(reinterpret_cast(amax_ptr), __float_as_uint(m)); + } } __global__ void bias_add_bf16_kernel( @@ -196,11 +246,21 @@ static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n, } static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out, - int64_t m, int64_t k, int64_t n, cudaStream_t stream) { + int64_t m, int64_t k, int64_t n, + const float* a_scale, const float* b_scale, + cudaStream_t stream) { std::lock_guard lock(g_mutex); set_layout(g_layout_a, k, n, k); // param A = rhs (op=T -> [N,K]) set_layout(g_layout_b, k, m, k); // param B = lhs (op=N -> [K,M]) set_layout(g_layout_c, n, m, n); // col-major [N,M] == row-major [M,N] + // Per-tensor FP32 scales applied inside the GEMM: + // D = alpha * A_SCALE * B_SCALE * A * B (alpha = 1). + TORCH_CHECK(cublasLtMatmulDescSetAttribute( + g_desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale, + sizeof(a_scale)) == CUBLAS_STATUS_SUCCESS); + TORCH_CHECK(cublasLtMatmulDescSetAttribute( + g_desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale, + sizeof(b_scale)) == CUBLAS_STATUS_SUCCESS); float alpha = 1.0f, beta = 0.0f; static AlgoCache cache; cublasLtMatmulAlgo_t algo; @@ -215,11 +275,22 @@ static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor ou "cublasLtMatmul failed: ", cublasLtGetStatusName(st)); } -torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w, - torch::Tensor bias) { +// --------------------------------------------------------------------------- +// Scaled FP8 linear forward: quantize x/w with per-tensor scales -> cublasLt +// GEMM (scales applied inside) -> bias in-place -> bf16 [..., N]. +// sx/sw: f32 scale tensors (device scalars); sx_inv/sw_inv: 1/scale. +// amax_x/amax_w: f32 buffers receiving max-abs of the quantized tensors. +// --------------------------------------------------------------------------- + +torch::Tensor fp8_linear_forward_scaled(torch::Tensor x, torch::Tensor w, + torch::Tensor bias, torch::Tensor sx, + torch::Tensor sw, torch::Tensor sx_inv, + torch::Tensor sw_inv, + torch::Tensor amax_x, + torch::Tensor amax_w) { TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required"); - TORCH_CHECK(x.dtype() == torch::kBFloat16, "x must be bf16"); - TORCH_CHECK(w.dtype() == torch::kBFloat16, "w must be bf16"); + TORCH_CHECK(x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16, + "x and w must be bf16"); const at::cuda::OptionalCUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); @@ -229,22 +300,30 @@ torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w, TORCH_CHECK(w_c.size(1) == k, "inner dim mismatch"); ensure_cublas_lt(); + const float* sx_ptr = sx.data_ptr(); + const float* sw_ptr = sw.data_ptr(); + const float* sxi_ptr = sx_inv.data_ptr(); + const float* swi_ptr = sw_inv.data_ptr(); + float* amax_x_ptr = amax_x.data_ptr(); + float* amax_w_ptr = amax_w.data_ptr(); + C10_CUDA_CHECK(cudaMemsetAsync(amax_x_ptr, 0, sizeof(float), stream.stream())); + C10_CUDA_CHECK(cudaMemsetAsync(amax_w_ptr, 0, sizeof(float), stream.stream())); + 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)); int64_t block = 256; - cast_bf16_to_fp8_kernel<<<(unsigned)((m * k + block - 1) / block), block, 0, stream>>>( - reinterpret_cast(x_c.data_ptr()), - reinterpret_cast<__nv_fp8_e4m3*>(x8.data_ptr()), m * k); - cast_bf16_to_fp8_kernel<<<(unsigned)((n * k + block - 1) / block), block, 0, stream>>>( - reinterpret_cast(w_c.data_ptr()), - reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), n * k); + quantize_kernel<__nv_fp8_e4m3> + <<<(unsigned)((m * k + block - 1) / block), block, 0, stream.stream()>>>( + reinterpret_cast(x_c.data_ptr()), sxi_ptr, + reinterpret_cast<__nv_fp8_e4m3*>(x8.data_ptr()), amax_x_ptr, m * k); + quantize_kernel<__nv_fp8_e4m3> + <<<(unsigned)((n * k + block - 1) / block), block, 0, stream.stream()>>>( + reinterpret_cast(w_c.data_ptr()), swi_ptr, + reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), amax_w_ptr, n * k); C10_CUDA_CHECK(cudaGetLastError()); - // A/B swap makes the col-major [N,M] storage directly represent the - // row-major output [M,N]. Keep this buffer as the public output so the - // bias path does not need a second allocation or a copy kernel. auto out = torch::empty({m, n}, x_c.options()); - fp8_gemm_into(x8, w8, out, m, k, n, stream.stream()); + fp8_gemm_into(x8, w8, out, m, k, n, sw_ptr, sx_ptr, stream.stream()); if (bias.defined() && bias.numel() > 0) { TORCH_CHECK(bias.scalar_type() == torch::kBFloat16 && bias.numel() == n, @@ -261,12 +340,15 @@ torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w, } // --------------------------------------------------------------------------- -// Fused FP8 linear backward: dX = g @ W, dW = g^T @ X, dB = sum(g). +// Scaled FP8 linear backward: dX = g @ W, dW = g^T @ X, dB = sum(g). +// Scales: g uses sg (immediate), w/x reuse the forward scales. // --------------------------------------------------------------------------- -std::tuple fp8_linear_backward( +std::tuple fp8_linear_backward_scaled( torch::Tensor g, torch::Tensor x, torch::Tensor w, - std::vector masks) { + std::vector masks, torch::Tensor sg, torch::Tensor sw, + torch::Tensor sx, torch::Tensor sg_inv, torch::Tensor sw_inv, + torch::Tensor sx_inv, torch::Tensor amax_g) { const at::cuda::OptionalCUDAGuard guard(g.device()); TORCH_CHECK(g.dtype() == torch::kBFloat16 && x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16, @@ -286,6 +368,15 @@ std::tuple fp8_linear_backward( auto grad_bias = torch::empty({0}, g_c.options().dtype(g.dtype())); ensure_cublas_lt(); + const float* sg_ptr = sg.data_ptr(); + const float* sw_ptr = sw.data_ptr(); + const float* sx_ptr = sx.data_ptr(); + const float* sgi_ptr = sg_inv.data_ptr(); + const float* swi_ptr = sw_inv.data_ptr(); + const float* sxi_ptr = sx_inv.data_ptr(); + float* amax_g_ptr = amax_g.data_ptr(); + C10_CUDA_CHECK(cudaMemsetAsync(amax_g_ptr, 0, sizeof(float), stream.stream())); + auto fp8_options = g_c.options().dtype(torch::kFloat8_e4m3fn); auto g8 = torch::empty({m, n}, fp8_options); auto gt8 = masks[1] ? torch::empty({n, m}, fp8_options) : torch::Tensor(); @@ -293,28 +384,36 @@ std::tuple fp8_linear_backward( auto xt8 = masks[1] ? torch::empty({k, m}, fp8_options) : torch::Tensor(); int64_t block = 256; - cast_bf16_to_fp8_kernel<<<(unsigned)((m * n + block - 1) / block), block, 0, stream>>>( - reinterpret_cast(g_c.data_ptr()), - reinterpret_cast<__nv_fp8_e4m3*>(g8.data_ptr()), m * n); + quantize_kernel<__nv_fp8_e4m3> + <<<(unsigned)((m * n + block - 1) / block), block, 0, stream.stream()>>>( + reinterpret_cast(g_c.data_ptr()), sgi_ptr, + reinterpret_cast<__nv_fp8_e4m3*>(g8.data_ptr()), amax_g_ptr, m * n); dim3 threads(32, 8); if (masks[0]) { dim3 blocks((k + 31) / 32, (n + 31) / 32); - transpose_cast_bf16_to_fp8_kernel<<>>( - reinterpret_cast(w_c.data_ptr()), - reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), n, k); - fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k, - stream.stream()); + transpose_quantize_kernel<__nv_fp8_e4m3> + <<>>( + reinterpret_cast(w_c.data_ptr()), swi_ptr, + reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), amax_g_ptr, + n, k); + fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k, sg_ptr, + sw_ptr, stream.stream()); } if (masks[1]) { dim3 g_blocks((n + 31) / 32, (m + 31) / 32); dim3 x_blocks((k + 31) / 32, (m + 31) / 32); - transpose_cast_bf16_to_fp8_kernel<<>>( - reinterpret_cast(g_c.data_ptr()), - reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), m, n); - transpose_cast_bf16_to_fp8_kernel<<>>( - reinterpret_cast(x_c.data_ptr()), - reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), m, k); - fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, stream.stream()); + transpose_quantize_kernel<__nv_fp8_e4m3> + <<>>( + reinterpret_cast(g_c.data_ptr()), sgi_ptr, + reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), amax_g_ptr, + m, n); + transpose_quantize_kernel<__nv_fp8_e4m3> + <<>>( + reinterpret_cast(x_c.data_ptr()), sxi_ptr, + reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), amax_g_ptr, + m, k); + fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, sg_ptr, sx_ptr, + stream.stream()); } C10_CUDA_CHECK(cudaGetLastError()); if (masks[2]) { @@ -327,12 +426,16 @@ std::tuple fp8_linear_backward( 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] -> bf16[M,N] (pre-scaled inputs)"); - m.def("fp8_linear_forward", &fp8_linear_forward, - py::arg("x"), py::arg("w"), py::arg("bias"), - "Fused FP8 linear forward: scale cast + cublasLt GEMM + bias " - "-> bf16, single call"); - m.def("fp8_linear_backward", &fp8_linear_backward, + m.def("fp8_linear_forward_scaled", &fp8_linear_forward_scaled, + py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("sx"), + py::arg("sw"), py::arg("sx_inv"), py::arg("sw_inv"), + py::arg("amax_x"), py::arg("amax_w"), + "Scaled FP8 linear forward: quantize with per-tensor scales + " + "cublasLt GEMM (scales applied inside) + bias -> bf16"); + m.def("fp8_linear_backward_scaled", &fp8_linear_backward_scaled, py::arg("g"), py::arg("x"), py::arg("w"), py::arg("masks"), - "Fused linear backward: dX = g*sw @ W, dW = (g*sx)^T @ X, " - "dB = sum(g), single call"); + py::arg("sg"), py::arg("sw"), py::arg("sx"), py::arg("sg_inv"), + py::arg("sw_inv"), py::arg("sx_inv"), py::arg("amax_g"), + "Scaled FP8 linear backward: dX = g*sw @ W, dW = (g*sx)^T @ X, " + "dB = sum(g)"); }