ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: vectorize tiled fp8 transpose quantize and arm amax via memset

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-08-31 14:24:51 +08:00
parent a1a1a6bf0f
commit 36e39496d4
2 changed files with 96 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/quantize.cuh
+91 -37
View File
@@ -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 <typename InT>
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<const __nv_bfloat162*>(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<const __half2*>(p));
f[0] = v.x;
f[1] = v.y;
}
};
template <>
@@ -67,6 +81,11 @@ struct quant_in_traits<float> {
#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 <FP8Format Fmt, typename InT>
__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<const InT*>(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<InT>::to_float(x[(int64_t)(r + j) * p.cols + c]);
local_amax = fmaxf(local_amax, fabsf(v));
q[j] = cvt_fp8<Fmt>(v * mult);
const InT* a = x + (int64_t)(r + j) * p.cols + c;
if (c + 1 < p.cols &&
(reinterpret_cast<uintptr_t>(a) & (kPairAlign - 1)) == 0) {
float f[2];
quant_in_traits<InT>::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<Fmt>(f[k] * mult);
}
} else {
const float v0 = quant_in_traits<InT>::to_float(a[0]);
local_amax = fmaxf(local_amax, fabsf(v0));
q[j][0] = cvt_fp8<Fmt>(v0 * mult);
if (c + 1 < p.cols) {
const float v1 = quant_in_traits<InT>::to_float(a[1]);
local_amax = fmaxf(local_amax, fabsf(v1));
q[j][1] = cvt_fp8<Fmt>(v1 * mult);
}
}
}
}
if (p.out_layout == 2) {
uint8_t* out = static_cast<uint8_t*>(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<unsigned short*>(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<uint8_t*>(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 <FP8Format Fmt, typename InT, bool Tiled = false>
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<Fmt, InT>
<<<grid, dim3(32, 8), 0, stream>>>(p);
fp8_quantize_tiled_kernel<Fmt, InT><<<grid, dim3(32, 8), 0, stream>>>(p);
} else {
constexpr int kThreads = 256;
constexpr int kVecElems = quant_in_traits<InT>::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<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
}
}