Fix #2093

[ikawrakow]

PR #2093 costed me the day. I was working on various DSA optimizations, and each time the new version was massively slower than what I thought we had on the main branch. Until I finally tested the main branch and got the same slow down.

The culprit is #2093. Not allocating the indexer cache in a layer had the effect of turning off DSA in that layer. I had missed that, and, not having redone the performance benchmarks (because painfully slow), didn't notice the performance degradation at context above 2048.

This PR fixes the unintended DSA removal in layers without own indexer cache.

[saood06]

I had missed that, and, not having redone the performance benchmarks (because painfully slow), didn't notice the performance degradation at context above 2048.

If it helps there is this small model that implements DSA. Obviously wouldn't have ran into or caught the bug in #2093 as it doesn't use IndexShare but it should be much faster to benchmark, and unlike random models you can actually validate coherency. Only other downside is you'd have to add support for it's architecture.

[sayap]

For non-unsloth quant that doesn't duplicate the indexer tensor to all layers, I think this PR is the first time that IndexShare actually takes effect with -dsa?

Testing with https://huggingface.co/sokann/GLM-5.2-GGUF-2.244bpw, the model doesn't seem to behave correctly with this PR. With a 16k prompt that consists of some log lines followed by a bunch of code, the model now responds with:

The code is a bit garbled/corrupted (it looks like it's been mangled through some kind of obfuscation, but the structure is recognizable).

[ikawrakow]

@sayap Can you post the request where you think the model does not behave correctly?

[ikawrakow]

@saood06 Thanks for trying to help. Does the model you linked to share indexer results between layers? Because if it does not, then no, testing with it wouldn't have caught the bug.

[sayap]

Can try this one:

{
  "model": "",
  "messages": [
    {
      "role": "user",
      "content": "Explain what this does:\n\n#include \"llama-sampling.h\"\n#include \"llama-vocab.h\"\n#include \"llama-grammar.h\"\n\n#include \"iqk/iqk_cpu_ops.h\"\n\n#include <algorithm>\n#include <cstring>\n#include <ctime>\n#include <cfloat>\n#include <numeric>\n#include <unordered_map>\n#include <fstream>\n\nstatic void llama_log_softmax(float * array, size_t size) {\n    float max_l = *std::max_element(array, array + size);\n    float sum = 0.f;\n    for (size_t i = 0; i < size; ++i) {\n        float p = expf(array[i] - max_l);\n        sum += p;\n        array[i] = p;\n    }\n\n    for (size_t i = 0; i < size; ++i) {\n        array[i] = logf(array[i] / sum);\n    }\n}\n\nvoid llama_set_rng_seed_impl(struct llama_sampling * smpl, uint32_t seed) {\n    if (seed == LLAMA_DEFAULT_SEED) {\n        seed = time(NULL);\n    }\n\n    smpl->rng.seed(seed);\n}\n\nstatic void llama_sort(llama_token_data_array * candidates, int32_t k) {\n    if (candidates->sorted || candidates->size < 2) {\n        return;\n    }\n    if (k < 0) {\n        k = candidates->size;\n    }\n    auto comp = [](const llama_token_data & a, const llama_token_data & b) {\n        return a.logit > b.logit;\n    };\n    if (k <= 1024) { //128) {\n        if (k == int(candidates->size)) {\n            std::sort(candidates->data, candidates->data + candidates->size, comp);\n        } else {\n            std::partial_sort(candidates->data, candidates->data + k, candidates->data + candidates->size, comp);\n        }\n    } else {\n        constexpr int   nbuckets     = 128;\n        constexpr float bucket_low   = -10.0f;\n        constexpr float bucket_high  =  10.0f;\n        constexpr float bucket_scale = nbuckets/(bucket_high - bucket_low);\n        constexpr float bucker_inter = -bucket_low * bucket_scale;\n\n        std::vector<int> bucket_idx(candidates->size);\n        std::vector<int> histo(nbuckets, 0);\n\n        for (int i = 0; i < (int)candidates->size; ++i) {\n            const float val = candidates->data[i].logit;\n            int ib = int(bucket_scale * val + bucker_inter); //nbuckets * (val - bucket_low) / (bucket_high - bucket_low);\n            ib = std::max(0, std::min(nbuckets-1, ib));\n            bucket_idx[i] = ib;\n            ++histo[ib];\n        }\n        int nhave = 0;\n        int ib = nbuckets - 1;\n        for ( ; ib >= 0; --ib) {\n            nhave += histo[ib];\n            if (nhave >= k) break;\n        }\n        std::vector<llama_token_data> tmp_tokens(nhave);\n        auto ptr = tmp_tokens.data();\n        std::vector<llama_token_data*> bucket_ptrs;\n        bucket_ptrs.reserve(nbuckets - ib);\n        for (int j = nbuckets - 1; j >= ib; --j) {\n            bucket_ptrs.push_back(ptr);\n            ptr += histo[j];\n        }\n        for (int i = 0; i < (int)candidates->size; ++i) {\n            int j = bucket_idx[i];\n            if (j >= ib) {\n                *bucket_ptrs[nbuckets-1-j]++ = candidates->data[i];\n            }\n        }\n\n        ptr = tmp_tokens.data();\n        int ndone = 0;\n        for (int j = nbuckets-1; j > ib; --j) {\n            std::sort(ptr, ptr + histo[j], comp);\n            ptr += histo[j];\n            ndone += histo[j];\n        }\n        std::partial_sort(ptr, ptr + k - ndone, ptr + histo[ib], comp);\n\n        std::memcpy(candidates->data, tmp_tokens.data(), k*sizeof(llama_token_data));\n\n    }\n    candidates->sorted = true;\n}\n\nvoid llama_sample_softmax_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, bool normalize) {\n    GGML_ASSERT(candidates->size > 0);\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    // Sort the logits in descending order if necessary\n    llama_sort(candidates, -1);\n\n    float max_l = candidates->data[0].logit;\n    float cum_sum = 0.0f;\n    for (size_t i = 0; i < candidates->size; ++i) {\n        float p = expf(candidates->data[i].logit - max_l);\n        candidates->data[i].p = p;\n        cum_sum += p;\n    }\n    if (normalize) {\n        for (size_t i = 0; i < candidates->size; ++i) {\n            candidates->data[i].p /= cum_sum;\n        }\n    }\n\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_top_k_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, int32_t k, size_t min_keep) {\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    if (k <= 0) {\n        k = candidates->size;\n    }\n\n    k = std::max(k, (int) min_keep);\n    k = std::min(k, (int) candidates->size);\n\n    llama_sort(candidates, k);\n\n    candidates->size = k;\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_top_p_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float p, size_t min_keep) {\n    if (p >= 1.0f) {\n        return;\n    }\n\n    llama_sample_softmax_impl(smpl, candidates);\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    // Compute the cumulative probabilities\n    float cum_sum = 0.0f;\n    size_t last_idx = candidates->size;\n\n    for (size_t i = 0; i < candidates->size; ++i) {\n        cum_sum += candidates->data[i].p;\n\n        // Check if the running sum is at least p or if we have kept at least min_keep tokens\n        // we set the last index to i+1 to indicate that the current iterate should be included in the set\n        if (cum_sum >= p && i + 1 >= min_keep) {\n            last_idx = i + 1;\n            break;\n        }\n    }\n\n    // Resize the output vector to keep only the top-p tokens\n    candidates->size = last_idx;\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_min_p_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float p, size_t min_keep) {\n    if (p <= 0.0f || !candidates->size) {\n        return;\n    }\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    bool min_p_applied = false;\n\n    // if the candidates aren't sorted, try the unsorted implementation first\n    if (!candidates->sorted) {\n        std::vector<llama_token_data> filtered_tokens;\n\n        float max_logit = -FLT_MAX;\n        for (size_t i = 0; i < candidates->size; ++i) {\n            max_logit = std::max(max_logit, candidates->data[i].logit);\n        }\n        const float min_logit = max_logit + logf(p); // min logit for p_i >= p * p_max\n\n        for (size_t i = 0; i < candidates->size; ++i) {\n            if (candidates->data[i].logit >= min_logit) {\n                filtered_tokens.push_back(candidates->data[i]);\n            }\n        }\n\n        // if we have enough values the operation was a success\n        if (filtered_tokens.size() >= min_keep) {\n            memcpy(candidates->data, filtered_tokens.data(), filtered_tokens.size()*sizeof(llama_token_data));\n            candidates->size = filtered_tokens.size();\n            min_p_applied = true;\n        }\n    }\n\n    // if the candidates are sorted or the unsorted implementation failed, use this implementation\n    if (!min_p_applied) {\n        // Sort the logits in descending order if needed\n        llama_sort(candidates, -1);\n\n        const float min_logit = candidates->data[0].logit + logf(p); // min logit for p_i >= p * p_max\n        size_t i = 1; // first token always matches\n\n        for (; i < candidates->size; ++i) {\n            if (candidates->data[i].logit < min_logit && i >= min_keep) {\n                break; // prob too small\n            }\n        }\n\n        // Resize the output vector to keep only the matching tokens\n        candidates->size = i;\n    }\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_tail_free_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float z, size_t min_keep) {\n    if (z >= 1.0f || candidates->size <= 2) {\n        return;\n    }\n\n    llama_sample_softmax_impl((struct llama_sampling *) nullptr, candidates);\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    // Compute the first and second derivatives\n    std::vector<float> first_derivatives(candidates->size - 1);\n    std::vector<float> second_derivatives(candidates->size - 2);\n\n    for (size_t i = 0; i < first_derivatives.size(); ++i) {\n        first_derivatives[i] = candidates->data[i].p - candidates->data[i + 1].p;\n    }\n    for (size_t i = 0; i < second_derivatives.size(); ++i) {\n        second_derivatives[i] = first_derivatives[i] - first_derivatives[i + 1];\n    }\n\n    // Calculate absolute value of second derivatives\n    for (size_t i = 0; i < second_derivatives.size(); ++i) {\n        second_derivatives[i] = std::abs(second_derivatives[i]);\n    }\n\n    // Normalize the second derivatives\n    {\n        const float second_derivatives_sum = std::accumulate(second_derivatives.begin(), second_derivatives.end(), 0.0f);\n\n        if (second_derivatives_sum > 1e-6f) {\n            for (float & value : second_derivatives) {\n                value /= second_derivatives_sum;\n            }\n        } else {\n            for (float & value : second_derivatives) {\n                value = 1.0f / second_derivatives.size();\n            }\n        }\n    }\n\n    if (min_keep < 1) min_keep = 1;\n    float cum_sum = 0.0f;\n    size_t last_idx = candidates->size;\n    for (size_t i = 0; i < second_derivatives.size(); ++i) {\n        cum_sum += second_derivatives[i];\n\n        // Check if the running sum is greater than z or if we have kept at least min_keep tokens\n        if (cum_sum > z && i >= min_keep) {\n            last_idx = i;\n            break;\n        }\n    }\n\n    // Resize the output vector to keep only the tokens above the tail location\n    candidates->size = last_idx;\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_typical_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float p, size_t min_keep) {\n    // Reference implementation:\n    // https://github.com/huggingface/transformers/compare/main...cimeister:typical-sampling:typical-pr\n    if (p >= 1.0f) {\n        return;\n    }\n\n    // Compute the softmax of logits and calculate entropy\n    llama_sample_softmax_impl((struct llama_sampling *) nullptr, candidates);\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    float entropy = 0.0f;\n    for (size_t i = 0; i < candidates->size; ++i) {\n        entropy += -candidates->data[i].p * logf(candidates->data[i].p);\n    }\n\n    // Compute the absolute difference between negative log probability and entropy for each candidate\n    std::vector<float> shifted_scores(candidates->size);\n    for (size_t i = 0; i < candidates->size; ++i) {\n        shifted_scores[i] = fabsf(-logf(candidates->data[i].p) - entropy);\n    }\n\n    // Sort tokens based on the shifted_scores and their corresponding indices\n    std::vector<size_t> indices(candidates->size);\n    std::iota(indices.begin(), indices.end(), 0);\n\n    std::sort(indices.begin(), indices.end(), [&](size_t a, size_t b) {\n        return shifted_scores[a] < shifted_scores[b];\n    });\n\n    // Compute the cumulative probabilities\n    float cum_sum = 0.0f;\n    size_t last_idx = indices.size();\n\n    for (size_t i = 0; i < indices.size(); ++i) {\n        size_t idx = indices[i];\n        cum_sum += candidates->data[idx].p;\n\n        // Check if the running sum is greater than typical or if we have kept at least min_keep tokens\n        if (cum_sum > p && i + 1 >= min_keep) {\n            last_idx = i + 1;\n            break;\n        }\n    }\n\n    // Resize the output vector to keep only the locally typical tokens\n    std::vector<llama_token_data> new_candidates(last_idx);\n    for (size_t i = 0; i < last_idx; ++i) {\n        size_t idx = indices[i];\n        new_candidates[i] = candidates->data[idx];\n    }\n\n    // Replace the data in candidates with the new_candidates data\n    std::copy(new_candidates.begin(), new_candidates.end(), candidates->data);\n    candidates->size = new_candidates.size();\n    candidates->sorted = false;\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_entropy_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float min_temp, float max_temp, float exponent_val) {\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    // no need to do anything if there is only one (or zero) candidates\n    if(candidates->size <= 1) {\n        return;\n    }\n\n    // Calculate maximum possible entropy\n    float max_entropy = -logf(1.0f / candidates->size);\n\n    llama_sample_softmax_impl((struct llama_sampling *) nullptr, candidates);\n\n    // Calculate entropy of the softmax probabilities\n    float entropy = 0.0f;\n    for (size_t i = 0; i < candidates->size; ++i) {\n        float prob = candidates->data[i].p;\n        if (prob > 0.0f) { // Ensure no log(0)\n            entropy -= prob * logf(prob);\n        }\n    }\n\n    // Normalize the entropy (max_entropy cannot be 0 here because we checked candidates->size != 1 above)\n    float normalized_entropy = entropy / max_entropy;\n\n    // Map the normalized entropy to the desired temperature range using the power function\n    float dyn_temp = min_temp + (max_temp - min_temp) * powf(normalized_entropy, exponent_val);\n\n#ifdef DEBUG\n    LLAMA_LOG_INFO(\"Your text maxtemp value is: %f\\n\", max_temp);\n    LLAMA_LOG_INFO(\"Entropy: %f\\n\", entropy);\n    LLAMA_LOG_INFO(\"Max Possible Entropy: %f\\n\", max_entropy);\n    LLAMA_LOG_INFO(\"Normalized Entropy: %f\\n\", normalized_entropy);\n    LLAMA_LOG_INFO(\"Exponent: %f\\n\", exponent_val);\n    LLAMA_LOG_INFO(\"Dynamic Temperature (dyn_temp): %f\\n\", dyn_temp);\n#endif\n\n    // Apply the dynamically calculated temperature scaling\n    for (size_t i = 0; i < candidates->size; ++i) {\n        candidates->data[i].logit /= dyn_temp;\n    }\n\n    // Re-compute softmax probabilities after scaling logits with dynamic temperature\n    double max_l_double = candidates->data[0].logit;\n    double cum_sum_double = 0.0;\n    for (size_t i = 0; i < candidates->size; ++i) {\n        double p = exp(candidates->data[i].logit - max_l_double);\n        candidates->data[i].p = p; // Store the scaled probability\n        cum_sum_double += p;\n    }\n    for (size_t i = 0; i < candidates->size; ++i) {\n        candidates->data[i].p /= cum_sum_double; // Re-normalize the probabilities\n    }\n\n#ifdef DEBUG\n    // Print the updated top 25 probabilities after temperature scaling\n    LLAMA_LOG_INFO(\"\\nUpdated Top 25 Probabilities After Dynamic Temperature Scaling (in percentages):\\n\");\n    for (size_t i = 0; i < 25 && i < candidates->size; ++i) {\n        LLAMA_LOG_INFO(\"Token %zu: %f%%\\n\", i + 1, candidates->data[i].p * 100.0f);\n    }\n#endif\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_temp_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float temp) {\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    for (size_t i = 0; i < candidates->size; ++i) {\n        candidates->data[i].logit /= temp;\n    }\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_xtc_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float probability, float threshold, size_t min_keep) {\n    if (probability <= 0 || threshold > 0.5f || candidates->size < 2) {\n        return;\n    }\n    GGML_ASSERT(smpl);\n    const int64_t t_start_sample_us = ggml_time_us();\n    if (probability < 1) {\n        std::uniform_real_distribution<float> distribution(0.0f, 1.0f);\n        float chance = distribution(smpl->rng);\n        if (chance > probability) return;\n    }\n\n    llama_sample_softmax_impl(nullptr, candidates);\n\n    int pos_last = 0;\n\n    for (size_t i = 0; i < candidates->size; ++i) {\n        if (candidates->data[i].p >= threshold) {\n            pos_last = i;\n        } else break;\n    }\n\n    if (candidates->size - pos_last >= min_keep && pos_last > 0) {\n        candidates->data += pos_last;\n        candidates->size -= pos_last;\n    }\n\n    smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    smpl->n_sample++;\n\n}\n\nvoid llama_sample_top_n_sigma_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float top_n_sigma) {\n\n    if (top_n_sigma <= 0.0f || candidates->size < 4) {\n        // top_n_sigma <= 0: disabled\n        // candidates->size < 4: no point in applying the transformation for fewer than 4 logits.\n        return;\n    }\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    double sum1 = 0, sum2 = 0;\n    float max = 0;\n    int count = 0;\n    for (int i = 0; i < (int)candidates->size; ++i) {\n        if (auto l = candidates->data[i].logit; l != -INFINITY) {\n            max = std::max(max, l);\n            ++count;\n            double dl = l;\n            sum1 += dl;\n            sum2 += dl*dl;\n        }\n    }\n    if (count < 4) {\n        return;\n    }\n    double dmean = sum1/count;\n    double dsigma2 = sum2/count - dmean*dmean;\n    if (dsigma2 <= 0) {\n        return;\n    }\n    float sigma = float(sqrt(dsigma2));\n\n    float thresh = max - top_n_sigma*sigma;\n    int n_use = 0;\n    for (int i = 0; i < (int)candidates->size; ++i) {\n        if (candidates->data[i].logit >= thresh) {\n            candidates->data[n_use++] = candidates->data[i];\n        }\n    }\n    if (n_use < (int)candidates->size) {\n        candidates->size = n_use;\n        llama_sample_softmax_impl(nullptr, candidates);\n    }\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n        smpl->n_sample++;\n    }\n}\n\n\nvoid llama_sample_repetition_penalties_impl(\n        struct llama_sampling * smpl,\n       llama_token_data_array * candidates,\n            const llama_token * last_tokens,\n                       size_t   penalty_last_n,\n                       float   penalty_repeat,\n                       float   penalty_freq,\n                       float   penalty_present) {\n    if (penalty_last_n == 0 || (penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f)) {\n        return;\n    }\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    // Create a frequency map to count occurrences of each token in last_tokens\n    std::unordered_map<llama_token, int> token_count;\n    for (size_t i = 0; i < penalty_last_n; ++i) {\n        token_count[last_tokens[i]]++;\n    }\n\n    // Apply frequency and presence penalties to the candidates\n    for (size_t i = 0; i < candidates->size; ++i) {\n        const auto token_iter = token_count.find(candidates->data[i].id);\n        if (token_iter == token_count.end()) {\n            continue;\n        }\n\n        const int count = token_iter->second;\n\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.\n        // This is common fix for this problem, which is to multiply by the penalty instead of dividing.\n        if (candidates->data[i].logit <= 0) {\n            candidates->data[i].logit *= penalty_repeat;\n        } else {\n            candidates->data[i].logit /= penalty_repeat;\n        }\n\n        candidates->data[i].logit -= float(count) * penalty_freq + float(count > 0) * penalty_present;\n    }\n\n    candidates->sorted = false;\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n}\n\nvoid llama_sample_apply_guidance_impl(\n        struct llama_sampling * smpl,\n                        float * logits,\n                        float * logits_guidance,\n                        float   scale) {\n    GGML_ASSERT(smpl);\n\n    const auto t_start_sample_us = ggml_time_us();\n    const auto n_vocab = smpl->n_vocab;\n\n    llama_log_softmax(logits, n_vocab);\n    llama_log_softmax(logits_guidance, n_vocab);\n\n    for (int i = 0; i < n_vocab; ++i) {\n              auto & l = logits[i];\n        const auto & g = logits_guidance[i];\n\n        l = scale * (l - g) + g;\n    }\n\n    smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n}\n\nllama_token llama_sample_token_mirostat_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float tau, float eta, int32_t m, float * mu) {\n    GGML_ASSERT(smpl);\n\n    const int32_t n_vocab = float(smpl->n_vocab);\n\n    int64_t t_start_sample_us = ggml_time_us();\n\n    llama_sample_softmax_impl((struct llama_sampling *) nullptr, candidates);\n\n    // Estimate s_hat using the most probable m tokens\n    float s_hat = 0.0;\n    float sum_ti_bi = 0.0;\n    float sum_ti_sq = 0.0;\n    for (size_t i = 0; i < size_t(m - 1) && i < candidates->size - 1; ++i) {\n        float t_i = logf(float(i + 2) / float(i + 1));\n        float b_i = logf(candidates->data[i].p / candidates->data[i + 1].p);\n        sum_ti_bi += t_i * b_i;\n        sum_ti_sq += t_i * t_i;\n    }\n    s_hat = sum_ti_bi / sum_ti_sq;\n\n    // Compute k from the estimated s_hat and target surprise value\n    float epsilon_hat = s_hat - 1;\n    float k = powf((epsilon_hat * powf(2, *mu)) / (1 - powf(n_vocab, -epsilon_hat)), 1 / s_hat);\n\n    // Sample the next word X using top-k sampling\n    llama_sample_top_k_impl((struct llama_sampling *) nullptr, candidates, int(k), 1);\n    smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    llama_token X = llama_sample_token_impl(smpl, candidates);\n    t_start_sample_us = ggml_time_us();\n\n    // Compute error as the difference between observed surprise and target surprise value\n    size_t X_idx = std::distance(candidates->data, std::find_if(candidates->data, candidates->data + candidates->size, [&](const llama_token_data & candidate) {\n        return candidate.id == X;\n    }));\n    float observed_surprise = -log2f(candidates->data[X_idx].p);\n    float e = observed_surprise - tau;\n\n    // Update mu using the learning rate and error\n    *mu = *mu - eta * e;\n\n    smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    return X;\n}\n\nllama_token llama_sample_token_mirostat_v2_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, float tau, float eta, float * mu) {\n    int64_t t_start_sample_us;\n    t_start_sample_us = ggml_time_us();\n\n    llama_sample_softmax_impl(smpl, candidates);\n\n    // Truncate the words with surprise values greater than mu\n    candidates->size = std::distance(candidates->data, std::find_if(candidates->data, candidates->data + candidates->size, [&](const llama_token_data & candidate) {\n        return -log2f(candidate.p) > *mu;\n    }));\n\n    if (candidates->size == 0) {\n        candidates->size = 1;\n    }\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n\n    // Normalize the probabilities of the remaining words\n    llama_sample_softmax_impl(smpl, candidates);\n\n    // Sample the next word X from the remaining words\n    llama_token X = llama_sample_token_impl(smpl, candidates);\n    t_start_sample_us = ggml_time_us();\n\n    // Compute error as the difference between observed surprise and target surprise value\n    size_t X_idx = std::distance(candidates->data, std::find_if(candidates->data, candidates->data + candidates->size, [&](const llama_token_data & candidate) {\n        return candidate.id == X;\n    }));\n    float observed_surprise = -log2f(candidates->data[X_idx].p);\n    float e = observed_surprise - tau;\n\n    // Update mu using the learning rate and error\n    *mu = *mu - eta * e;\n\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    }\n    return X;\n}\n\nllama_token llama_sample_token_greedy_impl(struct llama_sampling * smpl, llama_token_data_array * candidates) {\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    // Find max element\n    auto * max_iter = std::max_element(candidates->data, candidates->data + candidates->size, [](const llama_token_data & a, const llama_token_data & b) {\n        return a.logit < b.logit;\n    });\n\n    llama_token result = max_iter->id;\n    if (smpl) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n        smpl->n_sample++;\n    }\n    return result;\n}\n\nllama_token llama_sample_token_with_rng_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, std::mt19937 & rng) {\n    GGML_ASSERT(smpl);\n\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    if (candidates->size < 2) {\n        smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n        smpl->n_sample++;\n        return candidates->data[0].id;\n    }\n\n    std::vector<float> probs(candidates->size);\n    probs[0] = candidates->data[0].logit;\n    float max = probs[0];\n    for (int j = 1; j < candidates->size; ++j) {\n        probs[j] = candidates->data[j].logit;\n        max = std::max(max, probs[j]);\n    }\n\n    float sump = 0;\n    for (int j = 0; j < candidates->size; ++j) {\n        float p = expf(probs[j] - max);\n        sump += p;\n        probs[j] = sump;\n    }\n    probs.back() += sump;\n\n    auto r = rng();\n    auto p = sump * r / rng.max();\n    auto iter = std::upper_bound(probs.begin(), probs.end(), p);\n    if (iter == probs.end()) {\n        LLAMA_LOG_ERROR(\"=============================== Failed to sample token\\n\");\n        std::ofstream out(\"probabilities.txt\");\n        out << \"candidates->size: \" << candidates->size << std::endl;\n        out << \"max  = \" << max << std::endl;\n        out << \"sump = \" << sump << std::endl;\n        out << \"r    = \" << r << std::endl;\n        out << \"probabilities:\\n\";\n        for (int j = 0; j < candidates->size; ++j) {\n            out << j << \"  \" << candidates->data[j].id << \"  \" << candidates->data[j].logit << \"  \" << probs[j] << std::endl;\n        }\n        out.flush();\n        out.close();\n        LLAMA_LOG_ERROR(\"Data has been stored in probabilities.txt\\n\");\n        LLAMA_LOG_ERROR(\"Create an issue with full log and attach probabilities.txt to the issue\\n\");\n        LLAMA_LOG_ERROR(\"\\n\\nCrashing now\\n\");\n        GGML_ABORT(\"Fatal error\");\n    }\n    GGML_ASSERT(iter != probs.end());\n    auto idx = std::distance(probs.begin(), iter);\n    auto id  = candidates->data[idx].id;\n\n    smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    smpl->n_sample++;\n\n    return id;\n}\n\nllama_token llama_sample_token_impl(struct llama_sampling * smpl, llama_token_data_array * candidates) {\n    return llama_sample_token_with_rng_impl(smpl, candidates, smpl->rng);\n}\n\n\n// DRY\n\n// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)\nstatic 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) {\n    for (llama_token token_id = 0; token_id < (llama_token)vocab.n_tokens(); token_id++) {\n        auto word = vocab.detokenize( { token_id }, true);\n        if (word.find(str) != std::string::npos) {\n            token_sequences.emplace(token_id, std::vector<llama_token>());\n        }\n        else {\n            size_t word_len = word.size(), str_len = str.size();\n            size_t pos = -1;\n            while ((pos = word.find(str[0], pos + 1)) != std::string::npos) {\n                bool match = true;\n                size_t i;\n                for (i = 1; i < str_len && i + pos < word_len; ++i) {\n                    if (word[pos + i] != str[i]) {\n                        match = false;\n                        break;\n                    }\n                }\n                if (match) {\n                    auto tokenization = vocab.tokenize(str.substr(i), false, false);\n                    //std::vector<llama_token> tokenization = llama_tokenize_internal(vocab, str.substr(i), false, false);\n                    if (max_tail_len >= 0 && tokenization.size() > (size_t)max_tail_len) {\n                        tokenization.resize(max_tail_len);\n                    }\n\n                    // Ensure we don't already have a duplicate matching tokenization\n                    auto its = token_sequences.equal_range(token_id);\n                    bool found = false;\n                    for (auto it = its.first; it != its.second; ++it) {\n                        if (tokenization == it->second) {\n                            found = true;\n                            break;\n                        }\n                    }\n                    if (!found) {\n                        token_sequences.emplace(token_id, tokenization);\n                    }\n                }\n            }\n        }\n    }\n}\n\nstatic const char* llama_sampler_dry_name(const struct llama_sampler* /*smpl*/) {\n    return \"dry\";\n}\n\n\n\n// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)\nvoid llama_sampler_dry_apply(struct llama_sampler_dry* smpl, llama_token_data_array* cur_p) {\n    if (smpl->dry_multiplier == 0.0f || smpl->dry_base < 1.0f || smpl->dry_penalty_last_n == 0) {\n        return;\n    }\n\n    int32_t effective_dry_penalty_last_n = (smpl->dry_penalty_last_n == -1) ? smpl->total_context_size : std::max(smpl->dry_penalty_last_n, 0);\n    int last_n_repeat = std::min(std::min((int)smpl->last_tokens.size(), effective_dry_penalty_last_n), smpl->total_context_size);\n\n    if (last_n_repeat <= smpl->dry_allowed_length) {\n        return;\n    }\n\n    smpl->dry_repeat_count.assign(last_n_repeat, 0);\n    smpl->dry_max_token_repeat.clear();\n\n    // Step 1: Look for restart sequences to limit the maximum repetition length.\n    // Work backwards through the context looking for any token that begins a restart sequence.\n    //\n    // The collection `restart_sequences` is a mapping from a \"head\" token to all \"tail\"\n    // sequences that together comprise a restart sequence. This allows us to quickly check\n    // whether each token is the head of a complete sequence. Most restart sequences are actually\n    // a single token, and for these the \"tail\" is an empty vector.\n    //\n    // If the token is a \"head\", test all restart sequences that begin with this token\n    // (there will often only be one sequence for each token, but if sequences like 'aaaq1' and\n    // 'aaa1' are used as restart strings, both could start with 'aaa' when tokenized). The\n    // longest matching sequence (if any) is used to limit the maximum repetition length.\n    //\n    // Note that in the case case of a short sequence contained in a longer one, this might fail to\n    // find the smallest value for `rep_limit`. For example, if 'amniotic' and 'ni' are both used as\n    // restart sequences, 'ni' will be found first, and since it's shorter it will fail to suppress\n    // 'otic'. This is a minor issue since fully contained restart sequences are likely to be rare.\n    //\n    // This is theoretically worst-case O(N^2) for arbitrary restart sequences, which is why we\n    // have already clamped the maximum tail sequence length when generating `restart_sequences`.\n    // With clamping, this scan is O(N) in the context length.\n\n    int rep_limit = last_n_repeat;\n    for (int i = 0; i < last_n_repeat; ++i) {\n        llama_token token = smpl->last_tokens.rat(i);\n        auto its = smpl->dry_processed_breakers.equal_range(token);\n        if (its.first == smpl->dry_processed_breakers.end()) {\n            continue;\n        }\n        int longest_match = -1;\n        for (auto it = its.first; it != its.second; ++it) {\n            // Note that (*it) does not contain the head character, so seq_len will be\n            // the restart sequence length minus 1.\n            // In the common case of a single-token restart sequence, (*it) will be empty\n            // and we will trivially match.\n            int seq_len = (int)it->second.size();\n            if (seq_len > longest_match && seq_len <= (int)i) {\n                bool match = true;\n                for (int offset = 0; offset < seq_len; ++offset) {\n                    // The -1 when indexing `last_tokens` is because we already matched the head.\n                    if (it->second[offset] != smpl->last_tokens.rat(i - offset - 1)) {\n                        match = false;\n                        break;\n                    }\n                }\n                if (match) {\n                    longest_match = seq_len;\n                }\n            }\n        }\n        if (longest_match >= 0) {\n            // We found a restart sequence starting `i` tokens from the end and continuing for\n            // `longest_match` tokens.\n            rep_limit = i - longest_match;\n            break;\n        }\n    }\n    if (rep_limit < smpl->dry_allowed_length) {\n        return;\n    }\n\n    // Step 2: Iterate in reverse over the last N tokens of the context, using the \"Z-algorithm\" (in\n    // the reverse direction) to efficiently compute the positions and lengths of suffixes appearing\n    // elsewhere in the context. We limit the suffix length to `rep_limit` to respect restart sequences.\n    //\n    // This algorithm is not currently documented on Wikipedia, but there is a clear description here:\n    // https://ivanyu.me/blog/2014/10/15/z-algorithm/\n    //\n    // The code below is adapted from the public domain implementation by the same author here:\n    // https://github.com/ivanyu/string-algorithms/blob/master/z_algorithm.py\n    //\n    // Example:\n    // Last N tokens: a b c c b c y a b c\n    // Repeat counts: 0 0 3 1 0 2 0 0 0 0\n    //                    ^\n    //   This `3` means that the last three tokens of the context (a b c) also appear here.\n    //\n    // This step is worst case O(N) since the Z-algorithm is linear, despite the appearance of nested\n    // for/while loops. This can be seen by observing that the `lt` and `rt` bounds are set after each\n    // repeated suffix is detected (i.e. after each while loop when n > 0). These bound variables\n    // ensure that the inner while loops only examine each token in the context once as the outer\n    // for loop iterates over the context.\n\n    {\n        const int last = last_n_repeat - 1;\n        int rt = 0, lt = 0;\n\n        for (int k = 1; k < last_n_repeat; ++k) {\n            if (k > rt) {\n                // If k is outside the current Z-box, do naive computation.\n                int n = 0;\n                while (n + k < last_n_repeat && smpl->last_tokens.rat(n) == smpl->last_tokens.rat(n + k)) {\n                    ++n;\n                }\n                smpl->dry_repeat_count[last - k] = std::min(n, rep_limit);\n                if (n > 0) {\n                    lt = k;\n                    rt = k + n - 1;\n                }\n            }\n            else {\n                // If k is inside the current Z-box, consider two cases.\n\n                int p = k - lt; // Pair index.\n                int right_part_len = rt - k + 1;\n\n                if (smpl->dry_repeat_count[last - p] < right_part_len) {\n                    int n = std::min(smpl->dry_repeat_count[last - p], rep_limit);\n                    smpl->dry_repeat_count[last - k] = n;\n                }\n                else {\n                    int i = rt + 1;\n                    while (i < last_n_repeat && smpl->last_tokens.rat(i) == smpl->last_tokens.rat(i - k)) {\n                        i += 1;\n                    }\n\n                    int n = std::min(i - k, rep_limit);\n                    smpl->dry_repeat_count[last - k] = n;\n                    lt = k;\n                    rt = i - 1;\n                }\n            }\n        }\n    }\n\n    // Step 3: Iterate over dry_repeat_count and last_tokens, examining the maximum repeat length\n    // that would be generated by emitting each new token that would extend a sequence.\n    //\n    // Following the same example as above:\n    // Last N tokens: a b c c b c y a b c\n    // Repeat counts: 0 0 3 1 0 2 0 0 0 0\n    //\n    // For each non-zero, look ahead one token. This token, if emitted, would extend the repetition.\n    // c: 3 -> 4 (from `a b c` to `a b c c`)\n    // b: 1 -> 2 (from `c` to `c b`)\n    // y: 2 -> 3 (from `b c` to `b c y`)\n\n    for (int i = 0; i < last_n_repeat - 1; ++i) {\n        int repeat_len = smpl->dry_repeat_count[i];\n        if (repeat_len >= smpl->dry_allowed_length) {\n            // This token ends a repeat, so the next token would continue one.\n            // By convention, the value of `repeat_len` only includes the tokens currently\n            // in the context, not the new token that would be added.\n            llama_token token = smpl->last_tokens.rat(last_n_repeat - 2 - i);\n            // Track the maximum sequence ending in this token.\n            const auto& it = smpl->dry_max_token_repeat.find(token);\n            if (it == smpl->dry_max_token_repeat.end() || it->second < repeat_len) {\n                smpl->dry_max_token_repeat[token] = repeat_len;\n            }\n        }\n    }\n\n    // Step 4: Apply logit penalties based on the maximum repeat length for relevant tokens.\n\n    // Prevent floating point overflow in `pow(penalty_base, exponent)` by clamping to `max_exponent`.\n    // Compute it from `penalty_base` and the approximate log of `std::numeric_limits<float>::max()`\n    const float FLOAT_MAX_LOG = 88.7228391f;\n    int max_exponent = 0;\n    if (smpl->dry_base > 1.000001f) {\n        max_exponent = FLOAT_MAX_LOG / std::log(smpl->dry_base);\n    }\n\n    for (size_t i = 0; i < cur_p->size; ++i) {\n        const auto& af_kvp = smpl->dry_max_token_repeat.find(cur_p->data[i].id);\n        if (af_kvp != smpl->dry_max_token_repeat.end()) {\n            // Check all sequence breakers starting with this token\n            auto range = smpl->dry_processed_breakers.equal_range(cur_p->data[i].id);\n            bool is_single_token_breaker = false;\n\n            for (auto it = range.first; it != range.second; ++it) {\n                if (it->second.empty()) {\n                    is_single_token_breaker = true;\n                    break;\n                }\n            }\n\n            // Apply penalty only if it's not a single-token sequence breaker\n            if (!is_single_token_breaker) {\n                int repeat_exp = af_kvp->second - smpl->dry_allowed_length;\n                if (max_exponent > 0 && repeat_exp > max_exponent) {\n                    repeat_exp = max_exponent;\n                }\n                float penalty = smpl->dry_multiplier * std::pow(smpl->dry_base, repeat_exp);\n                cur_p->data[i].logit -= penalty;\n            }\n        }\n    }\n\n    cur_p->sorted = false;\n}\n\n\n\nstruct llama_sampler_dry* llama_sampler_init_dry_impl(const struct llama_vocab& vocab, int32_t context_size, 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) {\n    int32_t effective_dry_penalty_last_n = (dry_penalty_last_n == -1) ? context_size : std::max(dry_penalty_last_n, 0);\n    std::unordered_multimap<llama_token, std::vector<llama_token>> processed_breakers;\n    const int MAX_CHAR_LEN = 40;\n    const int MAX_SEQ_LEN = 20;\n\n    const bool dry_enabled = (dry_multiplier != 0.0f && dry_base >= 1.0f && dry_penalty_last_n != 0);\n\n    if (dry_enabled && seq_breakers != nullptr && num_breakers > 0) {\n        // Process sequence breakers\n        for (size_t i = 0; i < num_breakers; ++i) {\n            if (seq_breakers[i] == nullptr || std::strlen(seq_breakers[i]) == 0) {\n                LLAMA_LOG_WARN(\"skipping null or empty DRY sequence breaker at index %zu\\n\", i);\n                continue;\n            }\n\n            std::string sequence_break(seq_breakers[i]);\n            if (sequence_break.empty()) {\n                LLAMA_LOG_WARN(\"skipping empty DRY sequence breaker\\n\");\n                continue;\n            }\n\n            if (sequence_break.size() > MAX_CHAR_LEN) {\n                LLAMA_LOG_WARN(\"truncating DRY sequence breaker to %d characters\\n\", MAX_CHAR_LEN);\n                sequence_break.resize(MAX_CHAR_LEN);\n            }\n\n            get_overlapping_token_sequences(vocab, sequence_break, processed_breakers, MAX_SEQ_LEN);\n        }\n    }\n\n    return  new llama_sampler_dry {\n            /* .total_context_size     = */ context_size,\n            /* .dry_multiplier         = */ dry_multiplier,\n            /* .dry_base               = */ dry_base,\n            /* .dry_allowed_length     = */ dry_allowed_length,\n            /* .dry_penalty_last_n     = */ dry_penalty_last_n,\n            /* .dry_processed_breakers = */ std::move(processed_breakers),\n            /* .dry_repeat_count       = */ dry_enabled ? std::vector<int>(effective_dry_penalty_last_n, 0) : std::vector<int>{},\n            /* .dry_max_token_repeat   = */ {},\n            /* .last_tokens            = */ dry_enabled ? ring_buffer<llama_token>(effective_dry_penalty_last_n) : ring_buffer<llama_token>(0),\n    };\n}\n\n\n// adaptive p\n\nvoid llama_review_adaptive_p_impl(llama_sampler_adaptive_p * adapt_p_ctx, const size_t n_unsent, const bool rewind_status) {\n    // LLAMA_LOG_DEBUG(\"%s: n_unsent = %zu, rewind_status = %s\\n\", __func__, n_unsent, rewind_status ? \"true\" : \"false\");\n    if (adapt_p_ctx->target < 0.0f) {\n        // LLAMA_LOG_DEBUG(\"%s: sampler disabled, target = %f\\n\", __func__, adapt_p_ctx->target);\n        return;\n    }\n\n    auto & history = adapt_p_ctx->history;\n    const size_t hsz = history.size();\n    const size_t hsz_next = 1 + n_unsent;\n    // LLAMA_LOG_DEBUG(\"%s: hsz = %zu, hsz_next = %zu\\n\", __func__, hsz, hsz_next);\n    if (hsz_next >= hsz >> 1) { return; }   // skip small update\n\n    if (!rewind_status) {\n        // sent results, overwrite old history\n        // LLAMA_LOG_DEBUG(\"%s: hsz = %zu, hsz_next = %zu\\n\", __func__, hsz, hsz_next);\n        // LLAMA_LOG_DEBUG(\"%s: history[hsz-1].first = %f\\n\", __func__, history[hsz-1].first);\n        const size_t hsz_diff = hsz - hsz_next;\n        for (int j = 0; j < hsz_next; ++j) {\n            history[j] = history[j + hsz_diff];\n        }\n    }\n    history.resize(hsz_next);\n    // LLAMA_LOG_DEBUG(\"%s: history[hsz_next-1].first = %f\\n\", __func__, history[hsz_next-1].first);\n}\n\nllama_token llama_sample_token_adaptive_p_impl(\n                  struct llama_sampling * smpl,\n                 llama_token_data_array * candidates,\n        struct llama_sampler_adaptive_p * adapt_p_ctx) {\n    GGML_ASSERT(candidates->size > 0);\n    const int64_t t_start_sample_us = ggml_time_us();\n\n    struct llama_sampler_adaptive_p * ctx = adapt_p_ctx;\n    ctx->cum_probs.resize(candidates->size);\n\n    // compute cumulative probability distribution\n    const float max_logit = ctx->max_xform_logit;\n    float cum_prob = 0.0f;\n    for (size_t i = 0; i < candidates->size; ++i) {\n        cum_prob += expf(candidates->data[i].logit - max_logit);\n        ctx->cum_probs[i] = cum_prob;\n    }\n    ctx->cum_probs.back() += 1.0f;  // safety margin in case rng() ~= rng.max()\n\n    // select first token whose cum_prob > target_cum_prob\n    const float target_cum_prob = cum_prob * (float)ctx->rng() / (float)ctx->rng.max();\n    auto iter = std::upper_bound(ctx->cum_probs.begin(), ctx->cum_probs.end(), target_cum_prob);\n    GGML_ASSERT(iter != ctx->cum_probs.end());\n    const size_t idx = std::distance(ctx->cum_probs.begin(), iter);\n    llama_token id = candidates->data[idx].id;\n    GGML_ASSERT(id < int(ctx->orig_prob.size()));\n\n    // update history\n    const float update_prob = ctx->updt_w_cur\n        ? candidates->data[idx].p / ctx->cum_cur_p\n        : ctx->orig_prob[id] / ctx->cum_orig_prob;\n    if (update_prob > 0) {\n        ctx->history.push_back({\n            ctx->decay * ctx->history.back().first + update_prob,   // weighted_sum\n            ctx->decay * ctx->history.back().second + 1.0f });      // total_weight\n    }\n\n    smpl->t_sample_us += ggml_time_us() - t_start_sample_us;\n    smpl->n_sample++;\n\n    return id;\n}\n\nvoid llama_sample_adaptive_p_impl(struct llama_sampling * ctx, llama_token_data_array * candidates,\n        struct llama_sampler_adaptive_p * adapt_p_ctx) {\n    if (adapt_p_ctx->target < 0.0f) {\n        // LLAMA_LOG_DEBUG(\"%s: sampler disabled, target = %f\\n\", __func__, adapt_p_ctx->target);\n        llama_sample_softmax_impl(nullptr, candidates);\n        return;\n    }\n\n    auto t_start = ggml_time_us();\n\n    // incomplete softmax because final division can be fused\n    float max_l = candidates->data[0].logit;\n    if (!candidates->sorted) {\n        for (size_t i = 1; i < candidates->size; ++i) {\n            max_l = std::max(max_l, candidates->data[i].logit);\n        }\n    }\n    float cum_sum = 0.0f;\n    for (size_t i = 0; i < candidates->size; ++i) {\n        const float prob = expf(candidates->data[i].logit - max_l);\n        candidates->data[i].p = prob;\n        cum_sum += prob;\n    }\n    adapt_p_ctx->cum_cur_p = cum_sum;\n\n    // compute adapted target probability\n    const float weighted_sum = adapt_p_ctx->history.back().first;\n    const float total_weight = adapt_p_ctx->history.back().second;\n    const float target = std::clamp(adapt_p_ctx->target, 0.0f, 1.0f);\n    const float adapted_target = std::clamp(total_weight == 0.0f\n        ? target\n        : 2.0f * target - (weighted_sum / total_weight),\n        0.0f, 1.0f);\n\n    // transformation constants\n    static constexpr float peak_logit_value = 5.0f;\n    static constexpr float inv_width = 1.0f / 0.3f;\n    static constexpr float sharpness = 10.0f;\n\n    const float fused_target = adapted_target * inv_width;\n    const float fused_width = inv_width / cum_sum;\n\n    // quadratic near target for finite differentiation, transitioning to linear decay in tails\n    // unbounded negative logits suppress far-from-target tokens after softmax\n    float max_logit = -INFINITY;\n    for (size_t i = 0; i < candidates->size; ++i) {\n        const float dist = std::abs(candidates->data[i].p * fused_width - fused_target);\n        const float logit = peak_logit_value - sharpness * dist * dist / (1.0f + dist);\n        candidates->data[i].logit = logit;\n        max_logit = std::max(max_logit, logit);\n    }\n    candidates->sorted = false;\n    adapt_p_ctx->max_xform_logit = max_logit;\n\n    ctx->t_sample_us += ggml_time_us() - t_start;\n}\n\nvoid llama_prep_adaptive_p_impl(\n              struct llama_sampling * smpl,\n             llama_token_data_array * candidates,\n    struct llama_sampler_adaptive_p * adapt_p_ctx) {\n    if (adapt_p_ctx->updt_w_cur     // update with current probability, original not needed\n        || (adapt_p_ctx->target < 0.0f)) {  // or disabled\n        return;\n    }\n    constexpr float kDelta = 30.0f; //16.6f;\n    auto t_start = ggml_time_us();\n    auto & orig_prob = adapt_p_ctx->orig_prob;\n    if (candidates->size != orig_prob.size() || candidates->sorted) {\n        LLAMA_LOG_ERROR(\"%s: this function must be called before any other sampler has been applied\\n\", __func__);\n        LLAMA_LOG_ERROR(\"%s: the sampler has been initialized with a vocabulary of %zu, but is being called with %zu candidates\\n\",\n                __func__, orig_prob.size(), candidates->size);\n        GGML_ABORT(\"Bad candidates in adaptive_p sampler\");\n    }\n\n    float max_logit = -INFINITY;\n    for (int j = 0; j < int(candidates->size); ++j) {\n        orig_prob[j] = candidates->data[j].logit;\n        max_logit = std::max(max_logit, orig_prob[j]);\n    }\n    adapt_p_ctx->cum_orig_prob = iqk_exp_with_thresh(orig_prob.size(), orig_prob.data(), max_logit, max_logit - kDelta);\n\n    if (smpl) smpl->t_sample_us += ggml_time_us() - t_start;\n}\n\nstruct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(int n_vocab,\n       const float target,\n       const float decay,\n        const bool updt_w_cur,\n    const uint32_t seed) {\n    GGML_ASSERT(n_vocab > 0);\n    const float clamped_decay = std::clamp(decay, 0.0f, 0.99f);\n    auto result = new llama_sampler_adaptive_p {\n        /* .target            = */ target,\n        /* .decay             = */ clamped_decay,\n        /* .updt_w_cur        = */ updt_w_cur,\n        /* .rng               = */ std::mt19937(seed),\n        /* .history           = */ {},\n        /* .orig_prob         = */ {},\n        /* .cum_orig_prob     = */ 1.0f,\n        /* .cum_cur_p         = */ 1.0f,\n        /* .max_xform_logit   = */ -INFINITY,\n        /* .cum_probs         = */ {},\n    };\n    result->history.push_back({\n        target / (1.0f - clamped_decay),    // weighted_sum\n        1.0f / (1.0f - clamped_decay) });   // total_weight\n    result->orig_prob.resize(n_vocab);\n    return result;\n}\n\n// grammar\n\nstruct llama_sampler_grammar {\n    const struct llama_vocab* vocab;\n\n    std::string grammar_str;\n    std::string grammar_root;\n\n    struct llama_grammar* grammar;\n};\n\nstatic const char* llama_sampler_grammar_name(const struct llama_sampler* /*smpl*/) {\n    return \"grammar\";\n}\n\nstatic void llama_sampler_grammar_accept_impl(struct llama_sampler* smpl, llama_token token) {\n    auto* ctx = (llama_sampler_grammar*)smpl->ctx;\n    if (ctx->grammar) {\n        llama_grammar_accept_impl(*ctx->grammar,ctx->vocab ,nullptr, token);\n    }\n}\n\nstatic void llama_sampler_grammar_apply(struct llama_sampler* smpl, llama_token_data_array* cur_p) {\n    auto* ctx = (llama_sampler_grammar*)smpl->ctx;\n    if (ctx->grammar) {\n        llama_grammar_sample_impl(ctx->grammar, ctx->vocab, nullptr, cur_p);\n    }\n}\n\nvoid llama_sampler_reset(struct llama_sampler* smpl) {\n    if (smpl->iface->reset) {\n        smpl->iface->reset(smpl);\n    }\n}\n\n// Fwd declare to break reset --> init_impl --> llama_sampler_grammar_i --> reset cycle.\nstatic struct llama_grammar* llama_sampler_init_grammar_impl(\n    const struct llama_vocab* vocab,\n    const char* grammar_str,\n    const char* grammar_root,\n    bool lazy,\n    const char** trigger_words,\n    size_t num_trigger_words,\n    const llama_token* trigger_tokens,\n    size_t num_trigger_tokens,\n    const char** trigger_patterns,\n    size_t num_trigger_patterns);\n\nstatic void llama_sampler_grammar_reset(struct llama_sampler* smpl) {\n    auto* ctx = (llama_sampler_grammar*)smpl->ctx;\n    if (!ctx->grammar) {\n        return;\n    }\n\n    std::vector<const char*>  trigger_patterns_c;\n    trigger_patterns_c.reserve(ctx->grammar->trigger_patterns.size());\n    for (auto& trigger_pattern : ctx->grammar->trigger_patterns) {\n        trigger_patterns_c.push_back(trigger_pattern.pattern.c_str());\n    }\n    auto* grammar_new = llama_grammar_init_impl(ctx->grammar->vocab, ctx->grammar_str.c_str(), ctx->grammar_root.c_str(),\n        ctx->grammar->lazy, trigger_patterns_c.data(), trigger_patterns_c.size(),\n        ctx->grammar->trigger_tokens.data(), ctx->grammar->trigger_tokens.size());\n\n    llama_grammar_free_impl(ctx->grammar);\n    ctx->grammar = grammar_new;\n}\n\n//static struct llama_sampler* llama_sampler_grammar_clone(const struct llama_sampler* smpl) {\n//    const auto* ctx = (const llama_sampler_grammar*)smpl->ctx;\n//\n//    auto* result = llama_sampler_init_grammar_impl(ctx->vocab, nullptr, nullptr, false, nullptr, 0, nullptr, 0);\n//\n//    // copy the state\n//    {\n//        auto* result_ctx = (llama_sampler_grammar*)result->ctx;\n//\n//        if (ctx->grammar) {\n//            result_ctx->grammar_str = ctx->grammar_str;\n//            result_ctx->grammar_root = ctx->grammar_root;\n//\n//            result_ctx->grammar = llama_grammar_copy_impl(ctx->grammar);\n//        }\n//    }\n//\n//    return result;\n//}\n\nstatic void llama_sampler_grammar_free(struct llama_sampler* smpl) {\n    const auto* ctx = (llama_sampler_grammar*)smpl->ctx;\n\n    if (ctx->grammar) {\n        llama_grammar_free_impl(ctx->grammar);\n    }\n\n    delete ctx;\n}\n\n// ?\n//static struct llama_sampler_i llama_sampler_grammar_i = {\n//    /* .name   = */ llama_sampler_grammar_name,\n//    /* .accept = */ llama_sampler_grammar_accept_impl,\n//    /* .apply  = */ llama_sampler_grammar_apply,\n//    /* .reset  = */ llama_sampler_grammar_reset,\n//    /* .clone  = */ NULL,\n//    /* .free   = */ llama_sampler_grammar_free,\n//};\n\nstruct llama_grammar* llama_sampler_init_grammar_impl(\n    const struct llama_vocab* vocab,\n    const char* grammar_str,\n    const char* grammar_root,\n    bool lazy,\n    const char** trigger_words,\n    size_t num_trigger_words,\n    const llama_token* trigger_tokens,\n    size_t num_trigger_tokens,\n    const char** trigger_patterns,\n    size_t num_trigger_patterns) {\n    // Huh? this is not used and leaks. auto* ctx = new llama_sampler_grammar;\n    struct llama_grammar* grammar;\n    if (grammar_str != nullptr && grammar_str[0] != '\\0') {\n        // TODO: remove trigger_words support.\n        if (trigger_words != nullptr && num_trigger_words > 0) {\n            GGML_ASSERT(trigger_patterns == nullptr && num_trigger_patterns == 0);\n            std::string trigger_pattern(\"[\\\\s\\\\S]*?(\");\n            for (size_t i = 0; i < num_trigger_words; ++i) {\n                static const std::regex special_chars(\"[.^$|()*+?\\\\[\\\\]{}\\\\\\\\]\");\n                if (i > 0) {\n                    trigger_pattern += \"|\";\n                }\n                trigger_pattern += std::regex_replace(trigger_words[i], special_chars, \"\\\\$0\");\n            }\n            trigger_pattern += \")[\\\\s\\\\S]*\";\n            auto trigger_pattern_c = trigger_pattern.c_str();\n            trigger_patterns = &trigger_pattern_c;\n            num_trigger_patterns = 1;\n        }\n        grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, trigger_patterns, num_trigger_patterns, trigger_tokens, num_trigger_tokens);\n        if (!grammar) {\n            return nullptr;\n        }\n    } else {\n        grammar = nullptr;\n    }\n    return grammar;\n}\n\nstruct llama_grammar* llama_sampler_init_grammar(\n    const struct llama_vocab* vocab,\n    const char* grammar_str,\n    const char* grammar_root) {\n    return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ false, nullptr, 0, nullptr, 0, nullptr, 0);\n}\n\nstruct llama_grammar* llama_sampler_init_grammar_lazy(\n    const struct llama_vocab* vocab,\n    const char* grammar_str,\n    const char* grammar_root,\n    const char** trigger_words,\n    size_t num_trigger_words,\n    const llama_token* trigger_tokens,\n    size_t num_trigger_tokens) {\n    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);\n}\n\nstruct llama_grammar* llama_sampler_init_grammar_lazy_patterns(\n    const struct llama_vocab* vocab,\n    const char* grammar_str,\n    const char* grammar_root,\n    const char** trigger_patterns,\n    size_t num_trigger_patterns,\n    const llama_token* trigger_tokens,\n    size_t num_trigger_tokens) {\n    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);\n}"
    }
  ],
  "seed": 666,
  "temperature": 1,
  "max_tokens": 4096,
  "cache_prompt": true,
  "chat_template_kwargs": {
    "enable_thinking": false
  },
  "stream": false
}

I pick src/llama-sampling.cpp, and ask the model to explain what it does.

Answer before this PR (first 2 paragraphs):

This C++ code is an implementation file for the sampling engine used in llama.cpp (the popular C/C++ inference engine for Large Language Models like LLaMA, Mistral, etc.).

When a language model generates text, it outputs a list of "logits" (raw scores) for every possible token (word/character) in its vocabulary. The code you provided is responsible for taking those raw scores and applying various mathematical strategies (samplers) to decide which token the AI will actually generate next.

Answer with this PR and -ictk q8_0 (first 2 paragraphs):

The provided text appears to be heavily corrupted or obfuscated C++ code, likely from a project named llama_grammar. The code contains numerous syntax errors, nonsensical variables, and incomplete functions, making it impossible to interpret the original intent.

However, the code contains several recognizable C++ keywords, standard library usage, and recognizable function names, such as llama_sort, llama_sample_softmax_impl, llama_sample_top_k_impl, ll implements (probably intended as "implements"), and a mix of standard C++ and custom llama-style classes (like llama_token_data_array).

Answer with this PR without setting -ictk (last 2 paragraphs):

In short, this is the "brain" of the text generation engine. It translates raw mathematical scores from the neural network into coherent text by allowing users to mix and match filters (like Top-P, Top-K, Mirostat, and Grammar rules) to precisely control how the AI writes.

(Note: The provided code contains some syntax errors—like incomplete if statements, missing brackets, and stray text inside the DRY and Adaptive P functions—suggesting it might be a work-in-progress draft, a partially reverse-engineered file, or a copy-paste with missing pieces.)

[saood06]

@saood06 Thanks for trying to help. Does the model you linked to share indexer results between layers? Because if it does not, then no, testing with it wouldn't have caught the bug.

Well like I said before it wouldn't have helped you catch this bug, as it doesn't share indexer results between layers but now that the bug has been caught it might help you iterate faster if you are still optimizing the non index sharing part of DSA. Just like how DeepSeek-V2-Lite was a smaller model that you used to experiment with MLA and the deepseek architecture.

[ikawrakow]

@sayap

Here is what I get on the latest main branch

{"choices":[{"finish_reason":"stop","index":0,"message":{"role":"assistant","content":"This C++ source code is the implementation of the **token sampling logic** for `llama.cpp`, the popular C/C++ inference engine for Large Language Models (LLMs) like LLaMA, Mistral, etc. \n\nWhen an LLM generates text, it doesn't predict the *next word* directly; it outputs a list of \"logits\" (raw scores) for every possible token in its vocabulary. The code in this file is responsible for taking those raw scores and applying various mathematical and algorithmic filters (samplers) to decide exactly which token is chosen next.\n\nHere is a breakdown of what the different parts of the code do:\n\n### 1. Core Utilities & Softmax\n*   **`llama_log_softmax` / `llama_sample_softmax_impl`**: Converts raw logits into a normalized probability distribution (where all probabilities add up to 1.0) using the softmax function.\n*   `llama_sort`: An optimized sorting function. It sorts tokens by their logit scores in descending order. Notably, it uses a bucket-sort optimization for large arrays (over 1024 tokens) to improve performance before falling back to standard sorting or partial sorting.\n\n### 2. Standard Sampling Filters\nThese functions truncate or modify the list of candidate tokens before the final choice is made:\n*   **`llama_sample_top_k_impl`**: Keeps only the top `K` highest-scoring tokens and discards the rest.\n*   **`lla_sample_top_p_impl`**: (Nucleus Sampling) Keeps the smallest set of tokens whose cumulative probability exceeds `p` (e.g., 0.9).\n*   **`llama_sample_min_p_impl`**: Discards any token whose probability is less than `p` times the probability of the most likely token.\n*   **`llama_sample_temp_impl`**: Divides all logits by a temperature value. Temp > 1 makes the model more random; Temp < 1 makes it more deterministic.\n*   `llama_sample_tail_free_impl`: An advanced sampler that looks at the second derivative of the probability distribution to cut off the \"tail\" of unlikely tokens.\n*   `llama_sample_typical_impl`: Typical sampling. It favors tokens that have a probability close to the expected entropy of the distribution, preventing the model from choosing \"weird\" tokens with high or extremely low probabilities.\n\n### 3. Advanced & Dynamic Samplers\n*   **`llama_sample_entropy_impl`**: Dynamic Temperature. It calculates the \"entropy\" (randomness) of the current distribution. If the model is very uncertain, it lowers the temperature. If it is certain, it raises the temperature.\n*   **`llama_sample_xtc_impl`**: XTC (Exclude Top Choices). With some probability, it intentionally discards the most likely tokens to force the model to be more creative and avoid repetitive loops.\n*   **`llama_sample_top_n_sigma_impl`**: Keeps only tokens whose logits are within $N$ standard deviations from the maximum logit. \n*   **`llama_sample_repetition_penalties_impl`**: Prevents the model from repeating itself. It looks at the last $N$ generated tokens and applies penalties to them based on how often they appeared (frequency penalty) and if they appeared at all (presence penalty).\n\n### 4. Guidance & Mirostat\n*   **`llama_sample_apply_guidance_impl`**: Used when running two models (a base model and a guidance model). It alters the logits of the main model to steer it toward or away from the guidance model's predictions.\n*   **`llama_sample_token_mirostat_impl` / `_v2_impl`**: Mirostat is an algorithm that actively tries to maintain a specific level of \"surprise\" (perplexity) during generation. It adjusts a variable ($\\mu$) on the fly to keep the text from becoming too repetitive or too random.\n\n### 5. Final Token Selection\n*   `llama_sample_token_greedy_impl`: Simply picks the single token with the highest logit (greedy decoding).\n*   **`llama_sample_token_with_rng_impl`**: Performs \"Multinomial Sampling.\" It generates a random number and uses it to pick a token based on a cumulative probability distribution (like spinning a roulette wheel weighted by the probabilities). It includes error handling that writes debug data to `probabilities.txt` if a sampling failure occurs.\n\n### 6. DRY (Don't Repeat Yourself) Sampler\n*   **`llama_sampler_dry_apply`**: A sophisticated anti-repetition sampler. Instead of just penalizing the last $N$ tokens, it uses the **Z-algorithm** (a string-matching algorithm) to find long repeating suffixes in the recent context. If it detects that a token is about to continue a long repetition, it applies an exponentially scaling penalty to crush that token's probability. It also supports \"sequence breakers\" (specific words or tokens that reset the repetition tracking).\n\n### 7. Adaptive P Sampler\n*   **`llama_sample_adaptive_p_impl`**: A custom sampler that maintains a history of past probability distributions. It dynamically adjusts a target probability threshold based on past context to stabilize the randomness of generation over time.\n\n### 8. Grammar Sampling\n*   **`llama_sampler_grammar_apply`**: Restricts the model's output to a specific formal grammar (e.g., forcing the output to be valid JSON, or a specific coding language). It removes tokens that would violate the grammar rules.\n*   **Lazy Grammar Initialization**: The functions at the bottom (`llama_sampler_init_grammar_lazy`, etc.) allow the grammar restrictions to remain dormant until a specific \"trigger\" word or pattern is generated by the model. Once triggered, the grammar rules take over. This is useful for instruct models where you want free text, but a specific JSON structure *inside* a specific block."}}],"created":1783514131,"model":"GLM-5.2-UD-Q4_K_M-00001-of-00011.gguf","object":"chat.completion","usage":{"completion_tokens":1222,"prompt_tokens":13929,"total_tokens":15151,"prompt_tokens_details":{"cached_tokens":0}},"id":"chatcmpl-5RF9I34xlO5zjvkTyNq691hIEDv1PyLG","timings":{"prompt_n":13929,"prompt_ms":151377.88,"prompt_per_token_ms":10.867821092684329,"prompt_per_second":92.0147646406463,"predicted_n":1222,"predicted_ms":117509.055,"predicted_per_token_ms":96.16125613747954,"predicted_per_second":10.399198597929326,"n_ctx":20480,"n_past":15150}}

Looks perfectly fine to me. The model is Q4_K_M from Unsloth.

[ikawrakow]

OK, I get the same garbage using -ictk q8_0. So, I guess, quantizing the indexer cache is not viable?

[sayap]

Did some more testing with the main branch using the 2.244bpw quant...

• With -ictk q8_0, bad output every time.
• Without setting -ictk, the model is still coherent until about 40k context. Then, it started to mix in random tokens here and there, and the model was aware of that:
  • "I made a mess with the edit due to the weird tokens. Let me look at the actual current state of the file around that region and fix it properly with a write or a clean edit. Let me read the region fully."
  • "The previous edit got corrupted with stray text. Let me view the damaged region and repair it cleanly."
  • "My bash invocations got garbled somehow. Let me just run it cleanly."

So -ictk q8_0 is definitely no good for this quant, and I will use it without -dsa for anything that requires >40k context.

EDIT: Oohh I just saw #2099, seems promising!

EDIT 2: #2099 fixes the issue with corrupted tokens at long context! -ictk q8_0 still gives bad output though