A high-throughput and memory-efficient inference and serving engine for LLMs
[Security] Fix info disclosure via int32 truncation in GGUF dequantize kernels (#44971)
Signed-off-by: jperezde <jperezde@redhat.com>
Juan Pérez de Algaba committed
Jun 11, 2026 at 10:05 UTC
f219788f91952827132fa4fdf916427cd20d225e
3 files changed
+33
-29
csrc/libtorch_stable/quantization/gguf/dequantize.cuh
+18
-18
@@ -78,8 +78,8 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in
78
}
79
80
template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
81
-static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y, const int k) {
82
- const int i = 2*(blockDim.x*blockIdx.x + threadIdx.x);
81
+static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k) {
82
+ const int64_t i = 2*((int64_t)blockDim.x*blockIdx.x + threadIdx.x);
83
84
if (i >= k) {
85
return;
@@ -435,91 +435,91 @@ static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst
435
}
436
437
template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
438
-static void dequantize_block_cuda(const void * __restrict__ vx, dst_t * __restrict__ y, const int k, cudaStream_t stream) {
439
- const int num_blocks = (k + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2*CUDA_DEQUANTIZE_BLOCK_SIZE);
438
+static void dequantize_block_cuda(const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k, cudaStream_t stream) {
439
+ const int64_t num_blocks = (k + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2*CUDA_DEQUANTIZE_BLOCK_SIZE);
440
dequantize_block<qk, qr, dequantize_kernel><<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(vx, y, k);
441
}
442
443
template<typename dst_t>
444
-static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
444
+static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
445
const int nb = k / QK_K;
446
dequantize_block_q2_K<<<nb, 64, 0, stream>>>(vx, y);
447
}
448
449
template<typename dst_t>
450
-static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
450
+static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
451
const int nb = k / QK_K;
452
dequantize_block_q3_K<<<nb, 64, 0, stream>>>(vx, y);
453
}
454
455
template<typename dst_t>
456
-static void dequantize_row_q4_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
456
+static void dequantize_row_q4_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
457
const int nb = k / QK_K;
458
dequantize_block_q4_K<<<nb, 32, 0, stream>>>(vx, y);
459
}
460
461
template<typename dst_t>
462
-static void dequantize_row_q5_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
462
+static void dequantize_row_q5_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
463
const int nb = k / QK_K;
464
dequantize_block_q5_K<<<nb, 64, 0, stream>>>(vx, y);
465
}
466
467
template<typename dst_t>
468
-static void dequantize_row_q6_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
468
+static void dequantize_row_q6_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
469
const int nb = k / QK_K;
470
dequantize_block_q6_K<<<nb, 64, 0, stream>>>(vx, y);
471
}
472
473
template<typename dst_t>
474
-static void dequantize_row_iq2_xxs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
474
+static void dequantize_row_iq2_xxs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
475
const int nb = k / QK_K;
476
dequantize_block_iq2_xxs<<<nb, 32, 0, stream>>>(vx, y);
477
}
478
479
template<typename dst_t>
480
-static void dequantize_row_iq2_xs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
480
+static void dequantize_row_iq2_xs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
481
const int nb = k / QK_K;
482
dequantize_block_iq2_xs<<<nb, 32, 0, stream>>>(vx, y);
483
}
484
485
template<typename dst_t>
486
-static void dequantize_row_iq2_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
486
+static void dequantize_row_iq2_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
487
const int nb = k / QK_K;
488
dequantize_block_iq2_s<<<nb, 32, 0, stream>>>(vx, y);
489
}
490
491
template<typename dst_t>
492
-static void dequantize_row_iq3_xxs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
492
+static void dequantize_row_iq3_xxs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
493
const int nb = k / QK_K;
494
dequantize_block_iq3_xxs<<<nb, 32, 0, stream>>>(vx, y);
495
}
496
497
template<typename dst_t>
498
-static void dequantize_row_iq3_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
498
+static void dequantize_row_iq3_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
499
const int nb = k / QK_K;
500
dequantize_block_iq3_s<<<nb, 32, 0, stream>>>(vx, y);
501
}
502
503
template<typename dst_t>
504
-static void dequantize_row_iq1_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
504
+static void dequantize_row_iq1_s_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
505
const int nb = k / QK_K;
506
dequantize_block_iq1_s<<<nb, 32, 0, stream>>>(vx, y);
507
}
508
509
template<typename dst_t>
510
-static void dequantize_row_iq1_m_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
510
+static void dequantize_row_iq1_m_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
511
const int nb = k / QK_K;
512
dequantize_block_iq1_m<<<nb, 32, 0, stream>>>(vx, y);
513
}
514
515
template<typename dst_t>
516
-static void dequantize_row_iq4_nl_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
516
+static void dequantize_row_iq4_nl_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
517
const int nb = (k + QK_K - 1) / QK_K;
518
dequantize_block_iq4_nl<<<nb, 32, 0, stream>>>(vx, y);
519
}
520
521
template<typename dst_t>
522
-static void dequantize_row_iq4_xs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
522
+static void dequantize_row_iq4_xs_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
523
const int nb = (k + QK_K - 1) / QK_K;
524
dequantize_block_iq4_xs<<<nb, 32, 0, stream>>>(vx, y);
525
}
csrc/libtorch_stable/quantization/gguf/ggml-common.h
+1
-1
@@ -1064,7 +1064,7 @@ typedef half dfloat; // dequantize float
1064
typedef half2 dfloat2;
1065
typedef void (*dequantize_kernel_t)(const void * vx, const int ib, const int iqs, dfloat2 & v);
1066
template<typename dst_t>
1067
-using to_cuda_ggml_t = void (*)(const void * __restrict__ x, dst_t * __restrict__ y, int k, cudaStream_t stream);
1067
+using to_cuda_ggml_t = void (*)(const void * __restrict__ x, dst_t * __restrict__ y, int64_t k, cudaStream_t stream);
1068
typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs);
1069
typedef void (*allocate_tiles_cuda_t)(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc);
1070
typedef void (*load_tiles_cuda_t)(
csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu
+14
-10
@@ -79,6 +79,7 @@ torch::stable::Tensor ggml_dequantize(
79
W.get_device_index());
80
auto dtype_ = dtype.value_or(torch::headeronly::ScalarType::Half);
81
auto DW = torch::stable::empty({m, n}, dtype_, std::nullopt, W.device());
82
+ torch::stable::fill_(DW, 0.0);
83
cudaStream_t stream = get_current_cuda_stream();
84
85
VLLM_STABLE_DISPATCH_FLOATING_TYPES(DW.scalar_type(), "ggml_dequantize", [&] {
@@ -93,13 +94,14 @@ torch::stable::Tensor ggml_mul_mat_vec_a8(
94
torch::stable::Tensor W, // quant weight
95
torch::stable::Tensor X, // input
96
int64_t type, int64_t row) {
96
- int col = X.sizes()[1];
97
- int vecs = X.sizes()[0];
98
- const int padded = (col + 512 - 1) / 512 * 512;
97
+ int64_t col = X.sizes()[1];
98
+ int64_t vecs = X.sizes()[0];
99
+ const int64_t padded = (col + 512 - 1) / 512 * 512;
100
const torch::stable::accelerator::DeviceGuard device_guard(
101
X.get_device_index());
102
auto Y = torch::stable::empty({vecs, row}, X.scalar_type(), std::nullopt,
103
W.device());
104
+ torch::stable::fill_(Y, 0.0);
105
cudaStream_t stream = get_current_cuda_stream();
106
auto quant_X = torch::stable::empty({vecs, padded / 32 * 9},
107
torch::headeronly::ScalarType::Int,
@@ -213,13 +215,14 @@ torch::stable::Tensor ggml_mul_mat_vec_a8(
215
torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W, // quant weight
216
torch::stable::Tensor X, // input
217
int64_t type, int64_t row) {
216
- int col = X.sizes()[1];
217
- int padded = (col + 512 - 1) / 512 * 512;
218
- int batch = X.sizes()[0];
218
+ int64_t col = X.sizes()[1];
219
+ int64_t padded = (col + 512 - 1) / 512 * 512;
220
+ int64_t batch = X.sizes()[0];
221
const torch::stable::accelerator::DeviceGuard device_guard(
222
X.get_device_index());
223
auto Y = torch::stable::empty({batch, row}, X.scalar_type(), std::nullopt,
224
W.device());
225
+ torch::stable::fill_(Y, 0.0);
226
cudaStream_t stream = get_current_cuda_stream();
227
auto quant_X = torch::stable::empty({batch, padded / 32 * 9},
228
torch::headeronly::ScalarType::Int,
@@ -291,12 +294,13 @@ torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X, // input
294
torch::stable::Tensor num_tokens_post_padded,
295
int64_t type, int64_t row, int64_t top_k,
296
int64_t tokens) {
294
- int col = X.sizes()[1];
295
- int padded = (col + 512 - 1) / 512 * 512;
297
+ int64_t col = X.sizes()[1];
298
+ int64_t padded = (col + 512 - 1) / 512 * 512;
299
const torch::stable::accelerator::DeviceGuard device_guard(
300
X.get_device_index());
301
auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(),
302
std::nullopt, W.device());
303
+ torch::stable::fill_(Y, 0.0);
304
cudaStream_t stream = get_current_cuda_stream();
305
auto quant_X = torch::stable::empty({tokens, padded / 32 * 9},
306
torch::headeronly::ScalarType::Int,
@@ -395,8 +399,8 @@ torch::stable::Tensor ggml_moe_a8_vec(
399
torch::stable::Tensor W, // expert weights
400
torch::stable::Tensor topk_ids, int64_t top_k, int64_t type, int64_t row,
401
int64_t tokens) {
398
- int col = X.sizes()[1];
399
- const int padded = (col + 512 - 1) / 512 * 512;
402
+ int64_t col = X.sizes()[1];
403
+ const int64_t padded = (col + 512 - 1) / 512 * 512;
404
const torch::stable::accelerator::DeviceGuard device_guard(
405
X.get_device_index());
406
auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(),