llama: move suppress_tokens handling to common/sampling (#26276)

* llama: move suppress_tokens handling to common/sampling

* address security issues

* rm has_logit_bias

diff --git a/common/common.h b/common/common.h
index e7c55ae925fc5783bdde5fa336767e0f097ec9db..14889193d9b570ce2afac4eda213187b7cb088ff 100644 (file)
--- a/common/common.h
+++ b/common/common.h
@@ -294,10 +294,6 @@ struct common_params_sampling {
 
     bool backend_sampling = false;
 
-    bool has_logit_bias() const {
-        return !logit_bias.empty();
-    }
-
     // print the parameters into a string
     std::string print() const;
 };

diff --git a/common/sampling.cpp b/common/sampling.cpp
index 7b241e34f77f24384be3f66c3121a37a61f5354c..256ac161e20f14f28ccde26967380eba28438c86 100644 (file)
--- a/common/sampling.cpp
+++ b/common/sampling.cpp
@@ -310,8 +310,19 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st
         }
     }
 
-    if (params.has_logit_bias()) {
-        samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), params.logit_bias.size(), params.logit_bias.data()));
+    // logit bias: user biases + model suppress tokens (-INFINITY)
+    {
+        std::vector<llama_logit_bias> merged = params.logit_bias;
+
+        int32_t n_suppress = 0;
+        const llama_token * suppress = llama_vocab_get_suppress_tokens(vocab, &n_suppress);
+        for (int32_t i = 0; i < n_suppress; ++i) {
+            merged.push_back({ suppress[i], -INFINITY });
+        }
+
+        if (!merged.empty()) {
+            samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), merged.size(), merged.data()));
+        }
     }
 
     if (params.mirostat == 0) {

diff --git a/include/llama.h b/include/llama.h
index 3c6d22be8999226bcdf23dbff280337e23d8d01b..a3ceecf5ef474c4bd467fac09f6f54d1c9e1fc63 100644 (file)
--- a/include/llama.h
+++ b/include/llama.h
@@ -1102,6 +1102,9 @@ extern "C" {
     LLAMA_API bool llama_vocab_get_add_eos(const struct llama_vocab * vocab);
     LLAMA_API bool llama_vocab_get_add_sep(const struct llama_vocab * vocab);
 
+    // model-specific suppress tokens (gguf key: tokenizer.ggml.suppress_tokens)
+    LLAMA_API const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens);
+
     LLAMA_API llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab);
     LLAMA_API llama_token llama_vocab_fim_suf(const struct llama_vocab * vocab);
     LLAMA_API llama_token llama_vocab_fim_mid(const struct llama_vocab * vocab);

diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp
index 9164a4dd888d215b5cd595a0036c26bd2474d3bc..443cd46408581314cece607cfd712367efdb811d 100644 (file)
--- a/src/llama-vocab.cpp
+++ b/src/llama-vocab.cpp
@@ -2578,7 +2578,14 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             if (suppress_idx != -1) {
                 const int n = gguf_get_arr_n(ctx, suppress_idx);
                 const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx);
-                suppress_tokens.assign(data, data + n);
+                // drop out-of-range ids
+                suppress_tokens.reserve(n);
+                for (int i = 0; i < n; ++i) {
+                    const int32_t id = data[i];
+                    if (id >= 0 && id < (int) id_to_token.size()) {
+                        suppress_tokens.push_back(id);
+                    }
+                }
             }
         }
 
@@ -4205,6 +4212,14 @@ bool llama_vocab_get_add_sep(const struct llama_vocab * vocab) {
     return vocab->get_add_sep();
 }
 
+const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens) {
+    const std::vector<llama_token> & tokens = vocab->get_suppress_tokens();
+    if (n_suppress_tokens) {
+        *n_suppress_tokens = (int32_t) tokens.size();
+    }
+    return tokens.data();
+}
+
 llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab) {
     return vocab->token_fim_pre();
 }

diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp
index 6a96979cebde316694177ab5f0da4cdff0534493..e44f423bdbc5573654a866ab705203ef55b06272 100644 (file)
--- a/src/models/gemma4.cpp
+++ b/src/models/gemma4.cpp
@@ -142,33 +142,6 @@ static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, in
                         idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
 }
 
-// TODO @ngxson : maybe improve this in the future
-class llm_graph_input_logits_bias : public llm_graph_input_i {
-public:
-    llm_graph_input_logits_bias(const llama_vocab & vocab) {
-        arr.resize(vocab.n_tokens(), 0.0f);
-        for (llama_token id : vocab.get_suppress_tokens()) {
-            if (0 <= id && id < (int32_t)vocab.n_tokens()) {
-                arr[id] = -INFINITY;
-            }
-        }
-    }
-    virtual ~llm_graph_input_logits_bias() = default;
-
-    void set_input(const llama_ubatch * /*ubatch*/) override {
-        const int64_t n_vocab = arr.size();
-        ggml_backend_tensor_set(logits_bias, arr.data(), 0, n_vocab*ggml_element_size(logits_bias));
-    }
-
-    bool can_reuse(const llm_graph_params & /*params*/) override {
-        return true;
-    }
-
-    ggml_tensor * logits_bias = nullptr; // F32 [n_vocab]
-
-    std::vector<float> arr;
-};
-
 llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) :
         llm_graph_context(params),
         model(model),
@@ -429,16 +402,6 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
         cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
     }
 
-    // apply logits bias if needed (e.g. for gemma4_unified patch)
-    // this is to mirror the suppress_tokens patch on transformers, to avoid model from outputing <image|> and <audio|> tokens (which is a known issue related to the checkpoint)
-    // TODO: maybe handle this inside the sampling system in the future
-    if (!model.vocab.get_suppress_tokens().empty()) {
-        auto inp_bias = std::make_unique<llm_graph_input_logits_bias>(model.vocab);
-        inp_bias->logits_bias = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, inp_bias->arr.size());
-        cur = ggml_add(ctx0, cur, inp_bias->logits_bias);
-        res->add_input(std::move(inp_bias));
-    }
-
     cb(cur, "result_output", -1);
     res->t_logits = cur;