Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests 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
ggml/src/ggml-cuda/mmq.cuh
+89 -18
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,