Commit fc9ce6b9d for llama.cpp
commit fc9ce6b9d52a8504edcb262abc92737c2289f96c
Author: Johannes Gäßler <johannesg@5d6.de>
Date: Thu Oct 8 16:41:49 2026 +0200
CUDA: fix MMQ out-of-bounds reads (#29953)
diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu
index 43603bb8a..027a7e60d 100644
--- a/ggml/src/ggml-cuda/mmq.cu
+++ b/ggml/src/ggml-cuda/mmq.cu
@@ -141,7 +141,10 @@ void ggml_cuda_mul_mat_q(
GGML_TENSOR_BINARY_OP_LOCALS;
cudaStream_t stream = ctx.stream();
- const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+
+ const int id = ggml_cuda_get_device();
+ const int cc = ggml_cuda_info().devices[id].cc;
+ const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
const size_t ts_src0 = ggml_type_size(src0->type);
const size_t ts_src1 = ggml_type_size(src1->type);
@@ -176,7 +179,7 @@ void ggml_cuda_mul_mat_q(
const int64_t s03 = src0->nb[3] / ts_src0;
const int64_t s3 = dst->nb[3] / ts_dst;
- const bool fallback = ne01 % 128 != 0;
+ const bool fallback = ggml_cuda_mmq_needs_fallback(ne01);
const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
@@ -184,9 +187,52 @@ void ggml_cuda_mul_mat_q(
const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ;
+ int J_best = 0;
+ int nthreads_best = 0;
+ {
+ int64_t ncols_opt = ne11;
+ if (ids) {
+ const int64_t n_expert_used = ids->ne[0];
+ ncols_opt = ne12;
+
+ // Each expert only sees ne12*n_expert_used/ne02 tokens on average.
+ // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
+ if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
+ ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
+ }
+ }
+
+ int ntiles_J_best = INT_MAX;
+
+ for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
+ const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(src0->type, J, fallback, cc, prec_src1);
+ if (config.type == GGML_TYPE_COUNT) {
+ continue;
+ }
+
+ if (mmq_get_nbytes_shared(config, cc) > smpbo) {
+ continue;
+ }
+
+ const int ntiles_x = (ncols_opt + config.J - 1) / config.J;
+
+ if (ntiles_x < ntiles_J_best) {
+ J_best = J;
+ nthreads_best = config.nthreads;
+ ntiles_J_best = ntiles_x;
+ }
+ }
+ }
+ GGML_ASSERT(J_best > 0);
+
+ // A tile of size J can read in at most J - 1 extra columns.
+ // For simplicity, round up the padding of a full tile to a multiple of the number of bytes that nthreads can load in parallel.
+ const size_t src1_load_chunk_size = nthreads_best * sizeof(int);
+ const size_t src1_q8_1_padding = ((J_best * sizeof(block_q8_1_mmq) + src1_load_chunk_size - 1) / src1_load_chunk_size)
+ * src1_load_chunk_size;
+
if (!ids) {
- 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);
+ const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding;
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) {
@@ -223,7 +269,7 @@ void ggml_cuda_mul_mat_q(
ne00, ne01, ne1, s01, ne11, s1,
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
- ne1, ne1};
+ ne1, J_best};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
return;
}
@@ -237,7 +283,7 @@ void ggml_cuda_mul_mat_q(
GGML_ASSERT(ne1 == n_expert_used);
ggml_cuda_pool_alloc<int32_t> ids_src1(ctx.pool(), ne_get_rows);
- ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows);
+ ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows + J_best-1); // Needs to be padded for unconditional memory access.
ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
// gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
@@ -254,8 +300,7 @@ void ggml_cuda_mul_mat_q(
CUDA_CHECK(cudaGetLastError());
}
- 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, ne12) * sizeof(block_q8_1_mmq);
+ const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding;
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) {
@@ -296,13 +341,6 @@ void ggml_cuda_mul_mat_q(
ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
const int64_t s13 = ne12*s12;
- // Each expert only sees ne12*n_expert_used/ne02 tokens on average.
- // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
- int64_t ncols_opt = ne12;
- if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
- ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
- }
-
// 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,
@@ -310,7 +348,7 @@ void ggml_cuda_mul_mat_q(
ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
ne02, ne02, s02, s12, s2,
ne03, ne13, s03, s13, s3,
- ne12, ncols_opt};
+ ne12, J_best};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
}
diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh
index 0cd31d916..198aab4b5 100644
--- a/ggml/src/ggml-cuda/mmq.cuh
+++ b/ggml/src/ggml-cuda/mmq.cuh
@@ -208,7 +208,7 @@ struct ggml_cuda_mmq_config {
static_assert((nthreads_) % 32 == 0 && (nthreads_) <= 512, "bad nthreads"); \
static_assert( (occupancy_) <= 8, "bad occupancy"); \
static_assert((I_) % 32 == 0, "bad I"); \
- static_assert((J_) % 8 == 0, "bad J"); \
+ static_assert((J_) % 8 == 0 && (J_) <= 128, "bad J"); \
static_assert((K_vram_) % 256 == 0, "bad K_vram"); \
return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \
} \
@@ -295,6 +295,8 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
GGML_UNUSED_VARS(type, J, fallback, prec_src1);
}
+// FIXME all of the host functions are missing prec_src1, this can lead to inconsitent behavior.
+
static __host__ int ggml_cuda_mmq_get_type(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).type;
}
@@ -369,15 +371,8 @@ static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, in
return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, prec_src1));
}
-static __host__ int ggml_cuda_mmq_get_J_max(const ggml_type type, const bool fallback, const int cc, const int64_t ne11) {
- int ret = std::min(ne11, int64_t(512));
- ret -= ret % 8;
- for (;ret > 0; ret -= 8) {
- if (ggml_cuda_mmq_get_config(type, ret, fallback, cc).type != GGML_TYPE_COUNT) {
- return ret;
- }
- }
- return ret;
+static __host__ bool ggml_cuda_mmq_needs_fallback(const int64_t nrows_x) {
+ return nrows_x % 128 != 0;
}
static constexpr __device__ int ggml_cuda_mmq_get_rows_per_warp(ggml_type type, int J, bool fallback) {
@@ -1390,7 +1385,7 @@ struct mmq_args {
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;
int64_t ncols_max;
- int64_t ncols_opt; // value to optimize the tile size against, launch grid still uses ncols_max
+ int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance.
};
static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
@@ -1484,32 +1479,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
template <ggml_type type, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
- const int id = ggml_cuda_get_device();
- const int cc = ggml_cuda_info().devices[id].cc;
- const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
-
- int J_best = 0;
- int ntiles_J_best = INT_MAX;
-
- for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
- const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1);
- if (config.type == GGML_TYPE_COUNT) {
- continue;
- }
-
- if (mmq_get_nbytes_shared(config, cc) > smpbo) {
- continue;
- }
-
- const int ntiles_x = (args.ncols_opt + config.J - 1) / config.J;
-
- if (ntiles_x < ntiles_J_best) {
- J_best = J;
- ntiles_J_best = ntiles_x;
- }
- }
-
- switch (J_best) {
+ switch (args.J_best) {
case 8:
launch_mul_mat_q<type, 8, fallback, prec_src1>(ctx, args, stream);
break;
@@ -1559,7 +1529,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
launch_mul_mat_q<type, 128, fallback, prec_src1>(ctx, args, stream);
break;
default:
- fprintf(stderr, "J_best=%d\n", J_best);
+ fprintf(stderr, "J_best=%d\n", args.J_best);
GGML_ABORT("fatal error");
break;
}
@@ -1567,11 +1537,11 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
template <ggml_type type, ggml_prec prec_src1 = GGML_PREC_Q8>
void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
- if (args.nrows_x % 128 == 0) {
- constexpr bool fallback = false;
+ if (ggml_cuda_mmq_needs_fallback(args.nrows_x)) {
+ constexpr bool fallback = true;
mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
} else {
- constexpr bool fallback = true;
+ constexpr bool fallback = false;
mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
}
}