perf: transpose-quantize backward operands to route all gemms nt
- quantize gains out_layout (0 row-major / 1 transposed / 2 single-read dual-write); modes 1/2 run a new 32x32 smem-tile transpose kernel - backward feeds g8/w8T and g8T/x8T to trans_b=True gemms, dropping the NN-swap and TT crosswise kernels from training; fp8 weights keep the swap fallback - a 64x64 tile variant tied on the real step mix and was reverted; noted in the kernel header Benchmark: NVIDIA L20, 1.2B model, full train step fwd+bwd+CE - M=8192: fp8 551.8 -> 532.2 ms, 1.21x -> 1.26x vs bf16; M=2048 0.90x -> 0.95x - kernel-level grad_x +3.7..12.4%, grad_w +13.8..20.8%; layouts byte-exact, fp8 tests 36/36
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user