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); } }