From 851ec9b4d5fff23a604037f0a4820078cb8c6901 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 2 Jun 2026 09:02:52 +0000 Subject: [PATCH] P3 WIP: fused RMSNorm + quantize kernel skeleton (not yet integrated) --- dsv4/kernels/cuda/fused_rmsnorm_quantize.cu | 241 ++++++++++++++++++++ 1 file changed, 241 insertions(+) create mode 100644 dsv4/kernels/cuda/fused_rmsnorm_quantize.cu diff --git a/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu b/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu new file mode 100644 index 00000000..bdd71bc9 --- /dev/null +++ b/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu @@ -0,0 +1,241 @@ +/* + * fused_rmsnorm_amax_quantize.cu + * + * Fused RMSNorm + amax + NVFP4 quantize. + * Combines 5+ kernel launches into 2: + * Kernel 1: compute RMS + row-wise amax → gsa (GPU buffer) + * Kernel 2: normalize + quantize using gsa from GPU buffer + * + * Eliminates: + * - 3-4 BF16 RMSNorm kernel launches (x^2, mean, rsqrt, multiply) + * - 1 amax kernel launch + * - Total: ~4-5 launches per RMSNorm site × ~366 sites = ~1500 launches/token + * + * Production: used for attention input norm, FFN input norm, q_a_norm, kv_norm. + */ + +#include +#include +#include +#include +#include + +// FP4 E2M1 look-up table (same as production) +__device__ __constant__ float FP4_LUT[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + +// ============================================================================ +// Kernel 1: Compute RMS + row-wise amax + gsa +// ============================================================================ +// Each block handles one row. Threadblock: (WARP_SIZE, 1, 1) per row. +// For hidden_size=7168: 7168/32 = 224 warps needed. Use multiple blocks. + +__global__ void rmsnorm_amax_gsa_kernel( + const __nv_bfloat16* __restrict__ x, // (M, N) BF16 row-major + const float* __restrict__ norm_weight, // (N,) FP32 + float* __restrict__ gsa_out, // (M,) FP32 — per-row gsa + float* __restrict__ rms_out, // (M,) FP32 — per-row RMS (for kernel 2) + const int M, + const int N, + const float eps, + const float divisor // gsa = amax / divisor +) { + const int row = blockIdx.x; + if (row >= M) return; + + const __nv_bfloat16* x_row = x + row * N; + + // Step 1: Compute RMS (sum of x^2) using warp reduction + float sum_sq = 0.0f; + float row_amax = 0.0f; + + for (int col = threadIdx.x; col < N; col += blockDim.x) { + float val = __bfloat162float(x_row[col]); + float val_sq = val * val; + sum_sq += val_sq; + float abs_val = fabsf(val); + if (abs_val > row_amax) row_amax = abs_val; + } + + // Warp-level reduction for sum_sq and amax + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); + row_amax = fmaxf(row_amax, __shfl_down_sync(0xFFFFFFFF, row_amax, offset)); + } + + // Block-level reduction (if blockDim.x > warpSize) + // Use shared memory for cross-warp reduction + __shared__ float s_sum_sq[32]; // max 32 warps + __shared__ float s_amax[32]; + + int lane = threadIdx.x % warpSize; + int warp_id = threadIdx.x / warpSize; + int num_warps = blockDim.x / warpSize; + + if (lane == 0) { + s_sum_sq[warp_id] = sum_sq; + s_amax[warp_id] = row_amax; + } + __syncthreads(); + + if (warp_id == 0) { + sum_sq = (lane < num_warps) ? s_sum_sq[lane] : 0.0f; + row_amax = (lane < num_warps) ? s_amax[lane] : 0.0f; + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); + row_amax = fmaxf(row_amax, __shfl_down_sync(0xFFFFFFFF, row_amax, offset)); + } + } + + if (threadIdx.x == 0) { + float rms = sqrtf(sum_sq / N + eps); + float gsa = row_amax / divisor; + gsa_out[row] = gsa; + rms_out[row] = rms; + } +} + +// ============================================================================ +// Kernel 2: RMSNorm + quantize using gsa from GPU buffer +// ============================================================================ +// Each block handles one row. Quantizes x * rsqrt(rms) * weight to NVFP4. + +__global__ void rmsnorm_quantize_nvfp4_kernel( + const __nv_bfloat16* __restrict__ x, // (M, N) BF16 row-major + const float* __restrict__ norm_weight, // (N,) FP32 + const float* __restrict__ gsa, // (M,) FP32 — per-row global scale + const float* __restrict__ rms, // (M,) FP32 — per-row RMS + uint8_t* __restrict__ x_fp4, // (M, N//2) FP4 packed + __nv_fp8_e4m3* __restrict__ x_sf, // (M, N//16) E4M3 block scales + const int M, + const int N, + const float eps +) { + const int row = blockIdx.x; + if (row >= M) return; + + const __nv_bfloat16* x_row = x + row * N; + float row_gsa = gsa[row]; + float row_rms = rms[row]; + float inv_rms = 1.0f / row_rms; // rsqrt(sum/N + eps) = 1/rms + + // Process 16 elements per thread (one microblock) + // Each thread writes 8 FP4 values (16 BF16 → 8 bytes packed) + // Block scales: 1 E4M3 per 16 BF16 values + + const int num_blocks = N / 16; // number of 16-element microblocks + const int blocks_per_thread = 1; + const int total_threads = blockDim.x; + + for (int blk = threadIdx.x; blk < num_blocks; blk += total_threads) { + int col = blk * 16; + + // Load 16 BF16 values, normalize, and compute block scale + float vals[16]; + float block_amax = 0.0f; + + for (int i = 0; i < 16; i++) { + float v = __bfloat162float(x_row[col + i]); + v = v * inv_rms * norm_weight[col + i]; // RMSNorm + vals[i] = v; + float av = fabsf(v); + if (av > block_amax) block_amax = av; + } + + // Compute block scale + // block_scale = amax / (gsa * 6.0) + // E4M3 range is [-448, 448], FP4 range is [-6, 6] + // Quantized value = val / (gsa * block_scale) + float block_scale = block_amax / (row_gsa * 6.0f); + + // Clamp block_scale to E4M3 range + if (block_scale > 448.0f) block_scale = 448.0f; + if (block_scale < 0.0f) block_scale = 0.0f; + + // Convert block_scale to E4M3 + __nv_fp8_e4m3 bs_e4m3; + bs_e4m3.__x = 0; + if (block_scale > 0.0f) { + // Simple float → E4M3 conversion + // E4M3: 1 sign + 4 exponent + 3 mantissa bits + // Range: 2^-9 * (1 + 0/8) to 2^4 * (1 + 7/8) ≈ 0.00195 to 30.0 + // But we use it as a scale factor, so range [0, 448] + // Actually, E4M3 max = 448.0, and we clamp to that above. + // Use CUDA's native conversion: + float bs_f = block_scale; + __nv_fp8_e4m3 tmp; + tmp = __nv_fp8_e4m3(bs_f); + bs_e4m3 = tmp; + } + + // Quantize 16 values to FP4 (8 bytes) + uint8_t packed = 0; + + for (int i = 0; i < 16; i += 2) { + // Two values per byte + float v0 = vals[i] / (row_gsa * block_scale); + float v1 = vals[i + 1] / (row_gsa * block_scale); + + // Clamp to FP4 range [-6, 6] + v0 = fmaxf(fminf(v0, 6.0f), -6.0f); + v1 = fmaxf(fminf(v1, 6.0f), -6.0f); + + // Find nearest FP4 value + auto quantize_fp4 = [](float v) -> uint8_t { + bool sign = v < 0; + float av = fabsf(v); + float best_err = 1e10f; + uint8_t best_idx = 0; + for (int k = 0; k < 8; k++) { + float err = fabsf(av - FP4_LUT[k]); + if (err < best_err) { + best_err = err; + best_idx = k; + } + } + return best_idx | (sign ? 0x08 : 0x00); + }; + + uint8_t lo = quantize_fp4(v0); + uint8_t hi = quantize_fp4(v1); + packed |= (lo | (hi << 4)) << (i / 2 * 8); // pack 2 nibbles per byte + } + + // Wait, this packing is wrong. Each byte packs 2 FP4 values (nibbles). + // 16 BF16 → 8 FP4 bytes → 16 FP4 nibbles. + // Actually: 16 BF16 values → 8 bytes, each byte has 2 FP4 nibbles. + + // Let me redo this properly: + uint8_t fp4_bytes[8] = {0}; + for (int i = 0; i < 16; i += 2) { + float v0 = vals[i] / (row_gsa * block_scale); + float v1 = vals[i + 1] / (row_gsa * block_scale); + v0 = fmaxf(fminf(v0, 6.0f), -6.0f); + v1 = fmaxf(fminf(v1, 6.0f), -6.0f); + + auto qf4 = [](float v) -> uint8_t { + bool sign = v < 0; + float av = fabsf(v); + float best_err = 1e10f; + uint8_t best_idx = 0; + for (int k = 0; k < 8; k++) { + float err = fabsf(av - FP4_LUT[k]); + if (err < best_err) { best_err = err; best_idx = k; } + } + return best_idx | (sign ? 0x08 : 0x00); + }; + + uint8_t lo = qf4(v0); + uint8_t hi = qf4(v1); + fp4_bytes[i / 2] = lo | (hi << 4); + } + + // Write 8 FP4 bytes + uint8_t* dst = x_fp4 + row * (N / 2) + blk * 8; + for (int i = 0; i < 8; i++) { + dst[i] = fp4_bytes[i]; + } + + // Write block scale + x_sf[row * (N / 16) + blk] = bs_e4m3; + } +}