From 0eddfa0b4c07d081487e57db0758a53aa9c02305 Mon Sep 17 00:00:00 2001 From: Gaurav Garg Date: Sun, 5 Jul 2026 02:06:52 +0530 Subject: [PATCH] Clamp the mask sum before converting it into the sampled index --- src/llama-sampler.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 0bb7484b1d..9175351754 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -1245,6 +1245,9 @@ static void llama_sampler_dist_backend_apply( struct ggml_tensor * idxf = ggml_sum(ctx, mask); ggml_set_name(idxf, "dist_index_f32"); + // Clamp to prevent out-of-bounds access when computing the index. + idxf = ggml_clamp(ctx, idxf, 1.0f, mask->ne[0]); + // Use ggml_scale_bias to scale the index value by -1 and then add the size // of the mask to that value so we get the correct index ((-1 * idxf) + n). struct ggml_tensor * idx = ggml_cast(ctx, ggml_scale_bias(ctx, idxf, -1.0f, mask->ne[0]), GGML_TYPE_I32);