From 36e39496d4e081a256cc22c63dcb74daae7036df Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 28 Aug 2026 22:39:02 +0800 Subject: [PATCH] perf: vectorize tiled fp8 transpose quantize and arm amax via memset - tiled transpose quantize becomes one 64x32-tile kernel: native pair loads (128B warp reads) with in-kernel scalar fallback at unaligned or ragged rows, so odd widths and misaligned bases no longer route to a separate kernel - the old 32x32 scalar tiled kernel and its launcher correctness branch are gone; grid sizing simplifies to 1 + total / (vec * threads) since both elementwise loops are grid-stride - quantize arms the amax buffer with cudaMemsetAsync instead of the zeros() fill kernel, dropping one tensor-op dispatch and kernel launch per call - byte-exact parity holds over 1404 golden records (13 shapes x 3 dtypes x 3 scales x 2 formats x 3 layouts x aligned/misaligned) and tests/extension passes 65/65 - elementwise quantize kernel left unchanged: 16B-store pairing, __ldcs streaming hints and amax tree reduction all measured neutral at its ~52% DRAM ceiling and were reverted Benchmark: L20 (sm_89), profiler kernel time with L2 flushed between calls. - transposed quantize (layout 1): 230 -> 294 GB/s on 2048x1536 (+28%), 245 -> 299 on 2048x1536 weights (+22%); dual-layout (layout 2) 248 -> 329 (+33%) on the same shapes - DRAM-saturated sizes (~10.6M elements) regress ~5% (404 -> 384 GB/s on 8192x1536), ~0.02% of a training step; accepted for the single-kernel shape after scalar-path and geometry variants both measured the same - amax init fill kernel 3.0us -> memset 0.9us; quantize call CPU wall 18.5 -> 13.4us on 128x1536 --- csrc/kernels/fp8/ops.cu | 6 +- csrc/kernels/fp8/quantize.cuh | 128 ++++++++++++++++++++++++---------- 2 files changed, 96 insertions(+), 38 deletions(-) diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index 0d2587d..21da867 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -120,7 +120,11 @@ py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt, auto input = x.contiguous(); auto out_opts = input.options().dtype( fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn); - auto amax = torch::zeros({1}, input.options().dtype(torch::kFloat32)); + // amax is reduced via atomicMax of non-negative values; a driver memset + // arms it cheaper than the zeros() fill kernel (one fewer tensor-op + // dispatch + kernel launch on every quantize call). + auto amax = torch::empty({1}, input.options().dtype(torch::kFloat32)); + cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream()); FP8QuantizeParams p; p.input_ptr = input.data_ptr(); diff --git a/csrc/kernels/fp8/quantize.cuh b/csrc/kernels/fp8/quantize.cuh index 1f70159..1d51189 100644 --- a/csrc/kernels/fp8/quantize.cuh +++ b/csrc/kernels/fp8/quantize.cuh @@ -15,8 +15,8 @@ namespace astrai { namespace fp8 { -// Input element type traits: one element -> float, and the unpack of one -// 16-byte load into kVecElems floats. +// Input element type traits: one element -> float, the unpack of one +// 16-byte load into kVecElems floats, and a native 2-element pair load. template struct quant_in_traits; @@ -37,6 +37,13 @@ struct quant_in_traits<__nv_bfloat16> { f[2 * j + 1] = p.y; } } + static __device__ __forceinline__ void load_pair(const __nv_bfloat16* p, + float* f) { + const float2 v = __bfloat1622float2( + *reinterpret_cast(p)); + f[0] = v.x; + f[1] = v.y; + } }; template <> @@ -55,6 +62,13 @@ struct quant_in_traits<__half> { f[2 * j + 1] = p.y; } } + static __device__ __forceinline__ void load_pair(const __half* p, + float* f) { + const float2 v = + __half22float2(*reinterpret_cast(p)); + f[0] = v.x; + f[1] = v.y; + } }; template <> @@ -67,6 +81,11 @@ struct quant_in_traits { #pragma unroll for (int j = 0; j < 4; ++j) f[j] = __uint_as_float(w[j]); } + static __device__ __forceinline__ void load_pair(const float* p, + float* f) { + f[0] = p[0]; + f[1] = p[1]; + } }; // One float -> one fp8 byte (round-nearest-even + satfinite). @@ -161,76 +180,111 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) { // Tiled transpose quantize (out_layout 1/2): reads the [rows][cols] 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 row-major -// copy too. A 32x32 tile stages through shared memory: loads and 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 -// read is a 32-byte broadcast segment. (A 64x64 split-half variant measured -// +21% L2-resident but -3..5% DRAM-bound; the real step mix ties, so the -// simpler tile stays.) +// copy too. 64x32 tiles, one native pair load per row (a full 128B warp +// read); rows whose pair is unaligned or ragged (odd widths, misaligned +// bases) fall back to element loads in place. Staging goes through a byte +// tile whose pitch keeps the store stride coprime with the 32 banks. +// (+25-35% over the former 32x32 scalar kernel on sub-4M tensors; ~5% +// slower once DRAM-saturated — accepted for the single-kernel shape.) template __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) { - constexpr int kTile = 32; - __shared__ uint8_t tile[kTile][kTile + 4]; + constexpr int kTileC = 64, kTileR = 32; + // 34B pitch: staging stride is 17 words (coprime with the 32 banks) so + // the pair-byte stores stay conflict-free, and the byte-wise consume + // reads still span distinct words. + __shared__ uint8_t tile[kTileC][kTileR + 2]; 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 r0 = blockIdx.y * kTileR; + const int c0 = blockIdx.x * kTileC; const int r = r0 + threadIdx.y * 4; - const int c = c0 + threadIdx.x; + const int c = c0 + threadIdx.x * 2; // cols even => the pair is in-bounds - uint8_t q[4]; + uint8_t q[4][2]; float local_amax = 0.0f; + // Vectorize the pair when both elements are in-bounds and the native + // 2-element load is aligned; odd widths, misaligned bases and ragged + // edges fall back to element loads row by row. + constexpr int kPairAlign = 2 * (int)sizeof(InT); #pragma unroll for (int j = 0; j < 4; ++j) { - q[j] = 0; + q[j][0] = 0; + q[j][1] = 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)); - q[j] = cvt_fp8(v * mult); + const InT* a = x + (int64_t)(r + j) * p.cols + c; + if (c + 1 < p.cols && + (reinterpret_cast(a) & (kPairAlign - 1)) == 0) { + float f[2]; + quant_in_traits::load_pair(a, f); +#pragma unroll + for (int k = 0; k < 2; ++k) { + local_amax = fmaxf(local_amax, fabsf(f[k])); + q[j][k] = cvt_fp8(f[k] * mult); + } + } else { + const float v0 = quant_in_traits::to_float(a[0]); + local_amax = fmaxf(local_amax, fabsf(v0)); + q[j][0] = cvt_fp8(v0 * mult); + if (c + 1 < p.cols) { + const float v1 = quant_in_traits::to_float(a[1]); + local_amax = fmaxf(local_amax, fabsf(v1)); + q[j][1] = cvt_fp8(v1 * mult); + } + } } } 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]; + if (r + j < p.rows && c < p.cols) { + uint8_t* o = out + (int64_t)(r + j) * p.cols + c; + const int64_t off = (int64_t)(r + j) * p.cols + c; + if (c + 1 < p.cols && (off & 1) == 0) + *reinterpret_cast(o) = + (unsigned short)(q[j][0] | (q[j][1] << 8)); + else { + o[0] = q[j][0]; + if (c + 1 < p.cols) o[1] = q[j][1]; + } + } } #pragma unroll - for (int j = 0; j < 4; ++j) tile[threadIdx.x][threadIdx.y * 4 + j] = q[j]; + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int k = 0; k < 2; ++k) + tile[threadIdx.x * 2 + k][threadIdx.y * 4 + j] = q[j][k]; __syncthreads(); - // Transposed scatter: output element (c, r) lives at c * rows + r; r - // tracks threadIdx.x so each warp writes one contiguous run. tile was - // written as tile[col][row], so input (r0+tx, c0+ty*4+j) reads back - // from tile[ty*4+j][tx]. + // Transposed scatter: output element (c, r) lives at c * rows + r; + // threadIdx.x tracks r so each warp writes one contiguous run. tile is + // [col][row]; warp y walks 8 columns, threads read down one column. 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; + for (int i = 0; i < 8; ++i) { + const int oc = c0 + threadIdx.y * 8 + i; 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]; + tile[threadIdx.y * 8 + i][threadIdx.x]; } if (p.amax) publish_amax<8>(p.amax, local_amax); } // Unified quantize launcher: Tiled selects the transpose kernel (out_layout -// 1/2) over the vectorized elementwise one. +// 1/2) over the vectorized elementwise one. The transpose kernel vectorizes +// pair loads in-kernel and falls back to scalar loads at unaligned/ragged +// rows, so the host side picks only the grid. template void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) { if constexpr (Tiled) { - const dim3 grid((p.cols + 31) / 32, (p.rows + 31) / 32); + const dim3 grid((p.cols + 63) / 64, (p.rows + 31) / 32); if (grid.x == 0 || grid.y == 0) return; - fp8_quantize_tiled_kernel - <<>>(p); + fp8_quantize_tiled_kernel<<>>(p); } else { constexpr int kThreads = 256; constexpr int kVecElems = quant_in_traits::kVecElems; - // One block per 256 vectors; at least one block so a tiny or - // misaligned tensor's scalar tail is still covered. - int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads; - if (blocks < 1) blocks = 1; + // Grid-stride loops: any grid >= 1 is correct; one block per 256 + // vectors plus the tail block covers tiny and misaligned tensors. + const int64_t blocks = 1 + p.total / (kVecElems * kThreads); fp8_quantize_kernel<<>>(p); } }