llama: move suppress_tokens handling to common/sampling

This commit is contained in:
Xuan Son Nguyen
2026-07-29 15:09:21 +02:00
parent caa596ab3f
commit b1f192ff4a
4 changed files with 27 additions and 39 deletions
+16 -2
View File
@@ -310,8 +310,22 @@ 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) {
const llama_token id = suppress[i];
if (0 <= id && id < llama_vocab_n_tokens(vocab)) {
merged.push_back({ id, -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) {
+3
View File
@@ -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);
+8
View File
@@ -4205,6 +4205,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();
}
-37
View File
@@ -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;