llama-model-saver.cpp
llama-model.cpp
llama-quant.cpp
- llama-sampling.cpp
+ llama-sampler.cpp
llama-vocab.cpp
unicode.cpp
unicode-data.cpp
{ LLM_ARCH_RND1, "rnd1" },
{ LLM_ARCH_PANGU_EMBED, "pangu-embedded" },
{ LLM_ARCH_MISTRAL3, "mistral3" },
- { LLM_ARCH_MIMO2, "mimo2" },
+ { LLM_ARCH_MIMO2, "mimo2" },
+ { LLM_ARCH_STEP35, "step35" },
{ LLM_ARCH_LLAMA_EMBED, "llama-embed" },
{ LLM_ARCH_MAINCODER, "maincoder" },
+ { LLM_ARCH_KIMI_LINEAR, "kimi-linear" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
{ LLM_KV_EXPERT_FEED_FORWARD_LENGTH, "%s.expert_feed_forward_length" },
{ LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, "%s.expert_shared_feed_forward_length" },
{ LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, "%s.expert_chunk_feed_forward_length" },
+ { LLM_KV_SWIGLU_CLAMP_EXP, "%s.swiglu_clamp_exp" },
+ { LLM_KV_SWIGLU_CLAMP_SHEXP, "%s.swiglu_clamp_shexp" },
{ LLM_KV_USE_PARALLEL_RESIDUAL, "%s.use_parallel_residual" },
{ LLM_KV_TENSOR_DATA_LAYOUT, "%s.tensor_data_layout" },
{ LLM_KV_EXPERT_COUNT, "%s.expert_count" },
{ LLM_KV_ATTENTION_KEY_LENGTH_MLA, "%s.attention.key_length_mla" },
{ LLM_KV_ATTENTION_VALUE_LENGTH_MLA, "%s.attention.value_length_mla" },
- { LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" },
- { LLM_KV_ROPE_DIMENSION_SECTIONS, "%s.rope.dimension_sections" },
- { LLM_KV_ROPE_FREQ_BASE, "%s.rope.freq_base" },
- { LLM_KV_ROPE_FREQ_BASE_SWA, "%s.rope.freq_base_swa" },
- { LLM_KV_ROPE_SCALE_LINEAR, "%s.rope.scale_linear" },
- { LLM_KV_ROPE_SCALING_TYPE, "%s.rope.scaling.type" },
- { LLM_KV_ROPE_SCALING_FACTOR, "%s.rope.scaling.factor" },
- { LLM_KV_ROPE_SCALING_ATTN_FACTOR, "%s.rope.scaling.attn_factor" },
- { LLM_KV_ROPE_SCALING_ORIG_CTX_LEN, "%s.rope.scaling.original_context_length" },
- { LLM_KV_ROPE_SCALING_FINETUNED, "%s.rope.scaling.finetuned" },
- { LLM_KV_ROPE_SCALING_YARN_LOG_MUL, "%s.rope.scaling.yarn_log_multiplier" },
- { LLM_KV_ROPE_SCALING_YARN_EXT_FACTOR, "%s.rope.scaling.yarn_ext_factor" },
- { LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR, "%s.rope.scaling.yarn_attn_factor" },
- { LLM_KV_ROPE_SCALING_YARN_BETA_FAST, "%s.rope.scaling.yarn_beta_fast" },
- { LLM_KV_ROPE_SCALING_YARN_BETA_SLOW, "%s.rope.scaling.yarn_beta_slow" },
+ { LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" },
+ { LLM_KV_ROPE_DIMENSION_SECTIONS, "%s.rope.dimension_sections" },
+ { LLM_KV_ROPE_FREQ_BASE, "%s.rope.freq_base" },
+ { LLM_KV_ROPE_FREQ_BASE_SWA, "%s.rope.freq_base_swa" },
+ { LLM_KV_ROPE_SCALE_LINEAR, "%s.rope.scale_linear" },
+ { LLM_KV_ROPE_SCALING_TYPE, "%s.rope.scaling.type" },
+ { LLM_KV_ROPE_SCALING_FACTOR, "%s.rope.scaling.factor" },
+ { LLM_KV_ROPE_SCALING_ATTN_FACTOR, "%s.rope.scaling.attn_factor" },
+ { LLM_KV_ROPE_SCALING_ORIG_CTX_LEN, "%s.rope.scaling.original_context_length" },
+ { LLM_KV_ROPE_SCALING_FINETUNED, "%s.rope.scaling.finetuned" },
+ { LLM_KV_ROPE_SCALING_YARN_LOG_MUL, "%s.rope.scaling.yarn_log_multiplier" },
+ { LLM_KV_ROPE_SCALING_YARN_EXT_FACTOR, "%s.rope.scaling.yarn_ext_factor" },
+ { LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR, "%s.rope.scaling.yarn_attn_factor" },
+ { LLM_KV_ROPE_SCALING_YARN_BETA_FAST, "%s.rope.scaling.yarn_beta_fast" },
+ { LLM_KV_ROPE_SCALING_YARN_BETA_SLOW, "%s.rope.scaling.yarn_beta_slow" },
{ LLM_KV_SPLIT_NO, "split.no" },
{ LLM_KV_SPLIT_COUNT, "split.count" },
{ LLM_KV_SSM_GROUP_COUNT, "%s.ssm.group_count" },
{ LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" },
+ { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" },
+
{ LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" },
{ LLM_KV_POSNET_EMBEDDING_LENGTH, "%s.posnet.embedding_length" },
{ LLM_TENSOR_SSM_DT_NORM, "blk.%d.ssm_dt_norm" },
{ LLM_TENSOR_SSM_B_NORM, "blk.%d.ssm_b_norm" },
{ LLM_TENSOR_SSM_C_NORM, "blk.%d.ssm_c_norm" },
+ { LLM_TENSOR_SSM_CONV1D_Q, "blk.%d.ssm_conv1d_q" },
+ { LLM_TENSOR_SSM_CONV1D_K, "blk.%d.ssm_conv1d_k" },
+ { LLM_TENSOR_SSM_CONV1D_V, "blk.%d.ssm_conv1d_v" },
+ { LLM_TENSOR_SSM_F_A, "blk.%d.ssm_f_a" },
+ { LLM_TENSOR_SSM_F_B, "blk.%d.ssm_f_b" },
+ { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" },
+ { LLM_TENSOR_SSM_G_A, "blk.%d.ssm_g_a" },
+ { LLM_TENSOR_SSM_G_B, "blk.%d.ssm_g_b" },
+ { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" },
{ LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" },
{ LLM_TENSOR_ATTN_KV_A_NORM, "blk.%d.attn_kv_a_norm" },
{ LLM_TENSOR_ATTN_Q_A, "blk.%d.attn_q_a" },
LLM_TENSOR_FFN_UP_EXPS,
LLM_TENSOR_FFN_EXP_PROBS_B,
};
+ case LLM_ARCH_STEP35:
+ return {
+ LLM_TENSOR_TOKEN_EMBD,
+ LLM_TENSOR_OUTPUT_NORM,
+ LLM_TENSOR_OUTPUT,
+ LLM_TENSOR_ROPE_FREQS,
+ LLM_TENSOR_ROPE_FACTORS_LONG,
+ LLM_TENSOR_ROPE_FACTORS_SHORT,
+ LLM_TENSOR_ATTN_NORM,
+ LLM_TENSOR_ATTN_Q,
+ LLM_TENSOR_ATTN_Q_NORM,
+ LLM_TENSOR_ATTN_K,
+ LLM_TENSOR_ATTN_K_NORM,
+ LLM_TENSOR_ATTN_V,
+ LLM_TENSOR_ATTN_GATE,
+ LLM_TENSOR_ATTN_OUT,
+ LLM_TENSOR_FFN_NORM,
+ LLM_TENSOR_FFN_GATE,
+ LLM_TENSOR_FFN_DOWN,
+ LLM_TENSOR_FFN_UP,
+ LLM_TENSOR_FFN_GATE_INP,
+ LLM_TENSOR_FFN_GATE_EXPS,
+ LLM_TENSOR_FFN_DOWN_EXPS,
+ LLM_TENSOR_FFN_UP_EXPS,
+ LLM_TENSOR_FFN_GATE_SHEXP,
+ LLM_TENSOR_FFN_UP_SHEXP,
+ LLM_TENSOR_FFN_DOWN_SHEXP,
+ LLM_TENSOR_FFN_EXP_PROBS_B,
+ };
case LLM_ARCH_GPTJ:
case LLM_ARCH_UNKNOWN:
return {
LLM_TENSOR_FFN_DOWN,
LLM_TENSOR_FFN_UP,
};
+ case LLM_ARCH_KIMI_LINEAR:
+ return {
+ LLM_TENSOR_TOKEN_EMBD,
+ LLM_TENSOR_OUTPUT_NORM,
+ LLM_TENSOR_OUTPUT,
+ LLM_TENSOR_ROPE_FREQS,
+ LLM_TENSOR_ATTN_NORM,
+ LLM_TENSOR_ATTN_Q,
+ LLM_TENSOR_ATTN_K,
+ LLM_TENSOR_ATTN_V,
+ LLM_TENSOR_ATTN_OUT,
+ LLM_TENSOR_FFN_NORM,
+ // Dense FFN (layer 0 only)
+ LLM_TENSOR_FFN_GATE,
+ LLM_TENSOR_FFN_DOWN,
+ LLM_TENSOR_FFN_UP,
+ // MoE FFN (layers 1+)
+ LLM_TENSOR_FFN_GATE_INP,
+ LLM_TENSOR_FFN_GATE_EXPS,
+ LLM_TENSOR_FFN_DOWN_EXPS,
+ LLM_TENSOR_FFN_UP_EXPS,
+ LLM_TENSOR_FFN_EXP_PROBS_B,
+ // Shared experts
+ LLM_TENSOR_FFN_GATE_SHEXP,
+ LLM_TENSOR_FFN_DOWN_SHEXP,
+ LLM_TENSOR_FFN_UP_SHEXP,
+ // KDA (using SSM_ enum prefix, keeping GGUF names for backward compat)
+ LLM_TENSOR_SSM_CONV1D_Q,
+ LLM_TENSOR_SSM_CONV1D_K,
+ LLM_TENSOR_SSM_CONV1D_V,
+ LLM_TENSOR_SSM_F_A,
+ LLM_TENSOR_SSM_F_B,
+ LLM_TENSOR_SSM_BETA,
+ LLM_TENSOR_SSM_A,
+ LLM_TENSOR_SSM_G_A,
+ LLM_TENSOR_SSM_G_B,
+ LLM_TENSOR_SSM_DT,
+ LLM_TENSOR_SSM_NORM,
+ // MLA
+ LLM_TENSOR_ATTN_Q_A,
+ LLM_TENSOR_ATTN_Q_B,
+ LLM_TENSOR_ATTN_Q_A_NORM,
+ LLM_TENSOR_ATTN_KV_A_MQA,
+ LLM_TENSOR_ATTN_KV_B,
+ LLM_TENSOR_ATTN_K_B,
+ LLM_TENSOR_ATTN_V_B,
+ LLM_TENSOR_ATTN_KV_A_NORM,
+ };
default:
GGML_ABORT("unknown architecture for tensor mapping");
}
{LLM_TENSOR_SSM_C_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_SSM_D, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_SSM_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ // Kimi KDA - Conv tensors are 4D [d_conv, 1, d_inner, 1], reshaped to 2D at runtime
+ {LLM_TENSOR_SSM_CONV1D_Q, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ {LLM_TENSOR_SSM_CONV1D_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ {LLM_TENSOR_SSM_CONV1D_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ {LLM_TENSOR_SSM_F_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_SSM_F_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_SSM_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_SSM_G_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_SSM_G_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_TIME_MIX_LERP_X, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_TIME_MIX_LN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_CHANNEL_MIX_LERP_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
case LLM_ARCH_NEMOTRON_H:
case LLM_ARCH_NEMOTRON_H_MOE:
case LLM_ARCH_QWEN3NEXT:
+ case LLM_ARCH_KIMI_LINEAR:
return true;
default:
return false;
LLM_ARCH_PANGU_EMBED,
LLM_ARCH_MISTRAL3,
LLM_ARCH_MIMO2,
+ LLM_ARCH_STEP35,
LLM_ARCH_LLAMA_EMBED,
LLM_ARCH_MAINCODER,
+ LLM_ARCH_KIMI_LINEAR,
LLM_ARCH_UNKNOWN,
};
LLM_KV_EXPERT_FEED_FORWARD_LENGTH,
LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH,
LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH,
+ LLM_KV_SWIGLU_CLAMP_EXP,
+ LLM_KV_SWIGLU_CLAMP_SHEXP,
LLM_KV_USE_PARALLEL_RESIDUAL,
LLM_KV_TENSOR_DATA_LAYOUT,
LLM_KV_EXPERT_COUNT,
LLM_KV_SSM_GROUP_COUNT,
LLM_KV_SSM_DT_B_C_RMS,
+ LLM_KV_KDA_HEAD_DIM,
+
LLM_KV_WKV_HEAD_SIZE,
LLM_KV_TOKENIZER_MODEL,
LLM_TENSOR_SSM_NORM,
LLM_TENSOR_SSM_OUT,
LLM_TENSOR_SSM_BETA_ALPHA, // qwen3next
+ // Kimi Linear KDA (using SSM_ prefix for consistency)
+ LLM_TENSOR_SSM_CONV1D_Q, // kimi: Q conv1d weight
+ LLM_TENSOR_SSM_CONV1D_K, // kimi: K conv1d weight
+ LLM_TENSOR_SSM_CONV1D_V, // kimi: V conv1d weight
+ LLM_TENSOR_SSM_F_A, // kimi: forget gate projection A
+ LLM_TENSOR_SSM_F_B, // kimi: forget gate projection B
+ LLM_TENSOR_SSM_BETA, // kimi: beta mixing coefficient
+ LLM_TENSOR_SSM_G_A, // kimi: output gate projection A
+ LLM_TENSOR_SSM_G_B, // kimi: output gate projection B
LLM_TENSOR_TIME_MIX_W0,
LLM_TENSOR_TIME_MIX_W1,
LLM_TENSOR_TIME_MIX_W2,
llm_chat_template tmpl,
const std::vector<const llama_chat_message *> & chat,
std::string & dest, bool add_ass) {
- // Taken from the research: https://github.com/ggerganov/llama.cpp/issues/5527
+ // Taken from the research: https://github.com/ggml-org/llama.cpp/issues/5527
std::stringstream ss;
if (tmpl == LLM_CHAT_TEMPLATE_CHATML) {
// chatml template
auto dev_type = ggml_backend_dev_type(ggml_backend_get_device(backend.get()));
if (dev_type == GGML_BACKEND_DEVICE_TYPE_CPU) {
// ignore CPU backend
+ // TODO: should we ignore ACCEL types too?
continue;
}
auto * dev = ggml_backend_get_device(backend.get());
llama_sampler_chain_n(sampler) > 0;
if (sampler && can_offload) {
- ggml_backend_buffer_type_t buft = ggml_backend_dev_buffer_type(model.dev_output());
- auto * host_buft = ggml_backend_dev_host_buffer_type(model.dev_output());
- if (host_buft) {
- buft = host_buft;
- }
+ auto * buft = ggml_backend_dev_buffer_type(model.dev_output());
sampler->iface->backend_init(sampler, buft);
//
uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
- if (model.arch == LLM_ARCH_QWEN3NEXT) {
+ if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR) {
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
}
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
#include "llama-impl.h"
#include "llama-vocab.h"
-#include "llama-sampling.h"
+#include "llama-sampler.h"
#include <cmath>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstring>
+#include <numeric>
+#include <sstream>
#include <unordered_set>
void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {
return res;
}
+// TODO: Hybrid input classes are a bit redundant.
+// Instead of creating a hybrid input, the graph can simply create 2 separate inputs.
+// Refactoring is required in the future.
+void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) {
+ mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
+
+ mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+
+ const int64_t n_rs = mctx->get_recr()->get_n_rs();
+
+ if (inp_rs->s_copy) {
+ GGML_ASSERT(ggml_backend_buffer_is_host(inp_rs->s_copy->buffer));
+ int32_t * data = (int32_t *) inp_rs->s_copy->data;
+
+ // assuming copy destinations ALWAYS happen ONLY on the cells between head and head+n
+ for (uint32_t i = 0; i < n_rs; ++i) {
+ data[i] = mctx->get_recr()->s_copy(i);
+ }
+ }
+}
+
+bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) {
+ const auto * mctx = static_cast<const llama_memory_hybrid_context *>(params.mctx);
+
+ this->mctx = mctx;
+
+ bool res = true;
+
+ res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens;
+
+ res &= inp_attn->self_kq_mask->ne[0] == mctx->get_attn()->get_n_kv();
+ res &= inp_attn->self_kq_mask->ne[1] == params.ubatch.n_tokens;
+
+ res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
+
+ res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
+ res &= inp_rs->s_copy_extra->ne[0] == mctx->get_recr()->get_n_rs() - params.ubatch.n_seqs;
+
+ res &= inp_rs->head == mctx->get_recr()->get_head();
+ res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z();
+
+ return res;
+}
+
void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) {
const auto * attn_ctx = mctx->get_attn();
switch (type_op) {
case LLM_FFN_SILU:
if (gate && type_gate == LLM_FFN_PAR) {
+ // Step35: HF clamps gate (after SiLU) and up before multiplication
+ if (arch == LLM_ARCH_STEP35 && il >= 0) {
+ const float limit = hparams.swiglu_clamp_shexp[il];
+ constexpr float eps = 1e-6f;
+ if (limit > eps) {
+ ggml_tensor * gate_act = ggml_silu(ctx0, cur);
+ cb(gate_act, "ffn_silu", il);
+ gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit);
+ cb(gate_act, "ffn_silu_clamped", il);
+
+ tmp = ggml_clamp(ctx0, tmp, -limit, limit);
+ cb(tmp, "ffn_up_clamped", il);
+
+ cur = ggml_mul(ctx0, gate_act, tmp);
+ cb(cur, "ffn_swiglu_limited", il);
+ type_gate = LLM_FFN_SEQ;
+ break;
+ }
+ }
+
cur = ggml_swiglu_split(ctx0, cur, tmp);
cb(cur, "ffn_swiglu", il);
type_gate = LLM_FFN_SEQ;
switch (type_op) {
case LLM_FFN_SILU:
if (gate_exps) {
+ // Step35: per-layer clamp for routed experts
+ if (arch == LLM_ARCH_STEP35 && il >= 0) {
+ const float limit = hparams.swiglu_clamp_exp[il];
+ constexpr float eps = 1e-6f;
+ if (limit > eps) {
+ ggml_tensor * gate_act = ggml_silu(ctx0, cur);
+ cb(gate_act, "ffn_moe_silu", il);
+ gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit);
+ cb(gate_act, "ffn_moe_silu_clamped", il);
+
+ up = ggml_clamp(ctx0, up, -limit, limit);
+ cb(up, "ffn_moe_up_clamped", il);
+
+ cur = ggml_mul(ctx0, gate_act, up);
+ cb(cur, "ffn_moe_swiglu_limited", il);
+ break;
+ }
+ }
+
cur = ggml_swiglu_split(ctx0, cur, up);
cb(cur, "ffn_moe_swiglu", il);
} else {
return (llm_graph_input_mem_hybrid *) res->add_input(std::move(inp));
}
+llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const {
+ const auto * mctx_cur = static_cast<const llama_memory_hybrid_context *>(mctx);
+
+ auto inp_rs = build_rs_inp_impl (ctx0, ubatch, mctx_cur->get_recr());
+ auto inp_attn = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn());
+
+ auto inp = std::make_unique<llm_graph_input_mem_hybrid_k>(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur);
+
+ return (llm_graph_input_mem_hybrid_k *) res->add_input(std::move(inp));
+}
+
llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa() const {
const auto * mctx_cur = static_cast<const llama_memory_hybrid_iswa_context *>(mctx);
return;
}
+ std::array<ggml_tensor *, 2> outs;
+ outs[0] = res->t_logits;
+
auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers);
res->add_input(std::move(inp_sampling));
// add a dummy row of logits
// this trick makes the graph static, regardless of which samplers are activated
// this is important in order to minimize graph reallocations
- // TODO: use `ggml_build_forward_select()` when available (https://github.com/ggml-org/llama.cpp/pull/18550)
ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0);
for (const auto & [seq_id, sampler] : samplers) {
const auto it = seq_to_logit_row.find(seq_id);
// inactive samplers always work on the first row
- const auto row_idx = seq_to_logit_row.find(seq_id) != seq_to_logit_row.end() ? it->second : 0;
+ const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0;
+ const int i_out = it != seq_to_logit_row.end() ? 1 : 0;
ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]);
ggml_format_name(logits_seq, "logits_seq_%d", seq_id);
if (data.sampled != nullptr) {
res->t_sampled[seq_id] = data.sampled;
- ggml_build_forward_expand(gf, data.sampled);
+ outs[1] = data.sampled;
+ ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.probs != nullptr) {
res->t_sampled_probs[seq_id] = data.probs;
- ggml_build_forward_expand(gf, data.probs);
+ outs[1] = data.probs;
+ ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.logits != nullptr) {
res->t_sampled_logits[seq_id] = data.logits;
- ggml_build_forward_expand(gf, data.logits);
+ outs[1] = data.logits;
+ ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.candidates != nullptr) {
res->t_candidates[seq_id] = data.candidates;
- ggml_build_forward_expand(gf, data.candidates);
+ outs[1] = data.candidates;
+ ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
}
const llama_memory_hybrid_context * mctx;
};
+class llm_graph_input_mem_hybrid_k : public llm_graph_input_i {
+public:
+ llm_graph_input_mem_hybrid_k(
+ const llama_cparams & cparams,
+ std::unique_ptr<llm_graph_input_attn_k> inp_attn,
+ std::unique_ptr<llm_graph_input_rs> inp_rs,
+ const llama_memory_hybrid_context * mctx) :
+ inp_attn(std::move(inp_attn)),
+ inp_rs(std::move(inp_rs)),
+ cparams(cparams),
+ mctx(mctx) { }
+ virtual ~llm_graph_input_mem_hybrid_k() = default;
+
+ void set_input(const llama_ubatch * ubatch) override;
+
+ bool can_reuse(const llm_graph_params & params) override;
+
+ std::unique_ptr<llm_graph_input_attn_k> inp_attn;
+ std::unique_ptr<llm_graph_input_rs> inp_rs;
+
+ llm_graph_input_attn_k * get_attn() const { return inp_attn.get(); }
+ llm_graph_input_rs * get_recr() const { return inp_rs.get(); }
+
+ const llama_cparams cparams;
+
+ const llama_memory_hybrid_context * mctx;
+};
+
class llm_graph_input_mem_hybrid_iswa : public llm_graph_input_i {
public:
llm_graph_input_mem_hybrid_iswa(
//
llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const;
+ llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const;
llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const;
return n_embd * (n_shortconv_l_cache - 1);
}
+ if (n_embd_head_kda != 0) {
+ // for Kimi KDA layers
+ // Conv state for Q, K, V: 3 * (d_conv - 1) * n_head * head_dim
+ const uint32_t d_inner = n_head() * n_embd_head_kda; // 32 * 128 = 4096
+ return 3 * (ssm_d_conv > 0 ? ssm_d_conv - 1 : 3) * d_inner;
+ }
+
// TODO: maybe support other convolution strides than 1
// NOTE: since the first column of the conv_state is shifted out each time, it's not actually needed
// Corresponds to Mamba's conv_states size
return n_embd * wkv_head_size;
}
+ if (n_embd_head_kda != 0) {
+ // for Kimi KDA layers
+ // Full recurrent state: head_dim * head_dim * n_head
+ // h tensor shape for delta attention: [head_dim, head_dim, n_head]
+ return n_embd_head_kda * n_embd_head_kda * n_head(); // 128 * 128 * 32 = 524288
+ }
+
// corresponds to Mamba's ssm_states size
return ssm_d_state * ssm_d_inner;
}
uint32_t ssm_dt_rank = 0;
uint32_t ssm_n_group = 0;
+ // for Kimi Linear KDA
+ uint32_t n_embd_head_kda = 0;
+
// for hybrid state space models
std::array<bool, LLAMA_MAX_LAYERS> recurrent_layer_arr;
uint32_t n_deepstack_layers = 0;
// needed by encoder-decoder models (e.g. T5, FLAN-T5)
- // ref: https://github.com/ggerganov/llama.cpp/pull/8141
+ // ref: https://github.com/ggml-org/llama.cpp/pull/8141
llama_token dec_start_token_id = LLAMA_TOKEN_NULL;
uint32_t dec_n_layer = 0;
enum llama_rope_type rope_type = LLAMA_ROPE_TYPE_NONE;
enum llama_rope_scaling_type rope_scaling_type_train = LLAMA_ROPE_SCALING_TYPE_NONE;
+
+ // Step35: optional per-layer clamps for (Swi)GLU
+ std::array<float, LLAMA_MAX_LAYERS> swiglu_clamp_exp; // clamping for expert FFN
+ std::array<float, LLAMA_MAX_LAYERS> swiglu_clamp_shexp; // shared expert
+
// this value n_pattern means that every nth layer is dense (i.e. non-SWA)
// dense_first means whether the pattern is start with a dense layer
// note that if n_pattern == 0, all layers are SWA
}
bool llama_kv_cache_iswa::get_can_shift() const {
- return kv_base->get_size() == kv_swa->get_size();
+ return kv_base->get_can_shift() &&
+ kv_swa->get_can_shift() &&
+ kv_base->get_size() == kv_swa->get_size();
}
void llama_kv_cache_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
}
bool llama_kv_cache::get_can_shift() const {
+ // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot.
+ if (model.arch == LLM_ARCH_STEP35) {
+ return false;
+ }
return true;
}
io.write(&v_trans, sizeof(v_trans));
io.write(&n_layer, sizeof(n_layer));
- std::vector<uint8_t> tmp_buf;
-
// Iterate and write all the keys first, each row is a cell
// Get whole range at a time
for (const auto & layer : layers) {
const uint64_t k_size_row = ggml_row_size(k->type, n_embd_k_gqa);
io.write(&k_size_row, sizeof(k_size_row));
- // Read each range of cells of k_size length each into tmp_buf and write out
+ // Read each range of cells of k_size length and write out
for (const auto & range : cr.data) {
const size_t range_size = range.second - range.first;
const size_t buf_size = range_size * k_size_row;
const uint64_t v_size_row = ggml_row_size(v->type, n_embd_v_gqa);
io.write(&v_size_row, sizeof(v_size_row));
- // Read each range of cells of v_size length each into tmp_buf and write out
+ // Read each range of cells of v_size length and write out
for (const auto & range : cr.data) {
const size_t range_size = range.second - range.first;
const size_t buf_size = range_size * v_size_row;
// For each row, we get the element values of each cell
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
- // Read each range of cells of v_size_el length each into tmp_buf and write out
+ // Read each range of cells of v_size_el length and write out
for (const auto & range : cr.data) {
const size_t range_size = range.second - range.first;
const size_t src_offset = (range.first + j * kv_size) * v_size_el;
io.write(&s_trans, sizeof(s_trans));
io.write(&n_layer, sizeof(n_layer));
- std::vector<uint8_t> tmp_buf;
-
- // Iterate and write all the keys first, each row is a cell
+ // Iterate and write all the R tensors first, each row is a cell
// Get whole range at a time
for (uint32_t il = 0; il < n_layer; ++il) {
// skip null layers (read_data will handle this by checking "r_l" and "s_l" for null)
if (r_l[il] == nullptr) continue;
- // Write key type
+ // Write R tensor type
const int32_t r_type_i = (int32_t)r_l[il]->type;
io.write(&r_type_i, sizeof(r_type_i));
- // Write row size of key
+ // Write row size of R tensor
const uint64_t r_size_row = ggml_row_size(r_l[il]->type, hparams.n_embd_r());
io.write(&r_size_row, sizeof(r_size_row));
- // Read each range of cells of k_size length each into tmp_buf and write out
+ // Write each range of cells of r_size_row length
for (const auto & range : cell_ranges) {
const size_t range_size = range.second - range.first;
const size_t buf_size = range_size * r_size_row;
// skip null layers (read_data will handle this by checking "r_l" and "s_l" for null)
if (s_l[il] == nullptr) continue;
- // Write value type
+ // Write S tensor type
const int32_t s_type_i = (int32_t)s_l[il]->type;
io.write(&s_type_i, sizeof(s_type_i));
- // Write row size of value
+ // Write row size of S tensor
const uint64_t s_size_row = ggml_row_size(s_l[il]->type, hparams.n_embd_s());
io.write(&s_size_row, sizeof(s_size_row));
- // Read each range of cells of s_size length each into tmp_buf and write out
+ // Write each range of S tensor rows
for (const auto & range : cell_ranges) {
const size_t range_size = range.second - range.first;
const size_t buf_size = range_size * s_size_row;
}
}
} else {
- // When v is transposed, we also need the element size and get the element ranges from each row
+ // When S tensor is transposed, we also need the element size and get the element ranges from each row
const uint32_t mem_size = size;
for (uint32_t il = 0; il < n_layer; ++il) {
// skip null layers (read_data will handle this by checking "r_l" and "s_l" for null)
const uint32_t n_embd_s = hparams.n_embd_s();
- // Write value type
+ // Write S tensor type
const int32_t s_type_i = (int32_t)s_l[il]->type;
io.write(&s_type_i, sizeof(s_type_i));
// For each row, we get the element values of each cell
for (uint32_t j = 0; j < n_embd_s; ++j) {
- // Read each range of cells of v_size_el length each into tmp_buf and write out
+ // Write each range of cells of s_size_el length
for (const auto & range : cell_ranges) {
const size_t range_size = range.second - range.first;
const size_t src_offset = (range.first + j * mem_size) * s_size_el;
case LLM_TYPE_21B_A3B: return "21B.A3B";
case LLM_TYPE_30B_A3B: return "30B.A3B";
case LLM_TYPE_31B_A3_5B: return "31B.A3.5B";
+ case LLM_TYPE_48B_A3B: return "48B.A3B";
case LLM_TYPE_80B_A3B: return "80B.A3B";
case LLM_TYPE_100B_A6B: return "100B.A6B";
case LLM_TYPE_102B_A12B: return "102B.A12B";
case LLM_TYPE_106B_A12B: return "106B.A12B";
+ case LLM_TYPE_196B_A11B: return "196B.A11B";
case LLM_TYPE_230B_A10B: return "230B.A10B";
case LLM_TYPE_235B_A22B: return "235B.A22B";
case LLM_TYPE_300B_A47B: return "300B.A47B";
std::fill(hparams.xielu_alpha_p.begin(), hparams.xielu_alpha_p.end(), 0.0f);
std::fill(hparams.xielu_beta.begin(), hparams.xielu_beta.end(), 0.0f);
std::fill(hparams.xielu_eps.begin(), hparams.xielu_eps.end(), 0.0f);
+ std::fill(hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.end(), 0.0f);
+ std::fill(hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.end(), 0.0f);
ml.get_key_or_arr(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, hparams.n_layer, false);
ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, hparams.n_layer, false);
default: type = LLM_TYPE_UNKNOWN;
}
} break;
+ case LLM_ARCH_KIMI_LINEAR:
+ {
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
+ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
+ ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
+ ml.get_key(LLM_KV_ROPE_DIMENSION_COUNT, hparams.n_rot);
+ ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
+ ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);
+
+ // MLA qk_rope_head_dim (for reference)
+ // qk_rope_head_dim = 64, qk_nope_head_dim = 128, qk_head_dim = 192
+
+ // Mark KDA layers as recurrent using n_head_kv pattern (like Jamba)
+ // Set n_head_kv = 0 for KDA layers (recurrent), n_head_kv = n_head for MLA layers (attention)
+ for (uint32_t i = 0; i < hparams.n_layer; ++i) {
+ hparams.recurrent_layer_arr[i] = hparams.n_head_kv(i) == 0; // KDA layers are recurrent
+ }
+
+ // MoE parameters - Kimi uses moe_intermediate_size = 1024
+ ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
+ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
+ ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
+
+ switch (hparams.n_layer) {
+ case 27: type = LLM_TYPE_48B_A3B; break; // Kimi-Linear-48B-A3B
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+ } break;
+ case LLM_ARCH_STEP35:
+ {
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+
+ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+
+ // MoE + SWA parameters
+ ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
+ ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
+
+ // Step35 uses sigmoid gating by default (if not set in GGUF)
+ if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
+ hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
+ }
+
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+ ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa);
+ ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer);
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer, false);
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer, false);
+
+ switch (hparams.n_layer) {
+ case 45: type = LLM_TYPE_196B_A11B; break;
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+ } break;
default: throw std::runtime_error("unsupported model architecture");
}
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
}
} break;
+ case LLM_ARCH_KIMI_LINEAR:
+ {
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ // output
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+ // Check for KDA specific tensors to determine layer type or if it's a mixed model
+ // Assuming KDA layer if KDA tensors are present
+
+ // KDA uses head_dim = 128 (from linear_attn_config.head_dim)
+ const int64_t n_embd_head_k_kda = hparams.n_embd_head_kda;
+ const int64_t n_embd_head_v_kda = hparams.n_embd_head_kda;
+ const int64_t ssm_d_conv = hparams.ssm_d_conv;
+
+ // Try loading KDA specific tensors (using SSM_ prefix)
+ // Conv1d weights: try 4D first, then 3D (quantization may remove trailing 1)
+ // 4D: [d_conv, 1, d_inner, 1], 3D: [d_conv, 1, d_inner]
+ layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, n_embd_head_k_kda * n_head, 1}, TENSOR_NOT_REQUIRED);
+ if (!layer.ssm_q_conv) {
+ layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, n_embd_head_k_kda * n_head}, TENSOR_NOT_REQUIRED);
+ }
+
+ if (layer.ssm_q_conv) {
+ // KDA Layer - Conv1d weights may be 3D or 4D
+ layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", i), {ssm_d_conv, 1, n_embd_head_k_kda * n_head, 1}, TENSOR_NOT_REQUIRED);
+ if (!layer.ssm_k_conv) {
+ layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", i), {ssm_d_conv, 1, n_embd_head_k_kda * n_head}, 0);
+ }
+ layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", i), {ssm_d_conv, 1, n_embd_head_v_kda * n_head, 1}, TENSOR_NOT_REQUIRED);
+ if (!layer.ssm_v_conv) {
+ layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", i), {ssm_d_conv, 1, n_embd_head_v_kda * n_head}, 0);
+ }
+
+ // q, k, v projections
+ // Python: q_proj, k_proj, v_proj
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k_kda * n_head}, 0);
+ layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_head_k_kda * n_head}, 0);
+ layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_head_v_kda * n_head}, 0);
+
+ // KDA specific projections
+ // f_a_proj, f_b_proj
+ layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, n_embd_head_k_kda}, 0); // head_dim
+ layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {n_embd_head_k_kda, n_embd_head_k_kda * n_head}, 0); // projection_size
+
+ // b_proj (beta mixing coefficient)
+ layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);
+
+ // A_log - Shape in GGUF: [1, num_heads, 1, 1] (4D) or [1, num_heads] (2D after quantization) Note: -exp(A_log) is applied in convert_hf_to_gguf.py
+ layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head, 1, 1}, TENSOR_NOT_REQUIRED);
+ if (!layer.ssm_a) {
+ layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head}, 0);
+ }
+
+ // dt_bias - shape [n_embd_head_k_kda * n_head] = [4096]
+ layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_embd_head_k_kda * n_head}, 0);
+
+ // g_a_proj, g_b_proj (output gate)
+ layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, n_embd_head_k_kda}, 0);
+ layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {n_embd_head_k_kda, n_embd_head_k_kda * n_head}, 0);
+
+ // o_norm (reusing SSM_NORM)
+ layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {n_embd_head_k_kda}, 0); // FusedRMSNormGated
+
+ // o_proj
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v_kda * n_head, n_embd}, 0);
+
+ } else {
+ // MLA Layer - use MLA-specific head dimensions
+ const int64_t q_lora_rank = hparams.n_lora_q;
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+
+ layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED);
+ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0);
+
+ if (layer.attn_q_a_norm) {
+ layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
+ layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k_mla}, 0);
+ } else {
+ // Kimi MLA without Q compression: wq = [n_embd, n_head * n_embd_head_k_mla]
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k_mla}, 0);
+ }
+
+ // Kimi: qk_rope_head_dim = 64 (actual RoPE dimension for MLA)
+ // Note: hparams.n_rot may be 72 (from conversion) but actual is 64
+ const int64_t qk_rope_head_dim = hparams.n_rot; // From config: qk_rope_head_dim
+ layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, 0);
+ // Support Legacy GGUFs that don't split wkv_b (MLA KV cache disabled)
+ layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), {kv_lora_rank, n_head * (n_embd_head_k_mla - qk_rope_head_dim + n_embd_head_v_mla)}, TENSOR_NOT_REQUIRED);
+ if (!layer.wkv_b) { // MLA KV cache enabled
+ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_k_mla - qk_rope_head_dim, kv_lora_rank, n_head}, 0);
+ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, 0);
+ }
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, 0);
+ }
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ // MoE intermediate size (different from dense FFN)
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+
+ // Kimi uses n_layer_dense_lead to determine which layers use dense FFN vs MoE
+ // first_k_dense_replace = 1 means layer 0 uses dense FFN, layers 1+ use MoE
+ if (i < (int) hparams.n_layer_dense_lead) {
+ // Dense FFN layer - use normal n_ff
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
+ } else {
+ // MoE layer - use n_ff_exp (1024) instead of n_ff (9216)
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
+
+ // Shared experts use moe_intermediate_size * num_shared_experts
+ // Kimi: shared_expert_intermediate_size = 1024 * 1 = 1024
+ // Tensors are 2D: [n_embd, n_ff_shexp] or [n_ff_shexp, n_embd]
+ const int64_t n_ff_shexp_actual = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1);
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp_actual}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp_actual, n_embd}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp_actual}, TENSOR_NOT_REQUIRED);
+
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
+ }
+ }
+ } break;
case LLM_ARCH_COGVLM:
{
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
}
} break;
+ case LLM_ARCH_STEP35:
+ {
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ // output
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
+
+ // STEP35 supports per-layer partial RoPE dims; rope factors are stored as a single shared tensor
+ // ("rope_freqs.weight") and ggml uses only the first (n_rot_l/2) entries per layer.
+ uint32_t n_rot_max = 0;
+ for (int i = 0; i < n_layer; ++i) {
+ n_rot_max = std::max(n_rot_max, hparams.n_rot);
+ }
+ if (n_rot_max == 0) {
+ n_rot_max = n_rot;
+ }
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ const uint32_t n_head_l = hparams.n_head(i);
+ const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i);
+ const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i);
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
+ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
+
+ // optional rope factors (llama3) / longrope tensors
+ if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) {
+ layer.rope_long = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+ layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+ } else {
+ layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+ }
+
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head_l}, 0);
+ layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0);
+ layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, 0);
+
+ // head-wise attention gate (Step35 self_attn.g_proj)
+ layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head_l}, TENSOR_NOT_REQUIRED);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ // dense MLP (leading dense blocks)
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
+
+ // MoE routed experts + selection bias (router_bias)
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
+
+ // shared expert MLP
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
+ }
+ } break;
case LLM_ARCH_MAINCODER:
{
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
{
llm = std::make_unique<llm_build_mimo2_iswa>(*this, params);
} break;
+ case LLM_ARCH_KIMI_LINEAR:
+ {
+ llm = std::make_unique<llm_build_kimi_linear>(*this, params);
+ } break;
+ case LLM_ARCH_STEP35:
+ {
+ llm = std::make_unique<llm_build_step35_iswa>(*this, params);
+ } break;
default:
GGML_ABORT("fatal error");
}
case LLM_ARCH_WAVTOKENIZER_DEC:
case LLM_ARCH_NEMOTRON_H:
case LLM_ARCH_NEMOTRON_H_MOE:
+ case LLM_ARCH_KIMI_LINEAR:
return LLAMA_ROPE_TYPE_NONE;
// use what we call a normal RoPE, operating on pairs of consecutive head values
case LLM_ARCH_AFMOE:
case LLM_ARCH_QWEN3NEXT:
case LLM_ARCH_MIMO2:
+ case LLM_ARCH_STEP35:
return LLAMA_ROPE_TYPE_NEOX;
case LLM_ARCH_QWEN2VL:
LLM_TYPE_21B_A3B, // Ernie MoE small
LLM_TYPE_30B_A3B,
LLM_TYPE_31B_A3_5B,
+ LLM_TYPE_48B_A3B, // Kimi Linear
LLM_TYPE_80B_A3B, // Qwen3 Next
LLM_TYPE_100B_A6B,
LLM_TYPE_102B_A12B, // Solar-Open
LLM_TYPE_106B_A12B, // GLM-4.5-Air
+ LLM_TYPE_196B_A11B, // Step3.5-Flash
LLM_TYPE_230B_A10B, // Minimax M2
LLM_TYPE_235B_A22B,
LLM_TYPE_300B_A47B, // Ernie MoE big
struct ggml_tensor * ffn_act_beta = nullptr;
struct ggml_tensor * ffn_act_eps = nullptr;
+ // Kimi Linear KDA (using ssm_ prefix for consistency)
+ // Note: ssm_dt_b already exists above (mamba bias), reused for Kimi dt_bias
+ struct ggml_tensor * ssm_q_conv = nullptr;
+ struct ggml_tensor * ssm_k_conv = nullptr;
+ struct ggml_tensor * ssm_v_conv = nullptr;
+ struct ggml_tensor * ssm_f_a = nullptr;
+ struct ggml_tensor * ssm_f_b = nullptr;
+ struct ggml_tensor * ssm_beta = nullptr;
+ struct ggml_tensor * ssm_g_a = nullptr;
+ struct ggml_tensor * ssm_g_b = nullptr;
+ struct ggml_tensor * ssm_o_norm = nullptr;
+
struct llama_layer_posnet posnet;
struct llama_layer_convnext convnext;
quantize &= name != LLM_TN(model.arch)(LLM_TENSOR_POS_EMBD, "weight");
quantize &= name != LLM_TN(model.arch)(LLM_TENSOR_TOKEN_TYPES, "weight");
- // do not quantize Mamba's small yet 2D weights
+ // do not quantize Mamba /Kimi's small conv1d weights
// NOTE: can't use LLM_TN here because the layer number is not known
- quantize &= name.find("ssm_conv1d.weight") == std::string::npos;
+ quantize &= name.find("ssm_conv1d") == std::string::npos;
quantize &= name.find("shortconv.conv.weight") == std::string::npos;
// do not quantize RWKV's small yet 2D weights
--- /dev/null
+#include "llama-sampler.h"
+
+#include "llama-impl.h"
+#include "llama-vocab.h"
+#include "llama-grammar.h"
+
+#include "ggml-cpp.h"
+
+#include <array>
+#include <algorithm>
+#include <cassert>
+#include <cfloat>
+#include <chrono>
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <ctime>
+#include <numeric>
+#include <random>
+#include <unordered_map>
+#include <stdexcept>
+
+// the ring buffer works similarly to std::deque, but with a fixed capacity
+template<typename T>
+struct ring_buffer {
+ ring_buffer(size_t cap) : capacity(cap), data(cap) {}
+
+ T & front() {
+ if (sz == 0) {
+ throw std::runtime_error("ring buffer is empty");
+ }
+ return data[first];
+ }
+
+ const T & front() const {
+ if (sz == 0) {
+ throw std::runtime_error("ring buffer is empty");
+ }
+ return data[first];
+ }
+
+ T & back() {
+ if (sz == 0) {
+ throw std::runtime_error("ring buffer is empty");
+ }
+ return data[pos];
+ }
+
+ const T & back() const {
+ if (sz == 0) {
+ throw std::runtime_error("ring buffer is empty");
+ }
+ return data[pos];
+ }
+
+ void push_back(const T & value) {
+ if (capacity == 0) {
+ throw std::runtime_error("ring buffer: capacity is zero");
+ }
+
+ if (sz == capacity) {
+ // advance the start when buffer is full
+ first = (first + 1) % capacity;
+ } else {
+ sz++;
+ }
+ data[pos] = value;
+ pos = (pos + 1) % capacity;
+ }
+
+ T pop_front() {
+ if (sz == 0) {
+ throw std::runtime_error("ring buffer is empty");
+ }
+ T value = data[first];
+ first = (first + 1) % capacity;
+ sz--;
+ return value;
+ }
+
+ //T & operator[](size_t i) {
+ // if (i >= sz) {
+ // throw std::runtime_error("ring buffer: index out of bounds");
+ // }
+ // return data[(first + i) % capacity];
+ //}
+
+ //const T & at(size_t i) const {
+ // if (i >= sz) {
+ // throw std::runtime_error("ring buffer: index out of bounds");
+ // }
+ // return data[(first + i) % capacity];
+ //}
+
+ const T & rat(size_t i) const {
+ if (i >= sz) {
+ throw std::runtime_error("ring buffer: index out of bounds");
+ }
+ return data[(first + sz - i - 1) % capacity];
+ }
+
+ std::vector<T> to_vector() const {
+ std::vector<T> result;
+ result.reserve(sz);
+ for (size_t i = 0; i < sz; i++) {
+ result.push_back(data[(first + i) % capacity]);
+ }
+ return result;
+ }
+
+ void clear() {
+ // here only reset the status of the buffer
+ sz = 0;
+ first = 0;
+ pos = 0;
+ }
+
+ bool empty() const {
+ return sz == 0;
+ }
+
+ size_t size() const {
+ return sz;
+ }
+
+ size_t capacity = 0;
+ size_t sz = 0;
+ size_t first = 0;
+ size_t pos = 0;
+
+ std::vector<T> data;
+};
+
+// writes result in res, does not mutate cur
+static void llama_token_data_array_partial_sort(const llama_token_data_array & cur, int npartial, std::vector<llama_token_data> & res) {
+ static const auto comp = [](const llama_token_data & a, const llama_token_data & b) {
+ return a.logit > b.logit;
+ };
+
+ constexpr int nbuckets = 128;
+ constexpr float bucket_low = -10.0f;
+ constexpr float bucket_high = 10.0f;
+ constexpr float bucket_scale = nbuckets/(bucket_high - bucket_low);
+ constexpr float bucket_inter = -bucket_low * bucket_scale;
+
+ std::vector<int> bucket_idx;
+ std::vector<int> histo(nbuckets, 0);
+
+ std::vector<llama_token_data*> bucket_ptrs;
+
+ bucket_idx.reserve(cur.size);
+
+ for (int i = 0; i < (int)cur.size; ++i) {
+ const float val = cur.data[i].logit;
+ int ib = int(bucket_scale * val + bucket_inter); //nbuckets * (val - bucket_low) / (bucket_high - bucket_low);
+ ib = std::max(0, std::min(nbuckets - 1, ib));
+ bucket_idx.push_back(ib);
+ ++histo[ib];
+ }
+ int nhave = 0;
+ int ib = nbuckets - 1;
+ for ( ; ib >= 0; --ib) {
+ nhave += histo[ib];
+ if (nhave >= npartial) {
+ break;
+ }
+ }
+ res.resize(nhave);
+ auto * ptr = res.data();
+ bucket_ptrs.reserve(nbuckets - ib);
+ for (int j = nbuckets - 1; j >= ib; --j) {
+ bucket_ptrs.push_back(ptr);
+ ptr += histo[j];
+ }
+ for (int i = 0; i < (int)cur.size; ++i) {
+ int j = bucket_idx[i];
+ if (j >= ib) {
+ *bucket_ptrs[nbuckets - 1 - j]++ = cur.data[i];
+ }
+ }
+
+ ptr = res.data();
+ int ndone = 0;
+ for (int j = nbuckets - 1; j > ib; --j) {
+ std::sort(ptr, ptr + histo[j], comp);
+ ptr += histo[j];
+ ndone += histo[j];
+ }
+ std::partial_sort(ptr, ptr + npartial - ndone, ptr + histo[ib], comp);
+}
+
+// reduces the size of cur_p to npartial, keeping only the top npartial elements
+static void llama_token_data_array_partial_sort_inplace(llama_token_data_array * cur_p, int npartial) {
+ static const auto comp = [](const llama_token_data & a, const llama_token_data & b) {
+ return a.logit > b.logit;
+ };
+
+ if (npartial <= 128) {
+ std::partial_sort(cur_p->data, cur_p->data + npartial, cur_p->data + cur_p->size, comp);
+
+ cur_p->size = npartial;
+ cur_p->sorted = true;
+
+ return;
+ }
+
+ std::vector<llama_token_data> tmp;
+
+ llama_token_data_array_partial_sort(*cur_p, npartial, tmp);
+
+ std::copy(tmp.data(), tmp.data() + npartial, cur_p->data);
+
+ cur_p->size = npartial;
+ cur_p->sorted = true;
+}
+
+static int llama_sample_dist(llama_token_data_array * cur_p, std::mt19937 & rng) {
+ // iterator for the probabilities
+#ifdef __GNUC__
+ #pragma GCC diagnostic push
+ #pragma GCC diagnostic ignored "-Wunused-local-typedefs"
+#endif
+
+ struct probs_iterator {
+ typedef std::input_iterator_tag iterator_category;
+ typedef float value_type;
+ typedef float * pointer;
+ typedef float & reference;
+ typedef ptrdiff_t difference_type;
+
+ const llama_token_data * data;
+
+ bool operator==(const probs_iterator & other) const { return data == other.data; }
+ bool operator!=(const probs_iterator & other) const { return data != other.data; }
+ const float & operator*() const { return data->p; }
+ probs_iterator & operator++() { ++data; return *this; }
+ probs_iterator operator++(int) { probs_iterator tmp = *this; ++data; return tmp; }
+ };
+
+#ifdef __GNUC__
+ #pragma GCC diagnostic pop
+#endif
+
+ std::discrete_distribution<int> dist(probs_iterator{cur_p->data}, probs_iterator{cur_p->data + cur_p->size});
+
+ return dist(rng);
+}
+
+/*
+static void llama_log_softmax(float * array, size_t size) {
+ float max_l = *std::max_element(array, array + size);
+ float sum = 0.f;
+ for (size_t i = 0; i < size; ++i) {
+ float p = expf(array[i] - max_l);
+ sum += p;
+ array[i] = p;
+ }
+
+ for (size_t i = 0; i < size; ++i) {
+ array[i] = logf(array[i] / sum);
+ }
+}
+*/
+
+static void llama_sampler_temp_impl(llama_token_data_array * cur_p, float temp) {
+ if (temp <= 0.0f) {
+ // find the token with the highest logit and set the rest to -inf
+ size_t max_i = 0;
+ float max_l = cur_p->data[0].logit;
+
+ for (size_t i = 1; i < cur_p->size; ++i) {
+ if (cur_p->data[i ].logit > max_l) {
+ cur_p->data[max_i].logit = -INFINITY;
+ max_i = i;
+ max_l = cur_p->data[i].logit;
+ } else {
+ cur_p->data[i].logit = -INFINITY;
+ }
+ }
+
+ return;
+ }
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cur_p->data[i].logit /= temp;
+ }
+}
+
+static void llama_sampler_softmax_impl(llama_token_data_array * cur_p, bool do_sort) {
+ GGML_ASSERT(cur_p->size > 0);
+
+ // Sort the logits in descending order if requested
+ if (do_sort && !cur_p->sorted) {
+ llama_token_data_array_partial_sort_inplace(cur_p, cur_p->size);
+ }
+
+ float max_l = cur_p->data[0].logit;
+ if (!cur_p->sorted) {
+ for (size_t i = 1; i < cur_p->size; ++i) {
+ max_l = std::max(max_l, cur_p->data[i].logit);
+ }
+ }
+
+ float cum_sum = 0.0f;
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ float p = expf(cur_p->data[i].logit - max_l);
+ cur_p->data[i].p = p;
+ cum_sum += p;
+ }
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cur_p->data[i].p /= cum_sum;
+ }
+}
+
+static void llama_sampler_top_k_impl(llama_token_data_array * cur_p, int32_t k) {
+ // if (k >= (int32_t)cur_p->size) {
+ // return;
+ // }
+
+ if (k <= 0) {
+ return;
+ }
+
+ k = std::min(k, (int) cur_p->size);
+
+ // Sort scores in descending order
+ if (!cur_p->sorted) {
+ llama_token_data_array_partial_sort_inplace(cur_p, k);
+ }
+
+ cur_p->size = k;
+}
+
+static uint32_t get_rng_seed(uint32_t seed) {
+ if (seed == LLAMA_DEFAULT_SEED) {
+ // use system clock if std::random_device is not a true RNG
+ static bool is_rd_prng = std::random_device().entropy() == 0;
+ if (is_rd_prng) {
+ return (uint32_t) std::chrono::system_clock::now().time_since_epoch().count();
+ }
+ std::random_device rd;
+ return rd();
+ }
+ return seed;
+}
+
+// llama_sampler API
+
+struct llama_sampler * llama_sampler_init(
+ struct llama_sampler_i * iface,
+ llama_sampler_context_t ctx) {
+ return new llama_sampler {
+ /* .iface = */ iface,
+ /* .ctx = */ ctx,
+ };
+}
+
+const char * llama_sampler_name(const struct llama_sampler * smpl) {
+ if (!smpl->iface) {
+ return "(null)";
+ }
+
+ return smpl->iface->name(smpl);
+}
+
+void llama_sampler_accept(struct llama_sampler * smpl, llama_token token) {
+ if (!smpl) {
+ return;
+ }
+
+ if (smpl->iface->accept) {
+ smpl->iface->accept(smpl, token);
+ }
+}
+
+void llama_sampler_apply(struct llama_sampler * smpl, struct llama_token_data_array * cur_p) {
+ if (!smpl) {
+ return;
+ }
+
+ GGML_ASSERT(smpl->iface->apply);
+ smpl->iface->apply(smpl, cur_p);
+}
+
+void llama_sampler_reset(struct llama_sampler * smpl) {
+ if (!smpl) {
+ return;
+ }
+
+ if (smpl->iface->reset) {
+ smpl->iface->reset(smpl);
+ }
+}
+
+struct llama_sampler * llama_sampler_clone(const struct llama_sampler * smpl) {
+ if (!smpl) {
+ return nullptr;
+ }
+
+ if (smpl->iface->clone) {
+ return smpl->iface->clone(smpl);
+ }
+
+ if (smpl->ctx == nullptr) {
+ return llama_sampler_init(
+ /* .iface = */ smpl->iface,
+ /* .ctx = */ nullptr
+ );
+ }
+
+ GGML_ABORT("the sampler does not support cloning");
+}
+
+void llama_sampler_free(struct llama_sampler * smpl) {
+ if (smpl == nullptr) {
+ return;
+ }
+
+ if (smpl->iface->free) {
+ smpl->iface->free(smpl);
+ }
+
+ delete smpl;
+}
+
+// empty sampler
+
+struct llama_sampler_empty {
+ const char * name;
+};
+
+static struct llama_sampler * llama_sampler_init_empty(const char * name);
+
+static const char * llama_sampler_empty_name(const struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_empty *) smpl->ctx;
+ return ctx->name;
+}
+
+static void llama_sampler_empty_accept(struct llama_sampler * smpl, llama_token token) {
+ GGML_UNUSED(smpl);
+ GGML_UNUSED(token);
+}
+
+static void llama_sampler_empty_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ GGML_UNUSED(smpl);
+ GGML_UNUSED(cur_p);
+}
+
+static void llama_sampler_empty_reset(struct llama_sampler * smpl) {
+ GGML_UNUSED(smpl);
+}
+
+static struct llama_sampler * llama_sampler_empty_clone(const struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_empty *) smpl->ctx;
+ return llama_sampler_init_empty(ctx->name);
+}
+
+static void llama_sampler_empty_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_empty *) smpl->ctx;
+}
+
+static bool llama_sampler_empty_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ GGML_UNUSED(smpl);
+ GGML_UNUSED(buft);
+
+ return true;
+}
+
+static void llama_sampler_empty_backend_accept(
+ struct llama_sampler * smpl,
+ ggml_context * ctx,
+ ggml_cgraph * gf,
+ struct ggml_tensor * selected_token) {
+ GGML_UNUSED(smpl);
+ GGML_UNUSED(ctx);
+ GGML_UNUSED(gf);
+ GGML_UNUSED(selected_token);
+}
+
+static void llama_sampler_empty_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ GGML_UNUSED(smpl);
+ GGML_UNUSED(ctx);
+ GGML_UNUSED(gf);
+ GGML_UNUSED(data);
+}
+
+static void llama_sampler_empty_backend_set_input(struct llama_sampler * smpl) {
+ GGML_UNUSED(smpl);
+}
+
+static struct llama_sampler_i llama_sampler_empty_i = {
+ /* .name = */ llama_sampler_empty_name,
+ /* .accept = */ llama_sampler_empty_accept,
+ /* .apply = */ llama_sampler_empty_apply,
+ /* .reset = */ llama_sampler_empty_reset,
+ /* .clone = */ llama_sampler_empty_clone,
+ /* .free = */ llama_sampler_empty_free,
+ /* .backend_init = */ llama_sampler_empty_backend_init,
+ /* .backend_accept = */ llama_sampler_empty_backend_accept,
+ /* .backend_apply = */ llama_sampler_empty_backend_apply,
+ /* .backend_set_input = */ llama_sampler_empty_backend_set_input,
+};
+
+struct llama_sampler * llama_sampler_init_empty(const char * name) {
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_empty_i,
+ /* .ctx = */ new llama_sampler_empty {
+ /* .name = */ name,
+ }
+ );
+}
+
+// common backend sampler functionality
+//
+// +name : means that the sampler is support and will run on the backend
+// -name : means that a ggml operator is not supported by the backend
+//
+struct llama_sampler_backend {
+ llama_sampler_backend(const char * name) : name(name), name_ext(name), is_init(false), support(false) {}
+
+ const char * get_name() {
+ if (!is_init) {
+ return name.c_str();
+ }
+
+ if (support) {
+ name_ext = "+" + name;
+ } else {
+ name_ext = "-" + name;
+ }
+
+ return name_ext.c_str();
+ }
+
+ void init(bool support) {
+ GGML_ASSERT(this->is_init == false);
+
+ this->is_init = true;
+ this->support = support;
+ }
+
+private:
+ std::string name;
+ std::string name_ext;
+
+ bool is_init;
+ bool support;
+};
+
+// check if all ggml ops used by the sampler are supported by the backend
+static bool llama_sampler_backend_support(
+ llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * device = ggml_backend_buft_get_device(buft);
+ if (!device) {
+ // CPU backend always supported
+ return true;
+ }
+
+ ggml_init_params params = {
+ /*.mem_size =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
+ /*.mem_buffer =*/ NULL,
+ /*.no_alloc =*/ true,
+ };
+
+ ggml_context_ptr ctx_ptr { ggml_init(params) };
+ if (!ctx_ptr) {
+ throw std::runtime_error(format("failed to create ggml context"));
+ }
+
+ ggml_context * ctx = ctx_ptr.get();
+
+ const int64_t n = 1024*1024;
+
+ llama_sampler_data data = {
+ /*.logits = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n),
+ /*.probs = */ nullptr,
+ /*.sampled = */ nullptr,
+ /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
+ };
+
+ ggml_cgraph * gf = ggml_new_graph(ctx);
+
+ smpl->iface->backend_apply(smpl, ctx, gf, &data);
+
+ if (data.logits) {
+ ggml_build_forward_expand(gf, data.logits);
+ }
+
+ if (data.probs) {
+ ggml_build_forward_expand(gf, data.probs);
+ }
+
+ if (data.sampled) {
+ ggml_build_forward_expand(gf, data.sampled);
+ }
+
+ if (data.candidates) {
+ ggml_build_forward_expand(gf, data.candidates);
+ }
+
+ for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
+ struct ggml_tensor * op = ggml_graph_node(gf, i);
+
+ if (!ggml_backend_dev_supports_op(device, op)) {
+ LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",
+ __func__, ggml_backend_dev_name(device), ggml_op_name(op->op), smpl->iface->name(smpl));
+
+ return false;
+ }
+ }
+
+ return true;
+}
+
+// sampler chain
+
+static const char * llama_sampler_chain_name(const struct llama_sampler * /*smpl*/) {
+ return "chain";
+}
+
+static void llama_sampler_chain_accept(struct llama_sampler * smpl, llama_token token) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ time_meas tm(chain->t_sample_us, chain->params.no_perf);
+
+ for (auto & smpl : chain->samplers) {
+ llama_sampler_accept(smpl.ptr, token);
+ }
+
+ chain->n_sample++;
+}
+
+static void llama_sampler_chain_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ time_meas tm(chain->t_sample_us, chain->params.no_perf);
+
+ bool is_backend = chain->is_init;
+
+ for (auto & smpl : chain->samplers) {
+ if (is_backend && smpl.is_backend) {
+ continue;
+ }
+
+ is_backend = false;
+
+ if (smpl.ptr->iface->apply == nullptr) {
+ continue;
+ }
+
+ llama_sampler_apply(smpl.ptr, cur_p);
+ }
+}
+
+static void llama_sampler_chain_reset(struct llama_sampler * smpl) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ for (auto & smpl : chain->samplers) {
+ llama_sampler_reset(smpl.ptr);
+ }
+}
+
+static struct llama_sampler * llama_sampler_chain_clone(const struct llama_sampler * smpl) {
+ const auto * chain_src = (const llama_sampler_chain *) smpl->ctx;
+
+ auto * result = llama_sampler_chain_init(chain_src->params);
+
+ for (const auto & smpl : chain_src->samplers) {
+ llama_sampler_chain_add(result, llama_sampler_clone(smpl.ptr));
+ }
+
+ return result;
+}
+
+static void llama_sampler_chain_free(struct llama_sampler * smpl) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ for (auto & smpl : chain->samplers) {
+ llama_sampler_free(smpl.ptr);
+ }
+
+ delete chain;
+}
+
+static bool llama_sampler_chain_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");
+
+ chain->is_init = true;
+
+ bool res = true;
+
+ for (auto & smpl : chain->samplers) {
+ bool res_cur = true;
+
+ // to be able to run a sampler on the backend, it has to:
+ // - have the .backend_init() API implemented
+ // - return true during .backend_init()
+ if (smpl.ptr->iface->backend_init) {
+ if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) {
+ res_cur = false;
+ }
+ } else {
+ res_cur = false;
+ }
+
+ smpl.is_backend = res_cur;
+
+ res = res && res_cur;
+ }
+
+ return res;
+}
+
+static void llama_sampler_chain_backend_accept(
+ struct llama_sampler * smpl,
+ ggml_context * ctx,
+ ggml_cgraph * gf,
+ struct ggml_tensor * selected_token) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ for (auto & smpl : chain->samplers) {
+ if (!smpl.is_backend) {
+ break;
+ }
+
+ if (smpl.ptr->iface->backend_accept) {
+ smpl.ptr->iface->backend_accept(smpl.ptr, ctx, gf, selected_token);
+ }
+ }
+}
+
+static void llama_sampler_chain_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ GGML_ASSERT(chain->is_init && "llama_sampler_chain_backend_init() not called");
+
+ for (auto & smpl : chain->samplers) {
+ if (!smpl.is_backend) {
+ break;
+ }
+
+ if (smpl.ptr->iface->backend_apply) {
+ smpl.ptr->iface->backend_apply(smpl.ptr, ctx, gf, data);
+ }
+ }
+}
+
+static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+ for (auto & smpl : chain->samplers) {
+ if (!smpl.is_backend) {
+ break;
+ }
+
+ if (smpl.ptr->iface->backend_set_input) {
+ smpl.ptr->iface->backend_set_input(smpl.ptr);
+ }
+ }
+}
+
+static struct llama_sampler_i llama_sampler_chain_i = {
+ /* .name = */ llama_sampler_chain_name,
+ /* .accept = */ llama_sampler_chain_accept,
+ /* .apply = */ llama_sampler_chain_apply,
+ /* .reset = */ llama_sampler_chain_reset,
+ /* .clone = */ llama_sampler_chain_clone,
+ /* .free = */ llama_sampler_chain_free,
+ /* .backend_init = */ llama_sampler_chain_backend_init,
+ /* .backend_accept = */ llama_sampler_chain_backend_accept,
+ /* .backend_apply = */ llama_sampler_chain_backend_apply,
+ /* .backend_set_input = */ llama_sampler_chain_backend_set_input,
+};
+
+struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) {
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_chain_i,
+ /* .ctx = */ new llama_sampler_chain {
+ /* .params = */ params,
+ /* .is_init = */ false,
+ /* .samplers = */ {},
+ /* .cur = */ {},
+ /* .t_sample_us = */ 0,
+ /* .n_sample = */ 0,
+ }
+ );
+}
+
+llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {
+ const llama_token sampled_token = llama_get_sampled_token_ith (ctx, idx);
+ const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx);
+ const float * sampled_logits = llama_get_sampled_logits_ith (ctx, idx);
+ const llama_token * sampled_ids = llama_get_sampled_candidates_ith(ctx, idx);
+
+ // If a backend sampler has already sampled a token, return it.
+ if (sampled_token != LLAMA_TOKEN_NULL) {
+ LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx);
+ return sampled_token;
+ }
+
+ const llama_model * model = llama_get_model(ctx);
+ const llama_vocab * vocab = llama_model_get_vocab(model);
+
+ const int n_vocab = llama_vocab_n_tokens(vocab);
+
+ // use pre-allocated buffer from chain if available, otherwise allocate locally
+ std::vector<llama_token_data> * cur_ptr;
+ std::vector<llama_token_data> cur_local;
+
+ if (smpl->iface == &llama_sampler_chain_i) {
+ auto * chain = (llama_sampler_chain *) smpl->ctx;
+ cur_ptr = &chain->cur;
+ } else {
+ cur_ptr = &cur_local;
+ }
+
+ auto & cur = *cur_ptr;
+
+ if (sampled_probs) {
+ const uint32_t sampled_probs_count = llama_get_sampled_probs_count_ith(ctx, idx);
+ cur.resize(sampled_probs_count);
+ for (uint32_t i = 0; i < sampled_probs_count; ++i) {
+ cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], sampled_probs[i]};
+ }
+ } else if (sampled_logits) {
+ const uint32_t sampled_logits_count = llama_get_sampled_logits_count_ith(ctx, idx);
+ cur.resize(sampled_logits_count);
+ for (llama_token i = 0; i < (int)sampled_logits_count; i++) {
+ cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], 0.0f};
+ }
+ } else {
+ const auto * logits = llama_get_logits_ith(ctx, idx);
+ GGML_ASSERT(logits != nullptr);
+ cur.resize(n_vocab);
+ for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
+ cur[token_id] = llama_token_data{token_id, logits[token_id], 0.0f};
+ }
+ }
+
+ llama_token_data_array cur_p = {
+ /* .data = */ cur.data(),
+ /* .size = */ cur.size(),
+ /* .selected = */ -1,
+ /* .sorted = */ false,
+ };
+
+ llama_sampler_apply(smpl, &cur_p);
+
+ GGML_ASSERT(cur_p.selected >= 0 && cur_p.selected < (int32_t) cur_p.size);
+
+ auto token = cur_p.data[cur_p.selected].id;
+
+ llama_sampler_accept(smpl, token);
+
+ return token;
+}
+
+
+void llama_sampler_chain_add(struct llama_sampler * chain, struct llama_sampler * smpl) {
+ auto * p = (llama_sampler_chain *) chain->ctx;
+ p->samplers.push_back({
+ /* .is_backend = */ false,
+ /* .ptr = */ smpl,
+ });
+}
+
+struct llama_sampler * llama_sampler_chain_get(struct llama_sampler * chain, int32_t i) {
+ if (chain == nullptr) {
+ return nullptr;
+ }
+
+ if (chain->iface != &llama_sampler_chain_i) {
+ return nullptr;
+ }
+
+ if (i == -1) {
+ return chain;
+ }
+
+ const auto * p = (const llama_sampler_chain *) chain->ctx;
+
+ if (i < 0 || (size_t) i >= p->samplers.size()) {
+ return nullptr;
+ }
+
+ return p->samplers[i].ptr;
+}
+
+struct llama_sampler * llama_sampler_chain_remove(struct llama_sampler * chain, int32_t i) {
+ auto * p = (llama_sampler_chain *) chain->ctx;
+
+ if (i < 0 || (size_t) i >= p->samplers.size()) {
+ return nullptr;
+ }
+
+ auto * result = p->samplers[i].ptr;
+ p->samplers.erase(p->samplers.begin() + i);
+
+ return result;
+}
+
+int llama_sampler_chain_n(const struct llama_sampler * chain) {
+ const auto * p = (const llama_sampler_chain *) chain->ctx;
+
+ return p->samplers.size();
+}
+
+//
+// samplers
+//
+
+// greedy
+
+struct llama_sampler_greedy : public llama_sampler_backend {
+};
+
+static const char * llama_sampler_greedy_name(const struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_greedy *) smpl->ctx;
+ return sctx->get_name();
+}
+
+static void llama_sampler_greedy_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_greedy *) smpl->ctx;
+ GGML_UNUSED(ctx);
+}
+
+static struct llama_sampler * llama_sampler_greedy_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_greedy *) smpl->ctx;
+ auto * result = llama_sampler_init_greedy();
+
+ // copy the state
+ {
+ auto * result_ctx = (llama_sampler_greedy *) result->ctx;
+
+ GGML_UNUSED(ctx);
+ GGML_UNUSED(result_ctx);
+ }
+
+ return result;
+}
+
+static void llama_sampler_greedy_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_greedy *) smpl->ctx;
+}
+
+static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_token_data_array * cur_p) {
+ cur_p->selected = 0;
+ for (size_t i = 1; i < cur_p->size; ++i) {
+ if (cur_p->data[i].logit > cur_p->data[cur_p->selected].logit) {
+ cur_p->selected = i;
+ }
+ }
+}
+
+static bool llama_sampler_greedy_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_greedy *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_greedy_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ GGML_UNUSED(gf);
+ GGML_UNUSED(smpl);
+
+ struct ggml_tensor * curl = ggml_argmax(ctx, data->logits);
+ ggml_set_name(curl, "greedy_argmax");
+
+ data->sampled = curl;
+}
+
+static struct llama_sampler_i llama_sampler_greedy_i = {
+ /* .name = */ llama_sampler_greedy_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_greedy_apply,
+ /* .reset = */ llama_sampler_greedy_reset,
+ /* .clone = */ llama_sampler_greedy_clone,
+ /* .free = */ llama_sampler_greedy_free,
+ /* .backend_init = */ llama_sampler_greedy_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_greedy_backend_apply,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_greedy() {
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_greedy_i,
+ /* .ctx = */ new llama_sampler_greedy {
+ ("greedy"),
+ }
+ );
+}
+
+// dist
+
+struct llama_sampler_dist : public llama_sampler_backend {
+ const uint32_t seed;
+ uint32_t seed_cur;
+
+ std::mt19937 rng;
+
+ ggml_tensor * inp_uniform;
+};
+
+static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_dist *) smpl->ctx;
+ return sctx->get_name();
+}
+
+static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_dist *) smpl->ctx;
+
+ // edge cases
+ if (cur_p->size == 0) {
+ cur_p->selected = -1;
+ return;
+ }
+
+ cur_p->selected = 0;
+
+ if (cur_p->size == 1) {
+ cur_p->data[0].p = 1.0f;
+ return;
+ }
+
+ // max logit for numerical stability
+ float max_l = cur_p->data[0].logit;
+ if (!cur_p->sorted) {
+ for (size_t i = 1; i < cur_p->size; ++i) {
+ max_l = std::max(max_l, cur_p->data[i].logit);
+ }
+ }
+
+ // apply softmax to obtain the probabilities
+ double sum_cum = 0.0f;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ float p = expf(cur_p->data[i].logit - max_l);
+ cur_p->data[i].p = p;
+ sum_cum += p;
+ }
+
+#if 1
+ // sample from the obtained probabilities and normalize the probs in a single pass
+ // this is ~3x faster on Mac with full gpt-oss vocab than the version below
+ //
+ std::uniform_real_distribution<double> dist(0.0f, 1.0f);
+ const double rnd = dist(ctx->rng);
+
+ double sum_run = 0.0f;
+ const double sum_tgt = sum_cum*rnd;
+
+ bool found = false;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ if (!found) {
+ // accumulate probs until we reach the target sum
+ sum_run += cur_p->data[i].p;
+ if (sum_run >= sum_tgt) {
+ cur_p->selected = i;
+ found = true;
+ }
+ }
+
+ // normalize probs
+ cur_p->data[i].p /= sum_cum;
+ }
+
+ // fallback to the last token (don't think this can happen)
+ assert(found);
+ if (!found) {
+ cur_p->selected = cur_p->size - 1;
+ }
+#else
+ // for clarity, this is the same as above but does one pass for normalization and one extra pass for sampling
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cur_p->data[i].p /= sum_cum;
+ }
+
+ cur_p->selected = llama_sample_dist(cur_p, ctx->rng);
+#endif
+}
+
+static void llama_sampler_dist_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_dist *) smpl->ctx;
+ ctx->seed_cur = get_rng_seed(ctx->seed);
+ ctx->rng.seed(ctx->seed_cur);
+}
+
+static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_dist *) smpl->ctx;
+ auto * result = llama_sampler_init_dist(ctx->seed);
+
+ // copy the state
+ {
+ auto * result_ctx = (llama_sampler_dist *) result->ctx;
+
+ result_ctx->rng = ctx->rng;
+ }
+
+ return result;
+}
+
+static void llama_sampler_dist_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_dist *) smpl->ctx;
+}
+
+static bool llama_sampler_dist_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_dist *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_dist_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ GGML_UNUSED(gf);
+
+ auto * sctx = (llama_sampler_dist *) smpl->ctx;
+
+ sctx->inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+ ggml_set_name (sctx->inp_uniform, "uniform");
+ ggml_set_input(sctx->inp_uniform);
+
+ struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits);
+ ggml_set_name(probs, "dist_probs");
+
+ struct ggml_tensor * cumsum = ggml_cumsum(ctx, probs);
+ ggml_set_name(cumsum, "dist_cumsum");
+
+ // The uniform tensor has a random value and we subtract this tensor with
+ // the cumsum tensor (the uniform tensor will be broadcasted by ggml_sub).
+ // Recall that each entry in cumsum is the cumulative probability up to that
+ // index so values stay negative while the cumulative total is below the
+ // random value, and become zero/positive once the threshold is crossed.
+ struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform);
+ ggml_set_name(diff, "dist_cumsum");
+
+ // The ggml_step function produces a tensor where entries are 1 if the
+ // corresponding entry in diff is > 0, and 0 otherwise. So all values up to
+ // the index where the cumulative probability exceeds the random value are 0,
+ // and all entries after that are 1.
+ struct ggml_tensor * mask = ggml_step(ctx, diff);
+ ggml_set_name(mask, "dist_mask");
+
+ // Taking the sum of the mask gives us the sum of elements after the threshold
+ // we are interested in.
+ struct ggml_tensor * idxf = ggml_sum(ctx, mask);
+ ggml_set_name(idxf, "dist_index_f32");
+
+ // 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);
+ ggml_set_name(idx, "dist_index_i32");
+
+ // Map back to original vocab ids if a candidates tensor is available.
+ struct ggml_tensor * sampled_token = idx;
+ if (data->candidates != nullptr) {
+ struct ggml_tensor * candidates = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates));
+
+ sampled_token = ggml_get_rows(ctx, candidates, idx);
+ ggml_set_name(sampled_token, "dist_sampled_token");
+ }
+
+ data->sampled = sampled_token;
+ data->probs = probs;
+}
+
+static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_dist *) smpl->ctx;
+
+ GGML_ASSERT(sctx->inp_uniform != nullptr);
+
+ // We sample in double precision and cast to float to match rnd numbers of
+ // llama_dampler_dist which uses double precision (sampling from
+ // std::uniform_real_distribution<double> and
+ // std::uniform_real_distribution<float> with same rng will produce
+ // different sequences).
+ std::uniform_real_distribution<double> dist(0.0f, 1.0f);
+ const float rnd = dist(sctx->rng);
+
+ ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float));
+}
+
+static struct llama_sampler_i llama_sampler_dist_i = {
+ /* .name = */ llama_sampler_dist_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_dist_apply,
+ /* .reset = */ llama_sampler_dist_reset,
+ /* .clone = */ llama_sampler_dist_clone,
+ /* .free = */ llama_sampler_dist_free,
+ /* .backend_init = */ llama_sampler_dist_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_dist_backend_apply,
+ /* .backend_set_input = */ llama_sampler_dist_backend_set_input,
+};
+
+struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
+ auto seed_cur = get_rng_seed(seed);
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_dist_i,
+ /* .ctx = */ new llama_sampler_dist {
+ ("dist"),
+ /* .seed = */ seed,
+ /* .seed_cur = */ seed_cur,
+ /* .rng = */ std::mt19937(seed_cur),
+ /* .inp_uniform = */ nullptr,
+ }
+ );
+}
+
+// top-k
+
+struct llama_sampler_top_k : public llama_sampler_backend {
+ const int32_t k;
+};
+
+static const char * llama_sampler_top_k_name(const struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_top_k *) smpl->ctx;
+ return sctx->get_name();
+}
+
+static void llama_sampler_top_k_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_top_k *) smpl->ctx;
+ llama_sampler_top_k_impl(cur_p, ctx->k);
+}
+
+static struct llama_sampler * llama_sampler_top_k_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_top_k *) smpl->ctx;
+ return llama_sampler_init_top_k(ctx->k);
+}
+
+static void llama_sampler_top_k_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_top_k *) smpl->ctx;
+}
+
+static bool llama_sampler_top_k_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_top_k *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_top_k_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ auto * sctx = (llama_sampler_top_k *) smpl->ctx;
+
+ struct ggml_tensor * top_k = ggml_top_k(ctx, data->logits, sctx->k);
+ ggml_set_name(top_k, "top_k");
+
+ if (data->candidates) {
+ struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]);
+ data->candidates = ggml_get_rows(ctx, candidates_rows, top_k);
+ data->candidates = ggml_reshape_1d(ctx, data->candidates, sctx->k);
+ ggml_set_name(data->candidates, "top_k_candidates");
+ } else {
+ data->candidates = top_k;
+ }
+
+ struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
+ struct ggml_tensor * top_k_rows = ggml_get_rows(ctx, logits_rows, top_k);
+ data->logits = ggml_reshape_1d(ctx, top_k_rows, sctx->k);
+ ggml_set_name(top_k_rows, "top_k_rows");
+
+ GGML_UNUSED(gf);
+}
+
+static struct llama_sampler_i llama_sampler_top_k_i = {
+ /* .name = */ llama_sampler_top_k_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_top_k_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_top_k_clone,
+ /* .free = */ llama_sampler_top_k_free,
+ /* .backend_init = */ llama_sampler_top_k_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_top_k_backend_apply,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_top_k(int32_t k) {
+ const bool is_empty = (k <= 0);
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?top-k");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_top_k_i,
+ /* .ctx = */ new llama_sampler_top_k {
+ ("top-k"),
+ /* .k = */ k,
+ }
+ );
+}
+
+// top-p
+
+struct llama_sampler_top_p : public llama_sampler_backend {
+ const float p;
+ const size_t min_keep;
+
+ std::vector<llama_token_data> buf_sort;
+};
+
+static const char * llama_sampler_top_p_name(const struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_top_p *) smpl->ctx;
+ return sctx->get_name();
+}
+
+static void llama_sampler_top_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_top_p *) smpl->ctx;
+
+ if (ctx->p >= 1.0f) {
+ return;
+ }
+
+ llama_sampler_softmax_impl(cur_p, false);
+
+ size_t k = cur_p->size;
+ auto * pdata = cur_p->data;
+
+ auto & buf_sort = ctx->buf_sort;
+
+ // if not sorted, try adaptive top-k sorting
+ if (!cur_p->sorted && cur_p->size > 1024) {
+ k = std::min<size_t>(256, cur_p->size);
+ llama_token_data_array_partial_sort(*cur_p, k, buf_sort);
+ pdata = buf_sort.data();
+ } else if (!cur_p->sorted) {
+ // small candidates -> sort inplace
+ llama_token_data_array_partial_sort_inplace(cur_p, k);
+ }
+
+ // Compute the cumulative probabilities
+ float cum_sum = 0.0f;
+ size_t last_idx = cur_p->size;
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cum_sum += pdata[i].p;
+
+ // Check if the running sum is at least p or if we have kept at least min_keep tokens
+ // we set the last index to i+1 to indicate that the current iterate should be included in the set
+ if (cum_sum >= ctx->p && i + 1 >= ctx->min_keep) {
+ last_idx = i + 1;
+ break;
+ }
+
+ // we exceeded the current top-k heuristic -> increase k and continue
+ if (!cur_p->sorted && i == k - 1) {
+ k = cur_p->size;
+ llama_token_data_array_partial_sort(*cur_p, k, buf_sort);
+ pdata = buf_sort.data();
+ }
+ }
+
+ // Resize the output vector to keep only the top-p tokens
+ if (!cur_p->sorted) {
+ std::copy(buf_sort.data(), buf_sort.data() + last_idx, cur_p->data);
+ cur_p->sorted = true;
+ }
+
+ cur_p->size = last_idx;
+}
+
+static struct llama_sampler * llama_sampler_top_p_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_top_p *) smpl->ctx;
+ return llama_sampler_init_top_p(ctx->p, ctx->min_keep);
+}
+
+static void llama_sampler_top_p_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_top_p *) smpl->ctx;
+}
+
+static bool llama_sampler_top_p_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_top_p *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_top_p_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ auto * sctx = (llama_sampler_top_p *) smpl->ctx;
+
+ auto ggml_sort = [ctx](struct ggml_tensor * a, struct ggml_tensor * b) {
+ GGML_ASSERT(ggml_nrows(a) == 1);
+ struct ggml_tensor * a_reshaped = ggml_reshape_2d(ctx, a, 1, a->ne[0]);
+ struct ggml_tensor * a_sorted = ggml_get_rows(ctx, a_reshaped, b);
+ return ggml_reshape_1d(ctx, a_sorted, a->ne[0]);
+ };
+
+ // Get the sorted logits in descending order.
+ struct ggml_tensor * sorted_idx = ggml_argsort(ctx, data->logits, GGML_SORT_ORDER_DESC);
+ ggml_set_name(sorted_idx, "top_p_sorted_idx");
+
+ // Do the sorting via reshape + get_rows
+ struct ggml_tensor * sorted_logits = ggml_sort(data->logits, sorted_idx);
+ ggml_set_name(sorted_logits, "top_p_sorted_logits");
+
+ struct ggml_tensor * softmax = ggml_soft_max(ctx, sorted_logits);
+ ggml_set_name(softmax, "top_p_softmax");
+
+ // If candidates are provided, sort them as well. Otherwise, set sorted indices as candidates.
+ if (data->candidates) {
+ data->candidates = ggml_sort(data->candidates, sorted_idx);
+ } else {
+ data->candidates = sorted_idx;
+ }
+ ggml_set_name(data->candidates, "top_p_candidates");
+
+ // Compute Cumulative Distribution Function (CDF) by means of GGML_OP_CUMSUM.
+ struct ggml_tensor * cdf = ggml_cumsum(ctx, softmax);
+ ggml_set_name(cdf, "top_p_cdf");
+
+ // Invert CDF and add top-p value so that ggml_step yields 1 for values we want to keep
+ struct ggml_tensor * cdf_scaled = ggml_scale_bias(ctx, cdf, -1.0f, sctx->p);
+ ggml_set_name(cdf_scaled, "top_p_cdf_scaled");
+
+ struct ggml_tensor * mask = ggml_step(ctx, cdf_scaled);
+ ggml_set_name(mask, "top_p_mask");
+
+ // Taking the sum of the mask gives us the sum of elements after the threshold
+ // we are interested in.
+ struct ggml_tensor * idxf = ggml_sum(ctx, mask);
+ ggml_set_name(idxf, "top_p_index_f32");
+
+ // prevent out-of-bounds access
+ idxf = ggml_clamp(ctx, idxf, 0.0f, mask->ne[0] - 1);
+
+ // construct ones tensor to set the value in the mask
+ struct ggml_tensor * ones = ggml_scale_bias(ctx, idxf, 0.0f, 1.0f);
+ ggml_set_name(ones, "top_p_ones");
+
+ // Make top-p inclusive (i.e. return all values such that cum_sum/cdf >= p)
+ struct ggml_tensor * mask_reshaped = ggml_reshape_2d(ctx, mask, 1, mask->ne[0]);
+
+ mask_reshaped = ggml_set_rows(ctx, mask_reshaped, ones, ggml_cast(ctx, idxf, GGML_TYPE_I32));
+ mask = ggml_reshape_1d(ctx, mask_reshaped, mask->ne[0]);
+
+ // Apply -INFINITY bias for masked-out tokens
+ // log(1) = 0 (keep), log(0) = -INF (discard)
+ struct ggml_tensor * top_p_bias = ggml_log(ctx, mask);
+ ggml_set_name(top_p_bias, "top_p_bias");
+
+ data->logits = ggml_add(ctx, sorted_logits, top_p_bias);
+ ggml_set_name(data->logits, "top_p_logits");
+
+ GGML_UNUSED(gf);
+}
+
+static struct llama_sampler_i llama_sampler_top_p_i = {
+ /* .name = */ llama_sampler_top_p_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_top_p_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_top_p_clone,
+ /* .free = */ llama_sampler_top_p_free,
+ /* .backend_init = */ llama_sampler_top_p_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_top_p_backend_apply,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) {
+ const bool is_empty = p >= 1.0f;
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?top-p");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_top_p_i,
+ /* .ctx = */ new llama_sampler_top_p {
+ ("top-p"),
+ /* .p = */ p,
+ /* .min_keep = */ min_keep,
+ /* .buf_sort = */ {},
+ }
+ );
+}
+
+// min-p
+
+struct llama_sampler_min_p : public llama_sampler_backend {
+ const float p;
+ const size_t min_keep;
+};
+
+static const char * llama_sampler_min_p_name(const struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_min_p *) smpl->ctx;
+ return sctx->get_name();
+}
+
+static void llama_sampler_min_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_min_p *) smpl->ctx;
+
+ if (ctx->p <= 0.0f || !cur_p->size) {
+ return;
+ }
+
+ bool min_p_applied = false;
+
+ // if the cur_p aren't sorted, try the unsorted implementation first
+ if (!cur_p->sorted) {
+ std::vector<llama_token_data> filtered_tokens;
+
+ float max_logit = -FLT_MAX;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ max_logit = std::max(max_logit, cur_p->data[i].logit);
+ }
+ const float min_logit = max_logit + logf(ctx->p); // min logit for p_i >= p * p_max
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ if (cur_p->data[i].logit >= min_logit) {
+ filtered_tokens.push_back(cur_p->data[i]);
+ }
+ }
+
+ // if we have enough values the operation was a success
+ if (!filtered_tokens.empty() && filtered_tokens.size() >= ctx->min_keep) {
+ std::copy(filtered_tokens.begin(), filtered_tokens.end(), cur_p->data);
+ cur_p->size = filtered_tokens.size();
+ min_p_applied = true;
+ }
+ }
+
+ // if the cur_p are sorted or the unsorted implementation failed, use this implementation
+ if (!min_p_applied) {
+ // Sort the logits in descending order
+ if (!cur_p->sorted) {
+ llama_token_data_array_partial_sort_inplace(cur_p, cur_p->size);
+ }
+
+ const float min_logit = cur_p->data[0].logit + logf(ctx->p); // min logit for p_i >= p * p_max
+ size_t i = 1; // first token always matches
+
+ for (; i < cur_p->size; ++i) {
+ if (cur_p->data[i].logit < min_logit && i >= ctx->min_keep) {
+ break; // prob too small
+ }
+ }
+
+ // Resize the output vector to keep only the matching tokens
+ cur_p->size = i;
+ }
+}
+
+static struct llama_sampler * llama_sampler_min_p_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_min_p *) smpl->ctx;
+ return llama_sampler_init_min_p(ctx->p, ctx->min_keep);
+}
+
+static void llama_sampler_min_p_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_min_p *) smpl->ctx;
+}
+
+static bool llama_sampler_min_p_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_min_p *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_min_p_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ auto * sctx = (llama_sampler_min_p *) smpl->ctx;
+
+ struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits);
+ ggml_set_name(max_idx, "max_idx");
+
+ struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
+ ggml_set_name(logits_rows, "logits_rows");
+
+ struct ggml_tensor * max_logit = ggml_get_rows(ctx, logits_rows, max_idx);
+ ggml_set_name(max_logit, "max_logit");
+
+ // Calculate the threshold value.
+ struct ggml_tensor * threshold = ggml_scale_bias(ctx, max_logit, 1.0f, logf(sctx->p));
+ ggml_set_name(threshold, "min_p_threshold");
+
+ // Subtract the threshold from logits.
+ struct ggml_tensor * sub = ggml_sub(ctx, data->logits, threshold);
+
+ // Create a mask where logits below the threshold are 0 (discard),
+ // and others are 1 (keep).
+ struct ggml_tensor * mask = ggml_step(ctx, sub);
+ ggml_set_name(mask, "min_p_mask");
+
+ // Apply -INFINITY bias for masked-out tokens
+ // log(1) = 0 (keep), log(0) = -INF (discard)
+ struct ggml_tensor * min_p_bias = ggml_log(ctx, mask);
+ ggml_set_name(min_p_bias, "min_p_bias");
+
+ data->logits = ggml_add(ctx, data->logits, min_p_bias);
+ ggml_set_name(data->logits, "min_p_logits");
+
+ GGML_UNUSED(gf);
+}
+
+static struct llama_sampler_i llama_sampler_min_p_i = {
+ /* .name = */ llama_sampler_min_p_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_min_p_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_min_p_clone,
+ /* .free = */ llama_sampler_min_p_free,
+ /* .backend_init = */ llama_sampler_min_p_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_min_p_backend_apply,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) {
+ const bool is_empty = (p <= 0.0f);
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?min-p");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_min_p_i,
+ /* .ctx = */ new llama_sampler_min_p {
+ ("min-p"),
+ /* .p = */ p,
+ /* .min_keep = */ min_keep,
+ }
+ );
+}
+
+// typical
+
+struct llama_sampler_typical {
+ const float p;
+ const size_t min_keep;
+};
+
+static const char * llama_sampler_typical_name(const struct llama_sampler * /*smpl*/) {
+ return "typical";
+}
+
+static void llama_sampler_typical_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_typical *) smpl->ctx;
+
+ // Reference implementation:
+ // https://github.com/huggingface/transformers/compare/main...cimeister:typical-sampling:typical-pr
+ if (ctx->p >= 1.0f) {
+ return;
+ }
+
+ // Compute the softmax of logits and calculate entropy
+ llama_sampler_softmax_impl(cur_p, true);
+
+ float entropy = 0.0f;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ entropy += -cur_p->data[i].p * logf(cur_p->data[i].p);
+ }
+
+ // Compute the absolute difference between negative log probability and entropy for each candidate
+ std::vector<float> shifted_scores;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ float shifted_score = fabsf(-logf(cur_p->data[i].p) - entropy);
+ shifted_scores.push_back(shifted_score);
+ }
+
+ // Sort tokens based on the shifted_scores and their corresponding indices
+ std::vector<size_t> indices(cur_p->size);
+ std::iota(indices.begin(), indices.end(), 0);
+
+ std::sort(indices.begin(), indices.end(), [&](size_t a, size_t b) {
+ return shifted_scores[a] < shifted_scores[b];
+ });
+
+ // Compute the cumulative probabilities
+ float cum_sum = 0.0f;
+ size_t last_idx = indices.size();
+
+ for (size_t i = 0; i < indices.size(); ++i) {
+ size_t idx = indices[i];
+ cum_sum += cur_p->data[idx].p;
+
+ // Check if the running sum is greater than typical or if we have kept at least min_keep tokens
+ if (cum_sum > ctx->p && (ctx->min_keep == 0 || i >= ctx->min_keep - 1)) {
+ last_idx = i + 1;
+ break;
+ }
+ }
+
+ // Resize the output vector to keep only the locally typical tokens
+ std::vector<llama_token_data> cur_p_new;
+ for (size_t i = 0; i < last_idx; ++i) {
+ size_t idx = indices[i];
+ cur_p_new.push_back(cur_p->data[idx]);
+ }
+
+ // Replace the data in cur_p with the cur_p_new data
+ std::copy(cur_p_new.begin(), cur_p_new.end(), cur_p->data);
+ cur_p->size = cur_p_new.size();
+ cur_p->sorted = false;
+}
+
+static struct llama_sampler * llama_sampler_typical_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_typical *) smpl->ctx;
+ return llama_sampler_init_typical(ctx->p, ctx->min_keep);
+}
+
+static void llama_sampler_typical_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_typical *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_typical_i = {
+ /* .name = */ llama_sampler_typical_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_typical_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_typical_clone,
+ /* .free = */ llama_sampler_typical_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) {
+ const bool is_empty = (p >= 1.0f);
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?typical");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_typical_i,
+ /* .ctx = */ new llama_sampler_typical {
+ /* .p = */ p,
+ /* .min_keep = */ min_keep,
+ }
+ );
+}
+
+// temp
+
+struct llama_sampler_temp : public llama_sampler_backend {
+ const float temp;
+};
+
+static const char * llama_sampler_temp_name(const struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_temp *) smpl->ctx;
+ return sctx->get_name();
+}
+
+static void llama_sampler_temp_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ const auto * ctx = (llama_sampler_temp *) smpl->ctx;
+
+ llama_sampler_temp_impl(cur_p, ctx->temp);
+}
+
+static struct llama_sampler * llama_sampler_temp_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_temp *) smpl->ctx;
+ return llama_sampler_init_temp(ctx->temp);
+}
+
+static void llama_sampler_temp_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_temp *) smpl->ctx;
+}
+
+static void llama_sampler_backend_temp_sampling(
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data,
+ float temp) {
+ if (temp <= 0.0f) {
+ // Find the most probable token index.
+ struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits);
+ ggml_set_name(max_idx, "temp_max_idx");
+
+ if (data->candidates) {
+ struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]);
+ data->candidates = ggml_get_rows(ctx, candidates_rows, max_idx);
+ } else {
+ data->candidates = max_idx;
+ }
+
+ struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
+ data->logits = ggml_get_rows(ctx, logits_rows, max_idx);
+
+ return;
+ }
+
+ data->logits = ggml_scale(ctx, data->logits, 1.0f / temp);
+
+ GGML_UNUSED(gf);
+}
+
+static bool llama_sampler_temp_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_temp *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_temp_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ auto * sctx = (llama_sampler_temp *) smpl->ctx;
+ llama_sampler_backend_temp_sampling(ctx, gf, data, sctx->temp);
+}
+
+static struct llama_sampler_i llama_sampler_temp_i = {
+ /* .name = */ llama_sampler_temp_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_temp_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_temp_clone,
+ /* .free = */ llama_sampler_temp_free,
+ /* .backend_init = */ llama_sampler_temp_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_temp_backend_apply,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_temp(float temp) {
+ const bool is_empty = temp == 1.0f;
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?temp");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_temp_i,
+ /* .ctx = */ new llama_sampler_temp {
+ ("temp"),
+ /*.temp = */ temp,
+ }
+ );
+}
+
+// temp-ext
+
+struct llama_sampler_temp_ext : public llama_sampler_backend {
+ const float temp;
+ const float delta;
+ const float exponent;
+};
+
+static const char * llama_sampler_temp_ext_name(const struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
+ return sctx->get_name();
+}
+
+static void llama_sampler_temp_ext_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_temp_ext *) smpl->ctx;
+ if (ctx->delta > 0) {
+ const float min_temp = std::max(0.0f, ctx->temp - ctx->delta);
+ const float max_temp = ctx->temp + ctx->delta;
+
+ float exponent_val = ctx->exponent;
+
+ // no need to do anything if there is only one (or zero) candidates
+ if (cur_p->size <= 1) {
+ return;
+ }
+
+ // Calculate maximum possible entropy
+ float max_entropy = -logf(1.0f / cur_p->size);
+
+ llama_sampler_softmax_impl(cur_p, true);
+
+ // Calculate entropy of the softmax probabilities
+ float entropy = 0.0f;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ float prob = cur_p->data[i].p;
+ if (prob > 0.0f) { // Ensure no log(0)
+ entropy -= prob * logf(prob);
+ }
+ }
+
+ // Normalize the entropy (max_entropy cannot be 0 here because we checked cur_p->size != 1 above)
+ float normalized_entropy = entropy / max_entropy;
+
+ // Map the normalized entropy to the desired temperature range using the power function
+ float dyn_temp = min_temp + (max_temp - min_temp) * powf(normalized_entropy, exponent_val);
+
+ #ifdef DEBUG
+ LLAMA_LOG_INFO("Your text maxtemp value is: %f\n", max_temp);
+ LLAMA_LOG_INFO("Entropy: %f\n", entropy);
+ LLAMA_LOG_INFO("Max Possible Entropy: %f\n", max_entropy);
+ LLAMA_LOG_INFO("Normalized Entropy: %f\n", normalized_entropy);
+ LLAMA_LOG_INFO("Exponent: %f\n", exponent_val);
+ LLAMA_LOG_INFO("Dynamic Temperature (dyn_temp): %f\n", dyn_temp);
+ #endif
+
+ // Apply the dynamically calculated temperature scaling
+ llama_sampler_temp_impl(cur_p, dyn_temp);
+
+ // Re-compute softmax probabilities after scaling logits with dynamic temperature
+ const double max_l_double = cur_p->data[0].logit;
+
+ double cum_sum_double = 0.0;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ double p = exp(cur_p->data[i].logit - max_l_double);
+ cur_p->data[i].p = p; // Store the scaled probability
+ cum_sum_double += p;
+ }
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cur_p->data[i].p /= cum_sum_double; // Re-normalize the probabilities
+ }
+
+ #ifdef DEBUG
+ // Print the updated top 25 probabilities after temperature scaling
+ LLAMA_LOG_INFO("\nUpdated Top 25 Probabilities After Dynamic Temperature Scaling (in percentages):\n");
+ for (size_t i = 0; i < 25 && i < cur_p->size; ++i) {
+ LLAMA_LOG_INFO("Token %zu: %f%%\n", i + 1, cur_p->data[i].p * 100.0f);
+ }
+ #endif
+ } else {
+ llama_sampler_temp_impl(cur_p, ctx->temp);
+ }
+}
+
+static struct llama_sampler * llama_sampler_temp_ext_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_temp_ext *) smpl->ctx;
+ return llama_sampler_init_temp_ext(ctx->temp, ctx->delta, ctx->exponent);
+}
+
+static void llama_sampler_temp_ext_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_temp_ext *) smpl->ctx;
+}
+
+static bool llama_sampler_temp_ext_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_temp_ext_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
+
+ // Revert to standard temperature scaling if delta or temp are non-positive.
+ if (sctx->delta <= 0.0f || sctx->temp <= 0.0f) {
+ llama_sampler_backend_temp_sampling(ctx, gf, data, sctx->temp);
+ return;
+ }
+
+ // Calculate min_temp, max_temp, and max_entropy.
+ const float min_temp = std::max(0.0f, sctx->temp - sctx->delta);
+ const float max_temp = sctx->temp + sctx->delta;
+ const float max_entropy = logf(data->logits->ne[0]);
+
+ // Calculate the probabilities.
+ struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits);
+ ggml_set_name(probs, "temp_ext_softmax_probs");
+
+ // Clamp probabilities to avoid log(0) which would give -inf
+ struct ggml_tensor * probs_clamped = ggml_clamp(ctx, probs, 1e-10f, 1.0f);
+ ggml_set_name(probs_clamped, "temp_ext_probs_clamped");
+
+ // Calculate the entropy, entropy = -Σ(p * log(p)).
+ struct ggml_tensor * log_probs = ggml_log(ctx, probs_clamped);
+ struct ggml_tensor * p_log_p = ggml_mul(ctx, probs_clamped, log_probs);
+ struct ggml_tensor * sum_p_log_p = ggml_sum(ctx, p_log_p);
+ struct ggml_tensor * entropy = ggml_scale(ctx, sum_p_log_p, -1.0f);
+ ggml_set_name(log_probs, "temp_ext_log_probs");
+ ggml_set_name(p_log_p, "temp_ext_p_log_p");
+ ggml_set_name(sum_p_log_p, "temp_ext_sum_p_log_p");
+ ggml_set_name(entropy, "temp_ext_entropy");
+
+ // Normalize the entropy, norm_entropy = entropy / max_entropy
+ struct ggml_tensor * norm_entropy = ggml_scale(ctx, entropy, 1.0f / max_entropy);
+ ggml_set_name(norm_entropy, "temp_ext_norm_entropy");
+
+ // Calculate the dynamic temperature:
+ // dyn_temp = min_temp + (max_temp - min_temp) * powf(normalized_entropy, exponent);
+ //
+ // Calculate powf(normalized_entropy, exponent) as
+ // norm_entropy^exponent = exp(exponent * log(norm_entropy))
+ struct ggml_tensor * log_norm_entropy = ggml_log(ctx, norm_entropy);
+ struct ggml_tensor * scaled_log = ggml_scale(ctx, log_norm_entropy, sctx->exponent);
+ struct ggml_tensor * pow_entropy = ggml_exp(ctx, scaled_log);
+ // With pow_entropy computed we can now compute dyn_temp, scaling by
+ // (max_temp - min_temp) and then adding min_temp.
+ struct ggml_tensor * dyn_temp = ggml_scale_bias(ctx, pow_entropy, max_temp - min_temp, min_temp);
+ ggml_set_name(log_norm_entropy, "temp_ext_log_norm_entropy");
+ ggml_set_name(scaled_log, "temp_ext_scaled_log");
+ ggml_set_name(pow_entropy, "temp_ext_pow_entropy");
+ ggml_set_name(dyn_temp, "temp_ext_dyn_temp");
+
+ // Scale the logits by the dynamic temperature
+ struct ggml_tensor * scaled_logits = ggml_div(ctx, data->logits, dyn_temp);
+ ggml_set_name(scaled_logits, "temp_ext_scaled_logits");
+
+ data->logits = scaled_logits;
+}
+
+static struct llama_sampler_i llama_sampler_temp_ext_i = {
+ /* .name = */ llama_sampler_temp_ext_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_temp_ext_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_temp_ext_clone,
+ /* .free = */ llama_sampler_temp_ext_free,
+ /* .backend_init = */ llama_sampler_temp_ext_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_temp_ext_backend_apply,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) {
+ const bool is_empty = temp == 1.0f && delta <= 0.0f;
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?temp-ext");
+ }
+
+ auto * res = llama_sampler_init(
+ /* .iface = */ &llama_sampler_temp_ext_i,
+ /* .ctx = */ new llama_sampler_temp_ext {
+ ("temp-ext"),
+ /* .temp = */ temp,
+ /* .delta = */ delta,
+ /* .exponent = */ exponent,
+ }
+ );
+
+ return res;
+}
+
+// xtc
+
+struct llama_sampler_xtc {
+ const float probability;
+ const float threshold;
+ const size_t min_keep;
+
+ const uint32_t seed;
+ uint32_t seed_cur;
+
+ std::mt19937 rng;
+};
+
+static const char * llama_sampler_xtc_name(const struct llama_sampler * /*smpl*/) {
+ return "xtc";
+}
+
+static void llama_sample_xtc_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_xtc *) smpl->ctx;
+
+ if (ctx->probability <= 0.0f
+ || ctx->threshold > 0.5f
+ || cur_p->size < 2) {
+ return;
+ }
+
+ std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
+ float chance = distribution(ctx->rng);
+ if (chance > ctx->probability) {
+ return;
+ }
+
+ llama_sampler_softmax_impl(cur_p, true);
+
+ int pos_last = 0;
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ if (cur_p->data[i].p >= ctx->threshold) {
+ pos_last = i;
+ } else {
+ break;
+ }
+ }
+
+ if (cur_p->size - pos_last >= ctx->min_keep && pos_last > 0) {
+ cur_p->data += pos_last;
+ cur_p->size -= pos_last;
+ }
+}
+
+static struct llama_sampler * llama_sampler_xtc_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_xtc *) smpl->ctx;
+ auto * result = llama_sampler_init_xtc(ctx->probability, ctx->threshold, ctx->min_keep, ctx->seed);
+
+ // copy the state
+ {
+ auto * result_ctx = (llama_sampler_xtc *) result->ctx;
+
+ result_ctx->rng = ctx->rng;
+ }
+
+ return result;
+}
+
+static void llama_sampler_xtc_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_xtc *) smpl->ctx;
+}
+
+static void llama_sampler_xtc_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_xtc *) smpl->ctx;
+ ctx->seed_cur = get_rng_seed(ctx->seed);
+ ctx->rng.seed(ctx->seed_cur);
+}
+
+static struct llama_sampler_i llama_sampler_xtc_i = {
+ /* .name = */ llama_sampler_xtc_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sample_xtc_apply,
+ /* .reset = */ llama_sampler_xtc_reset,
+ /* .clone = */ llama_sampler_xtc_clone,
+ /* .free = */ llama_sampler_xtc_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) {
+ const bool is_empty = (p <= 0.0f || t > 0.5f);
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?xtc");
+ }
+
+ const auto seed_cur = get_rng_seed(seed);
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_xtc_i,
+ /* .ctx = */ new llama_sampler_xtc {
+ /* .probability = */ p,
+ /* .threshold = */ t,
+ /* .min_keep = */ min_keep,
+ /* .seed = */ seed,
+ /* .seed_cur = */ seed_cur,
+ /* .rng = */ std::mt19937(seed_cur),
+ }
+ );
+}
+
+// mirostat
+
+struct llama_sampler_mirostat {
+ const int32_t n_vocab;
+
+ const uint32_t seed;
+ uint32_t seed_cur;
+
+ const float tau;
+ const float eta;
+
+ const int32_t m;
+
+ float mu;
+
+ std::mt19937 rng;
+};
+
+static const char * llama_sampler_mirostat_name(const struct llama_sampler * /*smpl*/) {
+ return "mirostat";
+}
+
+static void llama_sampler_mirostat_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_mirostat *) smpl->ctx;
+
+ llama_sampler_softmax_impl(cur_p, true);
+
+ // Estimate s_hat using the most probable m tokens
+ float s_hat = 0.0;
+ float sum_ti_bi = 0.0;
+ float sum_ti_sq = 0.0;
+ for (size_t i = 0; i < size_t(ctx->m - 1) && i < cur_p->size - 1; ++i) {
+ float t_i = logf(float(i + 2) / float(i + 1));
+ float b_i = logf(cur_p->data[i].p / cur_p->data[i + 1].p);
+ sum_ti_bi += t_i * b_i;
+ sum_ti_sq += t_i * t_i;
+ }
+ s_hat = sum_ti_bi / sum_ti_sq;
+
+ // Compute k from the estimated s_hat and target surprise value
+ float epsilon_hat = s_hat - 1;
+ float k = powf((epsilon_hat * powf(2, ctx->mu)) / (1 - powf(ctx->n_vocab, -epsilon_hat)), 1 / s_hat);
+
+ llama_sampler_top_k_impl(cur_p, std::max(int(k), 1));
+
+ llama_sampler_softmax_impl(cur_p, true);
+
+ const int idx = llama_sample_dist(cur_p, ctx->rng);
+
+ cur_p->selected = idx;
+
+ float observed_surprise = -log2f(cur_p->data[idx].p);
+ float e = observed_surprise - ctx->tau;
+
+ // Update mu using the learning rate and error
+ ctx->mu = ctx->mu - ctx->eta * e;
+}
+
+static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_mirostat *) smpl->ctx;
+ auto * result = llama_sampler_init_mirostat(ctx->n_vocab, ctx->seed, ctx->tau, ctx->eta, ctx->m);
+
+ // copy the state
+ {
+ auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx;
+
+ result_ctx->mu = ctx->mu;
+ result_ctx->rng = ctx->rng;
+ }
+
+ return result;
+}
+
+static void llama_sampler_mirostat_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_mirostat *) smpl->ctx;
+ ctx->mu = 2.0f*ctx->tau;
+ ctx->seed_cur = get_rng_seed(ctx->seed);
+ ctx->rng.seed(ctx->seed_cur);
+}
+
+static void llama_sampler_mirostat_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_mirostat *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_mirostat_i = {
+ /* .name = */ llama_sampler_mirostat_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_mirostat_apply,
+ /* .reset = */ llama_sampler_mirostat_reset,
+ /* .clone = */ llama_sampler_mirostat_clone,
+ /* .free = */ llama_sampler_mirostat_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) {
+ const auto seed_cur = get_rng_seed(seed);
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_mirostat_i,
+ /* .ctx = */ new llama_sampler_mirostat {
+ /* .n_vocab = */ n_vocab,
+ /* .seed = */ seed,
+ /* .seed_cur = */ seed_cur,
+ /* .tau = */ tau,
+ /* .eta = */ eta,
+ /* .m = */ m,
+ /* .mu = */ 2.0f*tau,
+ /* .rng = */ std::mt19937(seed_cur),
+ }
+ );
+}
+
+// mirostat v2
+
+struct llama_sampler_mirostat_v2 {
+ const uint32_t seed;
+ uint32_t seed_cur;
+
+ const float tau;
+ const float eta;
+
+ float mu;
+
+ std::mt19937 rng;
+};
+
+static const char * llama_sampler_mirostat_v2_name(const struct llama_sampler * /*smpl*/) {
+ return "mirostat-v2";
+}
+
+static void llama_sampler_mirostat_v2_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_mirostat_v2 *) smpl->ctx;
+
+ llama_sampler_softmax_impl(cur_p, true);
+
+ // Truncate the words with surprise values greater than mu
+ cur_p->size = std::distance(cur_p->data, std::find_if(cur_p->data, cur_p->data + cur_p->size, [&](const llama_token_data & candidate) {
+ return -log2f(candidate.p) > ctx->mu;
+ }));
+
+ if (cur_p->size == 0) {
+ cur_p->size = 1;
+ }
+
+ // Normalize the probabilities of the remaining words
+ llama_sampler_softmax_impl(cur_p, true);
+
+ const int idx = llama_sample_dist(cur_p, ctx->rng);
+
+ cur_p->selected = idx;
+
+ float observed_surprise = -log2f(cur_p->data[idx].p);
+ float e = observed_surprise - ctx->tau;
+
+ // Update mu using the learning rate and error
+ ctx->mu = ctx->mu - ctx->eta * e;
+}
+
+static void llama_sampler_mirostat_v2_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_mirostat_v2 *) smpl->ctx;
+ ctx->mu = 2.0f*ctx->tau;
+ ctx->seed_cur = get_rng_seed(ctx->seed);
+ ctx->rng.seed(ctx->seed_cur);
+}
+
+static struct llama_sampler * llama_sampler_mirostat_v2_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_mirostat_v2 *) smpl->ctx;
+
+ auto * result = llama_sampler_init_mirostat_v2(ctx->seed, ctx->tau, ctx->eta);
+
+ // copy the state
+ {
+ auto * result_ctx = (llama_sampler_mirostat_v2 *) result->ctx;
+
+ result_ctx->mu = ctx->mu;
+ result_ctx->rng = ctx->rng;
+ }
+
+ return result;
+}
+
+static void llama_sampler_mirostat_v2_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_mirostat_v2 *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_mirostat_v2_i = {
+ /* .name = */ llama_sampler_mirostat_v2_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_mirostat_v2_apply,
+ /* .reset = */ llama_sampler_mirostat_v2_reset,
+ /* .clone = */ llama_sampler_mirostat_v2_clone,
+ /* .free = */ llama_sampler_mirostat_v2_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) {
+ auto seed_cur = get_rng_seed(seed);
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_mirostat_v2_i,
+ /* .ctx = */ new llama_sampler_mirostat_v2 {
+ /* .seed = */ seed,
+ /* .seed_cur = */ seed_cur,
+ /* .tau = */ tau,
+ /* .eta = */ eta,
+ /* .mu = */ 2.0f*tau,
+ /* .rng = */ std::mt19937(seed_cur),
+ }
+ );
+}
+
+// grammar
+
+struct llama_sampler_grammar {
+ const struct llama_vocab * vocab;
+
+ std::string grammar_str;
+ std::string grammar_root;
+
+ struct llama_grammar * grammar;
+};
+
+static const char * llama_sampler_grammar_name(const struct llama_sampler * /*smpl*/) {
+ return "grammar";
+}
+
+static void llama_sampler_grammar_accept_impl(struct llama_sampler * smpl, llama_token token) {
+ auto * ctx = (llama_sampler_grammar *) smpl->ctx;
+ if (ctx->grammar) {
+ llama_grammar_accept_impl(*ctx->grammar, token);
+ }
+}
+
+static void llama_sampler_grammar_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_grammar *) smpl->ctx;
+ if (ctx->grammar) {
+ llama_grammar_apply_impl(*ctx->grammar, cur_p);
+ }
+}
+
+// Fwd declare to break reset --> init_impl --> llama_sampler_grammar_i --> reset cycle.
+static struct llama_sampler * llama_sampler_init_grammar_impl(
+ const struct llama_vocab * vocab,
+ const char * grammar_str,
+ const char * grammar_root,
+ bool lazy,
+ const char ** trigger_words,
+ size_t num_trigger_words,
+ const llama_token * trigger_tokens,
+ size_t num_trigger_tokens,
+ const char ** trigger_patterns,
+ size_t num_trigger_patterns);
+
+static void llama_sampler_grammar_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_grammar *) smpl->ctx;
+ if (!ctx->grammar) {
+ return;
+ }
+
+ std::vector<const char *> trigger_patterns_c;
+ trigger_patterns_c.reserve(ctx->grammar->trigger_patterns.size());
+ for (auto & trigger_pattern : ctx->grammar->trigger_patterns) {
+ trigger_patterns_c.push_back(trigger_pattern.pattern.c_str());
+ }
+
+ auto * grammar_new = llama_grammar_init_impl(ctx->grammar->vocab, ctx->grammar_str.c_str(), ctx->grammar_root.c_str(),
+ ctx->grammar->lazy, trigger_patterns_c.data(), trigger_patterns_c.size(),
+ ctx->grammar->trigger_tokens.data(), ctx->grammar->trigger_tokens.size());
+
+ llama_grammar_free_impl(ctx->grammar);
+ ctx->grammar = grammar_new;
+}
+
+static struct llama_sampler * llama_sampler_grammar_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_grammar *) smpl->ctx;
+
+ auto * result = llama_sampler_init_grammar_impl(ctx->vocab, nullptr, nullptr, false, nullptr, 0, nullptr, 0, nullptr, 0);
+ GGML_ASSERT(result);
+
+ // copy the state
+ {
+ auto * result_ctx = (llama_sampler_grammar *) result->ctx;
+
+ if (ctx->grammar) {
+ result_ctx->grammar_str = ctx->grammar_str;
+ result_ctx->grammar_root = ctx->grammar_root;
+
+ result_ctx->grammar = llama_grammar_clone_impl(*ctx->grammar);
+ }
+ }
+
+ return result;
+}
+
+static void llama_sampler_grammar_free(struct llama_sampler * smpl) {
+ const auto * ctx = (llama_sampler_grammar *) smpl->ctx;
+
+ if (ctx->grammar) {
+ llama_grammar_free_impl(ctx->grammar);
+ }
+
+ delete ctx;
+}
+
+static struct llama_sampler_i llama_sampler_grammar_i = {
+ /* .name = */ llama_sampler_grammar_name,
+ /* .accept = */ llama_sampler_grammar_accept_impl,
+ /* .apply = */ llama_sampler_grammar_apply,
+ /* .reset = */ llama_sampler_grammar_reset,
+ /* .clone = */ llama_sampler_grammar_clone,
+ /* .free = */ llama_sampler_grammar_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+static struct llama_sampler * llama_sampler_init_grammar_impl(
+ const struct llama_vocab * vocab,
+ const char * grammar_str,
+ const char * grammar_root,
+ bool lazy,
+ const char ** trigger_words,
+ size_t num_trigger_words,
+ const llama_token * trigger_tokens,
+ size_t num_trigger_tokens,
+ const char ** trigger_patterns,
+ size_t num_trigger_patterns) {
+ auto * ctx = new llama_sampler_grammar;
+
+ if (grammar_str != nullptr && grammar_str[0] != '\0') {
+ std::string trigger_pattern;
+ llama_grammar * grammar = nullptr;
+ // TODO: remove trigger_words support.
+ if (trigger_words != nullptr && num_trigger_words > 0) {
+ GGML_ASSERT(trigger_patterns == nullptr && num_trigger_patterns == 0);
+ trigger_pattern = "[\\s\\S]*?(";
+ for (size_t i = 0; i < num_trigger_words; ++i) {
+ static const std::regex special_chars("[.^$|()*+?\\[\\]{}\\\\]");
+ if (i > 0) {
+ trigger_pattern += "|";
+ }
+ trigger_pattern += std::regex_replace(trigger_words[i], special_chars, "\\$0");
+ }
+ trigger_pattern += ")[\\s\\S]*";
+
+ std::array<const char *, 1> tmp_trigger_patterns = { trigger_pattern.c_str() };
+ grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, tmp_trigger_patterns.data(), tmp_trigger_patterns.size(), trigger_tokens, num_trigger_tokens);
+ } else {
+ grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, trigger_patterns, num_trigger_patterns, trigger_tokens, num_trigger_tokens);
+ }
+ *ctx = {
+ /* .vocab = */ vocab,
+ /* .grammar_str = */ grammar_str,
+ /* .grammar_root = */ grammar_root,
+ /* .grammar = */ grammar,
+ };
+ if (!ctx->grammar) {
+ delete ctx;
+ return nullptr;
+ }
+ } else {
+ *ctx = {
+ /* .vocab = */ vocab,
+ /* .grammar_str = */ {},
+ /* .grammar_root = */ {},
+ /* .grammar = */ nullptr,
+ };
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_grammar_i,
+ /* .ctx = */ ctx
+ );
+}
+
+struct llama_sampler * llama_sampler_init_grammar(
+ const struct llama_vocab * vocab,
+ const char * grammar_str,
+ const char * grammar_root) {
+ return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ false, nullptr, 0, nullptr, 0, nullptr, 0);
+}
+
+struct llama_sampler * llama_sampler_init_grammar_lazy(
+ const struct llama_vocab * vocab,
+ const char * grammar_str,
+ const char * grammar_root,
+ const char ** trigger_words,
+ size_t num_trigger_words,
+ const llama_token * trigger_tokens,
+ size_t num_trigger_tokens) {
+ return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ true, trigger_words, num_trigger_words, trigger_tokens, num_trigger_tokens, nullptr, 0);
+}
+
+struct llama_sampler * llama_sampler_init_grammar_lazy_patterns(
+ const struct llama_vocab * vocab,
+ const char * grammar_str,
+ const char * grammar_root,
+ const char ** trigger_patterns,
+ size_t num_trigger_patterns,
+ const llama_token * trigger_tokens,
+ size_t num_trigger_tokens) {
+ return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ true, nullptr, 0, trigger_tokens, num_trigger_tokens, trigger_patterns, num_trigger_patterns);
+}
+
+// penalties
+
+struct llama_sampler_penalties {
+ const int32_t penalty_last_n;
+ const float penalty_repeat;
+ const float penalty_freq;
+ const float penalty_present;
+
+ ring_buffer<llama_token> prev;
+
+ // a frequency map to count token occurrences
+ std::unordered_map<llama_token, int> token_count;
+};
+
+static const char * llama_sampler_penalties_name(const struct llama_sampler * /*smpl*/) {
+ return "penalties";
+}
+
+static void llama_sampler_penalties_accept(struct llama_sampler * smpl, llama_token token) {
+ auto * ctx = (llama_sampler_penalties *) smpl->ctx;
+ if (ctx->penalty_last_n == 0) {
+ return;
+ }
+
+ ctx->token_count[token]++;
+
+ // if the ring buffer is full, remove the oldest token
+ if (ctx->prev.size() >= (size_t) ctx->penalty_last_n) {
+ const auto old = ctx->prev.front();
+
+ ctx->token_count[old]--;
+ if (ctx->token_count[old] == 0) {
+ ctx->token_count.erase(old);
+ }
+ }
+
+ ctx->prev.push_back(token);
+
+#if 0
+ // sanity check
+ std::unordered_map<llama_token, int> tmp;
+ for (int i = 0; i < std::min<int>(ctx->penalty_last_n, ctx->prev.size()); ++i) {
+ tmp[ctx->prev.rat(i)]++;
+ }
+
+ assert(ctx->token_count == tmp);
+#endif
+}
+
+static void llama_sampler_penalties_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_penalties *) smpl->ctx;
+
+ if ((ctx->penalty_last_n == 0) ||
+ (ctx->penalty_repeat == 1.0f && ctx->penalty_freq == 0.0f && ctx->penalty_present == 0.0f)) {
+ return;
+ }
+
+ // Apply frequency and presence penalties to the cur_p
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ const auto token_iter = ctx->token_count.find(cur_p->data[i].id);
+ if (token_iter == ctx->token_count.end()) {
+ continue;
+ }
+
+ const int count = token_iter->second;
+
+ assert(count > 0 && count <= ctx->penalty_last_n);
+
+ // The academic publication that described this technique actually just only divided, but that would cause tokens with negative logits to become more likely, which is obviously wrong.
+ // This is common fix for this problem, which is to multiply by the penalty instead of dividing.
+ if (cur_p->data[i].logit <= 0) {
+ cur_p->data[i].logit *= ctx->penalty_repeat;
+ } else {
+ cur_p->data[i].logit /= ctx->penalty_repeat;
+ }
+
+ cur_p->data[i].logit -= float(count) * ctx->penalty_freq + float(count > 0) * ctx->penalty_present;
+ }
+
+ cur_p->sorted = false;
+}
+
+static void llama_sampler_penalties_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_penalties *) smpl->ctx;
+ ctx->prev.clear();
+ ctx->token_count.clear();
+}
+
+static struct llama_sampler * llama_sampler_penalties_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_penalties *) smpl->ctx;
+ auto * result = llama_sampler_init_penalties(
+ ctx->penalty_last_n,
+ ctx->penalty_repeat,
+ ctx->penalty_freq,
+ ctx->penalty_present);
+
+ // copy the state
+ {
+ auto * result_ctx = (llama_sampler_penalties *) result->ctx;
+
+ result_ctx->prev = ctx->prev;
+ }
+
+ return result;
+}
+
+static void llama_sampler_penalties_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_penalties *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_penalties_i = {
+ /* .name = */ llama_sampler_penalties_name,
+ /* .accept = */ llama_sampler_penalties_accept,
+ /* .apply = */ llama_sampler_penalties_apply,
+ /* .reset = */ llama_sampler_penalties_reset,
+ /* .clone = */ llama_sampler_penalties_clone,
+ /* .free = */ llama_sampler_penalties_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_penalties(
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present) {
+ penalty_last_n = std::max(penalty_last_n, 0);
+
+ const bool is_empty = (penalty_last_n == 0 || (penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f));
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?penalties");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_penalties_i,
+ /* .ctx = */ new llama_sampler_penalties {
+ /* .penalty_last_n = */ penalty_last_n,
+ /* .penalty_repeat = */ penalty_repeat,
+ /* .penalty_freq = */ penalty_freq,
+ /* .penalty_present = */ penalty_present,
+ /* .prev = */ ring_buffer<llama_token>(penalty_last_n),
+ /* .token_count = */ {},
+ }
+ );
+}
+
+// top-n-sigma
+
+struct llama_sampler_top_n_sigma {
+ const float n;
+};
+
+static const char * llama_sampler_top_n_sigma_name(const struct llama_sampler * /*smpl*/) {
+ return "top-n-sigma";
+}
+
+static void llama_sampler_top_n_sigma_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_top_n_sigma *) smpl->ctx;
+
+ if (ctx->n <= 0.0f || cur_p->size <= 1) {
+ return;
+ }
+
+ // find max logit and calculate mean
+ float max = cur_p->data[0].logit;
+ float logits_sum = 0;
+ size_t valid_count = 0;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ // Only count non-negative infinity values
+ if (cur_p->data[i].logit != -INFINITY) {
+ max = std::max(max, cur_p->data[i].logit);
+ logits_sum += cur_p->data[i].logit;
+ valid_count++;
+ }
+ }
+ float mean = valid_count > 0 ? logits_sum/valid_count : 0;
+
+ // calculate standard deviation
+ float acc = 0;
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ // Skip -infinity in std calculation
+ if (cur_p->data[i].logit != -INFINITY) {
+ acc += pow(cur_p->data[i].logit - mean, 2);
+ }
+ }
+ float std = valid_count > 0 ? sqrt(acc/valid_count) : 0;
+
+ // apply mask
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ if (cur_p->data[i].logit < max - (ctx->n * std)) {
+ cur_p->data[i].logit = -INFINITY;
+ }
+ }
+
+ llama_sampler_softmax_impl(cur_p, true);
+}
+
+static struct llama_sampler * llama_sampler_top_n_sigma_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_top_n_sigma *) smpl->ctx;
+ return llama_sampler_init_top_n_sigma(ctx->n);
+}
+
+static void llama_sampler_top_n_sigma_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_top_n_sigma *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_top_n_sigma_i = {
+ /* .name = */ llama_sampler_top_n_sigma_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_top_n_sigma_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_top_n_sigma_clone,
+ /* .free = */ llama_sampler_top_n_sigma_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
+ const bool is_empty = (n <= 0.0f);
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?top-n-sigma");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_top_n_sigma_i,
+ /* .ctx = */ new llama_sampler_top_n_sigma {
+ /* .n = */ n,
+ }
+ );
+}
+
+// DRY
+
+struct llama_sampler_dry {
+ int32_t total_context_size;
+
+ const float dry_multiplier;
+ const float dry_base;
+ const int32_t dry_allowed_length;
+ const int32_t dry_penalty_last_n;
+
+ std::unordered_multimap<llama_token, std::vector<llama_token>> dry_processed_breakers;
+ std::vector<int> dry_repeat_count;
+ std::unordered_map<llama_token, int> dry_max_token_repeat;
+ ring_buffer<llama_token> last_tokens;
+};
+
+// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
+static void get_overlapping_token_sequences(const llama_vocab & vocab, const std::string& str, std::unordered_multimap<llama_token, std::vector<llama_token>>& token_sequences, int max_tail_len = -1) {
+ for (llama_token token_id = 0; token_id < (llama_token) vocab.n_tokens(); token_id++) {
+ std::string word = vocab.detokenize({token_id}, true);
+ if (word.find(str) != std::string::npos) {
+ token_sequences.emplace(token_id, std::vector<llama_token>());
+ } else {
+ size_t word_len = word.size();
+ size_t str_len = str.size();
+ size_t pos = -1;
+ while ((pos = word.find(str[0], pos + 1)) != std::string::npos) {
+ bool match = true;
+ size_t i;
+ for (i = 1; i < str_len && i + pos < word_len; ++i) {
+ if (word[pos + i] != str[i]) {
+ match = false;
+ break;
+ }
+ }
+ if (match) {
+ std::vector<llama_token> tokenization = vocab.tokenize(str.substr(i), false, false);
+ if (max_tail_len >= 0 && tokenization.size() > (size_t)max_tail_len) {
+ tokenization.resize(max_tail_len);
+ }
+
+ // Ensure we don't already have a duplicate matching tokenization
+ auto its = token_sequences.equal_range(token_id);
+ bool found = false;
+ for (auto it = its.first; it != its.second; ++it) {
+ if (tokenization == it->second) {
+ found = true;
+ break;
+ }
+ }
+ if (!found) {
+ token_sequences.emplace(token_id, tokenization);
+ }
+ }
+ }
+ }
+ }
+}
+
+static const char * llama_sampler_dry_name(const struct llama_sampler * /*smpl*/) {
+ return "dry";
+}
+
+static void llama_sampler_dry_accept(struct llama_sampler * smpl, llama_token token) {
+ auto * ctx = (llama_sampler_dry *) smpl->ctx;
+ if (ctx->dry_multiplier == 0.0f || ctx->dry_base < 1.0f || ctx->dry_penalty_last_n == 0) {
+ return;
+ }
+
+ ctx->last_tokens.push_back(token);
+}
+
+// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
+static void llama_sampler_dry_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_dry *) smpl->ctx;
+
+ if (ctx->dry_multiplier == 0.0f || ctx->dry_base < 1.0f || ctx->dry_penalty_last_n == 0) {
+ return;
+ }
+
+ int32_t effective_dry_penalty_last_n = (ctx->dry_penalty_last_n == -1) ? ctx->total_context_size : std::max(ctx->dry_penalty_last_n, 0);
+ int last_n_repeat = std::min(std::min((int)ctx->last_tokens.size(), effective_dry_penalty_last_n), ctx->total_context_size);
+
+ if (last_n_repeat <= ctx->dry_allowed_length) {
+ return;
+ }
+
+ ctx->dry_repeat_count.assign(last_n_repeat, 0);
+ ctx->dry_max_token_repeat.clear();
+
+ // Step 1: Look for restart sequences to limit the maximum repetition length.
+ // Work backwards through the context looking for any token that begins a restart sequence.
+ //
+ // The collection `restart_sequences` is a mapping from a "head" token to all "tail"
+ // sequences that together comprise a restart sequence. This allows us to quickly check
+ // whether each token is the head of a complete sequence. Most restart sequences are actually
+ // a single token, and for these the "tail" is an empty vector.
+ //
+ // If the token is a "head", test all restart sequences that begin with this token
+ // (there will often only be one sequence for each token, but if sequences like 'aaaq1' and
+ // 'aaa1' are used as restart strings, both could start with 'aaa' when tokenized). The
+ // longest matching sequence (if any) is used to limit the maximum repetition length.
+ //
+ // Note that in the case case of a short sequence contained in a longer one, this might fail to
+ // find the smallest value for `rep_limit`. For example, if 'amniotic' and 'ni' are both used as
+ // restart sequences, 'ni' will be found first, and since it's shorter it will fail to suppress
+ // 'otic'. This is a minor issue since fully contained restart sequences are likely to be rare.
+ //
+ // This is theoretically worst-case O(N^2) for arbitrary restart sequences, which is why we
+ // have already clamped the maximum tail sequence length when generating `restart_sequences`.
+ // With clamping, this scan is O(N) in the context length.
+
+ int rep_limit = last_n_repeat;
+ for (int i = 0; i < last_n_repeat; ++i) {
+ llama_token token = ctx->last_tokens.rat(i);
+ auto its = ctx->dry_processed_breakers.equal_range(token);
+ if (its.first == ctx->dry_processed_breakers.end()) {
+ continue;
+ }
+ int longest_match = -1;
+ for (auto it = its.first; it != its.second; ++it) {
+ // Note that (*it) does not contain the head character, so seq_len will be
+ // the restart sequence length minus 1.
+ // In the common case of a single-token restart sequence, (*it) will be empty
+ // and we will trivially match.
+ int seq_len = (int)it->second.size();
+ if (seq_len > longest_match && seq_len <= (int)i) {
+ bool match = true;
+ for (int offset = 0; offset < seq_len; ++offset) {
+ // The -1 when indexing `last_tokens` is because we already matched the head.
+ if (it->second[offset] != ctx->last_tokens.rat(i - offset - 1)) {
+ match = false;
+ break;
+ }
+ }
+ if (match) {
+ longest_match = seq_len;
+ }
+ }
+ }
+ if (longest_match >= 0) {
+ // We found a restart sequence starting `i` tokens from the end and continuing for
+ // `longest_match` tokens.
+ rep_limit = i - longest_match;
+ break;
+ }
+ }
+ if (rep_limit < ctx->dry_allowed_length) {
+ return;
+ }
+
+ // Step 2: Iterate in reverse over the last N tokens of the context, using the "Z-algorithm" (in
+ // the reverse direction) to efficiently compute the positions and lengths of suffixes appearing
+ // elsewhere in the context. We limit the suffix length to `rep_limit` to respect restart sequences.
+ //
+ // This algorithm is not currently documented on Wikipedia, but there is a clear description here:
+ // https://ivanyu.me/blog/2014/10/15/z-algorithm/
+ //
+ // The code below is adapted from the public domain implementation by the same author here:
+ // https://github.com/ivanyu/string-algorithms/blob/master/z_algorithm.py
+ //
+ // Example:
+ // Last N tokens: a b c c b c y a b c
+ // Repeat counts: 0 0 3 1 0 2 0 0 0 0
+ // ^
+ // This `3` means that the last three tokens of the context (a b c) also appear here.
+ //
+ // This step is worst case O(N) since the Z-algorithm is linear, despite the appearance of nested
+ // for/while loops. This can be seen by observing that the `lt` and `rt` bounds are set after each
+ // repeated suffix is detected (i.e. after each while loop when n > 0). These bound variables
+ // ensure that the inner while loops only examine each token in the context once as the outer
+ // for loop iterates over the context.
+
+ {
+ const int last = last_n_repeat - 1;
+
+ int rt = 0;
+ int lt = 0;
+
+ for (int k = 1; k < last_n_repeat; ++k) {
+ if (k > rt) {
+ // If k is outside the current Z-box, do naive computation.
+ int n = 0;
+ while (n + k < last_n_repeat && ctx->last_tokens.rat(n) == ctx->last_tokens.rat(n+k)) {
+ ++n;
+ }
+ ctx->dry_repeat_count[last - k] = std::min(n, rep_limit);
+ if (n > 0) {
+ lt = k;
+ rt = k + n - 1;
+ }
+ } else {
+ // If k is inside the current Z-box, consider two cases.
+
+ int p = k - lt; // Pair index.
+ int right_part_len = rt - k + 1;
+
+ if (ctx->dry_repeat_count[last - p] < right_part_len) {
+ int n = std::min(ctx->dry_repeat_count[last - p], rep_limit);
+ ctx->dry_repeat_count[last - k] = n;
+ } else {
+ int i = rt + 1;
+ while (i < last_n_repeat && ctx->last_tokens.rat(i) == ctx->last_tokens.rat(i - k)) {
+ i += 1;
+ }
+
+ int n = std::min(i - k, rep_limit);
+ ctx->dry_repeat_count[last - k] = n;
+ lt = k;
+ rt = i - 1;
+ }
+ }
+ }
+ }
+
+ // Step 3: Iterate over dry_repeat_count and last_tokens, examining the maximum repeat length
+ // that would be generated by emitting each new token that would extend a sequence.
+ //
+ // Following the same example as above:
+ // Last N tokens: a b c c b c y a b c
+ // Repeat counts: 0 0 3 1 0 2 0 0 0 0
+ //
+ // For each non-zero, look ahead one token. This token, if emitted, would extend the repetition.
+ // c: 3 -> 4 (from `a b c` to `a b c c`)
+ // b: 1 -> 2 (from `c` to `c b`)
+ // y: 2 -> 3 (from `b c` to `b c y`)
+
+ for (int i = 0; i < last_n_repeat - 1; ++i) {
+ int repeat_len = ctx->dry_repeat_count[i];
+ if (repeat_len >= ctx->dry_allowed_length) {
+ // This token ends a repeat, so the next token would continue one.
+ // By convention, the value of `repeat_len` only includes the tokens currently
+ // in the context, not the new token that would be added.
+ llama_token token = ctx->last_tokens.rat(last_n_repeat - 2 - i);
+ // Track the maximum sequence ending in this token.
+ const auto& it = ctx->dry_max_token_repeat.find(token);
+ if (it == ctx->dry_max_token_repeat.end() || it->second < repeat_len) {
+ ctx->dry_max_token_repeat[token] = repeat_len;
+ }
+ }
+ }
+
+ // Step 4: Apply logit penalties based on the maximum repeat length for relevant tokens.
+
+ // Prevent floating point overflow in `pow(penalty_base, exponent)` by clamping to `max_exponent`.
+ // Compute it from `penalty_base` and the approximate log of `std::numeric_limits<float>::max()`
+ const float FLOAT_MAX_LOG = 88.7228391f;
+ int max_exponent = 0;
+ if (ctx->dry_base > 1.000001f) {
+ max_exponent = FLOAT_MAX_LOG / std::log(ctx->dry_base);
+ }
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ const auto& af_kvp = ctx->dry_max_token_repeat.find(cur_p->data[i].id);
+ if (af_kvp != ctx->dry_max_token_repeat.end()) {
+ // Check all sequence breakers starting with this token
+ auto range = ctx->dry_processed_breakers.equal_range(cur_p->data[i].id);
+ bool is_single_token_breaker = false;
+
+ for (auto it = range.first; it != range.second; ++it) {
+ if (it->second.empty()) {
+ is_single_token_breaker = true;
+ break;
+ }
+ }
+
+ // Apply penalty only if it's not a single-token sequence breaker
+ if (!is_single_token_breaker) {
+ int repeat_exp = af_kvp->second - ctx->dry_allowed_length;
+ if (max_exponent > 0 && repeat_exp > max_exponent) {
+ repeat_exp = max_exponent;
+ }
+ float penalty = ctx->dry_multiplier * std::pow(ctx->dry_base, repeat_exp);
+ cur_p->data[i].logit -= penalty;
+ }
+ }
+ }
+
+ cur_p->sorted = false;
+}
+
+static void llama_sampler_dry_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_dry *) smpl->ctx;
+ ctx->last_tokens.clear();
+ ctx->dry_repeat_count.clear();
+ ctx->dry_max_token_repeat.clear();
+}
+
+static struct llama_sampler * llama_sampler_dry_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (llama_sampler_dry *) smpl->ctx;
+
+ llama_vocab dummy_vocab;
+
+ // dummy vocab is passed because it is only needed for raw sequence breaker processing, which we have already done and will simply be copying
+ auto * result = llama_sampler_init_dry(&dummy_vocab, ctx->total_context_size, ctx->dry_multiplier, ctx->dry_base, ctx->dry_allowed_length, ctx->dry_penalty_last_n, NULL, 0);
+
+ // Copy the state, including the processed breakers
+ {
+ auto * result_ctx = (llama_sampler_dry *) result->ctx;
+ result_ctx->dry_processed_breakers = ctx->dry_processed_breakers;
+ result_ctx->dry_repeat_count = ctx->dry_repeat_count;
+ result_ctx->dry_max_token_repeat = ctx->dry_max_token_repeat;
+ result_ctx->last_tokens = ctx->last_tokens;
+ }
+
+ return result;
+}
+
+static void llama_sampler_dry_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_dry *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_dry_i = {
+ /* .name = */ llama_sampler_dry_name,
+ /* .accept = */ llama_sampler_dry_accept,
+ /* .apply = */ llama_sampler_dry_apply,
+ /* .reset = */ llama_sampler_dry_reset,
+ /* .clone = */ llama_sampler_dry_clone,
+ /* .free = */ llama_sampler_dry_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, int32_t n_ctx_train, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) {
+ int32_t effective_dry_penalty_last_n = (dry_penalty_last_n == -1) ? n_ctx_train : std::max(dry_penalty_last_n, 0);
+ std::unordered_multimap<llama_token, std::vector<llama_token>> processed_breakers;
+ const int MAX_CHAR_LEN = 40;
+ const int MAX_SEQ_LEN = 20;
+
+ const bool dry_enabled = (dry_multiplier != 0.0f && dry_base >= 1.0f && dry_penalty_last_n != 0);
+
+ if (!dry_enabled) {
+ return llama_sampler_init_empty("?dry");
+ }
+
+ if (dry_enabled && seq_breakers != nullptr && num_breakers > 0) {
+ // Process sequence breakers
+ for (size_t i = 0; i < num_breakers; ++i) {
+ if (seq_breakers[i] == nullptr || std::strlen(seq_breakers[i]) == 0) {
+ LLAMA_LOG_WARN("skipping null or empty DRY sequence breaker at index %zu\n", i);
+ continue;
+ }
+
+ std::string sequence_break(seq_breakers[i]);
+ if (sequence_break.empty()) {
+ LLAMA_LOG_WARN("skipping empty DRY sequence breaker\n");
+ continue;
+ }
+
+ if (sequence_break.size() > MAX_CHAR_LEN) {
+ LLAMA_LOG_WARN("truncating DRY sequence breaker to %d characters\n", MAX_CHAR_LEN);
+ sequence_break.resize(MAX_CHAR_LEN);
+ }
+
+ get_overlapping_token_sequences(*vocab, sequence_break, processed_breakers, MAX_SEQ_LEN);
+ }
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_dry_i,
+ /* .ctx = */ new llama_sampler_dry {
+ /* .total_context_size = */ n_ctx_train,
+ /* .dry_multiplier = */ dry_multiplier,
+ /* .dry_base = */ dry_base,
+ /* .dry_allowed_length = */ dry_allowed_length,
+ /* .dry_penalty_last_n = */ dry_penalty_last_n,
+ /* .dry_processed_breakers = */ std::move(processed_breakers),
+ /* .dry_repeat_count = */ dry_enabled ? std::vector<int>(effective_dry_penalty_last_n, 0) : std::vector<int>{},
+ /* .dry_max_token_repeat = */ {},
+ /* .last_tokens = */ dry_enabled ? ring_buffer<llama_token>(effective_dry_penalty_last_n) : ring_buffer<llama_token>(0),
+ }
+ );
+}
+
+// wrapper for test-sampling.cpp
+struct llama_sampler * llama_sampler_init_dry_testing(int32_t context_size, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const std::vector<std::vector<llama_token>>& seq_breakers) {
+ llama_vocab dummy_vocab;
+ auto * result = llama_sampler_init_dry(&dummy_vocab, context_size, dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, NULL, 0);
+ auto * ctx = (llama_sampler_dry *) result->ctx;
+
+ // Process the token-based sequence breakers
+ ctx->dry_processed_breakers.clear();
+ if (seq_breakers.empty()) {
+ LLAMA_LOG_WARN("empty DRY sequence breakers list in llama_sampler_init_dry_testing\n");
+ } else {
+ for (const auto& breaker : seq_breakers) {
+ if (breaker.empty()) {
+ LLAMA_LOG_WARN("skipping DRY empty sequence breaker\n");
+ continue;
+ }
+ llama_token head_token = breaker[0];
+ std::vector<llama_token> tail_tokens(breaker.begin() + 1, breaker.end());
+ ctx->dry_processed_breakers.emplace(head_token, std::move(tail_tokens));
+ }
+
+ if (ctx->dry_processed_breakers.empty()) {
+ LLAMA_LOG_WARN("no valid DRY sequence breakers processed in llama_sampler_init_dry_testing\n");
+ }
+ }
+
+ return result;
+}
+
+// adaptive-p sampler state
+//
+// maintains an exponential moving average of the *ORIGINAL* probabilities
+// of selected tokens, used to compute an adapted target at each sampling step.
+//
+// see llama.h for a full description of the sampler
+//
+// ref: https://github.com/ggml-org/llama.cpp/pull/17927
+//
+struct llama_sampler_adaptive_p {
+ const float target; // target probability (0.0 - 1.0; negative = disabled)
+ const float decay; // EMA decay; history ~= 1/(1-decay) tokens (0.0 - 0.99)
+ const uint32_t seed; // original RNG seed
+ uint32_t seed_cur; // actual RNG seed
+ std::mt19937 rng; // RNG state
+ float weighted_sum; // sum(p_i * decay^i)
+ float total_weight; // sum(decay^i), converges to 1/(1-decay)
+ std::vector<float> original_probs; // pre-transform probs, cached for EMA update
+ llama_token pending_token_id; // token ID of selected token
+ int32_t pending_token_idx; // index of orig. prob. of selected token in original_probs
+};
+
+// adaptive probability transformation constants
+static constexpr float DISTRIBUTION_WIDTH = 0.3f;
+static constexpr float PEAK_LOGIT_VALUE = 5.0f;
+static constexpr float SHARPNESS = 10.0f;
+static constexpr float INV_WIDTH = 1.0f / DISTRIBUTION_WIDTH;
+
+static const char * llama_sampler_adaptive_p_name(const struct llama_sampler * /*smpl*/) {
+ return "adaptive-p";
+}
+
+static void llama_sampler_adaptive_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
+
+ llama_sampler_softmax_impl(cur_p, false);
+
+ if (ctx->target < 0.0f) {
+ // at negative target values, adaptive-p is no-op
+ // we simply sample from the existing distribution
+ cur_p->selected = llama_sample_dist(cur_p, ctx->rng);
+ return;
+ }
+
+ // store the original probabilities
+ ctx->original_probs.resize(cur_p->size);
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ ctx->original_probs[i] = cur_p->data[i].p;
+ }
+
+ // using the EMA, compute the adapted target probability for the current sampling step
+ auto target = std::clamp(ctx->target, 0.0f, 1.0f);
+ float adapted_target = std::clamp(
+ ctx->total_weight == 0.0f ? target : 2.0f * target - (ctx->weighted_sum / ctx->total_weight),
+ 0.0f, 1.0f
+ );
+
+ // adaptive probability transform
+ //
+ // quadratic near target for fine differentiation, transitioning to linear decay in the
+ // tails. unbounded negative logits ensure proper suppression of far-from-target tokens
+ // after the softmax.
+ //
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ if (cur_p->data[i].logit == -INFINITY) {
+ // don't transform logits that are -INFINITY
+ // (as masked out by e.g. min-p and top-p when using backend sampling)
+ continue;
+ }
+ float dist = std::abs((cur_p->data[i].p - adapted_target) * INV_WIDTH);
+ cur_p->data[i].logit = PEAK_LOGIT_VALUE - SHARPNESS * dist * dist / (1.0f + dist);
+ }
+
+ // softmax and sample from the transformed distribution
+ llama_sampler_softmax_impl(cur_p, false);
+ const int idx = llama_sample_dist(cur_p, ctx->rng);
+ cur_p->selected = idx;
+
+ // store the selected token ID for acceptance later
+ ctx->pending_token_id = cur_p->data[idx].id;
+ ctx->pending_token_idx = idx;
+}
+
+static void llama_sampler_adaptive_p_accept(struct llama_sampler * smpl, llama_token token) {
+ auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
+ if (ctx->pending_token_id == token) {
+ GGML_ASSERT(ctx->pending_token_id != LLAMA_TOKEN_NULL);
+ GGML_ASSERT(ctx->pending_token_idx != -1);
+ // update EMA with the original probability of the selected token
+ ctx->weighted_sum = ctx->original_probs[ctx->pending_token_idx] + ctx->decay * ctx->weighted_sum;
+ ctx->total_weight = 1.0f + ctx->decay * ctx->total_weight;
+ }
+ ctx->pending_token_id = LLAMA_TOKEN_NULL;
+ ctx->pending_token_idx = -1;
+}
+
+static void llama_sampler_adaptive_p_reset(struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
+ // ctx->target and ctx->decay never change after init, so it's safe to keep them as is.
+ // original_probs is completely overwritten on every call to _apply.
+ // so we only need to reset the EMA state and pending token.
+ ctx->weighted_sum = ctx->target / (1.0f - ctx->decay);
+ ctx->total_weight = 1.0f / (1.0f - ctx->decay);
+ ctx->pending_token_id = LLAMA_TOKEN_NULL;
+ ctx->pending_token_idx = -1;
+ ctx->seed_cur = get_rng_seed(ctx->seed);
+ ctx->rng.seed(ctx->seed_cur);
+}
+
+static struct llama_sampler * llama_sampler_adaptive_p_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_adaptive_p *) smpl->ctx;
+ auto * result = llama_sampler_init_adaptive_p(ctx->target, ctx->decay, ctx->seed);
+ auto * result_ctx = (llama_sampler_adaptive_p *) result->ctx;
+
+ // copy everything (target, decay, seed, and RNG are already set)
+ result_ctx->weighted_sum = ctx->weighted_sum;
+ result_ctx->total_weight = ctx->total_weight;
+ result_ctx->pending_token_id = ctx->pending_token_id;
+ result_ctx->pending_token_idx = ctx->pending_token_idx;
+
+ return result;
+}
+
+static void llama_sampler_adaptive_p_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_adaptive_p *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_adaptive_p_i = {
+ /* .name = */ llama_sampler_adaptive_p_name,
+ /* .accept = */ llama_sampler_adaptive_p_accept,
+ /* .apply = */ llama_sampler_adaptive_p_apply,
+ /* .reset = */ llama_sampler_adaptive_p_reset,
+ /* .clone = */ llama_sampler_adaptive_p_clone,
+ /* .free = */ llama_sampler_adaptive_p_free,
+ /* .backend_init = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_adaptive_p(
+ float target,
+ float decay,
+ uint32_t seed
+) {
+ auto seed_cur = get_rng_seed(seed);
+ float clamped_decay = std::clamp(decay, 0.0f, 0.99f);
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_adaptive_p_i,
+ /* .ctx = */ new llama_sampler_adaptive_p {
+ /* .target = */ target,
+ /* .decay = */ clamped_decay,
+ /* .seed = */ seed,
+ /* .seed_cur = */ seed_cur,
+ /* .rng = */ std::mt19937(seed_cur),
+ /* .weighted_sum = */ target / (1.0f - clamped_decay),
+ /* .total_weight = */ 1.0f / (1.0f - clamped_decay),
+ /* .original_probs = */ {},
+ /* .pending_token_id = */ LLAMA_TOKEN_NULL,
+ /* .pending_token_idx = */ -1
+ }
+ );
+}
+
+// logit-bias
+
+struct llama_sampler_logit_bias : public llama_sampler_backend {
+ const int32_t n_vocab;
+
+ const std::vector<llama_logit_bias> logit_bias;
+
+ std::vector<llama_logit_bias> to_search;
+
+ struct ggml_tensor * inp_logit_bias;
+ struct ggml_tensor * inp_logit_idxs;
+};
+
+static const char * llama_sampler_logit_bias_name(const struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_logit_bias *) smpl->ctx;
+ return ctx->get_name();
+}
+
+static void llama_sampler_logit_bias_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_logit_bias *) smpl->ctx;
+
+ if (ctx->logit_bias.empty()) {
+ return;
+ }
+
+ ctx->to_search.clear();
+
+ // update the candidates that have not been shuffled in the vocabulary (i.e. idx == id)
+ for (const auto & lb : ctx->logit_bias) {
+ if (lb.token >= 0 && cur_p->size > (size_t) lb.token && cur_p->data[lb.token].id == lb.token) {
+ cur_p->data[lb.token].logit += lb.bias;
+ } else {
+ ctx->to_search.push_back(lb);
+ }
+ }
+
+ if (ctx->to_search.empty()) {
+ return;
+ }
+
+ // search for the remaining candidates that were not found in the previous step
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ for (const auto & lb : ctx->to_search) {
+ if (cur_p->data[i].id == lb.token) {
+ cur_p->data[i].logit += lb.bias;
+ break;
+ }
+ }
+ }
+}
+
+static struct llama_sampler * llama_sampler_logit_bias_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_logit_bias *) smpl->ctx;
+ return llama_sampler_init_logit_bias(ctx->n_vocab, ctx->logit_bias.size(), ctx->logit_bias.data());
+}
+
+static void llama_sampler_logit_bias_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_logit_bias *) smpl->ctx;
+}
+
+static void llama_sampler_logit_bias_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ GGML_UNUSED(gf);
+ GGML_UNUSED(ctx);
+
+ auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
+ if (sctx->logit_bias.empty()) {
+ return;
+ }
+
+ const size_t n = sctx->logit_bias.size();
+
+ sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
+ ggml_set_name(sctx->inp_logit_bias, "logit_bias");
+ ggml_set_input(sctx->inp_logit_bias);
+
+ sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
+ ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
+ ggml_set_input(sctx->inp_logit_idxs);
+
+ ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f);
+
+ cur = ggml_reshape_2d(ctx, cur, 1, ggml_nelements(cur));
+ cur = ggml_set_rows(ctx, cur, sctx->inp_logit_bias, sctx->inp_logit_idxs);
+ cur = ggml_reshape_1d(ctx, cur, ggml_nelements(cur));
+
+ data->logits = ggml_add(ctx, data->logits, cur);
+}
+
+static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
+ if (sctx->logit_bias.empty()) {
+ return;
+ }
+
+ GGML_ASSERT(sctx->inp_logit_bias != nullptr);
+ GGML_ASSERT(sctx->inp_logit_idxs != nullptr);
+
+ const size_t n = sctx->logit_bias.size();
+
+ std::vector<float> data_logit_bias(n, 0.0f);
+ std::vector<int32_t> data_logit_idxs(n, 0);
+ for (size_t i = 0; i < n; ++i) {
+ const auto & lb = sctx->logit_bias[i];
+ GGML_ASSERT(lb.token >= 0 && lb.token < (int32_t) sctx->n_vocab);
+ data_logit_bias[i] = lb.bias;
+ data_logit_idxs[i] = lb.token;
+ }
+
+ ggml_backend_tensor_set(sctx->inp_logit_bias, data_logit_bias.data(), 0, ggml_nbytes(sctx->inp_logit_bias));
+ ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs));
+}
+
+static bool llama_sampler_logit_bias_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ GGML_UNUSED(buft);
+
+ auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
+
+ sctx->init(true);
+
+ if (sctx->logit_bias.empty()) {
+ return true;
+ }
+
+ return true;
+}
+
+static struct llama_sampler_i llama_sampler_logit_bias_i = {
+ /* .name = */ llama_sampler_logit_bias_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_logit_bias_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_logit_bias_clone,
+ /* .free = */ llama_sampler_logit_bias_free,
+ /* .backend_init = */ llama_sampler_logit_bias_backend_init,
+ /* .backend_accept = */ nullptr,
+ /* .backend_apply = */ llama_sampler_logit_bias_backend_apply,
+ /* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input,
+};
+
+struct llama_sampler * llama_sampler_init_logit_bias(
+ int32_t n_vocab,
+ int32_t n_logit_bias,
+ const llama_logit_bias * logit_bias) {
+ const bool is_empty = n_logit_bias <= 0;
+
+ if (is_empty) {
+ return llama_sampler_init_empty("?logit-bias");
+ }
+
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_logit_bias_i,
+ /* .ctx = */ new llama_sampler_logit_bias {
+ ("logit-bias"),
+ /* .n_vocab = */ n_vocab,
+ /* .logit_bias = */ std::vector<llama_logit_bias>(logit_bias, logit_bias + n_logit_bias),
+ /* .to_search = */ {},
+ /* .inp_logit_bias = */ nullptr,
+ /* .inp_logit_idxs = */ nullptr,
+ }
+ );
+}
+
+// infill
+
+//#define GGML_DEBUG_SAMPLER_INFILL
+
+struct llama_sampler_infill {
+ const struct llama_vocab * vocab;
+
+ std::vector<char> buf0;
+ std::vector<char> buf1;
+};
+
+static const char * llama_sampler_infill_name(const struct llama_sampler * /*smpl*/) {
+ return "infill";
+}
+
+static void llama_sampler_infill_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+ auto * ctx = (llama_sampler_infill *) smpl->ctx;
+
+ llama_sampler_softmax_impl(cur_p, true);
+
+#if defined(GGML_DEBUG_SAMPLER_INFILL)
+#define LOG_DBG_CUR LLAMA_LOG_DEBUG
+#else
+#define LOG_DBG_CUR(...)
+#endif
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
+ }
+
+ float p_txt_sum = 0.0f;
+ float p_eog_sum = 0.0f;
+
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ if (ctx->vocab->is_eog(cur_p->data[i].id)) {
+ p_eog_sum += cur_p->data[i].p;
+ } else {
+ p_txt_sum += cur_p->data[i].p;
+ }
+ }
+
+ const float rat = p_eog_sum == 0.0 ? INFINITY : p_txt_sum / p_eog_sum; GGML_UNUSED(rat);
+
+ LOG_DBG_CUR("%s: p_txt_sum = %.2f, p_eog_sum = %.2f, rat = %.2f, n = %zu\n", __func__, p_txt_sum, p_eog_sum, rat, cur_p->size);
+
+ if (3*p_eog_sum*cur_p->size > p_txt_sum) {
+ LOG_DBG_CUR("%s: the ratio p_txt/p_eog = %.2f is too low -> sampling EOG\n", __func__, p_txt_sum/p_eog_sum);
+
+ // keep just the EOG tokens
+ const auto size_org = cur_p->size;
+
+ cur_p->size = 0;
+
+ float p_sum = 0.0f;
+
+ for (size_t i = 0; i < size_org; ++i) {
+ if (ctx->vocab->is_eog(cur_p->data[i].id)) {
+ p_sum += cur_p->data[i].p;
+
+ cur_p->data[cur_p->size++] = cur_p->data[i];
+ }
+ }
+
+ // normalize probs
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cur_p->data[i].p /= p_sum;
+ }
+
+ return;
+ }
+
+ size_t n_combined = 0; GGML_UNUSED(n_combined);
+
+ // combine tokens with common prefix
+ for (size_t i0 = 0; i0 < cur_p->size; ++i0) {
+ for (size_t i1 = 0; i1 < cur_p->size; ++i1) {
+ if (cur_p->data[i0].logit == -INFINITY) {
+ break;
+ }
+
+ if (i0 == i1 || cur_p->data[i1].logit == -INFINITY) {
+ continue;
+ }
+
+ int len0 = ctx->vocab->token_to_piece(cur_p->data[i0].id, ctx->buf0.data(), ctx->buf0.size(), 0, false);
+ if (len0 < 0) {
+ ctx->buf0.resize(len0);
+ len0 = ctx->vocab->token_to_piece(cur_p->data[i0].id, ctx->buf0.data(), ctx->buf0.size(), 0, false);
+ assert(len0 > 0);
+ }
+
+ int len1 = ctx->vocab->token_to_piece(cur_p->data[i1].id, ctx->buf1.data(), ctx->buf1.size(), 0, false);
+ if (len1 < 0) {
+ ctx->buf1.resize(len1);
+ len1 = ctx->vocab->token_to_piece(cur_p->data[i1].id, ctx->buf1.data(), ctx->buf1.size(), 0, false);
+ assert(len1 > 0);
+ }
+
+ // token i0 is a prefix of token i1
+ if (len0 > 0 && len0 <= len1 && memcmp(ctx->buf0.data(), ctx->buf1.data(), len0) == 0) {
+ int dst = i0;
+ int src = i1;
+
+ // merge into the token with higher probability
+ if (cur_p->data[i1].p > cur_p->data[i0].p) {
+ std::swap(dst, src);
+ }
+
+ cur_p->data[dst].p += cur_p->data[src].p;
+ cur_p->data[src].logit = -INFINITY;
+ cur_p->data[src].p = 0.0f;
+
+ n_combined++;
+ }
+ }
+ }
+
+ size_t n_non_eog = 0;
+
+ size_t size_org = cur_p->size;
+
+ float p_sum = 0.0f;
+ float thold = 0.2f;
+
+ cur_p->size = 0;
+
+ LOG_DBG_CUR("%s: n_combined = %zu, applying thold = %.3f\n", __func__, n_combined, thold);
+
+ for (size_t i = 0; i < size_org; ++i) {
+ const bool is_eog = ctx->vocab->is_eog(cur_p->data[i].id);
+
+ if (cur_p->data[i].p < thold && !is_eog) {
+ continue;
+ }
+
+ if (!is_eog) {
+ ++n_non_eog;
+ }
+
+ p_sum += cur_p->data[i].p;
+
+ // keep this token
+ cur_p->data[cur_p->size++] = cur_p->data[i];
+ }
+
+ LOG_DBG_CUR("%s: n_non_eog = %zu\n", __func__, n_non_eog);
+
+ // if no non-EOG tokens are left -> reduce cur_p to single EOT token
+ if (n_non_eog == 0) {
+ cur_p->size = 1;
+ cur_p->data[0].id = ctx->vocab->token_eot();
+ if (cur_p->data[0].id == LLAMA_TOKEN_NULL) {
+ cur_p->data[0].id = ctx->vocab->token_eos();
+ }
+ cur_p->data[0].logit = 1.0f;
+
+ GGML_ASSERT(cur_p->data[0].id != LLAMA_TOKEN_NULL);
+
+ return;
+ }
+
+ // normalize probs
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cur_p->data[i].p /= p_sum;
+
+ LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
+ }
+
+ size_org = cur_p->size;
+ p_sum = 0.0f;
+ thold = 1.0/(n_non_eog + 1);
+
+ cur_p->size = 0;
+
+ LOG_DBG_CUR("%s: applying thold = %.3f\n", __func__, thold);
+
+ for (size_t i = 0; i < size_org; ++i) {
+ const bool is_eog = ctx->vocab->is_eog(cur_p->data[i].id);
+
+ if (cur_p->data[i].p < thold && !is_eog) {
+ continue;
+ }
+
+ p_sum += cur_p->data[i].p;
+
+ cur_p->data[cur_p->size++] = cur_p->data[i];
+ }
+
+ // normalize probs
+ for (size_t i = 0; i < cur_p->size; ++i) {
+ cur_p->data[i].p /= p_sum;
+
+ LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
+ }
+
+#undef LOG_DBG_CUR
+}
+
+static struct llama_sampler * llama_sampler_infill_clone(const struct llama_sampler * smpl) {
+ const auto * ctx = (const llama_sampler_infill *) smpl->ctx;
+ return llama_sampler_init_infill(ctx->vocab);
+}
+
+static void llama_sampler_infill_free(struct llama_sampler * smpl) {
+ delete (llama_sampler_infill *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_infill_i = {
+ /* .name = */ llama_sampler_infill_name,
+ /* .accept = */ nullptr,
+ /* .apply = */ llama_sampler_infill_apply,
+ /* .reset = */ nullptr,
+ /* .clone = */ llama_sampler_infill_clone,
+ /* .free = */ llama_sampler_infill_free,
+ /* .backend_apply = */ nullptr,
+ /* .backend_accept = */ nullptr,
+ /* .backend_set_input = */ nullptr,
+ /* .backend_init = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
+ return llama_sampler_init(
+ /* .iface = */ &llama_sampler_infill_i,
+ /* .ctx = */ new llama_sampler_infill {
+ /* .vocab = */ vocab,
+ /* .buf0 = */ std::vector<char>(512),
+ /* .buf1 = */ std::vector<char>(512),
+ }
+ );
+}
+
+// utils
+
+uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) {
+ if (smpl->iface == &llama_sampler_dist_i) {
+ return ((const llama_sampler_dist *) smpl->ctx)->seed_cur;
+ }
+
+ if (smpl->iface == &llama_sampler_mirostat_i) {
+ return ((const llama_sampler_mirostat *) smpl->ctx)->seed_cur;
+ }
+
+ if (smpl->iface == &llama_sampler_mirostat_v2_i) {
+ return ((const llama_sampler_mirostat_v2 *) smpl->ctx)->seed_cur;
+ }
+
+ if (smpl->iface == &llama_sampler_chain_i) {
+ const auto * ctx = (const llama_sampler_chain *) smpl->ctx;
+ for (auto it = ctx->samplers.rbegin(); it != ctx->samplers.rend(); ++it) {
+ const uint32_t seed = llama_sampler_get_seed(it->ptr);
+ if (seed != LLAMA_DEFAULT_SEED) {
+ return seed;
+ }
+ }
+ }
+
+ return LLAMA_DEFAULT_SEED;
+}
+
+// perf
+
+struct llama_perf_sampler_data llama_perf_sampler(const struct llama_sampler * chain) {
+ struct llama_perf_sampler_data data = {};
+
+ if (chain == nullptr || chain->iface != &llama_sampler_chain_i) {
+ GGML_ABORT("%s: invalid sampler passed - requires a sampler created with llama_sampler_chain_init()\n", __func__);
+ }
+
+ const auto * ctx = (const struct llama_sampler_chain *) chain->ctx;
+
+ data.t_sample_ms = 1e-3 * ctx->t_sample_us;
+ data.n_sample = std::max(0, ctx->n_sample);
+
+ return data;
+}
+
+void llama_perf_sampler_print(const struct llama_sampler * chain) {
+ const auto data = llama_perf_sampler(chain);
+
+ LLAMA_LOG_INFO("%s: samplers time = %10.2f ms / %5d runs\n", __func__, data.t_sample_ms, data.n_sample);
+}
+
+void llama_perf_sampler_reset(struct llama_sampler * chain) {
+ if (chain == nullptr || chain->iface != &llama_sampler_chain_i) {
+ GGML_ABORT("%s: invalid sampler passed - requires a sampler created with llama_sampler_chain_init()\n", __func__);
+ }
+
+ auto * ctx = (struct llama_sampler_chain *) chain->ctx;
+
+ ctx->t_sample_us = 0;
+ ctx->n_sample = 0;
+}
--- /dev/null
+#pragma once
+
+#include "llama.h"
+
+#include <vector>
+
+struct llama_vocab;
+struct llama_grammar;
+
+// sampler chain
+
+struct llama_sampler_chain {
+ llama_sampler_chain_params params;
+
+ // has .backend_init() been called?
+ bool is_init = false;
+
+ struct info {
+ bool is_backend;
+
+ llama_sampler * ptr;
+ };
+
+ std::vector<info> samplers;
+
+ // pre-allocated buffer for llama_sampler_sample to avoid repeated allocations
+ std::vector<llama_token_data> cur;
+
+ // timing
+
+ mutable int64_t t_sample_us;
+
+ mutable int32_t n_sample;
+};
+
+struct llama_sampler * llama_sampler_init_dry_testing(
+ int32_t context_size,
+ float dry_multiplier,
+ float dry_base,
+ int32_t dry_allowed_length,
+ int32_t dry_penalty_last_n,
+ const std::vector<std::vector<llama_token>> & seq_breakers);
+++ /dev/null
-#include "llama-sampling.h"
-
-#include "llama-impl.h"
-#include "llama-vocab.h"
-#include "llama-grammar.h"
-
-#include "ggml-cpp.h"
-
-#include <array>
-#include <algorithm>
-#include <cassert>
-#include <cfloat>
-#include <chrono>
-#include <cmath>
-#include <cstdlib>
-#include <cstring>
-#include <ctime>
-#include <numeric>
-#include <random>
-#include <unordered_map>
-#include <stdexcept>
-
-// the ring buffer works similarly to std::deque, but with a fixed capacity
-template<typename T>
-struct ring_buffer {
- ring_buffer(size_t cap) : capacity(cap), data(cap) {}
-
- T & front() {
- if (sz == 0) {
- throw std::runtime_error("ring buffer is empty");
- }
- return data[first];
- }
-
- const T & front() const {
- if (sz == 0) {
- throw std::runtime_error("ring buffer is empty");
- }
- return data[first];
- }
-
- T & back() {
- if (sz == 0) {
- throw std::runtime_error("ring buffer is empty");
- }
- return data[pos];
- }
-
- const T & back() const {
- if (sz == 0) {
- throw std::runtime_error("ring buffer is empty");
- }
- return data[pos];
- }
-
- void push_back(const T & value) {
- if (capacity == 0) {
- throw std::runtime_error("ring buffer: capacity is zero");
- }
-
- if (sz == capacity) {
- // advance the start when buffer is full
- first = (first + 1) % capacity;
- } else {
- sz++;
- }
- data[pos] = value;
- pos = (pos + 1) % capacity;
- }
-
- T pop_front() {
- if (sz == 0) {
- throw std::runtime_error("ring buffer is empty");
- }
- T value = data[first];
- first = (first + 1) % capacity;
- sz--;
- return value;
- }
-
- //T & operator[](size_t i) {
- // if (i >= sz) {
- // throw std::runtime_error("ring buffer: index out of bounds");
- // }
- // return data[(first + i) % capacity];
- //}
-
- //const T & at(size_t i) const {
- // if (i >= sz) {
- // throw std::runtime_error("ring buffer: index out of bounds");
- // }
- // return data[(first + i) % capacity];
- //}
-
- const T & rat(size_t i) const {
- if (i >= sz) {
- throw std::runtime_error("ring buffer: index out of bounds");
- }
- return data[(first + sz - i - 1) % capacity];
- }
-
- std::vector<T> to_vector() const {
- std::vector<T> result;
- result.reserve(sz);
- for (size_t i = 0; i < sz; i++) {
- result.push_back(data[(first + i) % capacity]);
- }
- return result;
- }
-
- void clear() {
- // here only reset the status of the buffer
- sz = 0;
- first = 0;
- pos = 0;
- }
-
- bool empty() const {
- return sz == 0;
- }
-
- size_t size() const {
- return sz;
- }
-
- size_t capacity = 0;
- size_t sz = 0;
- size_t first = 0;
- size_t pos = 0;
-
- std::vector<T> data;
-};
-
-// writes result in res, does not mutate cur
-static void llama_token_data_array_partial_sort(const llama_token_data_array & cur, int npartial, std::vector<llama_token_data> & res) {
- static const auto comp = [](const llama_token_data & a, const llama_token_data & b) {
- return a.logit > b.logit;
- };
-
- constexpr int nbuckets = 128;
- constexpr float bucket_low = -10.0f;
- constexpr float bucket_high = 10.0f;
- constexpr float bucket_scale = nbuckets/(bucket_high - bucket_low);
- constexpr float bucket_inter = -bucket_low * bucket_scale;
-
- std::vector<int> bucket_idx;
- std::vector<int> histo(nbuckets, 0);
-
- std::vector<llama_token_data*> bucket_ptrs;
-
- bucket_idx.reserve(cur.size);
-
- for (int i = 0; i < (int)cur.size; ++i) {
- const float val = cur.data[i].logit;
- int ib = int(bucket_scale * val + bucket_inter); //nbuckets * (val - bucket_low) / (bucket_high - bucket_low);
- ib = std::max(0, std::min(nbuckets - 1, ib));
- bucket_idx.push_back(ib);
- ++histo[ib];
- }
- int nhave = 0;
- int ib = nbuckets - 1;
- for ( ; ib >= 0; --ib) {
- nhave += histo[ib];
- if (nhave >= npartial) {
- break;
- }
- }
- res.resize(nhave);
- auto * ptr = res.data();
- bucket_ptrs.reserve(nbuckets - ib);
- for (int j = nbuckets - 1; j >= ib; --j) {
- bucket_ptrs.push_back(ptr);
- ptr += histo[j];
- }
- for (int i = 0; i < (int)cur.size; ++i) {
- int j = bucket_idx[i];
- if (j >= ib) {
- *bucket_ptrs[nbuckets - 1 - j]++ = cur.data[i];
- }
- }
-
- ptr = res.data();
- int ndone = 0;
- for (int j = nbuckets - 1; j > ib; --j) {
- std::sort(ptr, ptr + histo[j], comp);
- ptr += histo[j];
- ndone += histo[j];
- }
- std::partial_sort(ptr, ptr + npartial - ndone, ptr + histo[ib], comp);
-}
-
-// reduces the size of cur_p to npartial, keeping only the top npartial elements
-static void llama_token_data_array_partial_sort_inplace(llama_token_data_array * cur_p, int npartial) {
- static const auto comp = [](const llama_token_data & a, const llama_token_data & b) {
- return a.logit > b.logit;
- };
-
- if (npartial <= 128) {
- std::partial_sort(cur_p->data, cur_p->data + npartial, cur_p->data + cur_p->size, comp);
-
- cur_p->size = npartial;
- cur_p->sorted = true;
-
- return;
- }
-
- std::vector<llama_token_data> tmp;
-
- llama_token_data_array_partial_sort(*cur_p, npartial, tmp);
-
- std::copy(tmp.data(), tmp.data() + npartial, cur_p->data);
-
- cur_p->size = npartial;
- cur_p->sorted = true;
-}
-
-static int llama_sample_dist(llama_token_data_array * cur_p, std::mt19937 & rng) {
- // iterator for the probabilities
-#ifdef __GNUC__
- #pragma GCC diagnostic push
- #pragma GCC diagnostic ignored "-Wunused-local-typedefs"
-#endif
-
- struct probs_iterator {
- typedef std::input_iterator_tag iterator_category;
- typedef float value_type;
- typedef float * pointer;
- typedef float & reference;
- typedef ptrdiff_t difference_type;
-
- const llama_token_data * data;
-
- bool operator==(const probs_iterator & other) const { return data == other.data; }
- bool operator!=(const probs_iterator & other) const { return data != other.data; }
- const float & operator*() const { return data->p; }
- probs_iterator & operator++() { ++data; return *this; }
- probs_iterator operator++(int) { probs_iterator tmp = *this; ++data; return tmp; }
- };
-
-#ifdef __GNUC__
- #pragma GCC diagnostic pop
-#endif
-
- std::discrete_distribution<int> dist(probs_iterator{cur_p->data}, probs_iterator{cur_p->data + cur_p->size});
-
- return dist(rng);
-}
-
-/*
-static void llama_log_softmax(float * array, size_t size) {
- float max_l = *std::max_element(array, array + size);
- float sum = 0.f;
- for (size_t i = 0; i < size; ++i) {
- float p = expf(array[i] - max_l);
- sum += p;
- array[i] = p;
- }
-
- for (size_t i = 0; i < size; ++i) {
- array[i] = logf(array[i] / sum);
- }
-}
-*/
-
-static void llama_sampler_temp_impl(llama_token_data_array * cur_p, float temp) {
- if (temp <= 0.0f) {
- // find the token with the highest logit and set the rest to -inf
- size_t max_i = 0;
- float max_l = cur_p->data[0].logit;
-
- for (size_t i = 1; i < cur_p->size; ++i) {
- if (cur_p->data[i ].logit > max_l) {
- cur_p->data[max_i].logit = -INFINITY;
- max_i = i;
- max_l = cur_p->data[i].logit;
- } else {
- cur_p->data[i].logit = -INFINITY;
- }
- }
-
- return;
- }
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- cur_p->data[i].logit /= temp;
- }
-}
-
-static void llama_sampler_softmax_impl(llama_token_data_array * cur_p, bool do_sort) {
- GGML_ASSERT(cur_p->size > 0);
-
- // Sort the logits in descending order if requested
- if (do_sort && !cur_p->sorted) {
- llama_token_data_array_partial_sort_inplace(cur_p, cur_p->size);
- }
-
- float max_l = cur_p->data[0].logit;
- if (!cur_p->sorted) {
- for (size_t i = 1; i < cur_p->size; ++i) {
- max_l = std::max(max_l, cur_p->data[i].logit);
- }
- }
-
- float cum_sum = 0.0f;
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- float p = expf(cur_p->data[i].logit - max_l);
- cur_p->data[i].p = p;
- cum_sum += p;
- }
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- cur_p->data[i].p /= cum_sum;
- }
-}
-
-static void llama_sampler_top_k_impl(llama_token_data_array * cur_p, int32_t k) {
- // if (k >= (int32_t)cur_p->size) {
- // return;
- // }
-
- if (k <= 0) {
- return;
- }
-
- k = std::min(k, (int) cur_p->size);
-
- // Sort scores in descending order
- if (!cur_p->sorted) {
- llama_token_data_array_partial_sort_inplace(cur_p, k);
- }
-
- cur_p->size = k;
-}
-
-static uint32_t get_rng_seed(uint32_t seed) {
- if (seed == LLAMA_DEFAULT_SEED) {
- // use system clock if std::random_device is not a true RNG
- static bool is_rd_prng = std::random_device().entropy() == 0;
- if (is_rd_prng) {
- return (uint32_t) std::chrono::system_clock::now().time_since_epoch().count();
- }
- std::random_device rd;
- return rd();
- }
- return seed;
-}
-
-// llama_sampler API
-
-struct llama_sampler * llama_sampler_init(
- struct llama_sampler_i * iface,
- llama_sampler_context_t ctx) {
- return new llama_sampler {
- /* .iface = */ iface,
- /* .ctx = */ ctx,
- };
-}
-
-const char * llama_sampler_name(const struct llama_sampler * smpl) {
- if (!smpl->iface) {
- return "(null)";
- }
-
- return smpl->iface->name(smpl);
-}
-
-void llama_sampler_accept(struct llama_sampler * smpl, llama_token token) {
- if (!smpl) {
- return;
- }
-
- if (smpl->iface->accept) {
- smpl->iface->accept(smpl, token);
- }
-}
-
-void llama_sampler_apply(struct llama_sampler * smpl, struct llama_token_data_array * cur_p) {
- if (!smpl) {
- return;
- }
-
- GGML_ASSERT(smpl->iface->apply);
- smpl->iface->apply(smpl, cur_p);
-}
-
-void llama_sampler_reset(struct llama_sampler * smpl) {
- if (!smpl) {
- return;
- }
-
- if (smpl->iface->reset) {
- smpl->iface->reset(smpl);
- }
-}
-
-struct llama_sampler * llama_sampler_clone(const struct llama_sampler * smpl) {
- if (!smpl) {
- return nullptr;
- }
-
- if (smpl->iface->clone) {
- return smpl->iface->clone(smpl);
- }
-
- if (smpl->ctx == nullptr) {
- return llama_sampler_init(
- /* .iface = */ smpl->iface,
- /* .ctx = */ nullptr
- );
- }
-
- GGML_ABORT("the sampler does not support cloning");
-}
-
-void llama_sampler_free(struct llama_sampler * smpl) {
- if (smpl == nullptr) {
- return;
- }
-
- if (smpl->iface->free) {
- smpl->iface->free(smpl);
- }
-
- delete smpl;
-}
-
-// empty sampler
-
-struct llama_sampler_empty {
- const char * name;
-};
-
-static struct llama_sampler * llama_sampler_init_empty(const char * name);
-
-static const char * llama_sampler_empty_name(const struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_empty *) smpl->ctx;
- return ctx->name;
-}
-
-static void llama_sampler_empty_accept(struct llama_sampler * smpl, llama_token token) {
- GGML_UNUSED(smpl);
- GGML_UNUSED(token);
-}
-
-static void llama_sampler_empty_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- GGML_UNUSED(smpl);
- GGML_UNUSED(cur_p);
-}
-
-static void llama_sampler_empty_reset(struct llama_sampler * smpl) {
- GGML_UNUSED(smpl);
-}
-
-static struct llama_sampler * llama_sampler_empty_clone(const struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_empty *) smpl->ctx;
- return llama_sampler_init_empty(ctx->name);
-}
-
-static void llama_sampler_empty_free(struct llama_sampler * smpl) {
- delete (llama_sampler_empty *) smpl->ctx;
-}
-
-static bool llama_sampler_empty_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- GGML_UNUSED(smpl);
- GGML_UNUSED(buft);
-
- return true;
-}
-
-static void llama_sampler_empty_backend_accept(
- struct llama_sampler * smpl,
- ggml_context * ctx,
- ggml_cgraph * gf,
- struct ggml_tensor * selected_token) {
- GGML_UNUSED(smpl);
- GGML_UNUSED(ctx);
- GGML_UNUSED(gf);
- GGML_UNUSED(selected_token);
-}
-
-static void llama_sampler_empty_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- GGML_UNUSED(smpl);
- GGML_UNUSED(ctx);
- GGML_UNUSED(gf);
- GGML_UNUSED(data);
-}
-
-static void llama_sampler_empty_backend_set_input(struct llama_sampler * smpl) {
- GGML_UNUSED(smpl);
-}
-
-static struct llama_sampler_i llama_sampler_empty_i = {
- /* .name = */ llama_sampler_empty_name,
- /* .accept = */ llama_sampler_empty_accept,
- /* .apply = */ llama_sampler_empty_apply,
- /* .reset = */ llama_sampler_empty_reset,
- /* .clone = */ llama_sampler_empty_clone,
- /* .free = */ llama_sampler_empty_free,
- /* .backend_init = */ llama_sampler_empty_backend_init,
- /* .backend_accept = */ llama_sampler_empty_backend_accept,
- /* .backend_apply = */ llama_sampler_empty_backend_apply,
- /* .backend_set_input = */ llama_sampler_empty_backend_set_input,
-};
-
-struct llama_sampler * llama_sampler_init_empty(const char * name) {
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_empty_i,
- /* .ctx = */ new llama_sampler_empty {
- /* .name = */ name,
- }
- );
-}
-
-// common backend sampler functionality
-//
-// +name : means that the sampler is support and will run on the backend
-// -name : means that a ggml operator is not supported by the backend
-//
-struct llama_sampler_backend {
- llama_sampler_backend(const char * name) : name(name), name_ext(name), is_init(false), support(false) {}
-
- const char * get_name() {
- if (!is_init) {
- return name.c_str();
- }
-
- if (support) {
- name_ext = "+" + name;
- } else {
- name_ext = "-" + name;
- }
-
- return name_ext.c_str();
- }
-
- void init(bool support) {
- GGML_ASSERT(this->is_init == false);
-
- this->is_init = true;
- this->support = support;
- }
-
-private:
- std::string name;
- std::string name_ext;
-
- bool is_init;
- bool support;
-};
-
-// check if all ggml ops used by the sampler are supported by the backend
-static bool llama_sampler_backend_support(
- llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * device = ggml_backend_buft_get_device(buft);
- if (!device) {
- // CPU backend always supported
- return true;
- }
-
- ggml_init_params params = {
- /*.mem_size =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
- /*.mem_buffer =*/ NULL,
- /*.no_alloc =*/ true,
- };
-
- ggml_context_ptr ctx_ptr { ggml_init(params) };
- if (!ctx_ptr) {
- throw std::runtime_error(format("failed to create ggml context"));
- }
-
- ggml_context * ctx = ctx_ptr.get();
-
- const int64_t n = 1024*1024;
-
- llama_sampler_data data = {
- /*.logits = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n),
- /*.probs = */ nullptr,
- /*.sampled = */ nullptr,
- /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
- };
-
- ggml_cgraph * gf = ggml_new_graph(ctx);
-
- smpl->iface->backend_apply(smpl, ctx, gf, &data);
-
- if (data.logits) {
- ggml_build_forward_expand(gf, data.logits);
- }
-
- if (data.probs) {
- ggml_build_forward_expand(gf, data.probs);
- }
-
- if (data.sampled) {
- ggml_build_forward_expand(gf, data.sampled);
- }
-
- if (data.candidates) {
- ggml_build_forward_expand(gf, data.candidates);
- }
-
- for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
- struct ggml_tensor * op = ggml_graph_node(gf, i);
-
- if (!ggml_backend_dev_supports_op(device, op)) {
- LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",
- __func__, ggml_backend_dev_name(device), ggml_op_name(op->op), smpl->iface->name(smpl));
-
- return false;
- }
- }
-
- return true;
-}
-
-// sampler chain
-
-static const char * llama_sampler_chain_name(const struct llama_sampler * /*smpl*/) {
- return "chain";
-}
-
-static void llama_sampler_chain_accept(struct llama_sampler * smpl, llama_token token) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- time_meas tm(chain->t_sample_us, chain->params.no_perf);
-
- for (auto & smpl : chain->samplers) {
- llama_sampler_accept(smpl.ptr, token);
- }
-
- chain->n_sample++;
-}
-
-static void llama_sampler_chain_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- time_meas tm(chain->t_sample_us, chain->params.no_perf);
-
- bool is_backend = chain->is_init;
-
- for (auto & smpl : chain->samplers) {
- if (is_backend && smpl.is_backend) {
- continue;
- }
-
- is_backend = false;
-
- if (smpl.ptr->iface->apply == nullptr) {
- continue;
- }
-
- llama_sampler_apply(smpl.ptr, cur_p);
- }
-}
-
-static void llama_sampler_chain_reset(struct llama_sampler * smpl) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- for (auto & smpl : chain->samplers) {
- llama_sampler_reset(smpl.ptr);
- }
-}
-
-static struct llama_sampler * llama_sampler_chain_clone(const struct llama_sampler * smpl) {
- const auto * chain_src = (const llama_sampler_chain *) smpl->ctx;
-
- auto * result = llama_sampler_chain_init(chain_src->params);
-
- for (const auto & smpl : chain_src->samplers) {
- llama_sampler_chain_add(result, llama_sampler_clone(smpl.ptr));
- }
-
- return result;
-}
-
-static void llama_sampler_chain_free(struct llama_sampler * smpl) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- for (auto & smpl : chain->samplers) {
- llama_sampler_free(smpl.ptr);
- }
-
- delete chain;
-}
-
-static bool llama_sampler_chain_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");
-
- chain->is_init = true;
-
- bool res = true;
-
- for (auto & smpl : chain->samplers) {
- bool res_cur = true;
-
- // to be able to run a sampler on the backend, it has to:
- // - have the .backend_init() API implemented
- // - return true during .backend_init()
- if (smpl.ptr->iface->backend_init) {
- if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) {
- res_cur = false;
- }
- } else {
- res_cur = false;
- }
-
- smpl.is_backend = res_cur;
-
- res = res && res_cur;
- }
-
- return res;
-}
-
-static void llama_sampler_chain_backend_accept(
- struct llama_sampler * smpl,
- ggml_context * ctx,
- ggml_cgraph * gf,
- struct ggml_tensor * selected_token) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- for (auto & smpl : chain->samplers) {
- if (!smpl.is_backend) {
- break;
- }
-
- if (smpl.ptr->iface->backend_accept) {
- smpl.ptr->iface->backend_accept(smpl.ptr, ctx, gf, selected_token);
- }
- }
-}
-
-static void llama_sampler_chain_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- GGML_ASSERT(chain->is_init && "llama_sampler_chain_backend_init() not called");
-
- for (auto & smpl : chain->samplers) {
- if (!smpl.is_backend) {
- break;
- }
-
- if (smpl.ptr->iface->backend_apply) {
- smpl.ptr->iface->backend_apply(smpl.ptr, ctx, gf, data);
- }
- }
-}
-
-static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
-
- for (auto & smpl : chain->samplers) {
- if (!smpl.is_backend) {
- break;
- }
-
- if (smpl.ptr->iface->backend_set_input) {
- smpl.ptr->iface->backend_set_input(smpl.ptr);
- }
- }
-}
-
-static struct llama_sampler_i llama_sampler_chain_i = {
- /* .name = */ llama_sampler_chain_name,
- /* .accept = */ llama_sampler_chain_accept,
- /* .apply = */ llama_sampler_chain_apply,
- /* .reset = */ llama_sampler_chain_reset,
- /* .clone = */ llama_sampler_chain_clone,
- /* .free = */ llama_sampler_chain_free,
- /* .backend_init = */ llama_sampler_chain_backend_init,
- /* .backend_accept = */ llama_sampler_chain_backend_accept,
- /* .backend_apply = */ llama_sampler_chain_backend_apply,
- /* .backend_set_input = */ llama_sampler_chain_backend_set_input,
-};
-
-struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) {
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_chain_i,
- /* .ctx = */ new llama_sampler_chain {
- /* .params = */ params,
- /* .is_init = */ false,
- /* .samplers = */ {},
- /* .cur = */ {},
- /* .t_sample_us = */ 0,
- /* .n_sample = */ 0,
- }
- );
-}
-
-llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {
- const llama_token sampled_token = llama_get_sampled_token_ith (ctx, idx);
- const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx);
- const float * sampled_logits = llama_get_sampled_logits_ith (ctx, idx);
- const llama_token * sampled_ids = llama_get_sampled_candidates_ith(ctx, idx);
-
- // If a backend sampler has already sampled a token, return it.
- if (sampled_token != LLAMA_TOKEN_NULL) {
- LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx);
- return sampled_token;
- }
-
- const llama_model * model = llama_get_model(ctx);
- const llama_vocab * vocab = llama_model_get_vocab(model);
-
- const int n_vocab = llama_vocab_n_tokens(vocab);
-
- // use pre-allocated buffer from chain if available, otherwise allocate locally
- std::vector<llama_token_data> * cur_ptr;
- std::vector<llama_token_data> cur_local;
-
- if (smpl->iface == &llama_sampler_chain_i) {
- auto * chain = (llama_sampler_chain *) smpl->ctx;
- cur_ptr = &chain->cur;
- } else {
- cur_ptr = &cur_local;
- }
-
- auto & cur = *cur_ptr;
-
- if (sampled_probs) {
- const uint32_t sampled_probs_count = llama_get_sampled_probs_count_ith(ctx, idx);
- cur.resize(sampled_probs_count);
- for (uint32_t i = 0; i < sampled_probs_count; ++i) {
- cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], sampled_probs[i]};
- }
- } else if (sampled_logits) {
- const uint32_t sampled_logits_count = llama_get_sampled_logits_count_ith(ctx, idx);
- cur.resize(sampled_logits_count);
- for (llama_token i = 0; i < (int)sampled_logits_count; i++) {
- cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], 0.0f};
- }
- } else {
- const auto * logits = llama_get_logits_ith(ctx, idx);
- GGML_ASSERT(logits != nullptr);
- cur.resize(n_vocab);
- for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
- cur[token_id] = llama_token_data{token_id, logits[token_id], 0.0f};
- }
- }
-
- llama_token_data_array cur_p = {
- /* .data = */ cur.data(),
- /* .size = */ cur.size(),
- /* .selected = */ -1,
- /* .sorted = */ false,
- };
-
- llama_sampler_apply(smpl, &cur_p);
-
- GGML_ASSERT(cur_p.selected >= 0 && cur_p.selected < (int32_t) cur_p.size);
-
- auto token = cur_p.data[cur_p.selected].id;
-
- llama_sampler_accept(smpl, token);
-
- return token;
-}
-
-
-void llama_sampler_chain_add(struct llama_sampler * chain, struct llama_sampler * smpl) {
- auto * p = (llama_sampler_chain *) chain->ctx;
- p->samplers.push_back({
- /* .is_backend = */ false,
- /* .ptr = */ smpl,
- });
-}
-
-struct llama_sampler * llama_sampler_chain_get(struct llama_sampler * chain, int32_t i) {
- if (chain == nullptr) {
- return nullptr;
- }
-
- if (chain->iface != &llama_sampler_chain_i) {
- return nullptr;
- }
-
- if (i == -1) {
- return chain;
- }
-
- const auto * p = (const llama_sampler_chain *) chain->ctx;
-
- if (i < 0 || (size_t) i >= p->samplers.size()) {
- return nullptr;
- }
-
- return p->samplers[i].ptr;
-}
-
-struct llama_sampler * llama_sampler_chain_remove(struct llama_sampler * chain, int32_t i) {
- auto * p = (llama_sampler_chain *) chain->ctx;
-
- if (i < 0 || (size_t) i >= p->samplers.size()) {
- return nullptr;
- }
-
- auto * result = p->samplers[i].ptr;
- p->samplers.erase(p->samplers.begin() + i);
-
- return result;
-}
-
-int llama_sampler_chain_n(const struct llama_sampler * chain) {
- const auto * p = (const llama_sampler_chain *) chain->ctx;
-
- return p->samplers.size();
-}
-
-//
-// samplers
-//
-
-// greedy
-
-struct llama_sampler_greedy : public llama_sampler_backend {
-};
-
-static const char * llama_sampler_greedy_name(const struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_greedy *) smpl->ctx;
- return sctx->get_name();
-}
-
-static void llama_sampler_greedy_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_greedy *) smpl->ctx;
- GGML_UNUSED(ctx);
-}
-
-static struct llama_sampler * llama_sampler_greedy_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_greedy *) smpl->ctx;
- auto * result = llama_sampler_init_greedy();
-
- // copy the state
- {
- auto * result_ctx = (llama_sampler_greedy *) result->ctx;
-
- GGML_UNUSED(ctx);
- GGML_UNUSED(result_ctx);
- }
-
- return result;
-}
-
-static void llama_sampler_greedy_free(struct llama_sampler * smpl) {
- delete (llama_sampler_greedy *) smpl->ctx;
-}
-
-static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_token_data_array * cur_p) {
- cur_p->selected = 0;
- for (size_t i = 1; i < cur_p->size; ++i) {
- if (cur_p->data[i].logit > cur_p->data[cur_p->selected].logit) {
- cur_p->selected = i;
- }
- }
-}
-
-static bool llama_sampler_greedy_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_greedy *) smpl->ctx;
-
- const bool res = llama_sampler_backend_support(smpl, buft);
-
- sctx->init(res);
-
- return res;
-}
-
-static void llama_sampler_greedy_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- GGML_UNUSED(gf);
- GGML_UNUSED(smpl);
-
- struct ggml_tensor * curl = ggml_argmax(ctx, data->logits);
- ggml_set_name(curl, "greedy_argmax");
-
- data->sampled = curl;
-}
-
-static struct llama_sampler_i llama_sampler_greedy_i = {
- /* .name = */ llama_sampler_greedy_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_greedy_apply,
- /* .reset = */ llama_sampler_greedy_reset,
- /* .clone = */ llama_sampler_greedy_clone,
- /* .free = */ llama_sampler_greedy_free,
- /* .backend_init = */ llama_sampler_greedy_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_greedy_backend_apply,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_greedy() {
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_greedy_i,
- /* .ctx = */ new llama_sampler_greedy {
- ("greedy"),
- }
- );
-}
-
-// dist
-
-struct llama_sampler_dist : public llama_sampler_backend {
- const uint32_t seed;
- uint32_t seed_cur;
-
- std::mt19937 rng;
-
- // backend input
- struct ggml_tensor * inp_uniform;
-
- ggml_context_ptr inp_ctx;
- ggml_backend_buffer_ptr inp_buf;
-};
-
-static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_dist *) smpl->ctx;
- return sctx->get_name();
-}
-
-static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_dist *) smpl->ctx;
-
- // edge cases
- if (cur_p->size == 0) {
- cur_p->selected = -1;
- return;
- }
-
- cur_p->selected = 0;
-
- if (cur_p->size == 1) {
- cur_p->data[0].p = 1.0f;
- return;
- }
-
- // max logit for numerical stability
- float max_l = cur_p->data[0].logit;
- if (!cur_p->sorted) {
- for (size_t i = 1; i < cur_p->size; ++i) {
- max_l = std::max(max_l, cur_p->data[i].logit);
- }
- }
-
- // apply softmax to obtain the probabilities
- double sum_cum = 0.0f;
- for (size_t i = 0; i < cur_p->size; ++i) {
- float p = expf(cur_p->data[i].logit - max_l);
- cur_p->data[i].p = p;
- sum_cum += p;
- }
-
-#if 1
- // sample from the obtained probabilities and normalize the probs in a single pass
- // this is ~3x faster on Mac with full gpt-oss vocab than the version below
- //
- std::uniform_real_distribution<double> dist(0.0f, 1.0f);
- const double rnd = dist(ctx->rng);
-
- double sum_run = 0.0f;
- const double sum_tgt = sum_cum*rnd;
-
- bool found = false;
- for (size_t i = 0; i < cur_p->size; ++i) {
- if (!found) {
- // accumulate probs until we reach the target sum
- sum_run += cur_p->data[i].p;
- if (sum_run >= sum_tgt) {
- cur_p->selected = i;
- found = true;
- }
- }
-
- // normalize probs
- cur_p->data[i].p /= sum_cum;
- }
-
- // fallback to the last token (don't think this can happen)
- assert(found);
- if (!found) {
- cur_p->selected = cur_p->size - 1;
- }
-#else
- // for clarity, this is the same as above but does one pass for normalization and one extra pass for sampling
- for (size_t i = 0; i < cur_p->size; ++i) {
- cur_p->data[i].p /= sum_cum;
- }
-
- cur_p->selected = llama_sample_dist(cur_p, ctx->rng);
-#endif
-}
-
-static void llama_sampler_dist_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_dist *) smpl->ctx;
- ctx->seed_cur = get_rng_seed(ctx->seed);
- ctx->rng.seed(ctx->seed_cur);
-}
-
-static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_dist *) smpl->ctx;
- auto * result = llama_sampler_init_dist(ctx->seed);
-
- // copy the state
- {
- auto * result_ctx = (llama_sampler_dist *) result->ctx;
-
- result_ctx->rng = ctx->rng;
- }
-
- return result;
-}
-
-static void llama_sampler_dist_free(struct llama_sampler * smpl) {
- delete (llama_sampler_dist *) smpl->ctx;
-}
-
-static bool llama_sampler_dist_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_dist *) smpl->ctx;
-
- // allocate inputs
- {
- ggml_init_params params = {
- /*.mem_size =*/ ggml_tensor_overhead(),
- /*.mem_buffer =*/ nullptr,
- /*.no_alloc =*/ true,
- };
-
- sctx->inp_ctx.reset(ggml_init(params));
-
- // Create the uniform random scalar input tensor. This will be set by
- // llama_sampler_dist_backend_set_input after this graph is built.
- sctx->inp_uniform = ggml_new_tensor_1d(sctx->inp_ctx.get(), GGML_TYPE_F32, 1);
- ggml_set_name (sctx->inp_uniform, "uniform");
- ggml_set_input(sctx->inp_uniform);
-
- // Allocate all tensors from our context to the backend
- sctx->inp_buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(sctx->inp_ctx.get(), buft));
-
- ggml_backend_buffer_clear(sctx->inp_buf.get(), 0);
- }
-
- const bool res = llama_sampler_backend_support(smpl, buft);
-
- sctx->init(res);
-
- if (!res) {
- sctx->inp_ctx.reset(nullptr);
- sctx->inp_buf.reset(nullptr);
- }
-
- return res;
-}
-
-static void llama_sampler_dist_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- GGML_UNUSED(gf);
- auto * sctx = (llama_sampler_dist *) smpl->ctx;
-
- struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits);
- ggml_set_name(probs, "dist_probs");
-
- struct ggml_tensor * cumsum = ggml_cumsum(ctx, probs);
- ggml_set_name(cumsum, "dist_cumsum");
-
- // The uniform tensor has a random value and we subtract this tensor with
- // the cumsum tensor (the uniform tensor will be broadcasted by ggml_sub).
- // Recall that each entry in cumsum is the cumulative probability up to that
- // index so values stay negative while the cumulative total is below the
- // random value, and become zero/positive once the threshold is crossed.
- struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform);
- ggml_set_name(diff, "dist_cumsum");
-
- // The ggml_step function produces a tensor where entries are 1 if the
- // corresponding entry in diff is > 0, and 0 otherwise. So all values up to
- // the index where the cumulative probability exceeds the random value are 0,
- // and all entries after that are 1.
- struct ggml_tensor * mask = ggml_step(ctx, diff);
- ggml_set_name(mask, "dist_mask");
-
- // Taking the sum of the mask gives us the sum of elements after the threshold
- // we are interested in.
- struct ggml_tensor * idxf = ggml_sum(ctx, mask);
- ggml_set_name(idxf, "dist_index_f32");
-
- // 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);
- ggml_set_name(idx, "dist_index_i32");
-
- // Map back to original vocab ids if a candidates tensor is available.
- struct ggml_tensor * sampled_token = idx;
- if (data->candidates != nullptr) {
- struct ggml_tensor * candidates = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates));
-
- sampled_token = ggml_get_rows(ctx, candidates, idx);
- ggml_set_name(sampled_token, "dist_sampled_token");
- }
-
- data->sampled = sampled_token;
- data->probs = probs;
-}
-
-static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_dist *) smpl->ctx;
- GGML_ASSERT(sctx->inp_uniform != nullptr);
-
- // We sample in double precision and cast to float to match rnd numbers of
- // llama_dampler_dist which uses double precision (sampling from
- // std::uniform_real_distribution<double> and
- // std::uniform_real_distribution<float> with same rng will produce
- // different sequences).
- std::uniform_real_distribution<double> dist(0.0f, 1.0f);
- const float rnd = dist(sctx->rng);
-
- ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float));
-}
-
-static struct llama_sampler_i llama_sampler_dist_i = {
- /* .name = */ llama_sampler_dist_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_dist_apply,
- /* .reset = */ llama_sampler_dist_reset,
- /* .clone = */ llama_sampler_dist_clone,
- /* .free = */ llama_sampler_dist_free,
- /* .backend_init = */ llama_sampler_dist_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_dist_backend_apply,
- /* .backend_set_input = */ llama_sampler_dist_backend_set_input,
-};
-
-struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
- auto seed_cur = get_rng_seed(seed);
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_dist_i,
- /* .ctx = */ new llama_sampler_dist {
- ("dist"),
- /* .seed = */ seed,
- /* .seed_cur = */ seed_cur,
- /* .rng = */ std::mt19937(seed_cur),
- /* .inp_uniform = */ nullptr,
- /* .inp_ctx = */ nullptr,
- /* .inp_buf = */ nullptr,
- }
- );
-}
-
-// top-k
-
-struct llama_sampler_top_k : public llama_sampler_backend {
- const int32_t k;
-};
-
-static const char * llama_sampler_top_k_name(const struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_top_k *) smpl->ctx;
- return sctx->get_name();
-}
-
-static void llama_sampler_top_k_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_top_k *) smpl->ctx;
- llama_sampler_top_k_impl(cur_p, ctx->k);
-}
-
-static struct llama_sampler * llama_sampler_top_k_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_top_k *) smpl->ctx;
- return llama_sampler_init_top_k(ctx->k);
-}
-
-static void llama_sampler_top_k_free(struct llama_sampler * smpl) {
- delete (llama_sampler_top_k *) smpl->ctx;
-}
-
-static bool llama_sampler_top_k_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_top_k *) smpl->ctx;
-
- const bool res = llama_sampler_backend_support(smpl, buft);
-
- sctx->init(res);
-
- return res;
-}
-
-static void llama_sampler_top_k_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- auto * sctx = (llama_sampler_top_k *) smpl->ctx;
-
- struct ggml_tensor * top_k = ggml_top_k(ctx, data->logits, sctx->k);
- ggml_set_name(top_k, "top_k");
-
- if (data->candidates) {
- struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]);
- data->candidates = ggml_get_rows(ctx, candidates_rows, top_k);
- data->candidates = ggml_reshape_1d(ctx, data->candidates, sctx->k);
- ggml_set_name(data->candidates, "top_k_candidates");
- } else {
- data->candidates = top_k;
- }
-
- struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
- struct ggml_tensor * top_k_rows = ggml_get_rows(ctx, logits_rows, top_k);
- data->logits = ggml_reshape_1d(ctx, top_k_rows, sctx->k);
- ggml_set_name(top_k_rows, "top_k_rows");
-
- GGML_UNUSED(gf);
-}
-
-static struct llama_sampler_i llama_sampler_top_k_i = {
- /* .name = */ llama_sampler_top_k_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_top_k_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_top_k_clone,
- /* .free = */ llama_sampler_top_k_free,
- /* .backend_init = */ llama_sampler_top_k_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_top_k_backend_apply,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_top_k(int32_t k) {
- const bool is_empty = (k <= 0);
-
- if (is_empty) {
- return llama_sampler_init_empty("?top-k");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_top_k_i,
- /* .ctx = */ new llama_sampler_top_k {
- ("top-k"),
- /* .k = */ k,
- }
- );
-}
-
-// top-p
-
-struct llama_sampler_top_p : public llama_sampler_backend {
- const float p;
- const size_t min_keep;
-
- std::vector<llama_token_data> buf_sort;
-};
-
-static const char * llama_sampler_top_p_name(const struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_top_p *) smpl->ctx;
- return sctx->get_name();
-}
-
-static void llama_sampler_top_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_top_p *) smpl->ctx;
-
- if (ctx->p >= 1.0f) {
- return;
- }
-
- llama_sampler_softmax_impl(cur_p, false);
-
- size_t k = cur_p->size;
- auto * pdata = cur_p->data;
-
- auto & buf_sort = ctx->buf_sort;
-
- // if not sorted, try adaptive top-k sorting
- if (!cur_p->sorted && cur_p->size > 1024) {
- k = std::min<size_t>(256, cur_p->size);
- llama_token_data_array_partial_sort(*cur_p, k, buf_sort);
- pdata = buf_sort.data();
- } else if (!cur_p->sorted) {
- // small candidates -> sort inplace
- llama_token_data_array_partial_sort_inplace(cur_p, k);
- }
-
- // Compute the cumulative probabilities
- float cum_sum = 0.0f;
- size_t last_idx = cur_p->size;
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- cum_sum += pdata[i].p;
-
- // Check if the running sum is at least p or if we have kept at least min_keep tokens
- // we set the last index to i+1 to indicate that the current iterate should be included in the set
- if (cum_sum >= ctx->p && i + 1 >= ctx->min_keep) {
- last_idx = i + 1;
- break;
- }
-
- // we exceeded the current top-k heuristic -> increase k and continue
- if (!cur_p->sorted && i == k - 1) {
- k = cur_p->size;
- llama_token_data_array_partial_sort(*cur_p, k, buf_sort);
- pdata = buf_sort.data();
- }
- }
-
- // Resize the output vector to keep only the top-p tokens
- if (!cur_p->sorted) {
- std::copy(buf_sort.data(), buf_sort.data() + last_idx, cur_p->data);
- cur_p->sorted = true;
- }
-
- cur_p->size = last_idx;
-}
-
-static struct llama_sampler * llama_sampler_top_p_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_top_p *) smpl->ctx;
- return llama_sampler_init_top_p(ctx->p, ctx->min_keep);
-}
-
-static void llama_sampler_top_p_free(struct llama_sampler * smpl) {
- delete (llama_sampler_top_p *) smpl->ctx;
-}
-
-static bool llama_sampler_top_p_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_top_p *) smpl->ctx;
-
- const bool res = llama_sampler_backend_support(smpl, buft);
-
- sctx->init(res);
-
- return res;
-}
-
-static void llama_sampler_top_p_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- auto * sctx = (llama_sampler_top_p *) smpl->ctx;
-
- auto ggml_sort = [ctx](struct ggml_tensor * a, struct ggml_tensor * b) {
- GGML_ASSERT(ggml_nrows(a) == 1);
- struct ggml_tensor * a_reshaped = ggml_reshape_2d(ctx, a, 1, a->ne[0]);
- struct ggml_tensor * a_sorted = ggml_get_rows(ctx, a_reshaped, b);
- return ggml_reshape_1d(ctx, a_sorted, a->ne[0]);
- };
-
- // Get the sorted logits in descending order.
- struct ggml_tensor * sorted_idx = ggml_argsort(ctx, data->logits, GGML_SORT_ORDER_DESC);
- ggml_set_name(sorted_idx, "top_p_sorted_idx");
-
- // Do the sorting via reshape + get_rows
- struct ggml_tensor * sorted_logits = ggml_sort(data->logits, sorted_idx);
- ggml_set_name(sorted_logits, "top_p_sorted_logits");
-
- struct ggml_tensor * softmax = ggml_soft_max(ctx, sorted_logits);
- ggml_set_name(softmax, "top_p_softmax");
-
- // If candidates are provided, sort them as well. Otherwise, set sorted indices as candidates.
- if (data->candidates) {
- data->candidates = ggml_sort(data->candidates, sorted_idx);
- } else {
- data->candidates = sorted_idx;
- }
- ggml_set_name(data->candidates, "top_p_candidates");
-
- // Compute Cumulative Distribution Function (CDF) by means of GGML_OP_CUMSUM.
- struct ggml_tensor * cdf = ggml_cumsum(ctx, softmax);
- ggml_set_name(cdf, "top_p_cdf");
-
- // Invert CDF and add top-p value so that ggml_step yields 1 for values we want to keep
- struct ggml_tensor * cdf_scaled = ggml_scale_bias(ctx, cdf, -1.0f, sctx->p);
- ggml_set_name(cdf_scaled, "top_p_cdf_scaled");
-
- struct ggml_tensor * mask = ggml_step(ctx, cdf_scaled);
- ggml_set_name(mask, "top_p_mask");
-
- // Taking the sum of the mask gives us the sum of elements after the threshold
- // we are interested in.
- struct ggml_tensor * idxf = ggml_sum(ctx, mask);
- ggml_set_name(idxf, "top_p_index_f32");
-
- // prevent out-of-bounds access
- idxf = ggml_clamp(ctx, idxf, 0.0f, mask->ne[0] - 1);
-
- // construct ones tensor to set the value in the mask
- struct ggml_tensor * ones = ggml_scale_bias(ctx, idxf, 0.0f, 1.0f);
- ggml_set_name(ones, "top_p_ones");
-
- // Make top-p inclusive (i.e. return all values such that cum_sum/cdf >= p)
- struct ggml_tensor * mask_reshaped = ggml_reshape_2d(ctx, mask, 1, mask->ne[0]);
-
- mask_reshaped = ggml_set_rows(ctx, mask_reshaped, ones, ggml_cast(ctx, idxf, GGML_TYPE_I32));
- mask = ggml_reshape_1d(ctx, mask_reshaped, mask->ne[0]);
-
- // Apply -INFINITY bias for masked-out tokens
- // log(1) = 0 (keep), log(0) = -INF (discard)
- struct ggml_tensor * top_p_bias = ggml_log(ctx, mask);
- ggml_set_name(top_p_bias, "top_p_bias");
-
- data->logits = ggml_add(ctx, sorted_logits, top_p_bias);
- ggml_set_name(data->logits, "top_p_logits");
-
- GGML_UNUSED(gf);
-}
-
-static struct llama_sampler_i llama_sampler_top_p_i = {
- /* .name = */ llama_sampler_top_p_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_top_p_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_top_p_clone,
- /* .free = */ llama_sampler_top_p_free,
- /* .backend_init = */ llama_sampler_top_p_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_top_p_backend_apply,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) {
- const bool is_empty = p >= 1.0f;
-
- if (is_empty) {
- return llama_sampler_init_empty("?top-p");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_top_p_i,
- /* .ctx = */ new llama_sampler_top_p {
- ("top-p"),
- /* .p = */ p,
- /* .min_keep = */ min_keep,
- /* .buf_sort = */ {},
- }
- );
-}
-
-// min-p
-
-struct llama_sampler_min_p : public llama_sampler_backend {
- const float p;
- const size_t min_keep;
-};
-
-static const char * llama_sampler_min_p_name(const struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_min_p *) smpl->ctx;
- return sctx->get_name();
-}
-
-static void llama_sampler_min_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_min_p *) smpl->ctx;
-
- if (ctx->p <= 0.0f || !cur_p->size) {
- return;
- }
-
- bool min_p_applied = false;
-
- // if the cur_p aren't sorted, try the unsorted implementation first
- if (!cur_p->sorted) {
- std::vector<llama_token_data> filtered_tokens;
-
- float max_logit = -FLT_MAX;
- for (size_t i = 0; i < cur_p->size; ++i) {
- max_logit = std::max(max_logit, cur_p->data[i].logit);
- }
- const float min_logit = max_logit + logf(ctx->p); // min logit for p_i >= p * p_max
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- if (cur_p->data[i].logit >= min_logit) {
- filtered_tokens.push_back(cur_p->data[i]);
- }
- }
-
- // if we have enough values the operation was a success
- if (!filtered_tokens.empty() && filtered_tokens.size() >= ctx->min_keep) {
- std::copy(filtered_tokens.begin(), filtered_tokens.end(), cur_p->data);
- cur_p->size = filtered_tokens.size();
- min_p_applied = true;
- }
- }
-
- // if the cur_p are sorted or the unsorted implementation failed, use this implementation
- if (!min_p_applied) {
- // Sort the logits in descending order
- if (!cur_p->sorted) {
- llama_token_data_array_partial_sort_inplace(cur_p, cur_p->size);
- }
-
- const float min_logit = cur_p->data[0].logit + logf(ctx->p); // min logit for p_i >= p * p_max
- size_t i = 1; // first token always matches
-
- for (; i < cur_p->size; ++i) {
- if (cur_p->data[i].logit < min_logit && i >= ctx->min_keep) {
- break; // prob too small
- }
- }
-
- // Resize the output vector to keep only the matching tokens
- cur_p->size = i;
- }
-}
-
-static struct llama_sampler * llama_sampler_min_p_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_min_p *) smpl->ctx;
- return llama_sampler_init_min_p(ctx->p, ctx->min_keep);
-}
-
-static void llama_sampler_min_p_free(struct llama_sampler * smpl) {
- delete (llama_sampler_min_p *) smpl->ctx;
-}
-
-static bool llama_sampler_min_p_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_min_p *) smpl->ctx;
-
- const bool res = llama_sampler_backend_support(smpl, buft);
-
- sctx->init(res);
-
- return res;
-}
-
-static void llama_sampler_min_p_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- auto * sctx = (llama_sampler_min_p *) smpl->ctx;
-
- struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits);
- ggml_set_name(max_idx, "max_idx");
-
- struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
- ggml_set_name(logits_rows, "logits_rows");
-
- struct ggml_tensor * max_logit = ggml_get_rows(ctx, logits_rows, max_idx);
- ggml_set_name(max_logit, "max_logit");
-
- // Calculate the threshold value.
- struct ggml_tensor * threshold = ggml_scale_bias(ctx, max_logit, 1.0f, logf(sctx->p));
- ggml_set_name(threshold, "min_p_threshold");
-
- // Subtract the threshold from logits.
- struct ggml_tensor * sub = ggml_sub(ctx, data->logits, threshold);
-
- // Create a mask where logits below the threshold are 0 (discard),
- // and others are 1 (keep).
- struct ggml_tensor * mask = ggml_step(ctx, sub);
- ggml_set_name(mask, "min_p_mask");
-
- // Apply -INFINITY bias for masked-out tokens
- // log(1) = 0 (keep), log(0) = -INF (discard)
- struct ggml_tensor * min_p_bias = ggml_log(ctx, mask);
- ggml_set_name(min_p_bias, "min_p_bias");
-
- data->logits = ggml_add(ctx, data->logits, min_p_bias);
- ggml_set_name(data->logits, "min_p_logits");
-
- GGML_UNUSED(gf);
-}
-
-static struct llama_sampler_i llama_sampler_min_p_i = {
- /* .name = */ llama_sampler_min_p_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_min_p_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_min_p_clone,
- /* .free = */ llama_sampler_min_p_free,
- /* .backend_init = */ llama_sampler_min_p_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_min_p_backend_apply,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) {
- const bool is_empty = (p <= 0.0f);
-
- if (is_empty) {
- return llama_sampler_init_empty("?min-p");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_min_p_i,
- /* .ctx = */ new llama_sampler_min_p {
- ("min-p"),
- /* .p = */ p,
- /* .min_keep = */ min_keep,
- }
- );
-}
-
-// typical
-
-struct llama_sampler_typical {
- const float p;
- const size_t min_keep;
-};
-
-static const char * llama_sampler_typical_name(const struct llama_sampler * /*smpl*/) {
- return "typical";
-}
-
-static void llama_sampler_typical_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_typical *) smpl->ctx;
-
- // Reference implementation:
- // https://github.com/huggingface/transformers/compare/main...cimeister:typical-sampling:typical-pr
- if (ctx->p >= 1.0f) {
- return;
- }
-
- // Compute the softmax of logits and calculate entropy
- llama_sampler_softmax_impl(cur_p, true);
-
- float entropy = 0.0f;
- for (size_t i = 0; i < cur_p->size; ++i) {
- entropy += -cur_p->data[i].p * logf(cur_p->data[i].p);
- }
-
- // Compute the absolute difference between negative log probability and entropy for each candidate
- std::vector<float> shifted_scores;
- for (size_t i = 0; i < cur_p->size; ++i) {
- float shifted_score = fabsf(-logf(cur_p->data[i].p) - entropy);
- shifted_scores.push_back(shifted_score);
- }
-
- // Sort tokens based on the shifted_scores and their corresponding indices
- std::vector<size_t> indices(cur_p->size);
- std::iota(indices.begin(), indices.end(), 0);
-
- std::sort(indices.begin(), indices.end(), [&](size_t a, size_t b) {
- return shifted_scores[a] < shifted_scores[b];
- });
-
- // Compute the cumulative probabilities
- float cum_sum = 0.0f;
- size_t last_idx = indices.size();
-
- for (size_t i = 0; i < indices.size(); ++i) {
- size_t idx = indices[i];
- cum_sum += cur_p->data[idx].p;
-
- // Check if the running sum is greater than typical or if we have kept at least min_keep tokens
- if (cum_sum > ctx->p && (ctx->min_keep == 0 || i >= ctx->min_keep - 1)) {
- last_idx = i + 1;
- break;
- }
- }
-
- // Resize the output vector to keep only the locally typical tokens
- std::vector<llama_token_data> cur_p_new;
- for (size_t i = 0; i < last_idx; ++i) {
- size_t idx = indices[i];
- cur_p_new.push_back(cur_p->data[idx]);
- }
-
- // Replace the data in cur_p with the cur_p_new data
- std::copy(cur_p_new.begin(), cur_p_new.end(), cur_p->data);
- cur_p->size = cur_p_new.size();
- cur_p->sorted = false;
-}
-
-static struct llama_sampler * llama_sampler_typical_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_typical *) smpl->ctx;
- return llama_sampler_init_typical(ctx->p, ctx->min_keep);
-}
-
-static void llama_sampler_typical_free(struct llama_sampler * smpl) {
- delete (llama_sampler_typical *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_typical_i = {
- /* .name = */ llama_sampler_typical_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_typical_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_typical_clone,
- /* .free = */ llama_sampler_typical_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) {
- const bool is_empty = (p >= 1.0f);
-
- if (is_empty) {
- return llama_sampler_init_empty("?typical");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_typical_i,
- /* .ctx = */ new llama_sampler_typical {
- /* .p = */ p,
- /* .min_keep = */ min_keep,
- }
- );
-}
-
-// temp
-
-struct llama_sampler_temp : public llama_sampler_backend {
- const float temp;
-};
-
-static const char * llama_sampler_temp_name(const struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_temp *) smpl->ctx;
- return sctx->get_name();
-}
-
-static void llama_sampler_temp_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- const auto * ctx = (llama_sampler_temp *) smpl->ctx;
-
- llama_sampler_temp_impl(cur_p, ctx->temp);
-}
-
-static struct llama_sampler * llama_sampler_temp_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_temp *) smpl->ctx;
- return llama_sampler_init_temp(ctx->temp);
-}
-
-static void llama_sampler_temp_free(struct llama_sampler * smpl) {
- delete (llama_sampler_temp *) smpl->ctx;
-}
-
-static void llama_sampler_backend_temp_sampling(
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data,
- float temp) {
- if (temp <= 0.0f) {
- // Find the most probable token index.
- struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits);
- ggml_set_name(max_idx, "temp_max_idx");
-
- if (data->candidates) {
- struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]);
- data->candidates = ggml_get_rows(ctx, candidates_rows, max_idx);
- } else {
- data->candidates = max_idx;
- }
-
- struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
- data->logits = ggml_get_rows(ctx, logits_rows, max_idx);
-
- return;
- }
-
- data->logits = ggml_scale(ctx, data->logits, 1.0f / temp);
-
- GGML_UNUSED(gf);
-}
-
-static bool llama_sampler_temp_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_temp *) smpl->ctx;
-
- const bool res = llama_sampler_backend_support(smpl, buft);
-
- sctx->init(res);
-
- return res;
-}
-
-static void llama_sampler_temp_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- auto * sctx = (llama_sampler_temp *) smpl->ctx;
- llama_sampler_backend_temp_sampling(ctx, gf, data, sctx->temp);
-}
-
-static struct llama_sampler_i llama_sampler_temp_i = {
- /* .name = */ llama_sampler_temp_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_temp_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_temp_clone,
- /* .free = */ llama_sampler_temp_free,
- /* .backend_init = */ llama_sampler_temp_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_temp_backend_apply,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_temp(float temp) {
- const bool is_empty = temp == 1.0f;
-
- if (is_empty) {
- return llama_sampler_init_empty("?temp");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_temp_i,
- /* .ctx = */ new llama_sampler_temp {
- ("temp"),
- /*.temp = */ temp,
- }
- );
-}
-
-// temp-ext
-
-struct llama_sampler_temp_ext : public llama_sampler_backend {
- const float temp;
- const float delta;
- const float exponent;
-};
-
-static const char * llama_sampler_temp_ext_name(const struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
- return sctx->get_name();
-}
-
-static void llama_sampler_temp_ext_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_temp_ext *) smpl->ctx;
- if (ctx->delta > 0) {
- const float min_temp = std::max(0.0f, ctx->temp - ctx->delta);
- const float max_temp = ctx->temp + ctx->delta;
-
- float exponent_val = ctx->exponent;
-
- // no need to do anything if there is only one (or zero) candidates
- if (cur_p->size <= 1) {
- return;
- }
-
- // Calculate maximum possible entropy
- float max_entropy = -logf(1.0f / cur_p->size);
-
- llama_sampler_softmax_impl(cur_p, true);
-
- // Calculate entropy of the softmax probabilities
- float entropy = 0.0f;
- for (size_t i = 0; i < cur_p->size; ++i) {
- float prob = cur_p->data[i].p;
- if (prob > 0.0f) { // Ensure no log(0)
- entropy -= prob * logf(prob);
- }
- }
-
- // Normalize the entropy (max_entropy cannot be 0 here because we checked cur_p->size != 1 above)
- float normalized_entropy = entropy / max_entropy;
-
- // Map the normalized entropy to the desired temperature range using the power function
- float dyn_temp = min_temp + (max_temp - min_temp) * powf(normalized_entropy, exponent_val);
-
- #ifdef DEBUG
- LLAMA_LOG_INFO("Your text maxtemp value is: %f\n", max_temp);
- LLAMA_LOG_INFO("Entropy: %f\n", entropy);
- LLAMA_LOG_INFO("Max Possible Entropy: %f\n", max_entropy);
- LLAMA_LOG_INFO("Normalized Entropy: %f\n", normalized_entropy);
- LLAMA_LOG_INFO("Exponent: %f\n", exponent_val);
- LLAMA_LOG_INFO("Dynamic Temperature (dyn_temp): %f\n", dyn_temp);
- #endif
-
- // Apply the dynamically calculated temperature scaling
- llama_sampler_temp_impl(cur_p, dyn_temp);
-
- // Re-compute softmax probabilities after scaling logits with dynamic temperature
- const double max_l_double = cur_p->data[0].logit;
-
- double cum_sum_double = 0.0;
- for (size_t i = 0; i < cur_p->size; ++i) {
- double p = exp(cur_p->data[i].logit - max_l_double);
- cur_p->data[i].p = p; // Store the scaled probability
- cum_sum_double += p;
- }
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- cur_p->data[i].p /= cum_sum_double; // Re-normalize the probabilities
- }
-
- #ifdef DEBUG
- // Print the updated top 25 probabilities after temperature scaling
- LLAMA_LOG_INFO("\nUpdated Top 25 Probabilities After Dynamic Temperature Scaling (in percentages):\n");
- for (size_t i = 0; i < 25 && i < cur_p->size; ++i) {
- LLAMA_LOG_INFO("Token %zu: %f%%\n", i + 1, cur_p->data[i].p * 100.0f);
- }
- #endif
- } else {
- llama_sampler_temp_impl(cur_p, ctx->temp);
- }
-}
-
-static struct llama_sampler * llama_sampler_temp_ext_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_temp_ext *) smpl->ctx;
- return llama_sampler_init_temp_ext(ctx->temp, ctx->delta, ctx->exponent);
-}
-
-static void llama_sampler_temp_ext_free(struct llama_sampler * smpl) {
- delete (llama_sampler_temp_ext *) smpl->ctx;
-}
-
-static bool llama_sampler_temp_ext_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
-
- const bool res = llama_sampler_backend_support(smpl, buft);
-
- sctx->init(res);
-
- return res;
-}
-
-static void llama_sampler_temp_ext_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
-
- // Revert to standard temperature scaling if delta or temp are non-positive.
- if (sctx->delta <= 0.0f || sctx->temp <= 0.0f) {
- llama_sampler_backend_temp_sampling(ctx, gf, data, sctx->temp);
- return;
- }
-
- // Calculate min_temp, max_temp, and max_entropy.
- const float min_temp = std::max(0.0f, sctx->temp - sctx->delta);
- const float max_temp = sctx->temp + sctx->delta;
- const float max_entropy = logf(data->logits->ne[0]);
-
- // Calculate the probabilities.
- struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits);
- ggml_set_name(probs, "temp_ext_softmax_probs");
-
- // Clamp probabilities to avoid log(0) which would give -inf
- struct ggml_tensor * probs_clamped = ggml_clamp(ctx, probs, 1e-10f, 1.0f);
- ggml_set_name(probs_clamped, "temp_ext_probs_clamped");
-
- // Calculate the entropy, entropy = -Σ(p * log(p)).
- struct ggml_tensor * log_probs = ggml_log(ctx, probs_clamped);
- struct ggml_tensor * p_log_p = ggml_mul(ctx, probs_clamped, log_probs);
- struct ggml_tensor * sum_p_log_p = ggml_sum(ctx, p_log_p);
- struct ggml_tensor * entropy = ggml_scale(ctx, sum_p_log_p, -1.0f);
- ggml_set_name(log_probs, "temp_ext_log_probs");
- ggml_set_name(p_log_p, "temp_ext_p_log_p");
- ggml_set_name(sum_p_log_p, "temp_ext_sum_p_log_p");
- ggml_set_name(entropy, "temp_ext_entropy");
-
- // Normalize the entropy, norm_entropy = entropy / max_entropy
- struct ggml_tensor * norm_entropy = ggml_scale(ctx, entropy, 1.0f / max_entropy);
- ggml_set_name(norm_entropy, "temp_ext_norm_entropy");
-
- // Calculate the dynamic temperature:
- // dyn_temp = min_temp + (max_temp - min_temp) * powf(normalized_entropy, exponent);
- //
- // Calculate powf(normalized_entropy, exponent) as
- // norm_entropy^exponent = exp(exponent * log(norm_entropy))
- struct ggml_tensor * log_norm_entropy = ggml_log(ctx, norm_entropy);
- struct ggml_tensor * scaled_log = ggml_scale(ctx, log_norm_entropy, sctx->exponent);
- struct ggml_tensor * pow_entropy = ggml_exp(ctx, scaled_log);
- // With pow_entropy computed we can now compute dyn_temp, scaling by
- // (max_temp - min_temp) and then adding min_temp.
- struct ggml_tensor * dyn_temp = ggml_scale_bias(ctx, pow_entropy, max_temp - min_temp, min_temp);
- ggml_set_name(log_norm_entropy, "temp_ext_log_norm_entropy");
- ggml_set_name(scaled_log, "temp_ext_scaled_log");
- ggml_set_name(pow_entropy, "temp_ext_pow_entropy");
- ggml_set_name(dyn_temp, "temp_ext_dyn_temp");
-
- // Scale the logits by the dynamic temperature
- struct ggml_tensor * scaled_logits = ggml_div(ctx, data->logits, dyn_temp);
- ggml_set_name(scaled_logits, "temp_ext_scaled_logits");
-
- data->logits = scaled_logits;
-}
-
-static struct llama_sampler_i llama_sampler_temp_ext_i = {
- /* .name = */ llama_sampler_temp_ext_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_temp_ext_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_temp_ext_clone,
- /* .free = */ llama_sampler_temp_ext_free,
- /* .backend_init = */ llama_sampler_temp_ext_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_temp_ext_backend_apply,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) {
- const bool is_empty = temp == 1.0f && delta <= 0.0f;
-
- if (is_empty) {
- return llama_sampler_init_empty("?temp-ext");
- }
-
- auto * res = llama_sampler_init(
- /* .iface = */ &llama_sampler_temp_ext_i,
- /* .ctx = */ new llama_sampler_temp_ext {
- ("temp-ext"),
- /* .temp = */ temp,
- /* .delta = */ delta,
- /* .exponent = */ exponent,
- }
- );
-
- return res;
-}
-
-// xtc
-
-struct llama_sampler_xtc {
- const float probability;
- const float threshold;
- const size_t min_keep;
-
- const uint32_t seed;
- uint32_t seed_cur;
-
- std::mt19937 rng;
-};
-
-static const char * llama_sampler_xtc_name(const struct llama_sampler * /*smpl*/) {
- return "xtc";
-}
-
-static void llama_sample_xtc_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_xtc *) smpl->ctx;
-
- if (ctx->probability <= 0.0f
- || ctx->threshold > 0.5f
- || cur_p->size < 2) {
- return;
- }
-
- std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
- float chance = distribution(ctx->rng);
- if (chance > ctx->probability) {
- return;
- }
-
- llama_sampler_softmax_impl(cur_p, true);
-
- int pos_last = 0;
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- if (cur_p->data[i].p >= ctx->threshold) {
- pos_last = i;
- } else {
- break;
- }
- }
-
- if (cur_p->size - pos_last >= ctx->min_keep && pos_last > 0) {
- cur_p->data += pos_last;
- cur_p->size -= pos_last;
- }
-}
-
-static struct llama_sampler * llama_sampler_xtc_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_xtc *) smpl->ctx;
- auto * result = llama_sampler_init_xtc(ctx->probability, ctx->threshold, ctx->min_keep, ctx->seed);
-
- // copy the state
- {
- auto * result_ctx = (llama_sampler_xtc *) result->ctx;
-
- result_ctx->rng = ctx->rng;
- }
-
- return result;
-}
-
-static void llama_sampler_xtc_free(struct llama_sampler * smpl) {
- delete (llama_sampler_xtc *) smpl->ctx;
-}
-
-static void llama_sampler_xtc_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_xtc *) smpl->ctx;
- ctx->seed_cur = get_rng_seed(ctx->seed);
- ctx->rng.seed(ctx->seed_cur);
-}
-
-static struct llama_sampler_i llama_sampler_xtc_i = {
- /* .name = */ llama_sampler_xtc_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sample_xtc_apply,
- /* .reset = */ llama_sampler_xtc_reset,
- /* .clone = */ llama_sampler_xtc_clone,
- /* .free = */ llama_sampler_xtc_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) {
- const bool is_empty = (p <= 0.0f || t > 0.5f);
-
- if (is_empty) {
- return llama_sampler_init_empty("?xtc");
- }
-
- const auto seed_cur = get_rng_seed(seed);
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_xtc_i,
- /* .ctx = */ new llama_sampler_xtc {
- /* .probability = */ p,
- /* .threshold = */ t,
- /* .min_keep = */ min_keep,
- /* .seed = */ seed,
- /* .seed_cur = */ seed_cur,
- /* .rng = */ std::mt19937(seed_cur),
- }
- );
-}
-
-// mirostat
-
-struct llama_sampler_mirostat {
- const int32_t n_vocab;
-
- const uint32_t seed;
- uint32_t seed_cur;
-
- const float tau;
- const float eta;
-
- const int32_t m;
-
- float mu;
-
- std::mt19937 rng;
-};
-
-static const char * llama_sampler_mirostat_name(const struct llama_sampler * /*smpl*/) {
- return "mirostat";
-}
-
-static void llama_sampler_mirostat_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_mirostat *) smpl->ctx;
-
- llama_sampler_softmax_impl(cur_p, true);
-
- // Estimate s_hat using the most probable m tokens
- float s_hat = 0.0;
- float sum_ti_bi = 0.0;
- float sum_ti_sq = 0.0;
- for (size_t i = 0; i < size_t(ctx->m - 1) && i < cur_p->size - 1; ++i) {
- float t_i = logf(float(i + 2) / float(i + 1));
- float b_i = logf(cur_p->data[i].p / cur_p->data[i + 1].p);
- sum_ti_bi += t_i * b_i;
- sum_ti_sq += t_i * t_i;
- }
- s_hat = sum_ti_bi / sum_ti_sq;
-
- // Compute k from the estimated s_hat and target surprise value
- float epsilon_hat = s_hat - 1;
- float k = powf((epsilon_hat * powf(2, ctx->mu)) / (1 - powf(ctx->n_vocab, -epsilon_hat)), 1 / s_hat);
-
- llama_sampler_top_k_impl(cur_p, std::max(int(k), 1));
-
- llama_sampler_softmax_impl(cur_p, true);
-
- const int idx = llama_sample_dist(cur_p, ctx->rng);
-
- cur_p->selected = idx;
-
- float observed_surprise = -log2f(cur_p->data[idx].p);
- float e = observed_surprise - ctx->tau;
-
- // Update mu using the learning rate and error
- ctx->mu = ctx->mu - ctx->eta * e;
-}
-
-static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_mirostat *) smpl->ctx;
- auto * result = llama_sampler_init_mirostat(ctx->n_vocab, ctx->seed, ctx->tau, ctx->eta, ctx->m);
-
- // copy the state
- {
- auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx;
-
- result_ctx->mu = ctx->mu;
- result_ctx->rng = ctx->rng;
- }
-
- return result;
-}
-
-static void llama_sampler_mirostat_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_mirostat *) smpl->ctx;
- ctx->mu = 2.0f*ctx->tau;
- ctx->seed_cur = get_rng_seed(ctx->seed);
- ctx->rng.seed(ctx->seed_cur);
-}
-
-static void llama_sampler_mirostat_free(struct llama_sampler * smpl) {
- delete (llama_sampler_mirostat *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_mirostat_i = {
- /* .name = */ llama_sampler_mirostat_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_mirostat_apply,
- /* .reset = */ llama_sampler_mirostat_reset,
- /* .clone = */ llama_sampler_mirostat_clone,
- /* .free = */ llama_sampler_mirostat_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) {
- const auto seed_cur = get_rng_seed(seed);
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_mirostat_i,
- /* .ctx = */ new llama_sampler_mirostat {
- /* .n_vocab = */ n_vocab,
- /* .seed = */ seed,
- /* .seed_cur = */ seed_cur,
- /* .tau = */ tau,
- /* .eta = */ eta,
- /* .m = */ m,
- /* .mu = */ 2.0f*tau,
- /* .rng = */ std::mt19937(seed_cur),
- }
- );
-}
-
-// mirostat v2
-
-struct llama_sampler_mirostat_v2 {
- const uint32_t seed;
- uint32_t seed_cur;
-
- const float tau;
- const float eta;
-
- float mu;
-
- std::mt19937 rng;
-};
-
-static const char * llama_sampler_mirostat_v2_name(const struct llama_sampler * /*smpl*/) {
- return "mirostat-v2";
-}
-
-static void llama_sampler_mirostat_v2_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_mirostat_v2 *) smpl->ctx;
-
- llama_sampler_softmax_impl(cur_p, true);
-
- // Truncate the words with surprise values greater than mu
- cur_p->size = std::distance(cur_p->data, std::find_if(cur_p->data, cur_p->data + cur_p->size, [&](const llama_token_data & candidate) {
- return -log2f(candidate.p) > ctx->mu;
- }));
-
- if (cur_p->size == 0) {
- cur_p->size = 1;
- }
-
- // Normalize the probabilities of the remaining words
- llama_sampler_softmax_impl(cur_p, true);
-
- const int idx = llama_sample_dist(cur_p, ctx->rng);
-
- cur_p->selected = idx;
-
- float observed_surprise = -log2f(cur_p->data[idx].p);
- float e = observed_surprise - ctx->tau;
-
- // Update mu using the learning rate and error
- ctx->mu = ctx->mu - ctx->eta * e;
-}
-
-static void llama_sampler_mirostat_v2_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_mirostat_v2 *) smpl->ctx;
- ctx->mu = 2.0f*ctx->tau;
- ctx->seed_cur = get_rng_seed(ctx->seed);
- ctx->rng.seed(ctx->seed_cur);
-}
-
-static struct llama_sampler * llama_sampler_mirostat_v2_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_mirostat_v2 *) smpl->ctx;
-
- auto * result = llama_sampler_init_mirostat_v2(ctx->seed, ctx->tau, ctx->eta);
-
- // copy the state
- {
- auto * result_ctx = (llama_sampler_mirostat_v2 *) result->ctx;
-
- result_ctx->mu = ctx->mu;
- result_ctx->rng = ctx->rng;
- }
-
- return result;
-}
-
-static void llama_sampler_mirostat_v2_free(struct llama_sampler * smpl) {
- delete (llama_sampler_mirostat_v2 *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_mirostat_v2_i = {
- /* .name = */ llama_sampler_mirostat_v2_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_mirostat_v2_apply,
- /* .reset = */ llama_sampler_mirostat_v2_reset,
- /* .clone = */ llama_sampler_mirostat_v2_clone,
- /* .free = */ llama_sampler_mirostat_v2_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) {
- auto seed_cur = get_rng_seed(seed);
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_mirostat_v2_i,
- /* .ctx = */ new llama_sampler_mirostat_v2 {
- /* .seed = */ seed,
- /* .seed_cur = */ seed_cur,
- /* .tau = */ tau,
- /* .eta = */ eta,
- /* .mu = */ 2.0f*tau,
- /* .rng = */ std::mt19937(seed_cur),
- }
- );
-}
-
-// grammar
-
-struct llama_sampler_grammar {
- const struct llama_vocab * vocab;
-
- std::string grammar_str;
- std::string grammar_root;
-
- struct llama_grammar * grammar;
-};
-
-static const char * llama_sampler_grammar_name(const struct llama_sampler * /*smpl*/) {
- return "grammar";
-}
-
-static void llama_sampler_grammar_accept_impl(struct llama_sampler * smpl, llama_token token) {
- auto * ctx = (llama_sampler_grammar *) smpl->ctx;
- if (ctx->grammar) {
- llama_grammar_accept_impl(*ctx->grammar, token);
- }
-}
-
-static void llama_sampler_grammar_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_grammar *) smpl->ctx;
- if (ctx->grammar) {
- llama_grammar_apply_impl(*ctx->grammar, cur_p);
- }
-}
-
-// Fwd declare to break reset --> init_impl --> llama_sampler_grammar_i --> reset cycle.
-static struct llama_sampler * llama_sampler_init_grammar_impl(
- const struct llama_vocab * vocab,
- const char * grammar_str,
- const char * grammar_root,
- bool lazy,
- const char ** trigger_words,
- size_t num_trigger_words,
- const llama_token * trigger_tokens,
- size_t num_trigger_tokens,
- const char ** trigger_patterns,
- size_t num_trigger_patterns);
-
-static void llama_sampler_grammar_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_grammar *) smpl->ctx;
- if (!ctx->grammar) {
- return;
- }
-
- std::vector<const char *> trigger_patterns_c;
- trigger_patterns_c.reserve(ctx->grammar->trigger_patterns.size());
- for (auto & trigger_pattern : ctx->grammar->trigger_patterns) {
- trigger_patterns_c.push_back(trigger_pattern.pattern.c_str());
- }
-
- auto * grammar_new = llama_grammar_init_impl(ctx->grammar->vocab, ctx->grammar_str.c_str(), ctx->grammar_root.c_str(),
- ctx->grammar->lazy, trigger_patterns_c.data(), trigger_patterns_c.size(),
- ctx->grammar->trigger_tokens.data(), ctx->grammar->trigger_tokens.size());
-
- llama_grammar_free_impl(ctx->grammar);
- ctx->grammar = grammar_new;
-}
-
-static struct llama_sampler * llama_sampler_grammar_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_grammar *) smpl->ctx;
-
- auto * result = llama_sampler_init_grammar_impl(ctx->vocab, nullptr, nullptr, false, nullptr, 0, nullptr, 0, nullptr, 0);
- GGML_ASSERT(result);
-
- // copy the state
- {
- auto * result_ctx = (llama_sampler_grammar *) result->ctx;
-
- if (ctx->grammar) {
- result_ctx->grammar_str = ctx->grammar_str;
- result_ctx->grammar_root = ctx->grammar_root;
-
- result_ctx->grammar = llama_grammar_clone_impl(*ctx->grammar);
- }
- }
-
- return result;
-}
-
-static void llama_sampler_grammar_free(struct llama_sampler * smpl) {
- const auto * ctx = (llama_sampler_grammar *) smpl->ctx;
-
- if (ctx->grammar) {
- llama_grammar_free_impl(ctx->grammar);
- }
-
- delete ctx;
-}
-
-static struct llama_sampler_i llama_sampler_grammar_i = {
- /* .name = */ llama_sampler_grammar_name,
- /* .accept = */ llama_sampler_grammar_accept_impl,
- /* .apply = */ llama_sampler_grammar_apply,
- /* .reset = */ llama_sampler_grammar_reset,
- /* .clone = */ llama_sampler_grammar_clone,
- /* .free = */ llama_sampler_grammar_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-static struct llama_sampler * llama_sampler_init_grammar_impl(
- const struct llama_vocab * vocab,
- const char * grammar_str,
- const char * grammar_root,
- bool lazy,
- const char ** trigger_words,
- size_t num_trigger_words,
- const llama_token * trigger_tokens,
- size_t num_trigger_tokens,
- const char ** trigger_patterns,
- size_t num_trigger_patterns) {
- auto * ctx = new llama_sampler_grammar;
-
- if (grammar_str != nullptr && grammar_str[0] != '\0') {
- std::string trigger_pattern;
- llama_grammar * grammar = nullptr;
- // TODO: remove trigger_words support.
- if (trigger_words != nullptr && num_trigger_words > 0) {
- GGML_ASSERT(trigger_patterns == nullptr && num_trigger_patterns == 0);
- trigger_pattern = "[\\s\\S]*?(";
- for (size_t i = 0; i < num_trigger_words; ++i) {
- static const std::regex special_chars("[.^$|()*+?\\[\\]{}\\\\]");
- if (i > 0) {
- trigger_pattern += "|";
- }
- trigger_pattern += std::regex_replace(trigger_words[i], special_chars, "\\$0");
- }
- trigger_pattern += ")[\\s\\S]*";
-
- std::array<const char *, 1> tmp_trigger_patterns = { trigger_pattern.c_str() };
- grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, tmp_trigger_patterns.data(), tmp_trigger_patterns.size(), trigger_tokens, num_trigger_tokens);
- } else {
- grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, trigger_patterns, num_trigger_patterns, trigger_tokens, num_trigger_tokens);
- }
- *ctx = {
- /* .vocab = */ vocab,
- /* .grammar_str = */ grammar_str,
- /* .grammar_root = */ grammar_root,
- /* .grammar = */ grammar,
- };
- if (!ctx->grammar) {
- delete ctx;
- return nullptr;
- }
- } else {
- *ctx = {
- /* .vocab = */ vocab,
- /* .grammar_str = */ {},
- /* .grammar_root = */ {},
- /* .grammar = */ nullptr,
- };
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_grammar_i,
- /* .ctx = */ ctx
- );
-}
-
-struct llama_sampler * llama_sampler_init_grammar(
- const struct llama_vocab * vocab,
- const char * grammar_str,
- const char * grammar_root) {
- return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ false, nullptr, 0, nullptr, 0, nullptr, 0);
-}
-
-struct llama_sampler * llama_sampler_init_grammar_lazy(
- const struct llama_vocab * vocab,
- const char * grammar_str,
- const char * grammar_root,
- const char ** trigger_words,
- size_t num_trigger_words,
- const llama_token * trigger_tokens,
- size_t num_trigger_tokens) {
- return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ true, trigger_words, num_trigger_words, trigger_tokens, num_trigger_tokens, nullptr, 0);
-}
-
-struct llama_sampler * llama_sampler_init_grammar_lazy_patterns(
- const struct llama_vocab * vocab,
- const char * grammar_str,
- const char * grammar_root,
- const char ** trigger_patterns,
- size_t num_trigger_patterns,
- const llama_token * trigger_tokens,
- size_t num_trigger_tokens) {
- return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ true, nullptr, 0, trigger_tokens, num_trigger_tokens, trigger_patterns, num_trigger_patterns);
-}
-
-// penalties
-
-struct llama_sampler_penalties {
- const int32_t penalty_last_n;
- const float penalty_repeat;
- const float penalty_freq;
- const float penalty_present;
-
- ring_buffer<llama_token> prev;
-
- // a frequency map to count token occurrences
- std::unordered_map<llama_token, int> token_count;
-};
-
-static const char * llama_sampler_penalties_name(const struct llama_sampler * /*smpl*/) {
- return "penalties";
-}
-
-static void llama_sampler_penalties_accept(struct llama_sampler * smpl, llama_token token) {
- auto * ctx = (llama_sampler_penalties *) smpl->ctx;
- if (ctx->penalty_last_n == 0) {
- return;
- }
-
- ctx->token_count[token]++;
-
- // if the ring buffer is full, remove the oldest token
- if (ctx->prev.size() >= (size_t) ctx->penalty_last_n) {
- const auto old = ctx->prev.front();
-
- ctx->token_count[old]--;
- if (ctx->token_count[old] == 0) {
- ctx->token_count.erase(old);
- }
- }
-
- ctx->prev.push_back(token);
-
-#if 0
- // sanity check
- std::unordered_map<llama_token, int> tmp;
- for (int i = 0; i < std::min<int>(ctx->penalty_last_n, ctx->prev.size()); ++i) {
- tmp[ctx->prev.rat(i)]++;
- }
-
- assert(ctx->token_count == tmp);
-#endif
-}
-
-static void llama_sampler_penalties_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_penalties *) smpl->ctx;
-
- if ((ctx->penalty_last_n == 0) ||
- (ctx->penalty_repeat == 1.0f && ctx->penalty_freq == 0.0f && ctx->penalty_present == 0.0f)) {
- return;
- }
-
- // Apply frequency and presence penalties to the cur_p
- for (size_t i = 0; i < cur_p->size; ++i) {
- const auto token_iter = ctx->token_count.find(cur_p->data[i].id);
- if (token_iter == ctx->token_count.end()) {
- continue;
- }
-
- const int count = token_iter->second;
-
- assert(count > 0 && count <= ctx->penalty_last_n);
-
- // The academic publication that described this technique actually just only divided, but that would cause tokens with negative logits to become more likely, which is obviously wrong.
- // This is common fix for this problem, which is to multiply by the penalty instead of dividing.
- if (cur_p->data[i].logit <= 0) {
- cur_p->data[i].logit *= ctx->penalty_repeat;
- } else {
- cur_p->data[i].logit /= ctx->penalty_repeat;
- }
-
- cur_p->data[i].logit -= float(count) * ctx->penalty_freq + float(count > 0) * ctx->penalty_present;
- }
-
- cur_p->sorted = false;
-}
-
-static void llama_sampler_penalties_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_penalties *) smpl->ctx;
- ctx->prev.clear();
- ctx->token_count.clear();
-}
-
-static struct llama_sampler * llama_sampler_penalties_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_penalties *) smpl->ctx;
- auto * result = llama_sampler_init_penalties(
- ctx->penalty_last_n,
- ctx->penalty_repeat,
- ctx->penalty_freq,
- ctx->penalty_present);
-
- // copy the state
- {
- auto * result_ctx = (llama_sampler_penalties *) result->ctx;
-
- result_ctx->prev = ctx->prev;
- }
-
- return result;
-}
-
-static void llama_sampler_penalties_free(struct llama_sampler * smpl) {
- delete (llama_sampler_penalties *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_penalties_i = {
- /* .name = */ llama_sampler_penalties_name,
- /* .accept = */ llama_sampler_penalties_accept,
- /* .apply = */ llama_sampler_penalties_apply,
- /* .reset = */ llama_sampler_penalties_reset,
- /* .clone = */ llama_sampler_penalties_clone,
- /* .free = */ llama_sampler_penalties_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_penalties(
- int32_t penalty_last_n,
- float penalty_repeat,
- float penalty_freq,
- float penalty_present) {
- penalty_last_n = std::max(penalty_last_n, 0);
-
- const bool is_empty = (penalty_last_n == 0 || (penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f));
-
- if (is_empty) {
- return llama_sampler_init_empty("?penalties");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_penalties_i,
- /* .ctx = */ new llama_sampler_penalties {
- /* .penalty_last_n = */ penalty_last_n,
- /* .penalty_repeat = */ penalty_repeat,
- /* .penalty_freq = */ penalty_freq,
- /* .penalty_present = */ penalty_present,
- /* .prev = */ ring_buffer<llama_token>(penalty_last_n),
- /* .token_count = */ {},
- }
- );
-}
-
-// top-n-sigma
-
-struct llama_sampler_top_n_sigma {
- const float n;
-};
-
-static const char * llama_sampler_top_n_sigma_name(const struct llama_sampler * /*smpl*/) {
- return "top-n-sigma";
-}
-
-static void llama_sampler_top_n_sigma_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_top_n_sigma *) smpl->ctx;
-
- if (ctx->n <= 0.0f || cur_p->size <= 1) {
- return;
- }
-
- // find max logit and calculate mean
- float max = cur_p->data[0].logit;
- float logits_sum = 0;
- size_t valid_count = 0;
- for (size_t i = 0; i < cur_p->size; ++i) {
- // Only count non-negative infinity values
- if (cur_p->data[i].logit != -INFINITY) {
- max = std::max(max, cur_p->data[i].logit);
- logits_sum += cur_p->data[i].logit;
- valid_count++;
- }
- }
- float mean = valid_count > 0 ? logits_sum/valid_count : 0;
-
- // calculate standard deviation
- float acc = 0;
- for (size_t i = 0; i < cur_p->size; ++i) {
- // Skip -infinity in std calculation
- if (cur_p->data[i].logit != -INFINITY) {
- acc += pow(cur_p->data[i].logit - mean, 2);
- }
- }
- float std = valid_count > 0 ? sqrt(acc/valid_count) : 0;
-
- // apply mask
- for (size_t i = 0; i < cur_p->size; ++i) {
- if (cur_p->data[i].logit < max - (ctx->n * std)) {
- cur_p->data[i].logit = -INFINITY;
- }
- }
-
- llama_sampler_softmax_impl(cur_p, true);
-}
-
-static struct llama_sampler * llama_sampler_top_n_sigma_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_top_n_sigma *) smpl->ctx;
- return llama_sampler_init_top_n_sigma(ctx->n);
-}
-
-static void llama_sampler_top_n_sigma_free(struct llama_sampler * smpl) {
- delete (llama_sampler_top_n_sigma *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_top_n_sigma_i = {
- /* .name = */ llama_sampler_top_n_sigma_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_top_n_sigma_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_top_n_sigma_clone,
- /* .free = */ llama_sampler_top_n_sigma_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
- const bool is_empty = (n <= 0.0f);
-
- if (is_empty) {
- return llama_sampler_init_empty("?top-n-sigma");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_top_n_sigma_i,
- /* .ctx = */ new llama_sampler_top_n_sigma {
- /* .n = */ n,
- }
- );
-}
-
-// DRY
-
-struct llama_sampler_dry {
- int32_t total_context_size;
-
- const float dry_multiplier;
- const float dry_base;
- const int32_t dry_allowed_length;
- const int32_t dry_penalty_last_n;
-
- std::unordered_multimap<llama_token, std::vector<llama_token>> dry_processed_breakers;
- std::vector<int> dry_repeat_count;
- std::unordered_map<llama_token, int> dry_max_token_repeat;
- ring_buffer<llama_token> last_tokens;
-};
-
-// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
-static void get_overlapping_token_sequences(const llama_vocab & vocab, const std::string& str, std::unordered_multimap<llama_token, std::vector<llama_token>>& token_sequences, int max_tail_len = -1) {
- for (llama_token token_id = 0; token_id < (llama_token) vocab.n_tokens(); token_id++) {
- std::string word = vocab.detokenize({token_id}, true);
- if (word.find(str) != std::string::npos) {
- token_sequences.emplace(token_id, std::vector<llama_token>());
- } else {
- size_t word_len = word.size();
- size_t str_len = str.size();
- size_t pos = -1;
- while ((pos = word.find(str[0], pos + 1)) != std::string::npos) {
- bool match = true;
- size_t i;
- for (i = 1; i < str_len && i + pos < word_len; ++i) {
- if (word[pos + i] != str[i]) {
- match = false;
- break;
- }
- }
- if (match) {
- std::vector<llama_token> tokenization = vocab.tokenize(str.substr(i), false, false);
- if (max_tail_len >= 0 && tokenization.size() > (size_t)max_tail_len) {
- tokenization.resize(max_tail_len);
- }
-
- // Ensure we don't already have a duplicate matching tokenization
- auto its = token_sequences.equal_range(token_id);
- bool found = false;
- for (auto it = its.first; it != its.second; ++it) {
- if (tokenization == it->second) {
- found = true;
- break;
- }
- }
- if (!found) {
- token_sequences.emplace(token_id, tokenization);
- }
- }
- }
- }
- }
-}
-
-static const char * llama_sampler_dry_name(const struct llama_sampler * /*smpl*/) {
- return "dry";
-}
-
-static void llama_sampler_dry_accept(struct llama_sampler * smpl, llama_token token) {
- auto * ctx = (llama_sampler_dry *) smpl->ctx;
- if (ctx->dry_multiplier == 0.0f || ctx->dry_base < 1.0f || ctx->dry_penalty_last_n == 0) {
- return;
- }
-
- ctx->last_tokens.push_back(token);
-}
-
-// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
-static void llama_sampler_dry_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_dry *) smpl->ctx;
-
- if (ctx->dry_multiplier == 0.0f || ctx->dry_base < 1.0f || ctx->dry_penalty_last_n == 0) {
- return;
- }
-
- int32_t effective_dry_penalty_last_n = (ctx->dry_penalty_last_n == -1) ? ctx->total_context_size : std::max(ctx->dry_penalty_last_n, 0);
- int last_n_repeat = std::min(std::min((int)ctx->last_tokens.size(), effective_dry_penalty_last_n), ctx->total_context_size);
-
- if (last_n_repeat <= ctx->dry_allowed_length) {
- return;
- }
-
- ctx->dry_repeat_count.assign(last_n_repeat, 0);
- ctx->dry_max_token_repeat.clear();
-
- // Step 1: Look for restart sequences to limit the maximum repetition length.
- // Work backwards through the context looking for any token that begins a restart sequence.
- //
- // The collection `restart_sequences` is a mapping from a "head" token to all "tail"
- // sequences that together comprise a restart sequence. This allows us to quickly check
- // whether each token is the head of a complete sequence. Most restart sequences are actually
- // a single token, and for these the "tail" is an empty vector.
- //
- // If the token is a "head", test all restart sequences that begin with this token
- // (there will often only be one sequence for each token, but if sequences like 'aaaq1' and
- // 'aaa1' are used as restart strings, both could start with 'aaa' when tokenized). The
- // longest matching sequence (if any) is used to limit the maximum repetition length.
- //
- // Note that in the case case of a short sequence contained in a longer one, this might fail to
- // find the smallest value for `rep_limit`. For example, if 'amniotic' and 'ni' are both used as
- // restart sequences, 'ni' will be found first, and since it's shorter it will fail to suppress
- // 'otic'. This is a minor issue since fully contained restart sequences are likely to be rare.
- //
- // This is theoretically worst-case O(N^2) for arbitrary restart sequences, which is why we
- // have already clamped the maximum tail sequence length when generating `restart_sequences`.
- // With clamping, this scan is O(N) in the context length.
-
- int rep_limit = last_n_repeat;
- for (int i = 0; i < last_n_repeat; ++i) {
- llama_token token = ctx->last_tokens.rat(i);
- auto its = ctx->dry_processed_breakers.equal_range(token);
- if (its.first == ctx->dry_processed_breakers.end()) {
- continue;
- }
- int longest_match = -1;
- for (auto it = its.first; it != its.second; ++it) {
- // Note that (*it) does not contain the head character, so seq_len will be
- // the restart sequence length minus 1.
- // In the common case of a single-token restart sequence, (*it) will be empty
- // and we will trivially match.
- int seq_len = (int)it->second.size();
- if (seq_len > longest_match && seq_len <= (int)i) {
- bool match = true;
- for (int offset = 0; offset < seq_len; ++offset) {
- // The -1 when indexing `last_tokens` is because we already matched the head.
- if (it->second[offset] != ctx->last_tokens.rat(i - offset - 1)) {
- match = false;
- break;
- }
- }
- if (match) {
- longest_match = seq_len;
- }
- }
- }
- if (longest_match >= 0) {
- // We found a restart sequence starting `i` tokens from the end and continuing for
- // `longest_match` tokens.
- rep_limit = i - longest_match;
- break;
- }
- }
- if (rep_limit < ctx->dry_allowed_length) {
- return;
- }
-
- // Step 2: Iterate in reverse over the last N tokens of the context, using the "Z-algorithm" (in
- // the reverse direction) to efficiently compute the positions and lengths of suffixes appearing
- // elsewhere in the context. We limit the suffix length to `rep_limit` to respect restart sequences.
- //
- // This algorithm is not currently documented on Wikipedia, but there is a clear description here:
- // https://ivanyu.me/blog/2014/10/15/z-algorithm/
- //
- // The code below is adapted from the public domain implementation by the same author here:
- // https://github.com/ivanyu/string-algorithms/blob/master/z_algorithm.py
- //
- // Example:
- // Last N tokens: a b c c b c y a b c
- // Repeat counts: 0 0 3 1 0 2 0 0 0 0
- // ^
- // This `3` means that the last three tokens of the context (a b c) also appear here.
- //
- // This step is worst case O(N) since the Z-algorithm is linear, despite the appearance of nested
- // for/while loops. This can be seen by observing that the `lt` and `rt` bounds are set after each
- // repeated suffix is detected (i.e. after each while loop when n > 0). These bound variables
- // ensure that the inner while loops only examine each token in the context once as the outer
- // for loop iterates over the context.
-
- {
- const int last = last_n_repeat - 1;
-
- int rt = 0;
- int lt = 0;
-
- for (int k = 1; k < last_n_repeat; ++k) {
- if (k > rt) {
- // If k is outside the current Z-box, do naive computation.
- int n = 0;
- while (n + k < last_n_repeat && ctx->last_tokens.rat(n) == ctx->last_tokens.rat(n+k)) {
- ++n;
- }
- ctx->dry_repeat_count[last - k] = std::min(n, rep_limit);
- if (n > 0) {
- lt = k;
- rt = k + n - 1;
- }
- } else {
- // If k is inside the current Z-box, consider two cases.
-
- int p = k - lt; // Pair index.
- int right_part_len = rt - k + 1;
-
- if (ctx->dry_repeat_count[last - p] < right_part_len) {
- int n = std::min(ctx->dry_repeat_count[last - p], rep_limit);
- ctx->dry_repeat_count[last - k] = n;
- } else {
- int i = rt + 1;
- while (i < last_n_repeat && ctx->last_tokens.rat(i) == ctx->last_tokens.rat(i - k)) {
- i += 1;
- }
-
- int n = std::min(i - k, rep_limit);
- ctx->dry_repeat_count[last - k] = n;
- lt = k;
- rt = i - 1;
- }
- }
- }
- }
-
- // Step 3: Iterate over dry_repeat_count and last_tokens, examining the maximum repeat length
- // that would be generated by emitting each new token that would extend a sequence.
- //
- // Following the same example as above:
- // Last N tokens: a b c c b c y a b c
- // Repeat counts: 0 0 3 1 0 2 0 0 0 0
- //
- // For each non-zero, look ahead one token. This token, if emitted, would extend the repetition.
- // c: 3 -> 4 (from `a b c` to `a b c c`)
- // b: 1 -> 2 (from `c` to `c b`)
- // y: 2 -> 3 (from `b c` to `b c y`)
-
- for (int i = 0; i < last_n_repeat - 1; ++i) {
- int repeat_len = ctx->dry_repeat_count[i];
- if (repeat_len >= ctx->dry_allowed_length) {
- // This token ends a repeat, so the next token would continue one.
- // By convention, the value of `repeat_len` only includes the tokens currently
- // in the context, not the new token that would be added.
- llama_token token = ctx->last_tokens.rat(last_n_repeat - 2 - i);
- // Track the maximum sequence ending in this token.
- const auto& it = ctx->dry_max_token_repeat.find(token);
- if (it == ctx->dry_max_token_repeat.end() || it->second < repeat_len) {
- ctx->dry_max_token_repeat[token] = repeat_len;
- }
- }
- }
-
- // Step 4: Apply logit penalties based on the maximum repeat length for relevant tokens.
-
- // Prevent floating point overflow in `pow(penalty_base, exponent)` by clamping to `max_exponent`.
- // Compute it from `penalty_base` and the approximate log of `std::numeric_limits<float>::max()`
- const float FLOAT_MAX_LOG = 88.7228391f;
- int max_exponent = 0;
- if (ctx->dry_base > 1.000001f) {
- max_exponent = FLOAT_MAX_LOG / std::log(ctx->dry_base);
- }
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- const auto& af_kvp = ctx->dry_max_token_repeat.find(cur_p->data[i].id);
- if (af_kvp != ctx->dry_max_token_repeat.end()) {
- // Check all sequence breakers starting with this token
- auto range = ctx->dry_processed_breakers.equal_range(cur_p->data[i].id);
- bool is_single_token_breaker = false;
-
- for (auto it = range.first; it != range.second; ++it) {
- if (it->second.empty()) {
- is_single_token_breaker = true;
- break;
- }
- }
-
- // Apply penalty only if it's not a single-token sequence breaker
- if (!is_single_token_breaker) {
- int repeat_exp = af_kvp->second - ctx->dry_allowed_length;
- if (max_exponent > 0 && repeat_exp > max_exponent) {
- repeat_exp = max_exponent;
- }
- float penalty = ctx->dry_multiplier * std::pow(ctx->dry_base, repeat_exp);
- cur_p->data[i].logit -= penalty;
- }
- }
- }
-
- cur_p->sorted = false;
-}
-
-static void llama_sampler_dry_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_dry *) smpl->ctx;
- ctx->last_tokens.clear();
- ctx->dry_repeat_count.clear();
- ctx->dry_max_token_repeat.clear();
-}
-
-static struct llama_sampler * llama_sampler_dry_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (llama_sampler_dry *) smpl->ctx;
-
- llama_vocab dummy_vocab;
-
- // dummy vocab is passed because it is only needed for raw sequence breaker processing, which we have already done and will simply be copying
- auto * result = llama_sampler_init_dry(&dummy_vocab, ctx->total_context_size, ctx->dry_multiplier, ctx->dry_base, ctx->dry_allowed_length, ctx->dry_penalty_last_n, NULL, 0);
-
- // Copy the state, including the processed breakers
- {
- auto * result_ctx = (llama_sampler_dry *) result->ctx;
- result_ctx->dry_processed_breakers = ctx->dry_processed_breakers;
- result_ctx->dry_repeat_count = ctx->dry_repeat_count;
- result_ctx->dry_max_token_repeat = ctx->dry_max_token_repeat;
- result_ctx->last_tokens = ctx->last_tokens;
- }
-
- return result;
-}
-
-static void llama_sampler_dry_free(struct llama_sampler * smpl) {
- delete (llama_sampler_dry *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_dry_i = {
- /* .name = */ llama_sampler_dry_name,
- /* .accept = */ llama_sampler_dry_accept,
- /* .apply = */ llama_sampler_dry_apply,
- /* .reset = */ llama_sampler_dry_reset,
- /* .clone = */ llama_sampler_dry_clone,
- /* .free = */ llama_sampler_dry_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, int32_t n_ctx_train, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) {
- int32_t effective_dry_penalty_last_n = (dry_penalty_last_n == -1) ? n_ctx_train : std::max(dry_penalty_last_n, 0);
- std::unordered_multimap<llama_token, std::vector<llama_token>> processed_breakers;
- const int MAX_CHAR_LEN = 40;
- const int MAX_SEQ_LEN = 20;
-
- const bool dry_enabled = (dry_multiplier != 0.0f && dry_base >= 1.0f && dry_penalty_last_n != 0);
-
- if (!dry_enabled) {
- return llama_sampler_init_empty("?dry");
- }
-
- if (dry_enabled && seq_breakers != nullptr && num_breakers > 0) {
- // Process sequence breakers
- for (size_t i = 0; i < num_breakers; ++i) {
- if (seq_breakers[i] == nullptr || std::strlen(seq_breakers[i]) == 0) {
- LLAMA_LOG_WARN("skipping null or empty DRY sequence breaker at index %zu\n", i);
- continue;
- }
-
- std::string sequence_break(seq_breakers[i]);
- if (sequence_break.empty()) {
- LLAMA_LOG_WARN("skipping empty DRY sequence breaker\n");
- continue;
- }
-
- if (sequence_break.size() > MAX_CHAR_LEN) {
- LLAMA_LOG_WARN("truncating DRY sequence breaker to %d characters\n", MAX_CHAR_LEN);
- sequence_break.resize(MAX_CHAR_LEN);
- }
-
- get_overlapping_token_sequences(*vocab, sequence_break, processed_breakers, MAX_SEQ_LEN);
- }
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_dry_i,
- /* .ctx = */ new llama_sampler_dry {
- /* .total_context_size = */ n_ctx_train,
- /* .dry_multiplier = */ dry_multiplier,
- /* .dry_base = */ dry_base,
- /* .dry_allowed_length = */ dry_allowed_length,
- /* .dry_penalty_last_n = */ dry_penalty_last_n,
- /* .dry_processed_breakers = */ std::move(processed_breakers),
- /* .dry_repeat_count = */ dry_enabled ? std::vector<int>(effective_dry_penalty_last_n, 0) : std::vector<int>{},
- /* .dry_max_token_repeat = */ {},
- /* .last_tokens = */ dry_enabled ? ring_buffer<llama_token>(effective_dry_penalty_last_n) : ring_buffer<llama_token>(0),
- }
- );
-}
-
-// wrapper for test-sampling.cpp
-struct llama_sampler * llama_sampler_init_dry_testing(int32_t context_size, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const std::vector<std::vector<llama_token>>& seq_breakers) {
- llama_vocab dummy_vocab;
- auto * result = llama_sampler_init_dry(&dummy_vocab, context_size, dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, NULL, 0);
- auto * ctx = (llama_sampler_dry *) result->ctx;
-
- // Process the token-based sequence breakers
- ctx->dry_processed_breakers.clear();
- if (seq_breakers.empty()) {
- LLAMA_LOG_WARN("empty DRY sequence breakers list in llama_sampler_init_dry_testing\n");
- } else {
- for (const auto& breaker : seq_breakers) {
- if (breaker.empty()) {
- LLAMA_LOG_WARN("skipping DRY empty sequence breaker\n");
- continue;
- }
- llama_token head_token = breaker[0];
- std::vector<llama_token> tail_tokens(breaker.begin() + 1, breaker.end());
- ctx->dry_processed_breakers.emplace(head_token, std::move(tail_tokens));
- }
-
- if (ctx->dry_processed_breakers.empty()) {
- LLAMA_LOG_WARN("no valid DRY sequence breakers processed in llama_sampler_init_dry_testing\n");
- }
- }
-
- return result;
-}
-
-// adaptive-p sampler state
-//
-// maintains an exponential moving average of the *ORIGINAL* probabilities
-// of selected tokens, used to compute an adapted target at each sampling step.
-//
-// see llama.h for a full description of the sampler
-//
-// ref: https://github.com/ggml-org/llama.cpp/pull/17927
-//
-struct llama_sampler_adaptive_p {
- const float target; // target probability (0.0 - 1.0; negative = disabled)
- const float decay; // EMA decay; history ~= 1/(1-decay) tokens (0.0 - 0.99)
- const uint32_t seed; // original RNG seed
- uint32_t seed_cur; // actual RNG seed
- std::mt19937 rng; // RNG state
- float weighted_sum; // sum(p_i * decay^i)
- float total_weight; // sum(decay^i), converges to 1/(1-decay)
- std::vector<float> original_probs; // pre-transform probs, cached for EMA update
- llama_token pending_token_id; // token ID of selected token
- int32_t pending_token_idx; // index of orig. prob. of selected token in original_probs
-};
-
-// adaptive probability transformation constants
-static constexpr float DISTRIBUTION_WIDTH = 0.3f;
-static constexpr float PEAK_LOGIT_VALUE = 5.0f;
-static constexpr float SHARPNESS = 10.0f;
-static constexpr float INV_WIDTH = 1.0f / DISTRIBUTION_WIDTH;
-
-static const char * llama_sampler_adaptive_p_name(const struct llama_sampler * /*smpl*/) {
- return "adaptive-p";
-}
-
-static void llama_sampler_adaptive_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
-
- llama_sampler_softmax_impl(cur_p, false);
-
- if (ctx->target < 0.0f) {
- // at negative target values, adaptive-p is no-op
- // we simply sample from the existing distribution
- cur_p->selected = llama_sample_dist(cur_p, ctx->rng);
- return;
- }
-
- // store the original probabilities
- ctx->original_probs.resize(cur_p->size);
- for (size_t i = 0; i < cur_p->size; ++i) {
- ctx->original_probs[i] = cur_p->data[i].p;
- }
-
- // using the EMA, compute the adapted target probability for the current sampling step
- auto target = std::clamp(ctx->target, 0.0f, 1.0f);
- float adapted_target = std::clamp(
- ctx->total_weight == 0.0f ? target : 2.0f * target - (ctx->weighted_sum / ctx->total_weight),
- 0.0f, 1.0f
- );
-
- // adaptive probability transform
- //
- // quadratic near target for fine differentiation, transitioning to linear decay in the
- // tails. unbounded negative logits ensure proper suppression of far-from-target tokens
- // after the softmax.
- //
- for (size_t i = 0; i < cur_p->size; ++i) {
- if (cur_p->data[i].logit == -INFINITY) {
- // don't transform logits that are -INFINITY
- // (as masked out by e.g. min-p and top-p when using backend sampling)
- continue;
- }
- float dist = std::abs((cur_p->data[i].p - adapted_target) * INV_WIDTH);
- cur_p->data[i].logit = PEAK_LOGIT_VALUE - SHARPNESS * dist * dist / (1.0f + dist);
- }
-
- // softmax and sample from the transformed distribution
- llama_sampler_softmax_impl(cur_p, false);
- const int idx = llama_sample_dist(cur_p, ctx->rng);
- cur_p->selected = idx;
-
- // store the selected token ID for acceptance later
- ctx->pending_token_id = cur_p->data[idx].id;
- ctx->pending_token_idx = idx;
-}
-
-static void llama_sampler_adaptive_p_accept(struct llama_sampler * smpl, llama_token token) {
- auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
- if (ctx->pending_token_id == token) {
- GGML_ASSERT(ctx->pending_token_id != LLAMA_TOKEN_NULL);
- GGML_ASSERT(ctx->pending_token_idx != -1);
- // update EMA with the original probability of the selected token
- ctx->weighted_sum = ctx->original_probs[ctx->pending_token_idx] + ctx->decay * ctx->weighted_sum;
- ctx->total_weight = 1.0f + ctx->decay * ctx->total_weight;
- }
- ctx->pending_token_id = LLAMA_TOKEN_NULL;
- ctx->pending_token_idx = -1;
-}
-
-static void llama_sampler_adaptive_p_reset(struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
- // ctx->target and ctx->decay never change after init, so it's safe to keep them as is.
- // original_probs is completely overwritten on every call to _apply.
- // so we only need to reset the EMA state and pending token.
- ctx->weighted_sum = ctx->target / (1.0f - ctx->decay);
- ctx->total_weight = 1.0f / (1.0f - ctx->decay);
- ctx->pending_token_id = LLAMA_TOKEN_NULL;
- ctx->pending_token_idx = -1;
- ctx->seed_cur = get_rng_seed(ctx->seed);
- ctx->rng.seed(ctx->seed_cur);
-}
-
-static struct llama_sampler * llama_sampler_adaptive_p_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_adaptive_p *) smpl->ctx;
- auto * result = llama_sampler_init_adaptive_p(ctx->target, ctx->decay, ctx->seed);
- auto * result_ctx = (llama_sampler_adaptive_p *) result->ctx;
-
- // copy everything (target, decay, seed, and RNG are already set)
- result_ctx->weighted_sum = ctx->weighted_sum;
- result_ctx->total_weight = ctx->total_weight;
- result_ctx->pending_token_id = ctx->pending_token_id;
- result_ctx->pending_token_idx = ctx->pending_token_idx;
-
- return result;
-}
-
-static void llama_sampler_adaptive_p_free(struct llama_sampler * smpl) {
- delete (llama_sampler_adaptive_p *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_adaptive_p_i = {
- /* .name = */ llama_sampler_adaptive_p_name,
- /* .accept = */ llama_sampler_adaptive_p_accept,
- /* .apply = */ llama_sampler_adaptive_p_apply,
- /* .reset = */ llama_sampler_adaptive_p_reset,
- /* .clone = */ llama_sampler_adaptive_p_clone,
- /* .free = */ llama_sampler_adaptive_p_free,
- /* .backend_init = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_adaptive_p(
- float target,
- float decay,
- uint32_t seed
-) {
- auto seed_cur = get_rng_seed(seed);
- float clamped_decay = std::clamp(decay, 0.0f, 0.99f);
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_adaptive_p_i,
- /* .ctx = */ new llama_sampler_adaptive_p {
- /* .target = */ target,
- /* .decay = */ clamped_decay,
- /* .seed = */ seed,
- /* .seed_cur = */ seed_cur,
- /* .rng = */ std::mt19937(seed_cur),
- /* .weighted_sum = */ target / (1.0f - clamped_decay),
- /* .total_weight = */ 1.0f / (1.0f - clamped_decay),
- /* .original_probs = */ {},
- /* .pending_token_id = */ LLAMA_TOKEN_NULL,
- /* .pending_token_idx = */ -1
- }
- );
-}
-
-// logit-bias
-
-struct llama_sampler_logit_bias : public llama_sampler_backend {
- const int32_t n_vocab;
-
- const std::vector<llama_logit_bias> logit_bias;
-
- std::vector<llama_logit_bias> to_search;
-
- struct ggml_tensor * inp_logit_bias;
- struct ggml_tensor * inp_logit_idxs;
-
- ggml_context_ptr inp_ctx;
- ggml_backend_buffer_ptr inp_buf;
-};
-
-static const char * llama_sampler_logit_bias_name(const struct llama_sampler * smpl) {
- auto * ctx = (llama_sampler_logit_bias *) smpl->ctx;
- return ctx->get_name();
-}
-
-static void llama_sampler_logit_bias_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_logit_bias *) smpl->ctx;
-
- if (ctx->logit_bias.empty()) {
- return;
- }
-
- ctx->to_search.clear();
-
- // update the candidates that have not been shuffled in the vocabulary (i.e. idx == id)
- for (const auto & lb : ctx->logit_bias) {
- if (lb.token >= 0 && cur_p->size > (size_t) lb.token && cur_p->data[lb.token].id == lb.token) {
- cur_p->data[lb.token].logit += lb.bias;
- } else {
- ctx->to_search.push_back(lb);
- }
- }
-
- if (ctx->to_search.empty()) {
- return;
- }
-
- // search for the remaining candidates that were not found in the previous step
- for (size_t i = 0; i < cur_p->size; ++i) {
- for (const auto & lb : ctx->to_search) {
- if (cur_p->data[i].id == lb.token) {
- cur_p->data[i].logit += lb.bias;
- break;
- }
- }
- }
-}
-
-static struct llama_sampler * llama_sampler_logit_bias_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_logit_bias *) smpl->ctx;
- return llama_sampler_init_logit_bias(ctx->n_vocab, ctx->logit_bias.size(), ctx->logit_bias.data());
-}
-
-static void llama_sampler_logit_bias_free(struct llama_sampler * smpl) {
- delete (llama_sampler_logit_bias *) smpl->ctx;
-}
-
-static void llama_sampler_logit_bias_backend_apply(
- struct llama_sampler * smpl,
- struct ggml_context * ctx,
- struct ggml_cgraph * gf,
- struct llama_sampler_data * data) {
- GGML_UNUSED(gf);
- GGML_UNUSED(ctx);
-
- auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
- if (sctx->logit_bias.empty()) {
- return;
- }
-
- ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f);
-
- cur = ggml_reshape_2d(ctx, cur, 1, ggml_nelements(cur));
- cur = ggml_set_rows(ctx, cur, sctx->inp_logit_bias, sctx->inp_logit_idxs);
- cur = ggml_reshape_1d(ctx, cur, ggml_nelements(cur));
-
- data->logits = ggml_add(ctx, data->logits, cur);
-}
-
-static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * smpl) {
- auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
- if (sctx->logit_bias.empty()) {
- return;
- }
-
- GGML_ASSERT(sctx->inp_logit_bias != nullptr);
- GGML_ASSERT(sctx->inp_logit_idxs != nullptr);
-
- const size_t n = sctx->logit_bias.size();
-
- std::vector<float> data_logit_bias(n, 0.0f);
- std::vector<int32_t> data_logit_idxs(n, 0);
- for (size_t i = 0; i < n; ++i) {
- const auto & lb = sctx->logit_bias[i];
- GGML_ASSERT(lb.token >= 0 && lb.token < (int32_t) sctx->n_vocab);
- data_logit_bias[i] = lb.bias;
- data_logit_idxs[i] = lb.token;
- }
-
- ggml_backend_tensor_set(sctx->inp_logit_bias, data_logit_bias.data(), 0, ggml_nbytes(sctx->inp_logit_bias));
- ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs));
-}
-
-static bool llama_sampler_logit_bias_backend_init(
- struct llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
-
- sctx->init(true);
-
- if (sctx->logit_bias.empty()) {
- return true;
- }
-
- ggml_init_params params = {
- /*.mem_size =*/ 2*ggml_tensor_overhead(),
- /*.mem_buffer =*/ nullptr,
- /*.no_alloc =*/ true,
- };
-
- sctx->inp_ctx.reset(ggml_init(params));
-
- const size_t n = sctx->logit_bias.size();
-
- sctx->inp_logit_bias = ggml_new_tensor_2d(sctx->inp_ctx.get(), GGML_TYPE_F32, 1, n);
- ggml_set_name(sctx->inp_logit_bias, "logit_bias");
- ggml_set_input(sctx->inp_logit_bias);
-
- sctx->inp_logit_idxs = ggml_new_tensor_1d(sctx->inp_ctx.get(), GGML_TYPE_I32, n);
- ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
- ggml_set_input(sctx->inp_logit_idxs);
-
- // Allocate all tensors from our context to the backend
- sctx->inp_buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(sctx->inp_ctx.get(), buft));
-
- ggml_backend_buffer_clear(sctx->inp_buf.get(), 0);
-
- return true;
-}
-
-static struct llama_sampler_i llama_sampler_logit_bias_i = {
- /* .name = */ llama_sampler_logit_bias_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_logit_bias_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_logit_bias_clone,
- /* .free = */ llama_sampler_logit_bias_free,
- /* .backend_init = */ llama_sampler_logit_bias_backend_init,
- /* .backend_accept = */ nullptr,
- /* .backend_apply = */ llama_sampler_logit_bias_backend_apply,
- /* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input,
-};
-
-struct llama_sampler * llama_sampler_init_logit_bias(
- int32_t n_vocab,
- int32_t n_logit_bias,
- const llama_logit_bias * logit_bias) {
- const bool is_empty = n_logit_bias <= 0;
-
- if (is_empty) {
- return llama_sampler_init_empty("?logit-bias");
- }
-
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_logit_bias_i,
- /* .ctx = */ new llama_sampler_logit_bias {
- ("logit-bias"),
- /* .n_vocab = */ n_vocab,
- /* .logit_bias = */ std::vector<llama_logit_bias>(logit_bias, logit_bias + n_logit_bias),
- /* .to_search = */ {},
- /* .inp_logit_bias = */ nullptr,
- /* .inp_logit_idxs = */ nullptr,
- /* .inp_ctx = */ nullptr,
- /* .inp_buf = */ nullptr,
- }
- );
-}
-
-// infill
-
-//#define GGML_DEBUG_SAMPLER_INFILL
-
-struct llama_sampler_infill {
- const struct llama_vocab * vocab;
-
- std::vector<char> buf0;
- std::vector<char> buf1;
-};
-
-static const char * llama_sampler_infill_name(const struct llama_sampler * /*smpl*/) {
- return "infill";
-}
-
-static void llama_sampler_infill_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
- auto * ctx = (llama_sampler_infill *) smpl->ctx;
-
- llama_sampler_softmax_impl(cur_p, true);
-
-#if defined(GGML_DEBUG_SAMPLER_INFILL)
-#define LOG_DBG_CUR LLAMA_LOG_DEBUG
-#else
-#define LOG_DBG_CUR(...)
-#endif
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
- }
-
- float p_txt_sum = 0.0f;
- float p_eog_sum = 0.0f;
-
- for (size_t i = 0; i < cur_p->size; ++i) {
- if (ctx->vocab->is_eog(cur_p->data[i].id)) {
- p_eog_sum += cur_p->data[i].p;
- } else {
- p_txt_sum += cur_p->data[i].p;
- }
- }
-
- const float rat = p_eog_sum == 0.0 ? INFINITY : p_txt_sum / p_eog_sum; GGML_UNUSED(rat);
-
- LOG_DBG_CUR("%s: p_txt_sum = %.2f, p_eog_sum = %.2f, rat = %.2f, n = %zu\n", __func__, p_txt_sum, p_eog_sum, rat, cur_p->size);
-
- if (3*p_eog_sum*cur_p->size > p_txt_sum) {
- LOG_DBG_CUR("%s: the ratio p_txt/p_eog = %.2f is too low -> sampling EOG\n", __func__, p_txt_sum/p_eog_sum);
-
- // keep just the EOG tokens
- const auto size_org = cur_p->size;
-
- cur_p->size = 0;
-
- float p_sum = 0.0f;
-
- for (size_t i = 0; i < size_org; ++i) {
- if (ctx->vocab->is_eog(cur_p->data[i].id)) {
- p_sum += cur_p->data[i].p;
-
- cur_p->data[cur_p->size++] = cur_p->data[i];
- }
- }
-
- // normalize probs
- for (size_t i = 0; i < cur_p->size; ++i) {
- cur_p->data[i].p /= p_sum;
- }
-
- return;
- }
-
- size_t n_combined = 0; GGML_UNUSED(n_combined);
-
- // combine tokens with common prefix
- for (size_t i0 = 0; i0 < cur_p->size; ++i0) {
- for (size_t i1 = 0; i1 < cur_p->size; ++i1) {
- if (cur_p->data[i0].logit == -INFINITY) {
- break;
- }
-
- if (i0 == i1 || cur_p->data[i1].logit == -INFINITY) {
- continue;
- }
-
- int len0 = ctx->vocab->token_to_piece(cur_p->data[i0].id, ctx->buf0.data(), ctx->buf0.size(), 0, false);
- if (len0 < 0) {
- ctx->buf0.resize(len0);
- len0 = ctx->vocab->token_to_piece(cur_p->data[i0].id, ctx->buf0.data(), ctx->buf0.size(), 0, false);
- assert(len0 > 0);
- }
-
- int len1 = ctx->vocab->token_to_piece(cur_p->data[i1].id, ctx->buf1.data(), ctx->buf1.size(), 0, false);
- if (len1 < 0) {
- ctx->buf1.resize(len1);
- len1 = ctx->vocab->token_to_piece(cur_p->data[i1].id, ctx->buf1.data(), ctx->buf1.size(), 0, false);
- assert(len1 > 0);
- }
-
- // token i0 is a prefix of token i1
- if (len0 > 0 && len0 <= len1 && memcmp(ctx->buf0.data(), ctx->buf1.data(), len0) == 0) {
- int dst = i0;
- int src = i1;
-
- // merge into the token with higher probability
- if (cur_p->data[i1].p > cur_p->data[i0].p) {
- std::swap(dst, src);
- }
-
- cur_p->data[dst].p += cur_p->data[src].p;
- cur_p->data[src].logit = -INFINITY;
- cur_p->data[src].p = 0.0f;
-
- n_combined++;
- }
- }
- }
-
- size_t n_non_eog = 0;
-
- size_t size_org = cur_p->size;
-
- float p_sum = 0.0f;
- float thold = 0.2f;
-
- cur_p->size = 0;
-
- LOG_DBG_CUR("%s: n_combined = %zu, applying thold = %.3f\n", __func__, n_combined, thold);
-
- for (size_t i = 0; i < size_org; ++i) {
- const bool is_eog = ctx->vocab->is_eog(cur_p->data[i].id);
-
- if (cur_p->data[i].p < thold && !is_eog) {
- continue;
- }
-
- if (!is_eog) {
- ++n_non_eog;
- }
-
- p_sum += cur_p->data[i].p;
-
- // keep this token
- cur_p->data[cur_p->size++] = cur_p->data[i];
- }
-
- LOG_DBG_CUR("%s: n_non_eog = %zu\n", __func__, n_non_eog);
-
- // if no non-EOG tokens are left -> reduce cur_p to single EOT token
- if (n_non_eog == 0) {
- cur_p->size = 1;
- cur_p->data[0].id = ctx->vocab->token_eot();
- if (cur_p->data[0].id == LLAMA_TOKEN_NULL) {
- cur_p->data[0].id = ctx->vocab->token_eos();
- }
- cur_p->data[0].logit = 1.0f;
-
- GGML_ASSERT(cur_p->data[0].id != LLAMA_TOKEN_NULL);
-
- return;
- }
-
- // normalize probs
- for (size_t i = 0; i < cur_p->size; ++i) {
- cur_p->data[i].p /= p_sum;
-
- LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
- }
-
- size_org = cur_p->size;
- p_sum = 0.0f;
- thold = 1.0/(n_non_eog + 1);
-
- cur_p->size = 0;
-
- LOG_DBG_CUR("%s: applying thold = %.3f\n", __func__, thold);
-
- for (size_t i = 0; i < size_org; ++i) {
- const bool is_eog = ctx->vocab->is_eog(cur_p->data[i].id);
-
- if (cur_p->data[i].p < thold && !is_eog) {
- continue;
- }
-
- p_sum += cur_p->data[i].p;
-
- cur_p->data[cur_p->size++] = cur_p->data[i];
- }
-
- // normalize probs
- for (size_t i = 0; i < cur_p->size; ++i) {
- cur_p->data[i].p /= p_sum;
-
- LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
- }
-
-#undef LOG_DBG_CUR
-}
-
-static struct llama_sampler * llama_sampler_infill_clone(const struct llama_sampler * smpl) {
- const auto * ctx = (const llama_sampler_infill *) smpl->ctx;
- return llama_sampler_init_infill(ctx->vocab);
-}
-
-static void llama_sampler_infill_free(struct llama_sampler * smpl) {
- delete (llama_sampler_infill *) smpl->ctx;
-}
-
-static struct llama_sampler_i llama_sampler_infill_i = {
- /* .name = */ llama_sampler_infill_name,
- /* .accept = */ nullptr,
- /* .apply = */ llama_sampler_infill_apply,
- /* .reset = */ nullptr,
- /* .clone = */ llama_sampler_infill_clone,
- /* .free = */ llama_sampler_infill_free,
- /* .backend_apply = */ nullptr,
- /* .backend_accept = */ nullptr,
- /* .backend_set_input = */ nullptr,
- /* .backend_init = */ nullptr,
-};
-
-struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
- return llama_sampler_init(
- /* .iface = */ &llama_sampler_infill_i,
- /* .ctx = */ new llama_sampler_infill {
- /* .vocab = */ vocab,
- /* .buf0 = */ std::vector<char>(512),
- /* .buf1 = */ std::vector<char>(512),
- }
- );
-}
-
-// utils
-
-uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) {
- if (smpl->iface == &llama_sampler_dist_i) {
- return ((const llama_sampler_dist *) smpl->ctx)->seed_cur;
- }
-
- if (smpl->iface == &llama_sampler_mirostat_i) {
- return ((const llama_sampler_mirostat *) smpl->ctx)->seed_cur;
- }
-
- if (smpl->iface == &llama_sampler_mirostat_v2_i) {
- return ((const llama_sampler_mirostat_v2 *) smpl->ctx)->seed_cur;
- }
-
- if (smpl->iface == &llama_sampler_chain_i) {
- const auto * ctx = (const llama_sampler_chain *) smpl->ctx;
- for (auto it = ctx->samplers.rbegin(); it != ctx->samplers.rend(); ++it) {
- const uint32_t seed = llama_sampler_get_seed(it->ptr);
- if (seed != LLAMA_DEFAULT_SEED) {
- return seed;
- }
- }
- }
-
- return LLAMA_DEFAULT_SEED;
-}
-
-// perf
-
-struct llama_perf_sampler_data llama_perf_sampler(const struct llama_sampler * chain) {
- struct llama_perf_sampler_data data = {};
-
- if (chain == nullptr || chain->iface != &llama_sampler_chain_i) {
- GGML_ABORT("%s: invalid sampler passed - requires a sampler created with llama_sampler_chain_init()\n", __func__);
- }
-
- const auto * ctx = (const struct llama_sampler_chain *) chain->ctx;
-
- data.t_sample_ms = 1e-3 * ctx->t_sample_us;
- data.n_sample = std::max(0, ctx->n_sample);
-
- return data;
-}
-
-void llama_perf_sampler_print(const struct llama_sampler * chain) {
- const auto data = llama_perf_sampler(chain);
-
- LLAMA_LOG_INFO("%s: samplers time = %10.2f ms / %5d runs\n", __func__, data.t_sample_ms, data.n_sample);
-}
-
-void llama_perf_sampler_reset(struct llama_sampler * chain) {
- if (chain == nullptr || chain->iface != &llama_sampler_chain_i) {
- GGML_ABORT("%s: invalid sampler passed - requires a sampler created with llama_sampler_chain_init()\n", __func__);
- }
-
- auto * ctx = (struct llama_sampler_chain *) chain->ctx;
-
- ctx->t_sample_us = 0;
- ctx->n_sample = 0;
-}
+++ /dev/null
-#pragma once
-
-// TODO: rename llama-sampling.h/.cpp to llama-sampler.h/.cpp ?
-
-#include "llama.h"
-
-#include <vector>
-
-struct llama_vocab;
-struct llama_grammar;
-
-// sampler chain
-
-struct llama_sampler_chain {
- llama_sampler_chain_params params;
-
- // has .backend_init() been called?
- bool is_init = false;
-
- struct info {
- bool is_backend;
-
- llama_sampler * ptr;
- };
-
- std::vector<info> samplers;
-
- // pre-allocated buffer for llama_sampler_sample to avoid repeated allocations
- std::vector<llama_token_data> cur;
-
- // timing
-
- mutable int64_t t_sample_us;
-
- mutable int32_t n_sample;
-};
-
-struct llama_sampler * llama_sampler_init_dry_testing(
- int32_t context_size,
- float dry_multiplier,
- float dry_base,
- int32_t dry_allowed_length,
- int32_t dry_penalty_last_n,
- const std::vector<std::vector<llama_token>> & seq_breakers);
//
// SPM tokenizer
// original implementation:
-// https://github.com/ggerganov/llama.cpp/commit/074bea2eb1f1349a0118239c4152914aecaa1be4
+// https://github.com/ggml-org/llama.cpp/commit/074bea2eb1f1349a0118239c4152914aecaa1be4
//
struct llm_bigram_spm {
// original regex from tokenizer.json
//"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
- // adapted: https://github.com/ggerganov/llama.cpp/pull/6920#issuecomment-2080233989
+ // adapted: https://github.com/ggml-org/llama.cpp/pull/6920#issuecomment-2080233989
"(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
};
break;
// read bpe merges and populate bpe ranks
const int merges_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_MERGES).c_str());
+ // Kimi-K2 uses custom tokenization without traditional BPE merges
+ const bool is_kimi_k2 = (tokenizer_pre == "kimi-k2");
+
if (merges_keyidx == -1) {
- throw std::runtime_error("cannot find tokenizer merges in model file\n");
- }
+ if (!is_kimi_k2) {
+ throw std::runtime_error("cannot find tokenizer merges in model file\n");
+ }
+ // Kimi-K2 doesn't need merges, skip
+ LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__);
+ } else {
+ const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
+ for (int i = 0; i < n_merges; i++) {
+ const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
+ //GGML_ASSERT(unicode_cpts_from_utf8(word).size() > 0);
- const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
- for (int i = 0; i < n_merges; i++) {
- const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
- //GGML_ASSERT(unicode_cpts_from_utf8(word).size() > 0);
+ std::string first;
+ std::string second;
- std::string first;
- std::string second;
+ const size_t pos = word.find(' ', 1);
- const size_t pos = word.find(' ', 1);
+ if (pos != std::string::npos) {
+ first = word.substr(0, pos);
+ second = word.substr(pos + 1);
+ }
- if (pos != std::string::npos) {
- first = word.substr(0, pos);
- second = word.substr(pos + 1);
+ bpe_ranks.emplace(std::make_pair(first, second), i);
}
-
- bpe_ranks.emplace(std::make_pair(first, second), i);
}
// default special tokens
|| t.first == "<|end_of_text|>" // granite
|| t.first == "<EOT>"
|| t.first == "_<EOT>"
+ || t.first == "[EOT]" // Kimi-K2
|| t.first == "<｜end▁of▁sentence｜>" // DeepSeek
|| t.first == "<end_of_utterance>" // smoldocling
) {
|| t.first == "<PRE>"
|| t.first == "▁<PRE>" // CodeLlama
|| t.first == "<|code_prefix|>" // GLM-4.5
+ || t.first == "<|prefix|>" // Falcon-H1-Tiny-Coder
) {
special_fim_pre_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
|| t.first == "<SUF>"
|| t.first == "▁<SUF>" // CodeLlama
|| t.first == "<|code_suffix|>" // GLM-4.5
+ || t.first == "<|suffix|>" // Falcon-H1-Tiny-Coder
) {
special_fim_suf_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
|| t.first == "<MID>"
|| t.first == "▁<MID>" // CodeLlama
|| t.first == "<|code_middle|>" // GLM-4.5
+ || t.first == "<|middle|>" // Falcon-H1-Tiny-Coder
) {
special_fim_mid_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
|| t.first == "<fim-pad>"
|| t.first == "<fim_pad>" // Granite
|| t.first == "<PAD>"
+ || t.first == "[PAD]" // Kimi-K2
) {
special_fim_pad_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
// maintain a list of tokens that cause end-of-generation
// this is currently determined based on the token text, which is obviously not ideal
- // ref: https://github.com/ggerganov/llama.cpp/issues/9606
+ // ref: https://github.com/ggml-org/llama.cpp/issues/9606
special_eog_ids.clear();
if (special_fim_pad_id != LLAMA_TOKEN_NULL && special_eog_ids.count(special_fim_pad_id) == 0) {
|| t.first == "<|eom_id|>"
|| t.first == "<EOT>"
|| t.first == "_<EOT>"
+ || t.first == "[EOT]" // Kimi-K2
+ || t.first == "[EOS]" // Kimi-K2
|| t.first == "<|end_of_text|>"
|| t.first == "<end_of_utterance>" // smoldocling
) {
}
int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t length, int32_t lstrip, bool special) const {
- // ref: https://github.com/ggerganov/llama.cpp/pull/7587#discussion_r1620983843
+ // ref: https://github.com/ggml-org/llama.cpp/pull/7587#discussion_r1620983843
static const int attr_special = LLAMA_TOKEN_ATTR_UNKNOWN | LLAMA_TOKEN_ATTR_CONTROL;
const llama_token_attr attr = token_get_attr(token);
if (!special && (attr & attr_special)) {
const uint32_t kv_lora_rank = hparams.n_lora_kv;
// We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
- // See https://github.com/ggerganov/llama.cpp/discussions/7416 for detailed explanation.
+ // See https://github.com/ggml-org/llama.cpp/discussions/7416 for detailed explanation.
// And also: https://github.com/ggml-org/llama.cpp/pull/17945 [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]
// first cancel the adjustment from llama_hparams::yarn_attn_factor_adjust to get the original attn_factor
--- /dev/null
+#include "models.h"
+#include "ggml.h"
+
+#define CHUNK_SIZE 64
+
+// Causal Conv1d function for Q,K,V
+// When qkv is 0, it is Q, 1 is K, 2 is V
+static ggml_tensor * causal_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, int64_t d_conv, int64_t head_dim, int64_t n_head, int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) {
+ const int64_t d_inner = head_dim * n_head;
+ const int64_t conv_state_size = (d_conv - 1) * d_inner;
+ const int64_t n_embd_r_total = 3 * conv_state_size; // Q + K + V
+
+ // conv_state_all is [n_embd_r_total, n_seqs], split into Q, K, V
+ // Each conv state is [(d_conv-1) * d_inner] per sequence, need to reshape to [d_conv-1, d_inner, n_seqs]
+ // Memory layout: for each seq, Q state is first conv_state_size elements, then K, then V
+ // conv_state_all has stride: nb[0] = element_size, nb[1] = n_embd_r_total * element_size
+ // View Q conv state: offset 0, size conv_state_size per seq
+ // conv_state_all is [n_embd_r_total, n_seqs] with memory layout:
+ // state[i + seq * n_embd_r_total] where i = conv_step + channel * (d_conv-1) + {0, conv_state_size, 2*conv_state_size} for Q/K/V
+ // We want [d_conv-1, d_inner, n_seqs] view:
+ // nb1 = (d_conv-1) * element_size (stride between channels)
+ // nb2 = n_embd_r_total * element_size (stride between seqs)
+ ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,
+ (d_conv - 1) * ggml_element_size(conv_state_all), // nb1: stride between channels
+ n_embd_r_total * ggml_element_size(conv_state_all), // nb2: stride between seqs
+ qkv * conv_state_size * ggml_element_size(conv_state_all));
+
+// Causal Conv1d function for Q,K,V
+// When qkv is 0, it is Q, 1 is K, 2 is V
+ // Step 1: Q, K, V projections -> [d_inner, n_tokens]
+ ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
+
+ // Reshape input: {d_inner, n_tokens} -> {d_inner, n_seq_tokens, n_seqs}
+ ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
+
+ // Concat Q conv state and current input: {d_conv-1 + n_seq_tokens, d_inner, n_seqs}
+ ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
+
+ // Save last (d_conv-1) columns back to Q conv state
+ ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
+ conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
+ ggml_build_forward_expand(gf,
+ ggml_cpy(ctx0, last_conv_x,
+ ggml_view_1d(ctx0, conv_states_all, conv_state_size * n_seqs,
+ (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+ // Reshape conv weight: GGUF [d_conv, 1, d_inner, 1] -> ggml_ssm_conv expects [d_conv, d_inner]
+ // GGUF stores as [d_conv, 1, d_inner, 1] with memory layout w[conv_step + channel * d_conv]
+ // vLLM stores as [d_inner, d_conv] with memory layout w[channel * d_conv + conv_step]
+ // ggml_ssm_conv computes: c[conv_step + channel * d_conv]
+ // GGUF layout: [d_conv, 1, d_inner] or [d_conv, 1, d_inner, 1] -> reshape to [d_conv, d_inner]
+ // Reshape conv weight from [d_conv, 1, d_inner, 1] to [d_conv, d_inner] for ggml_ssm_conv
+ ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
+
+ // Apply conv1d
+ // ggml_ssm_conv output: {d_inner, n_seq_tokens, n_seqs}
+ ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
+ // Reshape to 2D for bias add: {d_inner, n_tokens}
+ Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
+ Xcur = ggml_silu(ctx0, Xcur);
+
+ return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
+}
+
+llm_build_kimi_linear::llm_build_kimi_linear(const llama_model & model, const llm_graph_params & params) :
+ llm_graph_context_mamba(params), model(model) {
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ inpL = build_inp_embd(model.tok_embd);
+ cb(inpL, "model.embed_tokens", -1);
+
+ // Note: Kimi MLA does NOT use RoPE (rotary_emb=None in vLLM)
+ // So we don't need inp_pos
+
+ auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr;
+ auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr;
+ auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr();
+ auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr;
+ auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr;
+
+ // Output ids for selecting which tokens to output
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ ggml_tensor * chunked_causal_mask =
+ ggml_tri(ctx0, ggml_fill_inplace(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, CHUNK_SIZE, CHUNK_SIZE), 1.0f),
+ GGML_TRI_TYPE_LOWER);
+
+ ggml_tensor * chunked_identity = ggml_diag(ctx0, ggml_fill_inplace(ctx0, ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, CHUNK_SIZE), 1.0f));
+ ggml_tensor * chunked_diag_mask = ggml_add(ctx0, chunked_causal_mask, chunked_identity);
+
+ ggml_build_forward_expand(gf, chunked_causal_mask);
+ ggml_build_forward_expand(gf, chunked_identity);
+ ggml_build_forward_expand(gf, chunked_diag_mask);
+
+ // Kimi dimension constants
+ const int64_t n_head = hparams.n_head();
+ const int64_t head_dim = hparams.n_embd_head_kda;
+ const int64_t d_conv = hparams.ssm_d_conv;
+ const int64_t d_inner = n_head * head_dim; // 32 * 128 = 4096
+ const int64_t n_seqs = ubatch.n_seqs;
+ const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+
+ // Verify batch consistency for recurrent layers
+ GGML_ASSERT(n_seqs != 0);
+ GGML_ASSERT(ubatch.equal_seqs());
+ GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
+
+ // MLA params
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ // qk_rope_head_dim = 64 (from Kimi config) which is hparams.n_rot
+ // Confirmed from tensor shape: wkv_a_mqa [2304, 576] = [n_embd, kv_lora_rank + qk_rope_head_dim]
+ const int64_t n_embd_head_qk_rope = hparams.n_rot; // config.qk_rope_head_dim
+ const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; // 192 - 64 = 128
+ // Attention scale for MLA
+ const float kq_scale_mla = 1.0f / sqrtf((float)n_embd_head_k_mla);
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+ ggml_tensor * inpSA = inpL;
+
+ // Attention Norm
+ cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ // Check layer type by checking which tensors exist
+ // KDA layers have ssm_a_log tensor, MLA layers have wkv_a_mqa tensor
+ bool is_kda = (layer.ssm_a != nullptr);
+ bool is_mla = (layer.wkv_a_mqa != nullptr);
+
+ if (is_kda) {
+ // === KDA Layer (Kimi Delta Attention) with Recurrent State ===
+ // Reference: vLLM kda.py
+ const auto * mctx_cur = inp_rs->mctx;
+ const auto kv_head = mctx_cur->get_head();
+
+ // Get conv states from r_l tensor (Q, K, V each have separate state)
+ ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
+ cb(conv_states_all, "conv_states_all", il);
+ ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
+ ggml_tensor * Qcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head);
+ ggml_tensor * Kcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head);
+ ggml_tensor * Vcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head);
+
+ // g1 = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias)
+ ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
+ ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a);
+ cb(g1, "g1 f_b(f_a(cur))", il);
+ g1 = ggml_add(ctx0, g1, layer.ssm_dt_b);
+ g1 = ggml_softplus(ctx0, g1);
+ g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head, n_tokens);
+
+ // A_log shape is [1, n_head] or [1, n_head, 1, 1], need to broadcast to [head_dim, n_head, n_tokens]. No need to -exp(a_log) because it was done in convert_hf_to_gguf.py
+ // Reshape to [1, n_head, 1] for broadcasting with g1 [head_dim, n_head, n_tokens]
+ ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1);
+ g1 = ggml_mul(ctx0, g1, A);
+ cb(g1, "kda_g1", il);
+
+ // Compute beta (mixing coefficient)
+ ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
+ beta = ggml_reshape_4d(ctx0, beta, n_head, 1, n_seq_tokens, n_seqs);
+ cb(beta, "kda_beta", il);
+
+ // Reshape for KDA recurrence
+ // {n_embd, n_tokens} -> {n_embd, n_seq_tokens, n_seqs}
+ cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
+
+ g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head, n_seq_tokens, n_seqs);
+
+ // Get SSM state and compute KDA recurrence using ggml_kda_scan
+ ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);
+ ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs);
+ state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs);
+ // Choose between build_kda_chunking and build_kda_recurrent based on n_tokens
+ std::pair<ggml_tensor *, ggml_tensor *> attn_out = n_seq_tokens == 1 ?
+ build_kda_autoregressive(Qcur, Kcur, Vcur, g1, beta, state, il) :
+ build_kda_chunking(Qcur, Kcur, Vcur, g1, beta, state, chunked_causal_mask, chunked_identity, chunked_diag_mask, il);
+
+ ggml_tensor * output = attn_out.first;
+ ggml_tensor * new_state = attn_out.second;
+ cb(output, "attn_output", il);
+ cb(new_state, "new_state", il);
+
+ // Update the recurrent states
+ ggml_build_forward_expand(gf,
+ ggml_cpy(ctx0, new_state,
+ ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs,
+ kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all))));
+
+ // Output gating g2 = g_b(g_a(x))
+ ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens * n_seqs);
+ ggml_tensor * g_a = ggml_mul_mat(ctx0, layer.ssm_g_a, cur_2d);
+ ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g_b, g_a);
+ cb(g2, "g2 g_b(g_a(cur_2d))", il);
+ g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head, n_seq_tokens * n_seqs);
+
+ // Apply o_norm with sigmoid gating
+ // Note: Kimi model uses sigmoid gating, not SiLU (despite FusedRMSNormGated default being swish)
+ // Formula: output = RMSNorm(x) * sigmoid(g)
+ ggml_tensor * attn_out_final = ggml_reshape_3d(ctx0, output, head_dim, n_head, n_seq_tokens * n_seqs);
+ ggml_tensor * normed = build_norm(attn_out_final, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
+ cb(normed, "kda_normed", il);
+ ggml_tensor * gate = ggml_sigmoid(ctx0, g2);
+ ggml_tensor * gated = ggml_mul(ctx0, normed, gate);
+
+ // Output projection
+ gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens);
+ cur = ggml_mul_mat(ctx0, layer.wo, gated);
+ cb(cur, "kda_out", il);
+
+ } else if (is_mla) {
+ // === MLA Layer (Multi-head Latent Attention) without KV Cache ===
+ // Reference: vLLM mla.py
+ // Step 1: Q projection and reshape
+ // vLLM Kimi: q = q_proj(hidden_states), then view as [n_tokens, n_head, qk_head_dim]
+ // Note: Kimi MLA does NOT use RoPE (rotary_emb=None in vLLM)
+ ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wq, cur);
+
+ // Step 2: KV compression
+ // kv_cmpr_pe = kv_a_proj_with_mqa(hidden_states) -> [kv_lora_rank + qk_rope_head_dim, n_tokens]
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
+
+ // Split: kv_cmpr = kv_lora[:kv_lora_rank], k_pe = kv_lora[kv_lora_rank:]
+ ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);
+ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
+ // Note: Kimi MLA does NOT apply RoPE (rotary_emb=None in vLLM)
+ // k_pe is used directly without RoPE
+ // Normalize kv_c
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+
+ if (layer.wk_b && layer.wv_b) { // MLA KV cache enabled
+ // extract q_nope
+ ggml_tensor * q_nope =
+ ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(Qcur->type, n_embd_head_k_mla),
+ ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0);
+ cb(q_nope, "q_nope", il);
+
+ // and {n_embd_head_qk_rope, n_head, n_tokens}
+ ggml_tensor * q_pe = ggml_view_3d(
+ ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(Qcur->type, n_embd_head_k_mla),
+ ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, ggml_row_size(Qcur->type, n_embd_head_qk_nope));
+ cb(q_pe, "q_pe", il);
+
+ // {n_embd_head_qk_nope, n_tokens, n_head}
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+ cb(q_nope, "q_nope_perm", il);
+
+ // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}
+ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
+ cb(q_nope_absorbed, "q_nope_absorbed", il);
+
+ // {kv_lora_rank, n_head, n_tokens}
+ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
+ cb(q_nope_absorbed, "q_nope_absorbed_perm", il);
+
+ // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}
+ // note: rope must go first for in-place context shifting in build_rope_shift()
+ Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
+ cb(Qcur, "Qcur", il);
+
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "kv_cmpr_reshape", il);
+
+ // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}
+ ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);
+ cb(Kcur, "Kcur", il);
+
+ // {kv_lora_rank, 1, n_tokens}
+ ggml_tensor * Vcur = kv_cmpr;
+ cb(Vcur, "Vcur", il);
+
+ cur = build_attn(inp_attn_k, layer.wo, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale_mla, il);
+ cb(cur, "mla_out", il);
+ } else { // MLA KV cache disabled. Fall back to MHA KV cache.
+ Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens);
+ cb(Qcur, "mla_Q", il);
+ // KV decompression: kv = kv_b_proj(kv_c_normed)
+ ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr);
+ const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla;
+
+ // Split kv into k_nope and v
+ ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,
+ ggml_row_size(kv->type, kv_per_head),
+ ggml_row_size(kv->type, kv_per_head * n_head), 0);
+ ggml_tensor * Vcur = ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens,
+ ggml_row_size(kv->type, kv_per_head),
+ ggml_row_size(kv->type, kv_per_head * n_head),
+ ggml_row_size(kv->type, n_embd_head_qk_nope));
+ Vcur = ggml_cont(ctx0, Vcur);
+ cb(Vcur, "mla_V", il);
+
+ // Concatenate k_nope + k_pe (broadcast k_pe to all heads)
+ // K = [k_nope, k_pe] where k_nope is [qk_nope_head_dim, n_head, n_tokens]
+ // and k_pe is [qk_rope_head_dim, 1, n_tokens] broadcast to all heads
+ // Need to broadcast k_pe from [qk_rope, 1, n_tokens] to [qk_rope, n_head, n_tokens]
+ ggml_tensor * k_pe_target = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens);
+ ggml_tensor * k_pe_repeated = ggml_repeat(ctx0, k_pe, k_pe_target);
+ ggml_tensor * Kcur = ggml_concat(ctx0, k_pe_repeated, k_nope, 0);
+ cb(Kcur, "mla_K", il);
+
+ // Direct softmax attention (with MHA KV cache)
+ // Use build_attn with inp_attn for proper mask handling
+ cur = build_attn(inp_attn_kv, layer.wo, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale_mla, il);
+ cb(cur, "mla_out", il);
+ }
+ } else {
+ // Unknown layer type - this should not happen
+ GGML_ABORT("Kimi layer is neither KDA nor MLA - missing required tensors");
+ }
+
+ // On last layer, select only the output tokens
+ if (il == n_layer - 1 && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+ }
+
+ // Residual
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "ffn_inp", il);
+
+ // FFN Norm
+ cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ if ((uint32_t) il < hparams.n_layer_dense_lead) {
+ // Dense FFN layer
+ cur = build_ffn(cur,
+ layer.ffn_up, NULL, NULL,
+ layer.ffn_gate, NULL, NULL,
+ layer.ffn_down, NULL, NULL,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+ } else {
+ // MoE layer
+ // Kimi uses moe_renormalize=True and routed_scaling_factor (stored as expert_weights_scale) = 2.446
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.ffn_down_exps,
+ layer.ffn_exp_probs_b,
+ hparams.n_expert,
+ hparams.n_expert_used,
+ LLM_FFN_SILU, true,
+ true, hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il);
+ cb(moe_out, "ffn_moe_out", il);
+
+ // Shared expert
+ {
+ ggml_tensor * ffn_shexp = build_ffn(cur,
+ layer.ffn_up_shexp, NULL, NULL,
+ layer.ffn_gate_shexp, NULL, NULL,
+ layer.ffn_down_shexp, NULL, NULL,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "ffn_out", il);
+ }
+ }
+ // Residual
+ cur = ggml_add(ctx0, cur, ffn_inp);
+
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ inpL = cur;
+ }
+ cur = inpL;
+
+ // Final Norm
+ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ // Output
+ cur = ggml_mul_mat(ctx0, model.output, cur);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
+
+/*
+ This is a ggml implementation of the naive_chunk_kda function of
+ https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/kda/naive.py
+*/
+std::pair<ggml_tensor *, ggml_tensor *> llm_build_kimi_linear::build_kda_chunking(
+ ggml_tensor * q,
+ ggml_tensor * k,
+ ggml_tensor * v,
+ ggml_tensor * gk,
+ ggml_tensor * beta,
+ ggml_tensor * state,
+ ggml_tensor * causal_mask,
+ ggml_tensor * identity,
+ ggml_tensor * diag_mask,
+ int il) {
+ GGML_ASSERT(ggml_is_contiguous(state));
+
+ const int64_t S_k = q->ne[0];
+ const int64_t H_k = q->ne[1];
+ const int64_t n_tokens = q->ne[2];
+ const int64_t n_seqs = q->ne[3];
+
+ const int64_t S_v = v->ne[0];
+ const int64_t H_v = v->ne[1];
+
+ GGML_ASSERT(v->ne[2] == n_tokens);
+ GGML_ASSERT(k->ne[2] == n_tokens);
+ GGML_ASSERT(gk->ne[0] == S_v && gk->ne[1] == H_v && gk->ne[2] == n_tokens && gk->ne[3] == n_seqs);
+ GGML_ASSERT(beta->ne[0] == H_v && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs);
+ GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v && state->ne[2] == H_v && state->ne[3] == n_seqs);
+
+ GGML_ASSERT(q->ne[0] == S_k && q->ne[1] == H_k && q->ne[2] == n_tokens && q->ne[3] == n_seqs);
+ GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_k && k->ne[2] == n_tokens && k->ne[3] == n_seqs);
+
+ GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case
+
+ // TODO: can this ever be false?
+ const bool use_qk_l2norm = true;
+
+ if (use_qk_l2norm) {
+ const float eps_norm = hparams.f_norm_rms_eps;
+
+ q = ggml_l2_norm(ctx0, q, eps_norm);
+ k = ggml_l2_norm(ctx0, k, eps_norm);
+ }
+
+ const float scale = 1.0f / sqrtf(S_v);
+
+ beta = ggml_sigmoid(ctx0, beta);
+
+ cb(q, "q_in", il);
+ cb(k, "k_in", il);
+ cb(v, "v_in", il);
+ cb(beta, "beta_in", il);
+ cb(gk, "gk_in", il);
+
+ q = ggml_cont_4d(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3), S_k, n_tokens, H_k, n_seqs);
+ k = ggml_cont_4d(ctx0, ggml_permute(ctx0, k, 0, 2, 1, 3), S_k, n_tokens, H_k, n_seqs);
+ v = ggml_cont_4d(ctx0, ggml_permute(ctx0, v, 0, 2, 1, 3), S_v, n_tokens, H_v, n_seqs);
+ gk = ggml_cont_4d(ctx0, ggml_permute(ctx0, gk, 0, 2, 1, 3), S_v, n_tokens, H_v, n_seqs);
+
+ beta = ggml_cont(ctx0, ggml_permute(ctx0, beta, 2, 0, 1, 3));
+ state = ggml_reshape_4d(ctx0, state, S_v, S_v, H_v, n_seqs);
+
+ cb(q, "q_perm", il);
+ cb(k, "k_perm", il);
+ cb(v, "v_perm", il);
+ cb(beta, "beta_perm", il);
+ cb(gk, "gk_perm", il);
+ cb(state, "state_in", il);
+
+ GGML_ASSERT(q->ne[1] == n_tokens && q->ne[0] == S_k && q->ne[2] == H_k && q->ne[3] == n_seqs);
+ GGML_ASSERT(k->ne[1] == n_tokens && k->ne[0] == S_k && k->ne[2] == H_k && k->ne[3] == n_seqs);
+ GGML_ASSERT(v->ne[1] == n_tokens && v->ne[0] == S_v && v->ne[2] == H_k && v->ne[3] == n_seqs);
+ GGML_ASSERT(beta->ne[1] == n_tokens && beta->ne[2] == H_k && beta->ne[0] == 1 && beta->ne[3] == n_seqs);
+
+ // Do padding
+ const int64_t chunk_size = CHUNK_SIZE;
+
+ const int64_t pad = (chunk_size - n_tokens % chunk_size) % chunk_size;
+ const int64_t n_chunks = (n_tokens + pad) / chunk_size;
+
+ q = ggml_pad(ctx0, q, 0, pad, 0, 0);
+ k = ggml_pad(ctx0, k, 0, pad, 0, 0);
+ v = ggml_pad(ctx0, v, 0, pad, 0, 0);
+ gk = ggml_pad(ctx0, gk, 0, pad, 0, 0);
+ beta = ggml_pad(ctx0, beta, 0, pad, 0, 0);
+
+ cb(q, "q_pad", il);
+ cb(k, "k_pad", il);
+ cb(v, "v_pad", il);
+ cb(beta, "beta_pad", il);
+ cb(gk, "gk_pad", il);
+
+ ggml_tensor * v_beta = ggml_mul(ctx0, v, beta);
+ ggml_tensor * k_beta = ggml_mul(ctx0, k, beta);
+
+ cb(v_beta, "v_beta", il);
+ cb(k_beta, "k_beta", il);
+
+ const int64_t HB = H_k * n_seqs;
+
+ q = ggml_cont_4d(ctx0, q, S_k, chunk_size, n_chunks, HB);
+ k = ggml_cont_4d(ctx0, k, S_k, chunk_size, n_chunks, HB);
+ k_beta = ggml_cont_4d(ctx0, k_beta, S_k, chunk_size, n_chunks, HB);
+ v = ggml_cont_4d(ctx0, v, S_v, chunk_size, n_chunks, HB);
+ v_beta = ggml_cont_4d(ctx0, v_beta, S_v, chunk_size, n_chunks, HB);
+
+ gk = ggml_cont_4d(ctx0, gk, S_k, chunk_size, n_chunks, HB);
+ beta = ggml_cont_4d(ctx0, beta, 1, chunk_size, n_chunks, HB);
+
+ // switch for cumsum
+ gk = ggml_cont_4d(ctx0, ggml_permute(ctx0, gk, 1, 0, 2, 3), chunk_size, S_k, n_chunks, HB);
+ cb(gk, "gk", il);
+ ggml_tensor * gk_cumsum = ggml_cumsum(ctx0, gk);
+ cb(gk_cumsum, "gk_cumsum", il);
+
+/*
+ Compute Akk and Aqk loop together
+ Akk loop:
+ for i in range(BT):
+ k_i = k[..., i, :] # k_i [B,H,NT,S]
+ g_i = g[..., i:i+1, :] # g_i [B,H,NT,1,S]
+ A[..., i] = torch.einsum('... c d, ... d -> ... c', k * (g - g_i).exp(), k_i)
+ Aqk loop:
+ for j in range(BT):
+ k_j = k[:, :, i, j]
+ g_j = g[:, :, i, j:j+1, :]
+ A[..., j] = torch.einsum('... c d, ... d -> ... c', q_i * (g_i - g_j).exp(), k_j)
+*/
+ const int64_t CHB = n_chunks * H_k * n_seqs;
+ ggml_tensor * gkcs_i = ggml_reshape_4d(ctx0, gk_cumsum, chunk_size, 1, S_k, CHB); // [chunk_size, 1, S_k, CHB]
+ ggml_tensor * gkcs_j = ggml_reshape_4d(ctx0, gkcs_i, 1, chunk_size, S_k, CHB); // [1, chunk_size, S_k, CHB]
+
+ ggml_tensor * gkcs_j_bc = ggml_repeat_4d(ctx0, gkcs_j, chunk_size, chunk_size, S_k, CHB); // [1, chunk_size, S_k, CHB] -> [chunk_size, chunk_size, S_k, CHB]
+ // decay_mask [chunk_size,chunk_size,S_k,CHB]
+ ggml_tensor * decay_mask = ggml_sub(ctx0, gkcs_j_bc, gkcs_i);
+ cb(decay_mask, "decay_mask", il);
+
+ decay_mask = ggml_mul(ctx0, decay_mask, diag_mask);
+ cb(decay_mask, "decay_masked", il);
+ decay_mask = ggml_exp(ctx0, decay_mask);
+ decay_mask = ggml_mul(ctx0, decay_mask, diag_mask);
+
+ // decay_mask [S_k,BT_j,BT_i,CHB] *Note* second and third chunk_sizes are switched
+ decay_mask = ggml_cont_4d(ctx0, ggml_permute(ctx0, decay_mask, 2, 1, 0, 3), S_k, chunk_size, chunk_size, CHB);
+
+ ggml_tensor * k_i = ggml_reshape_4d(ctx0, k, S_k, chunk_size, 1, CHB);
+ ggml_tensor * k_j = ggml_reshape_4d(ctx0, k, S_k, 1, chunk_size, CHB);
+ ggml_tensor * q_i = ggml_reshape_4d(ctx0, q, S_k, chunk_size, 1, CHB);
+
+ ggml_tensor * decay_k_i = ggml_mul(ctx0, decay_mask, k_i);
+ ggml_tensor * decay_q_i = ggml_mul(ctx0, decay_mask, q_i);
+
+ // decay_k_i [S.BT,BT,CHB] @ k_j [S,1,BT,CHB] = Akk [BT,1,BT,CHB]
+ ggml_tensor * Akk = ggml_mul_mat(ctx0, decay_k_i, k_j);
+ ggml_tensor * Aqk = ggml_mul_mat(ctx0, decay_q_i, k_j);
+ Akk = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_reshape_4d(ctx0, Akk, chunk_size, chunk_size, n_chunks, HB)));
+ Aqk = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_reshape_4d(ctx0, Aqk, chunk_size, chunk_size, n_chunks, HB)));
+ cb(Akk, "Akk", il);
+ cb(Aqk, "Aqk", il);
+
+ Akk = ggml_mul(ctx0, Akk, beta);
+ Akk = ggml_neg(ctx0, ggml_mul(ctx0, Akk, causal_mask));
+ cb(Akk, "attn_pre_solve", il);
+
+ Aqk = ggml_mul(ctx0, Aqk, diag_mask);
+ Aqk = ggml_scale(ctx0, Aqk, scale); // scale q
+ cb(Aqk, "Aqk_masked", il);
+
+ // for i in range(1, chunk_size):
+ // row = attn[..., i, :i].clone()
+ // sub = attn[..., :i, :i].clone()
+ // attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
+ // attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+ //
+ // We reduce this to a linear triangular solve: AX = B, where B = attn, A = I - tril(A)
+ ggml_tensor * attn_lower = ggml_mul(ctx0, Akk, causal_mask);
+ ggml_tensor * lhs = ggml_sub(ctx0, ggml_repeat(ctx0, identity, attn_lower), attn_lower);
+
+ ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, Akk, true, true, false);
+ Akk = ggml_mul(ctx0, lin_solve, causal_mask);
+ Akk = ggml_add(ctx0, Akk, identity);
+
+ cb(Akk, "attn_solved", il);
+
+ // switch back for downstream
+ gk_cumsum = ggml_cont_4d(ctx0, ggml_permute(ctx0, gk_cumsum, 1, 0, 2, 3), S_k, chunk_size, n_chunks, HB);
+ ggml_tensor * gkexp = ggml_exp(ctx0, gk_cumsum);
+ cb(gk_cumsum, "gk_cumsum", il);
+
+ // u = (A*beta[..., None, :]) @ v aka U_[t]
+ ggml_tensor * vb = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta)), Akk);
+
+ ggml_tensor * kbeta_gkexp = ggml_mul(ctx0, k_beta, gkexp);
+ cb(kbeta_gkexp, "kbeta_gkexp", il);
+
+ ggml_tensor * k_cumdecay = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, kbeta_gkexp)), Akk);
+ cb(k_cumdecay, "k_cumdecay", il);
+
+ ggml_tensor * core_attn_out = nullptr;
+ ggml_tensor * new_state = ggml_dup(ctx0, state);
+
+ cb(new_state, "new_state", il);
+
+ for (int64_t chunk = 0; chunk < n_chunks; chunk++) {
+// extract one chunk worth of data
+ auto chunkify = [=](ggml_tensor * t) {
+ return ggml_cont(ctx0, ggml_view_4d(ctx0, t, t->ne[0], chunk_size, 1, t->ne[3],
+ t->nb[1], t->nb[2], t->nb[3], t->nb[2] * chunk));
+ };
+ auto chunkify_A = [=](ggml_tensor * t) {
+ return ggml_cont(ctx0, ggml_view_4d(ctx0, t, chunk_size, chunk_size, 1, t->ne[3],
+ t->nb[1], t->nb[2], t->nb[3], t->nb[2] * chunk));
+ };
+
+
+// k [S,BT,NT,H*B] => k_chunk [S,BT,1,H*B]
+ ggml_tensor * k_chunk = chunkify(k);
+ ggml_tensor * q_chunk = chunkify(q);
+ ggml_tensor * vb_chunk = chunkify(vb);
+
+// gk_cumsum [S,BT,NT,H*B] => gk_cs_chunk [S,BT,1,H*B]
+ ggml_tensor * gk_cs_chunk = chunkify(gk_cumsum);
+ ggml_tensor * k_cumdecay_chunk = chunkify(k_cumdecay);
+ ggml_tensor * gkexp_chunk = ggml_exp(ctx0, gk_cs_chunk);
+ ggml_tensor * Aqk_chunk = chunkify_A(Aqk);
+
+ ggml_tensor * state_t = ggml_cont_4d(ctx0, ggml_permute(ctx0, new_state, 1, 0, 2, 3), S_v, S_v, 1, H_v * n_seqs);
+
+ // new_state [S,S,1,H*B] k_cumdecay_chunk [S,BT,1,H*B]
+ // v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state or W_[t] @ S_[t]
+ ggml_tensor * v_prime = ggml_mul_mat(ctx0, state_t, k_cumdecay_chunk);
+
+ // v_new = v_i - v_prime or U_[t] - W_[t]*S_[t]
+ ggml_tensor * v_new = ggml_sub(ctx0, ggml_repeat(ctx0, vb_chunk, v_prime), v_prime);
+ ggml_tensor * v_new_t = ggml_cont(ctx0, ggml_transpose(ctx0, v_new));
+
+ // q_chunk [S,BT,1,H*B] gkexp_chunk [S,BT,1,H*B]
+ // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
+ // or Gamma_[t]*Q_]t] @ S
+ ggml_tensor * q_gk_exp = ggml_mul(ctx0, q_chunk, gkexp_chunk);
+ ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state_t, q_gk_exp);
+ attn_inter = ggml_scale(ctx0, attn_inter, scale); // scale q
+
+ // v_new_t [S,BT,1,H*B] Aqk [BT,BT,1,H*B]
+ // core_attn_out[:, :, i] = attn_inter + attn @ v_new or A' @ (U_[t] - W_[t]*S_[t])
+ ggml_tensor * v_attn = ggml_mul_mat(ctx0, v_new_t, Aqk_chunk);
+
+ // o[:, :, i] = (q_i * g_i.exp()) @ S + A @ v_i
+ ggml_tensor * core_attn_out_chunk = ggml_add(ctx0, attn_inter, v_attn);
+
+ core_attn_out = core_attn_out == nullptr ? core_attn_out_chunk : ggml_concat(ctx0, core_attn_out, core_attn_out_chunk, 1);
+
+ ggml_tensor * gk_cum_last =
+ ggml_cont(ctx0, ggml_view_4d(ctx0, gk_cs_chunk, gk_cs_chunk->ne[0], 1, gk_cs_chunk->ne[2], gk_cs_chunk->ne[3],
+ gk_cs_chunk->nb[1], gk_cs_chunk->nb[2], gk_cs_chunk->nb[3],
+ gk_cs_chunk->nb[1] * (gk_cs_chunk->ne[1] - 1)));
+
+ ggml_tensor * gkexp_last = ggml_exp(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, gk_cum_last)));
+
+ ggml_tensor * gk_diff = ggml_neg(ctx0, ggml_sub(ctx0, gk_cs_chunk, gk_cum_last));
+
+ ggml_tensor * gk_diff_exp = ggml_exp(ctx0, gk_diff);
+
+ ggml_tensor * key_gkdiff = ggml_mul(ctx0, k_chunk, gk_diff_exp);
+
+ // rearrange((g_i[:,:,-1:] - g_i).exp()*k_i, 'b h c k -> b h k c') @ (U_[t] - W_[t] @ S)
+ ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, v_new_t, ggml_cont(ctx0, ggml_transpose(ctx0, key_gkdiff)));
+
+ new_state = ggml_add(ctx0,
+ ggml_mul(ctx0, new_state, ggml_reshape_4d(ctx0, gkexp_last, gkexp_last->ne[0], gkexp_last->ne[1], H_v, n_seqs)),
+ ggml_reshape_4d(ctx0, kgdmulvnew, kgdmulvnew->ne[0], kgdmulvnew->ne[1], H_v, n_seqs));
+ }
+
+ core_attn_out = ggml_cont_4d(ctx0, core_attn_out, S_v, chunk_size * n_chunks, H_v, n_seqs);
+
+ // truncate padded tokens
+ ggml_tensor * output_tokens = ggml_view_4d(ctx0, core_attn_out,
+ S_v, n_tokens, H_v, n_seqs,
+ ggml_row_size(core_attn_out->type, S_v),
+ ggml_row_size(core_attn_out->type, S_v * chunk_size * n_chunks),
+ ggml_row_size(core_attn_out->type, S_v * chunk_size * n_chunks * H_v), 0);
+ output_tokens = ggml_cont(ctx0, output_tokens);
+ // permute back to (S_v, H_v, n_tokens, n_seqs)
+ output_tokens = ggml_permute(ctx0, output_tokens, 0, 2, 1, 3);
+ output_tokens = ggml_cont(ctx0, output_tokens);
+
+ cb(new_state, "output_state", il);
+
+ return {output_tokens, new_state};
+}
+
+std::pair<ggml_tensor *, ggml_tensor *> llm_build_kimi_linear::build_kda_autoregressive(
+ ggml_tensor * q,
+ ggml_tensor * k,
+ ggml_tensor * v,
+ ggml_tensor * gk,
+ ggml_tensor * beta,
+ ggml_tensor * state,
+ int il) {
+ GGML_ASSERT(ggml_is_contiguous(v));
+ GGML_ASSERT(ggml_is_contiguous(gk));
+
+ const int64_t S_k = q->ne[0];
+ const int64_t H_k = q->ne[1];
+ const int64_t n_tokens = q->ne[2];
+ const int64_t n_seqs = q->ne[3];
+
+ const int64_t S_v = v->ne[0];
+ const int64_t H_v = v->ne[1];
+
+ GGML_ASSERT(n_tokens == 1);
+ GGML_ASSERT(v->ne[2] == n_tokens);
+ GGML_ASSERT(k->ne[2] == n_tokens);
+ GGML_ASSERT(gk->ne[0] == S_k && gk->ne[1] == H_k && gk->ne[2] == n_tokens && gk->ne[3] == n_seqs);
+ GGML_ASSERT(beta->ne[0] == H_v && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs);
+ GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_k && state->ne[2] == H_v && state->ne[3] == n_seqs);
+
+ GGML_ASSERT(q->ne[0] == S_k && q->ne[1] == H_k && q->ne[2] == n_tokens && q->ne[3] == n_seqs);
+ GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_k && k->ne[2] == n_tokens && k->ne[3] == n_seqs);
+
+ GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case
+
+ const float eps_norm = hparams.f_norm_rms_eps;
+
+ q = ggml_l2_norm(ctx0, q, eps_norm);
+ k = ggml_l2_norm(ctx0, k, eps_norm);
+
+ const float scale = 1.0f / sqrtf(S_v);
+
+ q = ggml_scale(ctx0, q, scale);
+ beta = ggml_sigmoid(ctx0, beta);
+
+ cb(q, "q_in", il);
+ cb(k, "k_in", il);
+ cb(v, "v_in", il);
+ cb(beta, "beta_in", il);
+ cb(gk, "gk_in", il);
+
+// g [H,1,B,1] g_t [1,H,B,1] => [1,1,H,B]
+// gk [S,H,1,B] => [S,1,H,B] gk_t [1,S,H,B]
+// beta [H,1,1,B] beta_t [1,H,1,B] => [1,1,H,B]
+ gk = ggml_reshape_4d(ctx0, gk, S_k, 1, H_k, n_seqs);
+ ggml_tensor * gk_t = ggml_cont(ctx0, ggml_transpose(ctx0, gk));
+ ggml_tensor * beta_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, beta), 1, 1, H_k, n_seqs);
+
+ // Apply exponential to gk_t
+ gk_t = ggml_exp(ctx0, gk_t);
+ // Apply the gated delta rule for the single timestep
+ // last_recurrent_state = last_recurrent_state * gk_t
+ // S = S * g_i[..., None].exp()
+ state = ggml_mul(ctx0, state, gk_t);
+
+ ggml_tensor * state_t = ggml_cont(ctx0, ggml_transpose(ctx0, state));
+
+// state [S,S,H,B] k [S,1,H,B] k_state [S_v,1,H,B]
+ k = ggml_reshape_4d(ctx0, k, S_k, 1, H_k, n_seqs);
+ ggml_tensor * k_state = ggml_mul_mat(ctx0, state_t, k);
+
+ // v_i - (k_i[..., None] * S).sum(-2)
+ v = ggml_reshape_4d(ctx0, v, S_v, 1, H_v, n_seqs);
+ ggml_tensor * v_diff = ggml_sub(ctx0, v, k_state);
+
+ // b_i[..., None] * k_i
+ ggml_tensor * k_beta = ggml_mul(ctx0, k, beta_t);
+
+ // S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
+ // v_diff_t [1,S_v,H,B] k_beta_t [1,S_k,H,B] state [S_v,S_k,H,B]
+ state = ggml_add(ctx0, state, ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_diff)), ggml_cont(ctx0, ggml_transpose(ctx0, k_beta))));
+
+ q = ggml_reshape_4d(ctx0, q, S_k, 1, H_k, n_seqs);
+ state_t = ggml_cont(ctx0, ggml_transpose(ctx0, state));
+ ggml_tensor * core_attn_out = ggml_mul_mat(ctx0, state_t, q);
+ // core_attn_out should be [S_v, 1, H_v, n_seqs] after this
+ cb(core_attn_out, "output_tokens", il);
+ cb(state, "new_state", il);
+
+ return {core_attn_out, state};
+}
+
llm_build_jamba(const llama_model & model, const llm_graph_params & params);
};
+struct llm_build_kimi_linear : public llm_graph_context_mamba {
+ llm_build_kimi_linear(const llama_model & model, const llm_graph_params & params);
+
+ std::pair<ggml_tensor *, ggml_tensor *> build_kda_autoregressive(
+ ggml_tensor * q,
+ ggml_tensor * k,
+ ggml_tensor * v,
+ ggml_tensor * gk,
+ ggml_tensor * beta,
+ ggml_tensor * state,
+ int il);
+
+ std::pair<ggml_tensor *, ggml_tensor *> build_kda_chunking(
+ ggml_tensor * q,
+ ggml_tensor * k,
+ ggml_tensor * v,
+ ggml_tensor * gk,
+ ggml_tensor * beta,
+ ggml_tensor * state,
+ ggml_tensor * causal_mask,
+ ggml_tensor * identity,
+ ggml_tensor * diag_mask,
+ int il);
+
+ const llama_model & model;
+};
+
struct llm_build_lfm2 : public llm_graph_context {
const llama_model & model;
llm_build_starcoder(const llama_model & model, const llm_graph_params & params);
};
+struct llm_build_step35_iswa : public llm_graph_context {
+ llm_build_step35_iswa(const llama_model & model, const llm_graph_params & params);
+};
+
struct llm_build_t5_dec : public llm_graph_context {
llm_build_t5_dec(const llama_model & model, const llm_graph_params & params);
};
ggml_tensor * Kcur = ggml_view_3d(ctx0, cur, n_embd_head, n_head_kv, n_tokens, cur->nb[1], cur->nb[2], cur->nb[1]*n_head);
cb(Kcur, "Kcur", il);
- ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_3d(ctx0, cur, n_embd_head, n_head_kv, n_tokens, cur->nb[1], cur->nb[2], cur->nb[1]*(n_head+n_head_kv)));
+ ggml_tensor * Vcur = ggml_view_3d(ctx0, cur, n_embd_head, n_head_kv, n_tokens, cur->nb[1], cur->nb[2], cur->nb[1]*(n_head+n_head_kv));
cb(Vcur, "Vcur", il);
Qcur = build_norm(Qcur,
cb(g_diff, "g_diff", il); // shape: (chunk_size, 1, n_chunks, H_v * n_seqs)
ggml_tensor * g_diff_exp = ggml_exp(ctx0, g_diff);
- ggml_tensor * key_gdiff = ggml_mul(ctx0, k, g_diff_exp);
+ ggml_tensor * g_diff_exp_t = ggml_reshape_4d(ctx0, g_diff_exp,
+ 1, chunk_size, n_chunks, g_diff_exp->ne[3]);
+
+ ggml_tensor * key_gdiff = ggml_mul(ctx0, k, g_diff_exp_t);
cb(key_gdiff, "key_gdiff", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs)
+ ggml_tensor * key_gdiff_t = ggml_cont(ctx0, ggml_transpose(ctx0, key_gdiff));
+ cb(key_gdiff_t, "key_gdiff_t", il); // shape: (chunk_size, S_k, n_chunks, H_v * n_seqs)
+
// state to be updated per chunk
ggml_tensor * new_state = state; // ggml_dup(ctx0, state);
: ggml_concat(ctx0, core_attn_out, core_attn_out_chunk, 2);
// kgdmulvnew = (key_gdiff).transpose(-1, -2) @ v_new
- ggml_tensor * k_gdiff = ggml_cont(ctx0, get_slice_2d(ctx0, key_gdiff, chunk));
+ ggml_tensor * k_gdiff_t = get_slice_2d(ctx0, key_gdiff_t, chunk);
//ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, k_gdiff, v_new); // this is slower on metal, why?
- ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, v_new_t, ggml_cont(ctx0, ggml_transpose(ctx0, k_gdiff)));
+ ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, v_new_t, k_gdiff_t);
// last_recurrent_state = last_recurrent_state * g_last + kgdmulvnew
ggml_tensor * gexp_last_chunk = ggml_cont(ctx0, get_slice_2d(ctx0, g_last_exp, chunk));
--- /dev/null
+#include "models.h"
+
+llm_build_step35_iswa::llm_build_step35_iswa(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ inpL = build_inp_embd(model.tok_embd);
+ ggml_tensor * inp_pos = build_inp_pos();
+ auto * inp_attn = build_attn_inp_kv_iswa();
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ for (int il = 0; il < n_layer; ++il) {
+ ggml_tensor * inpSA = inpL;
+
+ const uint32_t n_head_l = hparams.n_head(il);
+ const uint32_t n_head_kv_l = hparams.n_head_kv(il);
+
+ const float freq_base_l = model.get_rope_freq_base(cparams, il);
+ const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
+
+ cur = inpL;
+
+ // dump pre-attn RMSNorm input to pinpoint layer boundary issues
+ cb(cur, "attn_norm_in", il);
+
+ // self-attention
+ {
+ cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+ ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
+ ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
+ ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
+
+ cb(Qcur, "Qcur", il);
+ cb(Kcur, "Kcur", il);
+ cb(Vcur, "Vcur", il);
+
+ Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens);
+ Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
+ Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
+
+ // Q/K per-head RMSNorm (Step35 q_norm / k_norm)
+ if (model.layers[il].attn_q_norm) {
+ Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Qcur, "Qcur_normed", il);
+ }
+ if (model.layers[il].attn_k_norm) {
+ Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Kcur, "Kcur_normed", il);
+ }
+
+ // RoPE (partial rotary factors per layer)
+ const bool is_swa = hparams.is_swa(il);
+ ggml_tensor * rope_factors = is_swa ? nullptr : model.get_rope_factors(cparams, il);
+ const int64_t n_rot_l = is_swa ? hparams.n_rot : (hparams.n_rot / 2);
+ Qcur = ggml_rope_ext(
+ ctx0, Qcur, inp_pos, rope_factors,
+ n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
+ ext_factor, attn_factor, beta_fast, beta_slow
+ );
+ Kcur = ggml_rope_ext(
+ ctx0, Kcur, inp_pos, rope_factors,
+ n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
+ ext_factor, attn_factor, beta_fast, beta_slow
+ );
+ cb(Qcur, "Qcur_pos", il);
+ cb(Kcur, "Kcur_pos", il);
+
+ const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
+ ggml_tensor * attn_out = build_attn(inp_attn,
+ nullptr, nullptr,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ cb(attn_out, "attn_out", il);
+ // head-wise attention gate: sigmoid(g_proj(x)) in torch
+ if (model.layers[il].wqkv_gate) {
+ ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, cur); // [n_head_l, n_tokens]
+ cb(gate, "attn_gate", il);
+
+ gate = ggml_sigmoid(ctx0, gate);
+ cb(gate, "attn_gate_sigmoid", il);
+
+ // reshape + broadcast to [n_embd_head_v, n_head_l, n_tokens]
+ ggml_tensor * attn_3d = ggml_reshape_3d(ctx0, attn_out, n_embd_head_v, n_head_l, n_tokens);
+ ggml_tensor * gate_3d = ggml_reshape_3d(ctx0, gate, 1, n_head_l, n_tokens);
+ cb(gate_3d, "attn_gate_3d", il);
+
+ attn_3d = ggml_mul(ctx0, attn_3d, gate_3d);
+ cb(attn_3d, "attn_gated_3d", il);
+
+ attn_out = ggml_reshape_2d(ctx0, attn_3d, n_embd_head_v * n_head_l, n_tokens);
+ cb(attn_out, "attn_gated", il);
+ }
+
+ // output projection
+ cur = build_lora_mm(model.layers[il].wo, attn_out);
+ cb(cur, "attn_proj", il);
+ }
+
+ if (il == n_layer - 1 && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "ffn_inp", il);
+
+ cur = build_norm(ffn_inp, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ // feed-forward
+ if (model.layers[il].ffn_gate_inp == nullptr) {
+ // dense MLP
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, model.layers[il].ffn_up_b, nullptr,
+ model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, nullptr,
+ model.layers[il].ffn_down, model.layers[il].ffn_down_b, nullptr,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+ } else {
+ // MoE routed experts
+ const bool norm_w = hparams.expert_weights_norm;
+ const float w_scale = hparams.expert_weights_scale;
+ const bool scale_w = w_scale != 0.0f;
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ model.layers[il].ffn_gate_inp,
+ model.layers[il].ffn_up_exps,
+ model.layers[il].ffn_gate_exps,
+ model.layers[il].ffn_down_exps,
+ model.layers[il].ffn_exp_probs_b,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU,
+ norm_w, scale_w, w_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il);
+ cb(moe_out, "ffn_moe_out", il);
+
+ // shared expert MLP (always added on MoE layers in Step35)
+ ggml_tensor * sh_out = build_ffn(cur,
+ model.layers[il].ffn_up_shexp, nullptr, nullptr,
+ model.layers[il].ffn_gate_shexp, nullptr, nullptr,
+ model.layers[il].ffn_down_shexp, nullptr, nullptr,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(sh_out, "ffn_shared_out", il);
+
+ cur = ggml_add(ctx0, moe_out, sh_out);
+ cb(cur, "ffn_out", il);
+ }
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ inpL = cur;
+ }
+
+ cur = inpL;
+
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ cur = build_lora_mm(model.output, cur);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
return bpe_offsets;
}
-// use std::wregex to split the text
-static std::vector<size_t> unicode_regex_split_stl(const std::wstring & wtext, const std::wstring & regex_expr, const std::vector<size_t> & offsets) {
- std::wregex expr(regex_expr, std::regex_constants::optimize | std::regex_constants::nosubs);
- std::vector<size_t> bpe_offsets; // store the offset of each word
- bpe_offsets.reserve(offsets.size()); // Reserve memory for the approximate size
- size_t start = 0;
- for (auto offset : offsets) {
- std::wcregex_iterator it(wtext.data() + start, wtext.data() + start + offset, expr);
- std::wcregex_iterator end;
-
- int64_t start_idx = 0;
- while (it != end) {
- std::wcmatch match = *it;
- if (match.position() > start_idx) {
- bpe_offsets.emplace_back(match.position() - start_idx);
- }
- bpe_offsets.emplace_back(match.length());
- start_idx = match.position() + match.length();
- ++it;
- }
-
- if (start_idx < (int64_t) offset) {
- bpe_offsets.emplace_back(offset - start_idx);
- }
- start += offset;
- }
-
- return bpe_offsets;
-}
-
-// use std::regex to split the text
-static std::vector<size_t> unicode_regex_split_stl(const std::string & text, const std::string & regex_expr, const std::vector<size_t> & offsets) {
- std::regex expr(regex_expr, std::regex_constants::optimize | std::regex_constants::nosubs);
+template <typename CharT>
+static std::vector<size_t> unicode_regex_split_stl(const std::basic_string<CharT> & text, const std::basic_string<CharT> & regex, const std::vector<size_t> & offsets) {
+ using BidirIt = typename std::basic_string<CharT>::const_iterator;
+#ifdef _MSC_VER
+ // Bypass bug in MSVC: https://github.com/ggml-org/llama.cpp/issues/17830
+ constexpr auto regex_flags = std::regex_constants::ECMAScript;
+#else
+ constexpr auto regex_flags = std::regex_constants::optimize | std::regex_constants::nosubs;
+#endif
+ std::basic_regex<CharT> expr(regex, regex_flags);
std::vector<size_t> bpe_offsets; // store the offset of each word
bpe_offsets.reserve(offsets.size()); // Reserve memory for the approximate size
size_t start = 0;
for (auto offset : offsets) {
- std::cregex_iterator it(text.data() + start, text.data() + start + offset, expr);
- std::cregex_iterator end;
+ std::regex_iterator<BidirIt> it(text.begin() + start, text.begin() + start + offset, expr);
+ std::regex_iterator<BidirIt> end;
int64_t start_idx = 0;
while (it != end) {
- std::cmatch match = *it;
+ std::match_results<BidirIt> match = *it;
if (match.position() > start_idx) {
bpe_offsets.emplace_back(match.position() - start_idx);
}
overview / pkg / ggml / sources / whisper.cpp / commitdiff