upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-23 11:10:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

CUDA: Improve NVFP4 W4A4 activation quantization (#25730)

Browse Source 
* 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
This commit is contained in:
Oliver Simons
2026-07-22 19:28:02 +02:00
committed by GitHub
parent 0278d8362d
commit 1a064ab092
5 changed files with 370 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
ggml/src/ggml-cuda/common.cuh
View File

@@ -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<uintptr_t>(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;

33
ggml/src/ggml-cuda/mmq.cu
View File

@@ -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<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
ggml_cuda_pool_alloc<float> 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<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
ggml_cuda_pool_alloc<float> 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,

107
ggml/src/ggml-cuda/mmq.cuh
View File

@@ -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 <ggml_type type, int J, bool fallback> 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 <ggml_type type, int J, bool fallback> 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 <ggml_type type, int J, bool fallback> static __device__ __forceinline_
template<ggml_type type, int J, bool fallback>
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 <ggml_type type, int J, bool fallback, bool fixup>
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<type, J, fallback, fixup>
(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<type, J, fallback, fixup>
(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<type, J, fallback, fixup>
(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<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
(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<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
(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,

339
ggml/src/ggml-cuda/quantize.cu
View File

@@ -1,6 +1,55 @@
#include "quantize.cuh"
#include <cstdint>
#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<uint8_t>(biased);
}
// scatter: grid over tokens, quantize once, write to all the token's compact rows
template <bool scatter>
template <bool scatter, bool use_aligned_float8>
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<const float8 *>(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<WARP_SIZE>(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<WARP_SIZE>(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<const float8 *>(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<const float8 *>(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<char2 *>(&q0_lo);
const char2 q0_hi_c = *reinterpret_cast<char2 *>(&q0_hi);
const char2 q1_lo_c = *reinterpret_cast<char2 *>(&q1_lo);
const char2 q1_hi_c = *reinterpret_cast<char2 *>(&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<uint32_t *>(yb->qs);
yqs[2 * sub + 0] = q0;
yqs[2 * sub + 1] = q1;
reinterpret_cast<uint8_t *>(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<uint32_t *>(yb->qs);
yqs[2 * sub + 0] = q0;
yqs[2 * sub + 1] = q1;
reinterpret_cast<uint8_t *>(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<uint32_t *>(yb->qs);
yqs[2 * sub + 0] = q0;
yqs[2 * sub + 1] = q1;
reinterpret_cast<uint8_t *>(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<true><<<num_blocks, block_size, 0, stream>>>(
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<true, true><<<num_blocks, block_size, 0, stream>>>(
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<true, false><<<num_blocks, block_size, 0, stream>>>(
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<false><<<num_blocks, block_size, 0, stream>>>(
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<false, true><<<num_blocks, block_size, 0, stream>>>(
x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
} else {
quantize_mmq_nvfp4<false, false><<<num_blocks, block_size, 0, stream>>>(
x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
}
} else {
GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);

4
ggml/src/ggml-cuda/quantize.cuh
View File

@@ -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,