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

refactor: dedupe fp8 kernel helpers and trim comments

Browse Source 
- merge the quantize launchers into one Tiled template; extract shared cvt_fp8/publish_amax helpers and replace the dtype x format ladder with two-level template dispatch
- fold the gemm interior/generic operand loads into one kInterior template and the fast/generic async loads into load_async<kFast>; Policy carries the smem budget
- compress kernel comments to the load-bearing invariants, dropping measured-number essays; Policy signature and kernel code unchanged

Benchmark: NVIDIA L20, 1.2B model train step fwd+bwd+CE
- M=8192: fp8 532.2 -> 530.4 ms (1.26x, noise); tests/extension 65 passed, quantize layouts byte-exact, NT routing diff 0.0
This commit is contained in:
ViperEkura
2026-08-28 16:45:21 +08:00
parent 8a353117ea
commit bf239d194c
4 changed files with 376 additions and 651 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/common.h
+33 -61
View File
@@ -11,30 +11,23 @@
namespace astrai {
namespace fp8 {
// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or
// E5M2 (gradient / large dynamic range, max 57344).
// Compile-time FP8 format: E4M3 (forward, max 448) or E5M2 (gradients,
// max 57344).
enum class FP8Format : int {
E4M3 = 0,
E5M2 = 1,
};
// Operand memory layouts as types (CUTLASS-style tags). The tag names the
// storage order of the raw buffer relative to the operand's canonical GEMM
// matrix — A is [M][K], B is [K][N]:
// A RowMajor = [M][K] storage (K-contiguous rows; the default)
// A ColMajor = [K][M] storage (M-contiguous; A^T)
// B RowMajor = [K][N] storage (N-contiguous; the plain a @ b operand)
// B ColMajor = [N][K] storage (K-contiguous; the nn.Linear weight layout)
// Empty tags: selection happens by type at compile time (see load_operand_tile).
// Operand storage tags (CUTLASS-style) relative to the canonical matrices
// A [M][K] / B [K][N]: A RowMajor = [M][K] (default), A ColMajor = [K][M],
// B RowMajor = [K][N], B ColMajor = [N][K] (the nn.Linear weight). Selection
// is by type at compile time (see gemm.cuh's stage loads).
struct RowMajor {};
struct ColMajor {};
// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
// warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128 CTA, or 32x32 on the
// cuBLAS-style 64x64 small CTA that lifts small-shape occupancy) and the
// cp.async pipeline depth.
// Compile-time tile configuration, mirroring KernelTraits in the attention
// kernels: CTA tile, warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128
// CTA, 32x32 on the 64x64 small CTA) and cp.async pipeline depth.
template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages,
int WarpM = 64, int WarpN = 32>
struct Fp8GemmTraits {
@@ -50,10 +43,8 @@ struct Fp8GemmTraits {
kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
// Derived launch geometry: WarpM x WarpN warp tiles tile the CTA. The
// shared-memory budget is layout-aware (crosswise operands add K-major
// staging + a canonical buffer), so it lives in Fp8GemmSmem in gemm.cuh
// together with the resident-CTA hint for __launch_bounds__.
// Derived geometry: warp tiles tile the CTA. The smem budget is
// layout-aware, so it lives in Fp8GemmSmem (gemm.cuh).
static constexpr int kWarpsM = BlockM / WarpM;
static constexpr int kWarpsN = BlockN / WarpN;
static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32;
@@ -63,73 +54,54 @@ struct Fp8GemmTraits {
"warp tile must be a multiple of the m16n8 MMA shape");
};
// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8
// with fused amax.
// Quantize-kernel parameter POD: float input -> FP8 with fused amax.
struct FP8QuantizeParams {
// Float input and FP8 output buffers; scale is the quantization
// multiplier (device scalar). amax (may be null) is zero-initialized by
// 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.
void* __restrict__ output_transposed_ptr = nullptr; // [cols][rows]
// Output layout: 0 = row-major only, 1 = transposed only, 2 = both from
// a single read. Modes 1/2 produce K-contiguous operands so crosswise
// consumers (backward grad_x / grad_w) route through the NT fast path.
int out_layout = 0;
const float* __restrict__ scale = nullptr;
float* __restrict__ amax = nullptr;
const float* __restrict__ scale = nullptr; // device multiplier
float* __restrict__ amax = nullptr; // raw-domain max out
// Element count (only the elementwise quantize kernel uses it); the
// tiled kernel views the same buffer as [rows][cols] row-major.
// Element count (elementwise kernel); 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
// through the pre-quantized GEMM kernels. Each kernel touches only the
// fields it needs; buffers are raw pointers packed by the torch binding.
// Pointer members default to null so optional paths cannot hold garbage.
// through the kernels; each kernel touches only the fields it needs.
struct FP8Params {
// Inputs: a/b are FP8 for the pre-quantized path. Scales are
// quantization steps (device scalars).
// Optional bf16 bias broadcast over output rows (fused into the epilogue
// before the bf16 rounding, so it adds in fp32 — one rounding fewer than
// the separate out + bias elementwise kernel it replaces). Null disables.
// FP8 operands + output; scales are quantization steps (device
// scalars). Optional bf16 bias fuses into the epilogue (fp32 add before
// the single bf16 rounding); null disables.
const void* __restrict__ a_ptr = nullptr;
const void* __restrict__ b_ptr = nullptr;
const void* __restrict__ bias_ptr = nullptr;
void* __restrict__ out_ptr = nullptr;
const float* __restrict__ scale = nullptr;
// Transposed-output mode (set by dispatch_fp8_gemm's swap for NN
// problems): the kernel computes E[N'][M'] over swapped operands and the
// epilogue scatters into the caller's [M][N] row-major buffer, so
// D[row][col] lives at out[col * p.m + row] — p.m/p.n are the swapped
// problem's dims and the D row stride is p.m. Zero in the plain
// orientation.
// NN-swap mode (canonicalize_gemm): the kernel computes the transposed
// problem and the epilogue scatters D[row][col] to out[col * p.m + row]
// in the caller's [M][N] buffer. Zero in the plain orientation.
int out_transposed = 0;
// Shapes. `int` covers every realistic LLM shape; the kernels promote
// to int64 for all pointer arithmetic.
int m, n, k;
int m, n, k; // int covers LLM shapes; kernels promote to int64
// Batched (bmm) geometry: grid.z slices step the operand/output pointers
// by these element strides (0 broadcasts the operand across batches).
// Batched (bmm) geometry: grid.z steps these element strides (0
// broadcasts the operand across batches).
int batch = 1;
int64_t a_batch_stride = 0;
int64_t b_batch_stride = 0;
int64_t out_batch_stride = 0;
// Physical leading dimensions (column count, i.e. row stride) of A and
// B. For a non-transposed operand the stride equals the contract dim;
// for a transposed operand it is the operand's own column count. The
// binding packs these so the kernel reads both buffers either naturally
// or transposed depending on the LayoutA/LayoutB tags (see gemm.cuh).
// Physical leading dims (row strides) of A and B; the binding packs
// them so the kernel reads each buffer naturally or transposed per the
// LayoutA/LayoutB tags.
int a_ld, b_ld;
};
csrc/kernels/fp8/gemm.cuh
+226 -434
View File
File diff suppressed because it is too large Load Diff
csrc/kernels/fp8/ops.cu
+36 -57
View File
@@ -6,7 +6,6 @@
#include <cstdint>
#include <mutex>
#include <tuple>
#include <unordered_map>
#include "../common/device.cuh"
@@ -46,32 +45,13 @@ void check_scale(const torch::Tensor& scale, const torch::Tensor& input) {
"scale must be a CUDA float32 scalar on the input device");
}
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) {
p.a_ptr = a;
p.b_ptr = b;
p.out_ptr = output;
p.scale = scale.data_ptr<float>();
p.m = static_cast<int>(m);
p.n = static_cast<int>(n);
p.k = static_cast<int>(k);
p.a_ld = static_cast<int>(a_ld);
p.b_ld = static_cast<int>(b_ld);
}
// Layout dispatch (the NN swap in canonicalize_gemm) and launch planning
// (plan_gemm/launch_plan) live in gemm.cuh behind fp8::gemm — pure CUDA,
// shared with the C test suite.
// Inner-layout resolution for one GEMM operand. The user flag names the
// math (0 = tensor's last two dims are [rows][contract], 1 = transposed);
// the storage may independently be a col-major view (.t() of a contiguous
// math (0 = last two dims are [rows][contract], 1 = transposed); the
// storage may independently be a col-major view (.t() of a contiguous
// buffer), which folds into the returned dispatch flag at zero copy — the
// kernel's LayoutA/LayoutB tags cover both storages. m/n/k derive from the
// user flag only; the fold never swaps them (see the layout table in
// gemm.cuh). Tensors whose inner dims are neither natural layout fall back
// to .contiguous().
// user flag only. Tensors whose inner dims are neither natural layout fall
// back to .contiguous().
bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
int64_t& batch_stride, torch::Tensor& storage) {
torch::Tensor t = t_in;
@@ -89,43 +69,36 @@ 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 <bool Tiled, FP8Format Fmt, typename InT>
void launch_one(const FP8QuantizeParams& p, cudaStream_t stream) {
if constexpr (Tiled)
launch_fp8_quantize_tiled<Fmt, InT>(p, stream);
else
launch_fp8_quantize<Fmt, InT>(p, stream);
// Dtype dispatch over the unified quantize launcher.
template <bool Tiled, FP8Format Fmt>
void launch_for_dtype(const torch::Tensor& x, const FP8QuantizeParams& p,
cudaStream_t stream) {
switch (x.scalar_type()) {
case torch::kHalf:
launch_fp8_quantize<Fmt, __half, Tiled>(p, stream);
break;
case torch::kFloat32:
launch_fp8_quantize<Fmt, float, Tiled>(p, stream);
break;
default:
launch_fp8_quantize<Fmt, __nv_bfloat16, Tiled>(p, stream);
}
}
template <bool Tiled>
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<Tiled, FP8Format::E5M2, __half>(p, stream);
else
launch_one<Tiled, FP8Format::E4M3, __half>(p, stream);
} else if (x.scalar_type() == torch::kFloat32) {
if (e5m2)
launch_one<Tiled, FP8Format::E5M2, float>(p, stream);
else
launch_one<Tiled, FP8Format::E4M3, float>(p, stream);
} else {
if (e5m2)
launch_one<Tiled, FP8Format::E5M2, __nv_bfloat16>(p, stream);
else
launch_one<Tiled, FP8Format::E4M3, __nv_bfloat16>(p, stream);
}
if (e5m2)
launch_for_dtype<Tiled, FP8Format::E5M2>(x, p, stream);
else
launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
}
} // namespace
// 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).
// Output-layout dispatch: 0 = [rows][cols] row-major (2-tuple return),
// 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations from a
// single read (3-tuple). Layouts 1/2 feed the NT GEMM fast path.
py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
int64_t layout) {
TORCH_CHECK(x.is_cuda(), "CUDA tensors required");
@@ -219,8 +192,15 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
? torch::empty({batch, m, n}, a.options().dtype(torch::kBFloat16))
: torch::empty({m, n}, a.options().dtype(torch::kBFloat16));
FP8Params p;
pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
m, n, k, a_ld, b_ld);
p.a_ptr = a_st.data_ptr();
p.b_ptr = b_st.data_ptr();
p.out_ptr = output.data_ptr();
p.scale = scale.data_ptr<float>();
p.m = static_cast<int>(m);
p.n = static_cast<int>(n);
p.k = static_cast<int>(k);
p.a_ld = static_cast<int>(a_ld);
p.b_ld = static_cast<int>(b_ld);
// Fused epilogue bias (bf16, broadcast over rows and batches). An
// undefined or 0-element tensor keeps the plain scaled output.
if (bias.defined() && bias.numel() > 0) {
@@ -243,9 +223,8 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
return output;
}
// mm_fp8 binding: Python None and an omitted argument both mean "no bias"
// (resolved to an undefined tensor here, so every Python layer can pass its
// bias argument through untouched instead of normalizing it host-side).
// mm_fp8 binding: Python None and an omitted argument both mean "no bias",
// so every Python layer can pass its bias argument through untouched.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
py::arg("fmt"), py::arg("layout") = 0);
csrc/kernels/fp8/quantize.cuh
+81 -99
View File
@@ -1,10 +1,7 @@
#pragma once
// FP8 quantize device code — pure CUDA, no torch. Any float input element
// type (bf16 / fp16 / fp32) converts to E4M3 or E5M2 with a fused amax over
// the raw (unscaled) values. Mirrors the GEMM file's split: kernels take the
// FP8QuantizeParams POD, formats and input types ride on template parameters,
// and the launcher is a plain function usable from both the torch binding and
// pure C tests.
// FP8 quantize device code — pure CUDA, no torch: kernels take the
// FP8QuantizeParams POD, format and input type ride on template parameters,
// and the launcher is shared by the torch binding and the C tests.
#include <cuda_bf16.h>
#include <cuda_fp16.h>
@@ -18,8 +15,8 @@
namespace astrai {
namespace fp8 {
// Input element type traits: one element -> float, and the vectorized
// unpack of one 16-byte load into kVecElems floats.
// Input element type traits: one element -> float, and the unpack of one
// 16-byte load into kVecElems floats.
template <typename InT>
struct quant_in_traits;
@@ -31,12 +28,13 @@ struct quant_in_traits<__nv_bfloat16> {
}
static __device__ __forceinline__ void load_vec(const uint4& raw,
float* f) {
const unsigned w[4] = {raw.x, raw.y, raw.z, raw.w};
const __nv_bfloat162* b2 =
reinterpret_cast<const __nv_bfloat162*>(&raw);
#pragma unroll
for (int j = 0; j < 4; ++j) {
f[2 * j] =
__bfloat162float(__ushort_as_bfloat16(w[j] & 0xffffu));
f[2 * j + 1] = __bfloat162float(__ushort_as_bfloat16(w[j] >> 16));
const float2 p = __bfloat1622float2(b2[j]);
f[2 * j] = p.x;
f[2 * j + 1] = p.y;
}
}
};
@@ -65,15 +63,22 @@ struct quant_in_traits<float> {
static __device__ __forceinline__ float to_float(float v) { return v; }
static __device__ __forceinline__ void load_vec(const uint4& raw,
float* f) {
f[0] = __uint_as_float(raw.x);
f[1] = __uint_as_float(raw.y);
f[2] = __uint_as_float(raw.z);
f[3] = __uint_as_float(raw.w);
const unsigned* w = reinterpret_cast<const unsigned*>(&raw);
#pragma unroll
for (int j = 0; j < 4; ++j) f[j] = __uint_as_float(w[j]);
}
};
// Convert one float pair to one packed fp8 pair. The stored bytes see
// value * mult (round-nearest-even + satfinite).
// One float -> one fp8 byte (round-nearest-even + satfinite).
template <FP8Format Fmt>
__device__ __forceinline__ uint8_t cvt_fp8(float v) {
if constexpr (Fmt == FP8Format::E5M2)
return __nv_fp8_e5m2(v).__x;
else
return __nv_fp8_e4m3(v).__x;
}
// One float pair -> one packed fp8x2 word (round-nearest-even + satfinite).
template <FP8Format Fmt>
__device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
constexpr __nv_fp8_interpretation_t kFmt =
@@ -82,23 +87,36 @@ __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
make_float2(a, b), __NV_SATFINITE, kFmt));
}
// Quantize kernel: float input -> FP8 (E4M3 or E5M2), fused amax over raw
// values.
// Block-wide amax reduce -> one atomic per block: warp-reduce, park one
// value per warp, thread 0 folds. kWarps must cover the block's warp count.
template <int kWarps>
__device__ __forceinline__ void publish_amax(float* amax, float v) {
v = warp_reduce_max(v);
__shared__ float slots[kWarps];
const int tid = threadIdx.y * blockDim.x + threadIdx.x;
if ((tid & 31) == 0) slots[tid >> 5] = v;
__syncthreads();
if (tid == 0) {
#pragma unroll
for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]);
atomic_max_float(amax, v);
}
}
// Elementwise quantize kernel (out_layout 0): vectorized 16B loads -> fp8
// stores, fused amax over raw values.
template <FP8Format Fmt, typename InT>
__global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
const float mult = *p.scale;
const auto* x = static_cast<const InT*>(p.input_ptr);
void* x8 = p.output_ptr;
float* amax = p.amax;
uint8_t* x8 = static_cast<uint8_t*>(p.output_ptr);
float local_amax = 0.0f;
const int64_t stride = (int64_t)blockDim.x * gridDim.x;
// Vectorized body: one 16B load -> kVecElems fp8 bytes per step (8
// elements for 16-bit inputs, 4 for fp32). Torch allocations are >=16B
// aligned and the binding passes freshly allocated contiguous buffers,
// so element 0 keeps the uint4 access natural; a misaligned base
// (contiguous view with an odd storage offset) falls back to the scalar
// loop below via total_vec = 0.
// One 16B load -> kVecElems bytes per step. Torch allocations are >=16B
// aligned, so element 0 keeps the uint4 access natural; a misaligned
// base (odd storage offset view) falls to the scalar tail via
// total_vec = 0.
constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
const bool aligned =
((reinterpret_cast<uintptr_t>(x) |
@@ -125,8 +143,7 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
packed[j] = (lo & 0xffffu) | (hi << 16);
}
if constexpr (kVecElems == 8)
reinterpret_cast<uint2*>(x8)[i] =
make_uint2(packed[0], packed[1]);
reinterpret_cast<uint2*>(x8)[i] = make_uint2(packed[0], packed[1]);
else
reinterpret_cast<unsigned*>(x8)[i] = packed[0];
}
@@ -136,51 +153,20 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
i < p.total; i += stride) {
const float v = quant_in_traits<InT>::to_float(x[i]);
local_amax = fmaxf(local_amax, fabsf(v));
if constexpr (Fmt == FP8Format::E5M2) {
reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] =
__nv_fp8_e5m2(v * mult);
} else {
reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] =
__nv_fp8_e4m3(v * mult);
}
}
if (amax) {
local_amax = warp_reduce_max(local_amax);
__shared__ float slots[32];
if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax;
__syncthreads();
if (threadIdx.x == 0) {
float v = 0.0f;
for (int w = 0; w < (blockDim.x >> 5); ++w)
v = fmaxf(v, slots[w]);
atomic_max_float(amax, v);
}
x8[i] = cvt_fp8<Fmt>(v * mult);
}
if (p.amax) publish_amax<8>(p.amax, local_amax);
}
template <FP8Format Fmt, typename InT>
void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
constexpr int kThreads = 256;
// One block per 256 vectors; at least one block so the scalar tail of a
// tiny / misaligned tensor is still covered.
constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
if (blocks < 1) blocks = 1;
fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(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.
// 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.)
template <FP8Format Fmt, typename InT>
__global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
constexpr int kTile = 32;
@@ -201,10 +187,7 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
const float v =
quant_in_traits<InT>::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;
q[j] = cvt_fp8<Fmt>(v * mult);
}
}
if (p.out_layout == 2) {
@@ -218,9 +201,9 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
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].
// 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].
uint8_t* out_t = static_cast<uint8_t*>(p.output_transposed_ptr);
#pragma unroll
for (int j = 0; j < 4; ++j) {
@@ -229,28 +212,27 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
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);
}
}
if (p.amax) publish_amax<8>(p.amax, local_amax);
}
template <FP8Format Fmt, typename InT>
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<Fmt, InT>
<<<grid, dim3(32, 8), 0, stream>>>(p);
// Unified quantize launcher: Tiled selects the transpose kernel (out_layout
// 1/2) over the vectorized elementwise one.
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);
if (grid.x == 0 || grid.y == 0) return;
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;
fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
}
}
} // namespace fp8