diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index fd2a26f..65194a6 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -407,11 +407,21 @@ class _LinearFp8(torch.autograd.Function): meta.g.seed(g2, fmt) sg = meta.g.scale.clone() sw, sx = _sw_fwd, _sx_fwd - g8, amax_g = quantize(g2, sg.reciprocal(), fmt) - x8, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt) - w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0] - grad_x = mm_fp8(g8, w8, sg * sw).reshape(x.shape) # g8[m,n] @ w8[n,k] - grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True) # g8.T @ x8 + # Backward GEMMs route through the NT fast path via transposed + # quantize outputs: g8 [m,n] with w8T [k,n] (trans_b=True) gives + # grad_x, g8T [n,m] with x8T [k,m] gives grad_w — no NN-swap or TT + # crosswise kernel in the training path. g is consumed in both + # orientations, so one dual-layout pass feeds both. + g8, g8T, amax_g = quantize(g2, sg.reciprocal(), fmt, layout=2) + x8T, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, layout=1) + if _is_fp8(w.dtype): + # Pre-quantized weight has no transposed copy: keep the swap + # path for grad_x (grad_w is unaffected). + grad_x = mm_fp8(g8, w, sg * sw).reshape(x.shape) + else: + w8T, _ = quantize(w, sw.reciprocal(), fmt, layout=1) + grad_x = mm_fp8(g8, w8T, sg * sw, trans_b=True).reshape(x.shape) + grad_w = mm_fp8(g8T, x8T, sg * sx, trans_b=True) # g8.T @ x8 # bias-free linears must not pay the column-sum # reduce: g2.sum(0) is another full read of the gradient. grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index 038c54c..abe2427 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -63,6 +63,71 @@ def _fp8_quantize_fake(x, scale, fmt): _QUANT_INPUT_DTYPES = (torch.bfloat16, torch.float16, torch.float32) +@custom_op("custom::fp8_quantize_t", mutates_args=()) +def fp8_quantize_t( + x: torch.Tensor, scale: torch.Tensor, fmt: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Transposed-output variant of fp8_quantize: returns ``(x8T, amax)`` + where ``x8T`` is the [cols][rows] row-major transpose of the quantized + input (the K-contiguous operand orientation for NT GEMMs).""" + + +@fp8_quantize_t.register_fake +def _fp8_quantize_t_fake(x, scale, fmt): + dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn + rows, cols = x.shape[-2], x.shape[-1] + return ( + torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype), + torch.empty(1, device=x.device, dtype=torch.float32), + ) + + +@fp8_quantize_t.register_kernel("cuda") +def _fp8_quantize_t_cuda(x, scale, fmt): + if x.dtype not in _QUANT_INPUT_DTYPES: + raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}") + return get_module("fp8_ops").quantize(x, scale, int(fmt), 1) + + +@fp8_quantize_t.register_kernel("cpu") +def _fp8_quantize_t_cpu(x, scale, fmt): + x8, amax = _fp8_quantize_cpu(x, scale, fmt) + return x8.transpose(-2, -1).contiguous(), amax + + +@custom_op("custom::fp8_quantize_dual", mutates_args=()) +def fp8_quantize_dual( + x: torch.Tensor, scale: torch.Tensor, fmt: int +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Dual-orientation quantize: one read of ``x`` produces both the + row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors + consumed by GEMMs on both orientations (backward ``g``).""" + + +@fp8_quantize_dual.register_fake +def _fp8_quantize_dual_fake(x, scale, fmt): + dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn + rows, cols = x.shape[-2], x.shape[-1] + return ( + torch.empty(x.shape, device=x.device, dtype=dtype), + torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype), + torch.empty(1, device=x.device, dtype=torch.float32), + ) + + +@fp8_quantize_dual.register_kernel("cuda") +def _fp8_quantize_dual_cuda(x, scale, fmt): + if x.dtype not in _QUANT_INPUT_DTYPES: + raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}") + return get_module("fp8_ops").quantize(x, scale, int(fmt), 2) + + +@fp8_quantize_dual.register_kernel("cpu") +def _fp8_quantize_dual_cpu(x, scale, fmt): + x8, amax = _fp8_quantize_cpu(x, scale, fmt) + return x8, x8.transpose(-2, -1).contiguous(), amax + + @fp8_quantize.register_kernel("cuda") def _fp8_quantize_cuda(x, scale, fmt): if x.dtype not in _QUANT_INPUT_DTYPES: @@ -125,13 +190,16 @@ def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None): def quantize( - x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3" -) -> Tuple[torch.Tensor, torch.Tensor]: - """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; returns - ``(x8, amax)``. + x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3", layout: int = 0 +) -> tuple: + """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax. ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects - E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. + E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. ``layout`` + picks the output orientation: 0 = row-major ``(x8, amax)``; 1 = + transposed ``[cols][rows]`` ``(x8T, amax)`` — the K-contiguous operand + orientation NT GEMMs want; 2 = both from one read ``(x8, x8T, amax)`` + (for tensors consumed in both orientations, e.g. backward ``g``). """ # Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a # 512-wide GEMM): real CUDA tensors of a supported dtype go straight to @@ -144,8 +212,12 @@ def quantize( and x.dtype in _QUANT_INPUT_DTYPES and fmt in _FMT_TO_INT ): - return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt]) - return fp8_quantize(x, scale, _fmt_int(fmt)) + return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt], layout) + if layout == 0: + return fp8_quantize(x, scale, _fmt_int(fmt)) + if layout == 1: + return fp8_quantize_t(x, scale, _fmt_int(fmt)) + return fp8_quantize_dual(x, scale, _fmt_int(fmt)) def mm_fp8( diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index 1b4326e..d09d6f9 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -71,12 +71,24 @@ struct FP8QuantizeParams { // the binding and receives the raw-domain absolute maximum. const void* __restrict__ input_ptr = nullptr; void* __restrict__ output_ptr = nullptr; + void* __restrict__ output_transposed_ptr = nullptr; + // Transposed-output destination ([cols][rows]); the output-layout modes: + // 0 = row-major only (output_ptr; the vectorized elementwise kernel) + // 1 = transposed only (output_transposed_ptr; the tiled kernel) + // 2 = both destinations in one read of the input (the tiled kernel) + // Modes 1/2 exist so crosswise-layout GEMM operands (NN grad_x, TT + // grad_w) can be produced K-contiguous instead, routing every training + // GEMM through the dual-congruous NT fast path. + int out_layout = 0; const float* __restrict__ scale = nullptr; float* __restrict__ amax = nullptr; - // Element count (only the elementwise quantize kernel uses it). + // Element count (only the elementwise quantize kernel uses it); the + // tiled kernel views the same buffer as [rows][cols] row-major. int total = 0; + int rows = 0; + int cols = 0; }; // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index 02bd106..ca8a2ab 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -46,16 +46,6 @@ void check_scale(const torch::Tensor& scale, const torch::Tensor& input) { "scale must be a CUDA float32 scalar on the input device"); } -void pack_quantize(FP8QuantizeParams& p, const void* input, void* output, - const torch::Tensor& scale, torch::Tensor& amax, - int64_t total) { - p.input_ptr = input; - p.output_ptr = output; - p.scale = scale.data_ptr(); - p.amax = amax.data_ptr(); - p.total = static_cast(total); -} - void pack_gemm(FP8Params& p, const void* a, const void* b, void* output, const torch::Tensor& scale, int64_t m, int64_t n, int64_t k, int64_t a_ld, int64_t b_ld) { @@ -99,11 +89,45 @@ bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld, return flag ^ col_major; } +// Dtype x format switch shared by both quantize kernels; Tiled selects the +// transpose kernel (out_layout 1/2) over the vectorized elementwise one. +template +void launch_one(const FP8QuantizeParams& p, cudaStream_t stream) { + if constexpr (Tiled) + launch_fp8_quantize_tiled(p, stream); + else + launch_fp8_quantize(p, stream); +} + +template +void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p, + bool e5m2, cudaStream_t stream) { + if (x.scalar_type() == torch::kHalf) { + if (e5m2) + launch_one(p, stream); + else + launch_one(p, stream); + } else if (x.scalar_type() == torch::kFloat32) { + if (e5m2) + launch_one(p, stream); + else + launch_one(p, stream); + } else { + if (e5m2) + launch_one(p, stream); + else + launch_one(p, stream); + } +} + } // namespace -std::tuple quantize(torch::Tensor x, - torch::Tensor scale, - int64_t fmt) { +// Output-layout dispatch: 0 = [rows][cols] row-major (the historic 2-tuple +// return), 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations +// from a single read of the input (3-tuple). Layouts 1/2 feed the NT GEMM +// fast path from crosswise consumers (backward grad_x / grad_w). +py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt, + int64_t layout) { TORCH_CHECK(x.is_cuda(), "CUDA tensors required"); TORCH_CHECK(x.scalar_type() == torch::kBFloat16 || x.scalar_type() == torch::kHalf || @@ -112,39 +136,44 @@ std::tuple quantize(torch::Tensor x, TORCH_CHECK(fmt == static_cast(FP8Format::E4M3) || fmt == static_cast(FP8Format::E5M2), "unsupported quantization type: expected E4M3 (0) or E5M2 (1)"); + TORCH_CHECK(layout >= 0 && layout <= 2, + "layout must be 0 (row-major), 1 (transposed) or 2 (both)"); + TORCH_CHECK(layout == 0 || x.dim() >= 2, + "transposed quantize layouts need a 2D+ tensor"); check_scale(scale, x); check_fp8_device(x); const at::cuda::OptionalCUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto input = x.contiguous(); - auto output = torch::empty_like( - input, input.options().dtype(fmt ? torch::kFloat8_e5m2 - : torch::kFloat8_e4m3fn)); + auto out_opts = input.options().dtype( + fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn); auto amax = torch::zeros({1}, input.options().dtype(torch::kFloat32)); + FP8QuantizeParams p; - pack_quantize(p, input.data_ptr(), output.data_ptr(), scale, amax, - input.numel()); - const bool e5m2 = fmt == static_cast(FP8Format::E5M2); - if (x.scalar_type() == torch::kHalf) { - if (e5m2) - launch_fp8_quantize(p, stream.stream()); - else - launch_fp8_quantize(p, stream.stream()); - } else if (x.scalar_type() == torch::kFloat32) { - if (e5m2) - launch_fp8_quantize(p, stream.stream()); - else - launch_fp8_quantize(p, stream.stream()); - } else { - if (e5m2) - launch_fp8_quantize( - p, stream.stream()); - else - launch_fp8_quantize( - p, stream.stream()); + p.input_ptr = input.data_ptr(); + p.scale = scale.data_ptr(); + p.amax = amax.data_ptr(); + p.total = static_cast(input.numel()); + p.out_layout = static_cast(layout); + p.rows = static_cast(input.size(-2)); + p.cols = static_cast(input.size(-1)); + torch::Tensor output, output_t; + if (layout == 0 || layout == 2) { + output = torch::empty_like(input, out_opts); + p.output_ptr = output.data_ptr(); } + if (layout >= 1) { + output_t = torch::empty({input.size(-1), input.size(-2)}, out_opts); + p.output_transposed_ptr = output_t.data_ptr(); + } + const bool e5m2 = fmt == static_cast(FP8Format::E5M2); + if (layout != 0) + launch_quantize_for(input, p, e5m2, stream.stream()); + else + launch_quantize_for(input, p, e5m2, stream.stream()); C10_CUDA_CHECK(cudaGetLastError()); - return {output, amax}; + if (layout == 2) return py::make_tuple(output, output_t, amax); + return py::make_tuple(layout == 1 ? output_t : output, amax); } torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale, @@ -219,7 +248,7 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale, // bias argument through untouched instead of normalizing it host-side). PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("quantize", &quantize, py::arg("x"), py::arg("scale"), - py::arg("fmt")); + py::arg("fmt"), py::arg("layout") = 0); m.def( "mm_fp8", [](torch::Tensor a, torch::Tensor b, torch::Tensor scale, diff --git a/csrc/kernels/fp8/quantize.cuh b/csrc/kernels/fp8/quantize.cuh index 0a56e5c..4642888 100644 --- a/csrc/kernels/fp8/quantize.cuh +++ b/csrc/kernels/fp8/quantize.cuh @@ -169,5 +169,89 @@ void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) { fp8_quantize_kernel<<>>(p); } +// Tiled transpose quantize (out_layout 1/2): reads the [rows][cols] +// row-major input once and writes the fp8 bytes transposed ([cols][rows], +// so the contract dim lands K-contiguous for NT GEMM operands) and, in +// mode 2, the plain row-major copy too. A 32x32 tile stages through shared +// memory: input-row-major loads and output writes both stay coalesced, and +// the byte-wide staging is conflict-free — the +4 pad makes the store +// stride 9 (words) coprime with the 32 banks and the load is a 32-byte +// broadcast segment. A 64x64 split-half variant (16 elems/thread, paired +// 2-byte scatter stores) measured +21% on L2-resident shapes but -3..5% +// on the DRAM-bound ones that carry the training traffic (occupancy and +// memory-level parallelism, not instruction count, gate the DRAM regime); +// weighted by the real step's mix the two tie, so the simpler tile stays. +template +__global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) { + constexpr int kTile = 32; + __shared__ uint8_t tile[kTile][kTile + 4]; + const float mult = *p.scale; + const auto* x = static_cast(p.input_ptr); + const int r0 = blockIdx.y * kTile; + const int c0 = blockIdx.x * kTile; + const int r = r0 + threadIdx.y * 4; + const int c = c0 + threadIdx.x; + + uint8_t q[4]; + float local_amax = 0.0f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + q[j] = 0; + if (r + j < p.rows && c < p.cols) { + const float v = + quant_in_traits::to_float(x[(int64_t)(r + j) * p.cols + c]); + local_amax = fmaxf(local_amax, fabsf(v)); + if constexpr (Fmt == FP8Format::E5M2) + q[j] = __nv_fp8_e5m2(v * mult).__x; + else + q[j] = __nv_fp8_e4m3(v * mult).__x; + } + } + if (p.out_layout == 2) { + uint8_t* out = static_cast(p.output_ptr); +#pragma unroll + for (int j = 0; j < 4; ++j) + if (r + j < p.rows && c < p.cols) + out[(int64_t)(r + j) * p.cols + c] = q[j]; + } +#pragma unroll + for (int j = 0; j < 4; ++j) tile[threadIdx.x][threadIdx.y * 4 + j] = q[j]; + __syncthreads(); + // Transposed scatter: output element (c, r) lives at c * rows + r; r + // tracks threadIdx.x so each warp writes one contiguous run. The read + // swaps the staging indices — tile[col][row] was written, so the value + // for input (r0+tx, c0+ty*4+j) sits at tile[ty*4+j][tx]. + uint8_t* out_t = static_cast(p.output_transposed_ptr); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int oc = c0 + threadIdx.y * 4 + j; + if (oc < p.cols && r0 + threadIdx.x < p.rows) + out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] = + tile[threadIdx.y * 4 + j][threadIdx.x]; + } + if (p.amax) { + local_amax = warp_reduce_max(local_amax); + __shared__ float slots[8]; + // blockDim.x is 32, so warp id == threadIdx.y; only complete warps + // exist (blockDim.y == 8). + if (threadIdx.x == 0) slots[threadIdx.y] = local_amax; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y == 0) { + float v = 0.0f; + for (int w = 0; w < (int)blockDim.y; ++w) v = fmaxf(v, slots[w]); + atomic_max_float(p.amax, v); + } + } +} + +template +void launch_fp8_quantize_tiled(const FP8QuantizeParams& p, + cudaStream_t stream) { + const dim3 grid((p.cols + 31) / 32, (p.rows + 31) / 32); + if (grid.x == 0 || grid.y == 0) return; + fp8_quantize_tiled_kernel + <<>>(p); +} + } // namespace fp8 } // namespace astrai