From 1a064ab0921238c1daa397d6f4a900ef33884de2 Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Wed, 22 Jul 2026 19:28:02 +0200 Subject: [PATCH] CUDA: Improve NVFP4 W4A4 activation quantization (#25730) * Squash history before conflict-resolution during rebase on master WIP commit Add 32-byte loads, restore per-block amax Use nvfp4x4 intrinsic when available Fuse per-channel amax and quantization kernels Do pointer arithmetic only once on x Remove unnecessary ternary in the load We assert on host side that ne00 is 64-aligned Add back scale-search, but optimize it with intrinsics Code cleanup Make scale in MMQ-epilogue NVFP4-specific/restrictive for now Remove unneeded include, add comment Fix trailing whitespace Guard __builtin_align__(32) struct to NVIDIA Seems like HIP doesn't have this available, see https://github.com/ggml-org/llama.cpp/actions/runs/29438651734/job/87431623001 * compiler massaging to avoid unnecessary LDCs * kvalues_mxfp4 -> kvalues_nvfp4 in quantize_mmq_nvfp4 * Always pass in src1_scale.ptr * Extract ggml_cuda_is_aligned helper --- ggml/src/ggml-cuda/common.cuh | 9 + ggml/src/ggml-cuda/mmq.cu | 33 +++- ggml/src/ggml-cuda/mmq.cuh | 107 ++++++++-- ggml/src/ggml-cuda/quantize.cu | 339 +++++++++++++++++++++++--------- ggml/src/ggml-cuda/quantize.cuh | 4 + 5 files changed, 370 insertions(+), 122 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index ba5f673c73..fa2dd26202 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -362,6 +362,15 @@ static bool blackwell_mma_available(const int cc) { ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN; } +// Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes. +static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) { + GGML_ASSERT(tensor != nullptr); + return (reinterpret_cast(tensor->data) % alignment) == 0 && + tensor->nb[1] % alignment == 0 && + tensor->nb[2] % alignment == 0 && + tensor->nb[3] % alignment == 0; +} + static constexpr __device__ int ggml_cuda_get_physical_warp_size() { #if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__)) return 64; diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 0dd68ebfb6..1f0747456d 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -130,14 +130,20 @@ void ggml_cuda_mul_mat_q( const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq); ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + ggml_cuda_pool_alloc src1_scale(ctx.pool()); + if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) { + src1_scale.alloc(ne13*ne12*ne11); + } { const int64_t s11 = src1->nb[1] / ts_src1; const int64_t s12 = src1->nb[2] / ts_src1; const int64_t s13 = src1->nb[3] / ts_src1; if (use_native_fp4) { + static constexpr size_t align_float8 = 32; + const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8); static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1)); - quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, + quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); } else { @@ -155,6 +161,7 @@ void ggml_cuda_mul_mat_q( const mmq_args args = { src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d, + src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr, ne00, ne01, ne1, s01, ne11, s1, ne02, ne12, s02, s12, s2, ne03, ne13, s03, s13, s3, @@ -192,6 +199,10 @@ void ggml_cuda_mul_mat_q( const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block + ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq); ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + ggml_cuda_pool_alloc src1_scale(ctx.pool()); + if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) { + src1_scale.alloc(ne12*n_expert_used); + } const int64_t ne11_flat = ne12*n_expert_used; const int64_t ne12_flat = 1; @@ -202,18 +213,19 @@ void ggml_cuda_mul_mat_q( const int64_t s12 = src1->nb[2] / ts_src1; const int64_t s13 = src1->nb[3] / ts_src1; - if (dedup_bcast) { - // quantize each token once, scatter its block to all n_expert_used slots - if (use_native_fp4) { - quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, + if (use_native_fp4) { + static constexpr size_t align_float8 = 32; + const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8); + if (dedup_bcast) { + quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); } else { - quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, - /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); + quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, + ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); } - } else if (use_native_fp4) { - quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, - ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); + } else if (dedup_bcast) { + quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, + /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); } else { quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); @@ -229,6 +241,7 @@ void ggml_cuda_mul_mat_q( // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid. const mmq_args args = { src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d, + src1_scale.ptr, ne00, ne01, ne_get_rows, s01, ne_get_rows, s1, ne02, ne02, s02, s12, s2, ne03, ne13, s03, s13, s3, diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index ad3d6a5a68..f25835aa37 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -13,7 +13,7 @@ typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride); typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00); typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted, - float * __restrict__ dst, const int stride, const int i_max, const int j_max); + float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max); enum mmq_q8_1_ds_layout { MMQ_Q8_1_DS_LAYOUT_D4, @@ -413,11 +413,13 @@ static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config template static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a( const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst, - const int stride, const int i_max, const int j_max) { + const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + const bool y_scale_used = y_scale != nullptr; + #pragma unroll for (int j0 = 0; j0 < J; j0 += nwarps) { const int j = j0 + threadIdx.y; @@ -434,7 +436,16 @@ template static __device__ __forceinline_ continue; } - dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + if constexpr (type == GGML_TYPE_NVFP4) { + if (y_scale_used) { + dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + } else { + dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + } + } else { + dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + GGML_UNUSED(y_scale_used); + } } } } @@ -442,7 +453,8 @@ template static __device__ __forceinline_ template static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst, - const int stride, const int i_max, const int j_max) { + const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) { + #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; #else @@ -457,6 +469,8 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I); + const bool y_scale_used = y_scale != nullptr; + #pragma unroll for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { #pragma unroll @@ -475,7 +489,16 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( continue; } - dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + if constexpr (type == GGML_TYPE_NVFP4) { + if (y_scale_used) { + dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]; + } else { + dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + } + } else { + dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + GGML_UNUSED(y_scale_used); + } } } } @@ -819,6 +842,7 @@ template static __device__ __forceinline__ void mul_mat_q_process_tile( const char * __restrict__ x, const int offset_x, const int * __restrict__ y, const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, + const float * __restrict__ y_scale, const int stride_row_x, const int ncols_y, const int stride_col_dst, const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { @@ -884,9 +908,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( } if (fixup) { - write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), I, I, J); + write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J); } else { - write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j); + write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); } } @@ -898,6 +922,7 @@ __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_g static __global__ void mul_mat_q( const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup, + const float * __restrict__ y_scale, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst, const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, @@ -943,8 +968,14 @@ static __global__ void mul_mat_q( int col_low = 0; int col_high = ncols_dst; int col_diff = ncols_dst; - int offset_y = wt*stride_sample_y + zt*stride_channel_y; - int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; + } else { + GGML_UNUSED(offset_y_scale); + } if (ids_dst) { col_low = expert_bounds[zt + 0]; @@ -953,6 +984,9 @@ static __global__ void mul_mat_q( offset_y = 0; offset_dst = 0; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = 0; + } if (jt*J >= col_diff) { return; @@ -974,6 +1008,11 @@ static __global__ void mul_mat_q( offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int)); offset_dst += it*I; + const float * y_scale_tile = nullptr; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale += col_low + jt*J; + y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; + } const int tile_x_max_i = nrows_x - it*I - 1; const int tile_y_max_j = col_diff - jt*J - 1; @@ -982,7 +1021,8 @@ static __global__ void mul_mat_q( constexpr bool fixup = false; mul_mat_q_process_tile - (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z); return; } @@ -1016,8 +1056,14 @@ static __global__ void mul_mat_q( int col_low = 0; int col_high = ncols_dst; int col_diff = ncols_dst; - int offset_y = wt*stride_sample_y + zt*stride_channel_y; - int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; + } else { + GGML_UNUSED(offset_y_scale); + } if (ids_dst) { col_low = expert_bounds[zt + 0]; @@ -1026,6 +1072,9 @@ static __global__ void mul_mat_q( offset_y = 0; offset_dst = 0; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = 0; + } if (jt*J >= col_diff) { kbc += blocks_per_ne00.z; @@ -1053,6 +1102,11 @@ static __global__ void mul_mat_q( offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int)); offset_dst += it*I; + const float * y_scale_tile = nullptr; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale += col_low + jt * J; + y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; + } const int tile_x_max_i = nrows_x - it*I - 1; const int tile_y_max_j = col_diff - jt*J - 1; @@ -1061,7 +1115,8 @@ static __global__ void mul_mat_q( constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. mul_mat_q_process_tile - (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); kbc += blocks_per_ne00.z; @@ -1090,8 +1145,14 @@ static __global__ void mul_mat_q( int col_low = 0; int col_high = ncols_dst; int col_diff = ncols_dst; - int offset_y = wt*stride_sample_y + zt*stride_channel_y; - int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; + } else { + GGML_UNUSED(offset_y_scale); + } if (ids_dst) { col_low = expert_bounds[zt + 0]; @@ -1100,6 +1161,9 @@ static __global__ void mul_mat_q( offset_y = 0; offset_dst = 0; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = 0; + } if (jt*J >= col_diff) { return; @@ -1122,6 +1186,11 @@ static __global__ void mul_mat_q( offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int)); offset_dst += it*I; + const float * y_scale_tile = nullptr; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale += col_low + jt * J; + y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr; + } const int tile_x_max_i = nrows_x - it*I - 1; const int tile_y_max_j = col_diff - jt*J - 1; @@ -1130,7 +1199,8 @@ static __global__ void mul_mat_q( constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. mul_mat_q_process_tile - (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); } @@ -1274,6 +1344,7 @@ static __global__ void mul_mat_q_stream_k_fixup( struct mmq_args { const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst; + const float * y_scale; int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst; int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst; int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst; @@ -1323,7 +1394,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) { mul_mat_q<<>> - (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, @@ -1352,7 +1423,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z); mul_mat_q<<>> - (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, + (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index af7fe105b6..2bd9b62623 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu @@ -1,6 +1,55 @@ #include "quantize.cuh" #include +#if defined(BLACKWELL_MMA_AVAILABLE) +// this maps to 256-bit loads in PTX on supported devices, +// and otherwise falls back to 2 128-bit loads +struct __builtin_align__(32) float8 { + float x; float y; float z; float w; + float p; float q; float r; float s; +}; +#endif + +#if CUDART_VERSION >= 12080 +static __device__ __forceinline__ float nvfp4_native_scale_error( + const float vals[QK_NVFP4_SUB], const float inv_col_scale, const float inv_scale, const float scale) { + const float scale_dequant = 2.0f * scale; + float err = 0.0f; + +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; k += 4) { + const float v0 = vals[k + 0] * inv_col_scale; + const float v1 = vals[k + 1] * inv_col_scale; + const float v2 = vals[k + 2] * inv_col_scale; + const float v3 = vals[k + 3] * inv_col_scale; + + const __nv_fp4x4_e2m1 q(make_float4(v0 * inv_scale, v1 * inv_scale, v2 * inv_scale, v3 * inv_scale)); + const __nv_fp4x4_storage_t q_storage = q.__x; + const __nv_fp4x2_storage_t q_lo = static_cast<__nv_fp4x2_storage_t>(q_storage); + const __nv_fp4x2_storage_t q_hi = static_cast<__nv_fp4x2_storage_t>(q_storage >> 8U); + + const __half2_raw hraw2_lo = __nv_cvt_fp4x2_to_halfraw2(q_lo, __NV_E2M1); + const __half2_raw hraw2_hi = __nv_cvt_fp4x2_to_halfraw2(q_hi, __NV_E2M1); + const __half2 h2_lo = static_cast<__half2>(hraw2_lo); + const __half2 h2_hi = static_cast<__half2>(hraw2_hi); + const float2 dq_lo = __half22float2(h2_lo); + const float2 dq_hi = __half22float2(h2_hi); + + const float err0 = fabsf(v0) - fabsf(dq_lo.x) * scale_dequant; + const float err1 = fabsf(v1) - fabsf(dq_lo.y) * scale_dequant; + const float err2 = fabsf(v2) - fabsf(dq_hi.x) * scale_dequant; + const float err3 = fabsf(v3) - fabsf(dq_hi.y) * scale_dequant; + + err = fmaf(err0, err0, err); + err = fmaf(err1, err1, err); + err = fmaf(err2, err2, err); + err = fmaf(err3, err3, err); + } + + return err; +} +#endif // CUDART_VERSION >= 12080 + __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1) static __global__ void quantize_q8_1( const float * x_ptr, void * vy_ptr, @@ -74,115 +123,209 @@ __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) { return static_cast(biased); } - // scatter: grid over tokens, quantize once, write to all the token's compact rows -template +template static __global__ void quantize_mmq_nvfp4( - const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, + const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, float * __restrict__ scale, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) { #if defined(BLACKWELL_MMA_AVAILABLE) - const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB; - if (i0_base >= ne0) { - return; - } - - const int64_t k_block = i0_base / QK_FP4_MMQ; const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ; - if (k_block >= blocks_per_col) { - return; - } - const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB; int64_t base_idx; if constexpr (scatter) { base_idx = (int64_t) blockIdx.x * s02; // one physical row per token } else { - const int64_t i2 = blockIdx.z % ne2; - const int64_t i3 = blockIdx.z / ne2; + const int64_t i2 = blockIdx.y % ne2; + const int64_t i3 = blockIdx.y / ne2; const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x; base_idx = i3 * s03 + i2 * s02 + i01 * s01; } + const float * __restrict__ x_row = x + base_idx; - float vals_raw[QK_NVFP4_SUB]; - float amax_raw = 0.0f; + float amax = 0.0f; + if constexpr (use_aligned_float8) { + for (int64_t i0 = 8 * threadIdx.x; i0 < ne00; i0 += 8 * blockDim.x) { + const float * x_base = x_row + i0; + const float8 v = reinterpret_cast(x_base)[0]; + amax = fmaxf(amax, fabsf(v.x)); + amax = fmaxf(amax, fabsf(v.y)); + amax = fmaxf(amax, fabsf(v.z)); + amax = fmaxf(amax, fabsf(v.w)); + amax = fmaxf(amax, fabsf(v.p)); + amax = fmaxf(amax, fabsf(v.q)); + amax = fmaxf(amax, fabsf(v.r)); + amax = fmaxf(amax, fabsf(v.s)); + } + } else { + for (int64_t i0 = threadIdx.x; i0 < ne00; i0 += blockDim.x) { + amax = fmaxf(amax, fabsf(x_row[i0])); + } + } + + amax = warp_reduce_max(amax); + + __shared__ float warp_amax[CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE]; + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + + if (lane == 0) { + warp_amax[warp] = amax; + } + __syncthreads(); + + if (warp == 0) { + amax = threadIdx.x < int(CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE) ? warp_amax[lane] : 0.0f; + amax = warp_reduce_max(amax); + if (lane == 0) { + warp_amax[0] = amax / (6.0f * 448.0f); + if constexpr (scatter) { #pragma unroll - for (int k = 0; k < QK_NVFP4_SUB; k++) { - const int64_t i00 = i0_base + k; - if (i00 < ne00) { - const float v = x[base_idx + i00]; - vals_raw[k] = v; - amax_raw = fmaxf(amax_raw, fabsf(v)); - } else { - vals_raw[k] = 0.0f; + for (int slot = 0; slot < n_expert_used; ++slot) { + const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; + scale[i] = warp_amax[0]; + } + } else { + scale[blockIdx.y * ne1 + blockIdx.x] = warp_amax[0]; + } } } - - static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2}; - const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_raw / 6.0f); - - float best_err = FLT_MAX; - uint8_t fp8_code = 0; - float subblock_scale = 0.0f; - -#pragma unroll // Check +/- 2 to find best code to reduce NVFP4 activation loss. Negligible overhead on Blackwell. - for (int i = 0; i < 5; i++) { - const int test_code = first_fp8_code + test_offsets[i]; - if (test_code < 0 || test_code > 0x7e) { - continue; - } - const uint8_t code = (uint8_t) test_code; - const float test_scale = ggml_cuda_ue4m3_to_fp32(code); - const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f; - float cur_err = 0.0f; -#pragma unroll - for (int k = 0; k < QK_NVFP4_SUB; ++k) { - const float v = vals_raw[k]; - const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale); - const float err_diff = fabsf(v) - fabsf(kvalues_mxfp4[q & 0x7]) * test_scale; - cur_err = fmaf(err_diff, err_diff, cur_err); - } - - if (cur_err < best_err) { - best_err = cur_err; - fp8_code = test_code; - subblock_scale = test_scale; - } - } - - const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f; - uint32_t q0 = 0; - uint32_t q1 = 0; -#pragma unroll // this is faster than the previous __nv_fp4x4_e2m1 - for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) { - q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 0], inv_scale) << (8 * k); - q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 8], inv_scale) << (8 * k + 4); - q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 4], inv_scale) << (8 * k); - q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4); - } + __syncthreads(); block_fp4_mmq * y = (block_fp4_mmq *) vy; - if constexpr (scatter) { + const int64_t n_subblocks = (ne0 + QK_NVFP4_SUB - 1) / QK_NVFP4_SUB; + + for (int64_t isb = threadIdx.x; isb < n_subblocks; isb += blockDim.x) { + const int64_t i0_base = isb * QK_NVFP4_SUB; + const int64_t k_block = i0_base / QK_FP4_MMQ; + const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB; + + const float row_scale = warp_amax[0]; + const float inv_col_scale = row_scale > 0.0f ? 1.0f / row_scale : 0.0f; + + float vals[QK_NVFP4_SUB]; + if constexpr (use_aligned_float8) { + const float * x_base = x_row + i0_base; + const float8 v0 = i0_base + 7 < ne00 ? reinterpret_cast(x_base)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + const float8 v1 = i0_base + 15 < ne00 ? reinterpret_cast(x_base + 8)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + vals[0] = v0.x; vals[1] = v0.y; vals[2] = v0.z; vals[3] = v0.w; + vals[4] = v0.p; vals[5] = v0.q; vals[6] = v0.r; vals[7] = v0.s; + vals[8] = v1.x; vals[9] = v1.y; vals[10] = v1.z; vals[11] = v1.w; + vals[12] = v1.p; vals[13] = v1.q; vals[14] = v1.r; vals[15] = v1.s; + } else { #pragma unroll - for (int slot = 0; slot < n_expert_used; ++slot) { - const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; - block_fp4_mmq * yb = y + (k_block * ne1 + i); + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + const int64_t i00 = i0_base + k; + vals[k] = i00 < ne00 ? x_row[i00] : 0.0f; + } + } + + uint32_t q0 = 0; + uint32_t q1 = 0; + + float amax_sub = 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + amax_sub = fmaxf(amax_sub, fabsf(vals[k] * inv_col_scale)); + } + + static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2 }; + const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_sub / 6.0f); + + uint8_t fp8_code = (uint8_t) first_fp8_code; + float subblock_scale = ggml_cuda_ue4m3_to_fp32(fp8_code); + float inv_scale_err = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f; +#if CUDART_VERSION >= 12080 + float best_err = nvfp4_native_scale_error(vals, inv_col_scale, inv_scale_err, subblock_scale); +#else + float best_err = 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + const float v = vals[k] * inv_col_scale; + const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, inv_scale_err); + const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * subblock_scale; + best_err = fmaf(err_diff, err_diff, best_err); + } +#endif // CUDART_VERSION >= 12080 + +#pragma unroll + for (int i = 1; i < 5; ++i) { + const int test_code = first_fp8_code + test_offsets[i]; + if (test_code < 0 || test_code > 0x7e) { + continue; + } + + const float test_scale = ggml_cuda_ue4m3_to_fp32((uint8_t) test_code); + const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f; +#if CUDART_VERSION >= 12080 + const float cur_err = nvfp4_native_scale_error(vals, inv_col_scale, test_inv_scale, test_scale); +#else + float cur_err = 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB; ++k) { + const float v = vals[k] * inv_col_scale; + const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale); + const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * test_scale; + cur_err = fmaf(err_diff, err_diff, cur_err); + } +#endif // CUDART_VERSION >= 12080 + + if (cur_err < best_err) { + best_err = cur_err; + fp8_code = (uint8_t) test_code; + subblock_scale = test_scale; + } + } +#if CUDART_VERSION >= 12080 + const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f; + const float s = inv_col_scale * inv_scale; + + __nv_fp4x4_e2m1 q0_lo(make_float4(vals[0] * s, vals[8] * s, vals[1] * s, vals[9] * s)); + __nv_fp4x4_e2m1 q0_hi(make_float4(vals[2] * s, vals[10] * s, vals[3] * s, vals[11] * s)); + __nv_fp4x4_e2m1 q1_lo(make_float4(vals[4] * s, vals[12] * s, vals[5] * s, vals[13] * s)); + __nv_fp4x4_e2m1 q1_hi(make_float4(vals[6] * s, vals[14] * s, vals[7] * s, vals[15] * s)); + + const char2 q0_lo_c = *reinterpret_cast(&q0_lo); + const char2 q0_hi_c = *reinterpret_cast(&q0_hi); + const char2 q1_lo_c = *reinterpret_cast(&q1_lo); + const char2 q1_hi_c = *reinterpret_cast(&q1_hi); + + q0 = uint32_t(uint8_t(q0_lo_c.x)) | (uint32_t(uint8_t(q0_lo_c.y)) << 8) | + (uint32_t(uint8_t(q0_hi_c.x)) << 16) | (uint32_t(uint8_t(q0_hi_c.y)) << 24); + q1 = uint32_t(uint8_t(q1_lo_c.x)) | (uint32_t(uint8_t(q1_lo_c.y)) << 8) | + (uint32_t(uint8_t(q1_hi_c.x)) << 16) | (uint32_t(uint8_t(q1_hi_c.y)) << 24); +#else + const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f; +#pragma unroll + for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) { + q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 0] * inv_col_scale, inv_scale)) << (8 * k); + q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 8] * inv_col_scale, inv_scale)) << (8 * k + 4); + q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 4] * inv_col_scale, inv_scale)) << (8 * k); + q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 12] * inv_col_scale, inv_scale)) << (8 * k + 4); + } +#endif // CUDART_VERSION >= 12080 + + if constexpr (scatter) { +#pragma unroll + for (int slot = 0; slot < n_expert_used; ++slot) { + const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; + block_fp4_mmq * yb = y + (k_block * ne1 + i); + uint32_t * yqs = reinterpret_cast(yb->qs); + yqs[2 * sub + 0] = q0; + yqs[2 * sub + 1] = q1; + reinterpret_cast(yb->d4)[sub] = fp8_code; + } + } else { + block_fp4_mmq * yb = y + (blockIdx.y * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x); uint32_t * yqs = reinterpret_cast(yb->qs); yqs[2 * sub + 0] = q0; yqs[2 * sub + 1] = q1; reinterpret_cast(yb->d4)[sub] = fp8_code; } - } else { - block_fp4_mmq * yb = y + (blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x); - uint32_t * yqs = reinterpret_cast(yb->qs); - yqs[2 * sub + 0] = q0; - yqs[2 * sub + 1] = q1; - reinterpret_cast(yb->d4)[sub] = fp8_code; } - GGML_UNUSED(n_expert_used); #else - GGML_UNUSED(n_expert_used); + GGML_UNUSED_VARS(x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, n_expert_used); NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only. #endif // defined(BLACKWELL_MMA_AVAILABLE) @@ -491,18 +634,22 @@ void quantize_scatter_mmq_q8_1_cuda( // scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row) void quantize_scatter_mmq_fp4_cuda( - const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0, + const float * x, const int32_t * ids_src1_inv, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8, const int64_t ne00, const int64_t stride_token, const int64_t ne0, const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) { GGML_ASSERT(ne0 > 0); if (type_src0 == GGML_TYPE_NVFP4) { + GGML_ASSERT(scale); GGML_ASSERT(ne00 % QK_NVFP4 == 0); - constexpr int nvfp4_block_size = 128; - const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size); - const dim3 block_size(nvfp4_block_size, 1, 1); - const dim3 num_blocks(n_tokens, block_num_y, 1); - quantize_mmq_nvfp4<<>>( - x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + const dim3 num_blocks(n_tokens, 1, 1); + if (use_aligned_float8) { + quantize_mmq_nvfp4<<>>( + x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used); + } else { + quantize_mmq_nvfp4<<>>( + x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used); + } } else { GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4); constexpr int nwarps = 8; @@ -516,20 +663,24 @@ void quantize_scatter_mmq_fp4_cuda( } void quantize_mmq_fp4_cuda( - const float * x, const int32_t * ids, void * vy, const ggml_type type_src0, + const float * x, const int32_t * ids, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4); GGML_ASSERT(ne0 > 0); if (type_src0 == GGML_TYPE_NVFP4) { + GGML_ASSERT(scale); GGML_ASSERT(ne00 % QK_NVFP4 == 0); - constexpr int nvfp4_block_size = 128; - const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size); - const dim3 block_size(nvfp4_block_size, 1, 1); - const dim3 num_blocks(ne1, block_num_y, ne2 * ne3); - quantize_mmq_nvfp4<<>>( - x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + const dim3 num_blocks(ne1, ne2 * ne3, 1); + if (use_aligned_float8) { + quantize_mmq_nvfp4<<>>( + x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + } else { + quantize_mmq_nvfp4<<>>( + x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0); + } } else { GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0); diff --git a/ggml/src/ggml-cuda/quantize.cuh b/ggml/src/ggml-cuda/quantize.cuh index ed47fa027e..5f08dcbfe3 100644 --- a/ggml/src/ggml-cuda/quantize.cuh +++ b/ggml/src/ggml-cuda/quantize.cuh @@ -29,7 +29,9 @@ void quantize_mmq_q8_1_cuda( void quantize_mmq_fp4_cuda(const float * x, const int32_t * ids, void * vy, + float * scale, ggml_type type_src0, + bool use_aligned_float8, int64_t ne00, int64_t s01, int64_t s02, @@ -44,7 +46,9 @@ void quantize_mmq_fp4_cuda(const float * x, void quantize_scatter_mmq_fp4_cuda(const float * x, const int32_t * ids_src1_inv, void * vy, + float * scale, ggml_type type_src0, + bool use_aligned_float8, int64_t ne00, int64_t stride_token, int64_t ne0,