Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 09:15:04 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into xsn/server_docker_isolate

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-03 18:52:40 +02:00
parent ea1abe6f43 e8e06f78e2
commit 7e565b10f7
29 changed files with 1810 additions and 437 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/arg.cpp
+18 -3
View File
@@ -27,6 +27,7 @@
#include <algorithm>
#include <cinttypes>
#include <climits>
#include <cmath>
#include <cstdarg>
#include <filesystem>
#include <fstream>
@@ -2036,7 +2037,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
{"--repeat-penalty"}, "N",
string_format("penalize repeat sequence of tokens (default: %.2f, 1.0 = disabled)", (double)params.sampling.penalty_repeat),
[](common_params & params, const std::string & value) {
params.sampling.penalty_repeat = std::stof(value);
const float penalty_repeat = std::stof(value);
if (!std::isfinite(penalty_repeat) ||
penalty_repeat <= 0.0f ||
!std::isfinite(1.0f/penalty_repeat)) {
throw std::runtime_error("error: repeat-penalty must be finite and greater than 0\n");
}
params.sampling.penalty_repeat = penalty_repeat;
params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT;
}
).set_sampling());
@@ -2044,14 +2051,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
{"--presence-penalty"}, "N",
string_format("repeat alpha presence penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_present),
[](common_params & params, const std::string & value) {
params.sampling.penalty_present = std::stof(value);
const float penalty_present = std::stof(value);
if (!std::isfinite(penalty_present)) {
throw std::runtime_error("error: presence-penalty must be finite\n");
}
params.sampling.penalty_present = penalty_present;
}
).set_sampling());
add_opt(common_arg(
{"--frequency-penalty"}, "N",
string_format("repeat alpha frequency penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_freq),
[](common_params & params, const std::string & value) {
params.sampling.penalty_freq = std::stof(value);
const float penalty_freq = std::stof(value);
if (!std::isfinite(penalty_freq)) {
throw std::runtime_error("error: frequency-penalty must be finite\n");
}
params.sampling.penalty_freq = penalty_freq;
}
).set_sampling());
add_opt(common_arg(
common/common.cpp
+2 -1
View File
@@ -1299,8 +1299,9 @@ common_init_result::common_init_result(common_params & params, bool model_only)
pimpl->samplers.resize(cparams.n_seq_max);
pimpl->samplers_seq_config.resize(cparams.n_seq_max);
const int32_t n_ctx = cparams.n_ctx > 0 ? (int32_t) cparams.n_ctx : llama_model_n_ctx_train(model);
for (int i = 0; i < (int) cparams.n_seq_max; ++i) {
pimpl->samplers[i].reset(common_sampler_init(model, params.sampling));
pimpl->samplers[i].reset(common_sampler_init(model, params.sampling, n_ctx));
pimpl->samplers_seq_config[i] = { i, common_sampler_get(pimpl->samplers[i].get()) };
}
common/sampling.cpp
+19 -2
View File
@@ -184,9 +184,26 @@ std::string common_params_sampling::print() const {
return std::string(result);
}
struct common_sampler * common_sampler_init(const struct llama_model * model, struct common_params_sampling & params) {
const llama_vocab * vocab = llama_model_get_vocab(model);
struct common_sampler * common_sampler_init(
const struct llama_model * model,
struct common_params_sampling & params,
int32_t n_ctx) {
if (!std::isfinite(params.penalty_repeat) ||
params.penalty_repeat <= 0.0f ||
!std::isfinite(1.0f/params.penalty_repeat)) {
throw std::invalid_argument("penalty_repeat must be finite and greater than 0");
}
if (!std::isfinite(params.penalty_freq)) {
throw std::invalid_argument("penalty_freq must be finite");
}
if (!std::isfinite(params.penalty_present)) {
throw std::invalid_argument("penalty_present must be finite");
}
if (params.penalty_last_n == -1) {
params.penalty_last_n = n_ctx > 0 ? n_ctx : llama_model_n_ctx_train(model);
}
const llama_vocab * vocab = llama_model_get_vocab(model);
llama_sampler_chain_params lparams = llama_sampler_chain_default_params();
lparams.no_perf = params.no_perf;
common/sampling.h
+4 -1
View File
@@ -37,7 +37,10 @@ struct common_sampler;
// llama_sampler API overloads
// note: can mutate params in some cases
struct common_sampler * common_sampler_init(const struct llama_model * model, struct common_params_sampling & params);
struct common_sampler * common_sampler_init(
const struct llama_model * model,
struct common_params_sampling & params,
int32_t n_ctx = 0);
void common_sampler_free(struct common_sampler * gsmpl);
ggml/src/ggml-backend.cpp
+57 -6
View File
@@ -765,8 +765,9 @@ struct ggml_backend_sched_split {
int backend_id;
int i_start;
int i_end;
struct ggml_tensor * inputs[GGML_SCHED_MAX_SPLIT_INPUTS];
struct ggml_tensor ** inputs;
int n_inputs;
int inputs_capacity;
// graph view of this split
struct ggml_cgraph graph;
};
@@ -805,8 +806,9 @@ struct ggml_backend_sched {
int cur_copy;
int next_copy;
ggml_backend_event_t events[GGML_SCHED_MAX_BACKENDS][GGML_SCHED_MAX_COPIES];
struct ggml_tensor * graph_inputs[GGML_SCHED_MAX_SPLIT_INPUTS];
struct ggml_tensor ** graph_inputs;
int n_graph_inputs;
int graph_inputs_capacity;
struct ggml_context * ctx;
@@ -832,6 +834,36 @@ struct ggml_backend_sched {
#define tensor_id_copy(id, backend_id, copy_id) sched->hv_tensor_copies[(id) * sched->n_backends * sched->n_copies + (backend_id) * sched->n_copies + (copy_id)]
#define tensor_copy(tensor, backend_id, copy_id) tensor_id_copy(hash_id(tensor), backend_id, copy_id)
static void ggml_backend_sched_split_inputs_grow(struct ggml_backend_sched_split * split) {
int new_cap = GGML_SCHED_MAX_SPLIT_INPUTS;
if (split->inputs_capacity > 0) {
new_cap = 2*split->inputs_capacity;
GGML_LOG_WARN("%s: increasing split inputs capacity from %d to %d\n", __func__, split->inputs_capacity, new_cap);
}
auto * pnew = (struct ggml_tensor **) realloc((void *) split->inputs, new_cap * sizeof(struct ggml_tensor *));
if (pnew == NULL) {
GGML_LOG_ERROR("%s: failed to allocate %zu bytes\n", __func__, new_cap * sizeof(struct ggml_tensor *));
GGML_ABORT("failed to grow split inputs container");
}
split->inputs = pnew;
split->inputs_capacity = new_cap;
}
static void ggml_backend_sched_graph_inputs_grow(ggml_backend_sched_t sched) {
int new_cap = GGML_SCHED_MAX_SPLIT_INPUTS;
if (sched->graph_inputs_capacity > 0) {
new_cap = 2*sched->graph_inputs_capacity;
GGML_LOG_WARN("%s: increasing graph inputs capacity from %d to %d\n", __func__, sched->graph_inputs_capacity, new_cap);
}
auto * pnew = (struct ggml_tensor **) realloc((void *) sched->graph_inputs, new_cap * sizeof(struct ggml_tensor *));
if (pnew == NULL) {
GGML_LOG_ERROR("%s: failed to allocate %zu bytes\n", __func__, new_cap * sizeof(struct ggml_tensor *));
GGML_ABORT("failed to grow graph inputs container");
}
sched->graph_inputs = pnew;
sched->graph_inputs_capacity = new_cap;
}
// returns the priority of the backend, lower id is higher priority
static int ggml_backend_sched_backend_id(ggml_backend_sched_t sched, ggml_backend_t backend) {
for (int i = 0; i < sched->n_backends; i++) {
@@ -1297,7 +1329,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
}
// check if the split has too many inputs
// FIXME: count the number of inputs instead of only checking when full
if (split->n_inputs == GGML_SCHED_MAX_SPLIT_INPUTS) {
if (split->n_inputs >= split->inputs_capacity) {
const size_t id = hash_id(src);
int src_backend_id = sched->hv_tensor_backend_ids[id];
bool supported = ggml_backend_sched_buffer_supported(sched, src, cur_backend_id);
@@ -1313,10 +1345,14 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
split->i_end = i;
i_split++;
if (i_split >= sched->splits_capacity) {
int old_cap = sched->splits_capacity;
sched->splits_capacity *= 2;
sched->splits = (ggml_backend_sched_split *)
realloc(sched->splits, sched->splits_capacity * sizeof(struct ggml_backend_sched_split));
GGML_ASSERT(sched->splits != NULL);
for (int k = old_cap; k < sched->splits_capacity; k++) {
memset(&sched->splits[k], 0, sizeof(struct ggml_backend_sched_split));
}
}
split = &sched->splits[i_split];
split->backend_id = node_backend_id;
@@ -1353,7 +1389,9 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
SET_CAUSE(tensor_copy, "4.cpy");
}
int n_graph_inputs = sched->n_graph_inputs++;
GGML_ASSERT(n_graph_inputs < GGML_SCHED_MAX_SPLIT_INPUTS);
if (n_graph_inputs >= sched->graph_inputs_capacity) {
ggml_backend_sched_graph_inputs_grow(sched);
}
sched->graph_inputs[n_graph_inputs] = src;
}
}
@@ -1373,7 +1411,9 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
SET_CAUSE(tensor_copy, "4.cpy");
}
int n_inputs = split->n_inputs++;
GGML_ASSERT(n_inputs < GGML_SCHED_MAX_SPLIT_INPUTS);
if (n_inputs >= split->inputs_capacity) {
ggml_backend_sched_split_inputs_grow(split);
}
split->inputs[n_inputs] = src;
}
node->src[j] = tensor_id_copy(src_id, cur_backend_id, sched->cur_copy);
@@ -1399,7 +1439,11 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
sched->prev_leaf_backend_ids = tmp;
}
int graph_size = std::max(graph->n_nodes, graph->n_leafs) + sched->n_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2*sched->n_copies;
int total_inputs = sched->n_graph_inputs;
for (int i = 0; i < sched->n_splits; i++) {
total_inputs += sched->splits[i].n_inputs;
}
int graph_size = std::max(graph->n_nodes, graph->n_leafs) + total_inputs * 2 * sched->n_copies;
// remember the actual graph_size for performing reallocation checks later [GGML_SCHED_DEBUG_REALLOC]
sched->debug_prev_graph_size = sched->debug_graph_size;
@@ -1782,6 +1826,9 @@ ggml_backend_sched_t ggml_backend_sched_new(
sched->splits = (ggml_backend_sched_split *) calloc(initial_splits_capacity, sizeof(sched->splits[0]));
sched->splits_capacity = initial_splits_capacity;
sched->graph_inputs_capacity = GGML_SCHED_MAX_SPLIT_INPUTS;
sched->graph_inputs = (struct ggml_tensor **) calloc(sched->graph_inputs_capacity, sizeof(struct ggml_tensor *));
for (int b = 0; b < n_backends; b++) {
sched->backends[b] = backends[b];
sched->bufts[b] = bufts ? bufts[b] : ggml_backend_get_default_buffer_type(backends[b]);
@@ -1814,7 +1861,11 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
ggml_gallocr_free(sched->galloc);
ggml_free(sched->ctx);
ggml_hash_set_free(&sched->hash_set);
for (int i = 0; i < sched->splits_capacity; i++) {
free(sched->splits[i].inputs);
}
free(sched->splits);
free(sched->graph_inputs);
free(sched->hv_tensor_backend_ids);
free(sched->hv_tensor_copies);
free(sched->node_backend_ids);
ggml/src/ggml-cuda/common.cuh
+2 -1
View File
@@ -627,7 +627,8 @@ template <typename T> struct block_reduce_policy<block_reduce_method::MAX, T> {
};
template <block_reduce_method reduce_method_t, const unsigned int block_size_template = 0, typename T>
static __device__ T block_reduce(T val, T * shared_vals) {
static __device__ T block_reduce(T val, [[maybe_unused]] T * shared_vals) {
// for multi-warp reductions, callers must not reuse shared_vals until all reads from this invocation have completed
val = block_reduce_policy<reduce_method_t, T>::reduce(val);
const unsigned int block_size = block_size_template == 0 ? blockDim.x : block_size_template;
if (block_size > WARP_SIZE) {
ggml/src/ggml-cuda/norm.cu
+2 -2
View File
@@ -64,7 +64,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr
tmp += xi * xi;
}
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum + 32);
const float variance = tmp / group_size;
const float scale = rsqrtf(variance + eps);
@@ -297,7 +297,7 @@ static void group_norm_f32_cuda(
group_norm_f32<WARP_SIZE><<<num_groups, block_dims, 0, stream>>>(x, dst, group_size, ne_elements, eps);
} else {
const dim3 block_dims(1024, 1, 1);
group_norm_f32<1024><<<num_groups, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream>>>(x, dst, group_size, ne_elements, eps);
group_norm_f32<1024><<<num_groups, block_dims, block_dims.x > WARP_SIZE ? 2 * 32 * sizeof(float): 0, stream>>>(x, dst, group_size, ne_elements, eps);
}
}
ggml/src/ggml-cuda/softmax.cu
+14 -6
View File
@@ -116,6 +116,11 @@ static __global__ void soft_max_f32(
vals[col] = val;
}
if (block_size > WARP_SIZE) {
// sync is needed as we reuse buf_iw across block_reduce invocations, see #26385
// for block_size <= WARP_SIZE, block_reduce does not access buf_iw
__syncthreads();
}
// find the sum of exps in the block
tmp = block_reduce<block_reduce_method::SUM, block_size_template>(tmp, buf_iw);
@@ -142,6 +147,8 @@ static __device__ void soft_max_f32_parallelize_cols_single_row(const float * __
float * __restrict__ dst,
float * __restrict__ tmp_maxs,
float * __restrict__ tmp_sums,
float * shared_vals_max,
float * shared_vals_sum,
const soft_max_params p) {
namespace cg = cooperative_groups;
@@ -154,7 +161,6 @@ static __device__ void soft_max_f32_parallelize_cols_single_row(const float * __
float local_vals[n_elem_per_thread] = { -INFINITY, -INFINITY, -INFINITY, -INFINITY };
float local_max = -INFINITY;
const int step_size = gridDim.x * blockDim.x;
__shared__ float shared_vals[32];
// Compute thread-local max
for (int col = col_start; col < p.ncols;) {
@@ -171,7 +177,7 @@ static __device__ void soft_max_f32_parallelize_cols_single_row(const float * __
}
// Compute CTA-level max
local_max = block_reduce<block_reduce_method::MAX>(local_max, shared_vals);
local_max = block_reduce<block_reduce_method::MAX>(local_max, shared_vals_max);
// Store CTA-level max to GMEM
if (tid == 0) {
@@ -186,7 +192,7 @@ static __device__ void soft_max_f32_parallelize_cols_single_row(const float * __
} else {
local_max = -INFINITY;
}
local_max = block_reduce<block_reduce_method::MAX>(local_max, shared_vals);
local_max = block_reduce<block_reduce_method::MAX>(local_max, shared_vals_max);
// Compute softmax dividends, accumulate divisor
float tmp_expf = 0.0f;
@@ -209,7 +215,7 @@ static __device__ void soft_max_f32_parallelize_cols_single_row(const float * __
}
// Reduce divisor within CTA
tmp_expf = block_reduce<block_reduce_method::SUM>(tmp_expf, shared_vals);
tmp_expf = block_reduce<block_reduce_method::SUM>(tmp_expf, shared_vals_sum);
// Store CTA-level sum to GMEM
if (tid == 0) {
@@ -223,7 +229,7 @@ static __device__ void soft_max_f32_parallelize_cols_single_row(const float * __
} else {
tmp_expf = 0.0f;
}
tmp_expf = block_reduce<block_reduce_method::SUM>(tmp_expf, shared_vals);
tmp_expf = block_reduce<block_reduce_method::SUM>(tmp_expf, shared_vals_sum);
// Divide dividend by global sum + store data
for (int col = col_start; col < p.ncols;) {
@@ -310,9 +316,11 @@ __launch_bounds__(8*WARP_SIZE, 1) static __global__ void soft_max_f32_paralleliz
// https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/device-callable-apis.html#grid-synchronization
// https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/device-callable-apis.html#class-cluster-group
{
__shared__ float shared_vals[2][32];
for (int rowx = 0; rowx < p.ne01 * p.ne02 * p.ne03; rowx++) {
soft_max_f32_parallelize_cols_single_row(x + int64_t(rowx) * p.ncols, dst + int64_t(rowx) * p.ncols, tmp_maxs,
tmp_sums, p);
tmp_sums, shared_vals[0], shared_vals[1], p);
}
}
ggml/src/ggml-opencl/ggml-opencl.cpp
+14 -6
View File
@@ -7065,7 +7065,7 @@ static inline bool use_flat_gemv_for_large_m_q4_K(const ggml_tensor *tensor) {
return tensor->ne[1] >= 32768 && tensor->ne[2] == 1 && tensor->ne[3] == 1;
}
static inline bool use_flat_gemv_for_large_m_q6_K(const ggml_tensor *tensor) {
static inline bool use_flat_gemv_for_large_m_q6_K(const ggml_backend_opencl_context *backend_ctx, const ggml_tensor *tensor) {
// gemv_noshuffle variant perf drops for large M, use flat variant for large M.
// threshold is well above typical hidden/FFN dims, but below typical vocab sizes.
// q6_K flat gemv is worse for smaller K; 2048 seems to be a reasonable threshold.
@@ -7083,7 +7083,15 @@ static inline bool use_flat_gemv_for_large_m_q6_K(const ggml_tensor *tensor) {
if ((tensor->ne[1] % 128 != 0) && tensor->ne[2] == 1 && tensor->ne[3] == 1) {
return true;
}
return tensor->ne[1] >= 32768 && tensor->ne[0] >= 2048 && tensor->ne[2] == 1 && tensor->ne[3] == 1;
// The gemv_noshuffle slowdown tracks TOTAL weight size, not ne0 alone; ne0 >= 2048 is a
// proxy for "large weight" that misses a narrow-hidden vocab-scale lm_head.
// Add a direct size escape so such weights also take the flat path, without changing
// which weights ne0 >= 2048 already routes there.
// The size escape is not taken on the A7X since its compiler miscompiles the flat K-quant GEMV
return tensor->ne[1] >= 32768
&& (tensor->ne[0] >= 2048 || (backend_ctx->adreno_gen != ADRENO_GPU_GEN::A7X && ggml_nbytes(tensor) >= (256ull << 20)))
&& tensor->ne[2] == 1 && tensor->ne[3] == 1;
}
static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
@@ -9403,7 +9411,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
cl_kernel kernel;
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
kernel = backend_ctx->kernel_convert_block_q6_K;
if (use_adreno_kernels(backend_ctx, tensor) && !use_flat_gemv_for_large_m_q6_K(tensor)) {
if (use_adreno_kernels(backend_ctx, tensor) && !use_flat_gemv_for_large_m_q6_K(backend_ctx, tensor)) {
kernel = backend_ctx->kernel_convert_block_q6_K_noshuffle;
}
#else
@@ -9436,7 +9444,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
tensor->extra = extra;
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
if (use_adreno_kernels(backend_ctx, tensor) && !use_flat_gemv_for_large_m_q6_K(tensor)) {
if (use_adreno_kernels(backend_ctx, tensor) && !use_flat_gemv_for_large_m_q6_K(backend_ctx, tensor)) {
cl_int M = tensor->ne[1]; // ne01
cl_int K = tensor->ne[0]; // ne00
@@ -10473,7 +10481,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
CL_CHECK(clReleaseMemObject(data_device));
return;
}
if (use_adreno_kernels(backend_ctx, tensor) && !use_flat_gemv_for_large_m_q6_K(tensor)) {
if (use_adreno_kernels(backend_ctx, tensor) && !use_flat_gemv_for_large_m_q6_K(backend_ctx, tensor)) {
static ggml_cl_buffer buf_trans_ql;
static ggml_cl_buffer buf_trans_qh;
static ggml_cl_buffer buf_trans_s;
@@ -18895,7 +18903,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
}
// q6_K x fp32
if (src0t == GGML_TYPE_Q6_K && src1t == GGML_TYPE_F32 && !use_flat_gemv_for_large_m_q6_K(src0)) {
if (src0t == GGML_TYPE_Q6_K && src1t == GGML_TYPE_F32 && !use_flat_gemv_for_large_m_q6_K(backend_ctx, src0)) {
ggml_cl_mul_mat_q6_K_f32_adreno(backend, src0, src1, dst);
return;
}
include/llama.h
+4 -3
View File
@@ -1256,6 +1256,7 @@ extern "C" {
struct ggml_tensor * probs;
struct ggml_tensor * sampled;
struct ggml_tensor * candidates;
int64_t n_vocab;
};
// user code can implement the interface below in order to create custom llama_sampler
@@ -1425,9 +1426,9 @@ extern "C" {
/// NOTE: Avoid using on the full vocabulary as searching for repeated tokens can become slow. For example, apply top-k or top-p sampling first.
LLAMA_API struct llama_sampler * llama_sampler_init_penalties(
int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty, -1 = context size)
float penalty_repeat, // 1.0 = disabled
float penalty_freq, // 0.0 = disabled
float penalty_present); // 0.0 = disabled
float penalty_repeat, // must be > 0.0, 1.0 = disabled
float penalty_freq, // must be finite, 0.0 = disabled
float penalty_present); // must be finite, 0.0 = disabled
/// @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982
LLAMA_API struct llama_sampler * llama_sampler_init_dry(
src/CMakeLists.txt
+1
View File
@@ -25,6 +25,7 @@ add_library(llama
llama-kv-cache.cpp
llama-kv-cache-iswa.cpp
llama-kv-cache-dsa.cpp
llama-kv-cache-msa.cpp
llama-kv-cache-dsv4.cpp
llama-memory.cpp
llama-memory-hybrid.cpp
src/llama-graph.cpp
+64
View File
@@ -8,6 +8,7 @@
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
#include "llama-kv-cache-dsa.h"
#include "llama-kv-cache-msa.h"
#include "llama-kv-cache-dsv4.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-iswa.h"
@@ -518,6 +519,40 @@ bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) {
return res;
}
llm_graph_input_attn_kv_msa::llm_graph_input_attn_kv_msa(
const llama_hparams & hparams,
const llama_cparams & cparams,
const llama_kv_cache_msa_context * mctx) :
llm_graph_input_attn_kv(hparams, cparams, mctx->get_base()),
mctx_msa(mctx) {
}
void llm_graph_input_attn_kv_msa::set_input(const llama_ubatch * ubatch) {
llm_graph_input_attn_kv::set_input(ubatch);
if (self_k_idxs_idx) {
mctx_msa->get_idx()->set_input_k_idxs(self_k_idxs_idx, ubatch);
}
}
bool llm_graph_input_attn_kv_msa::can_reuse(const llm_graph_params & params) {
mctx_msa = static_cast<const llama_kv_cache_msa_context *>(params.mctx);
// the parent class operates on the base cache context
this->mctx = mctx_msa->get_base();
bool res = true;
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
if (self_k_idxs_idx) {
res &= self_k_idxs_idx->ne[0] == params.ubatch.n_tokens;
}
res &= can_reuse_kq_mask(self_kq_mask, this->mctx, params.ubatch, params.cparams);
return res;
}
void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
mctx->get_mla()->set_input_k_idxs(self_k_idxs_mla, ubatch);
@@ -3187,6 +3222,34 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
return (llm_graph_input_attn_k_dsa *) res->add_input(std::move(inp));
}
llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_enabled) const {
const auto * mctx_cur = static_cast<const llama_kv_cache_msa_context *>(mctx);
auto inp = std::make_unique<llm_graph_input_attn_kv_msa>(hparams, cparams, mctx_cur);
const auto * mctx_base = mctx_cur->get_base();
const auto * mctx_idx = mctx_cur->get_idx();
{
GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA");
inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch);
inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_base, ubatch, cparams);
inp->self_kq_mask_cnv = inp->self_kq_mask;
}
inp->self_k_rot = mctx_base->build_input_k_rot(ctx0);
inp->self_v_rot = mctx_base->build_input_v_rot(ctx0);
if (msa_enabled) {
inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch);
}
return (llm_graph_input_attn_kv_msa *) res->add_input(std::move(inp));
}
// TODO: maybe separate the inner implementation into a separate function
// like with the non-sliding window equivalent
// once sliding-window hybrid caches are a thing.
@@ -3620,6 +3683,7 @@ void llm_graph_context::build_sampling() const {
/*.probs =*/ nullptr,
/*.sampled =*/ nullptr,
/*.candidates =*/ nullptr,
/*.n_vocab =*/ logits_seq->ne[0],
};
assert(sampler->iface->backend_apply);
src/llama-graph.h
+23
View File
@@ -23,6 +23,7 @@ struct llama_memory_context_i;
class llama_kv_cache_context;
class llama_kv_cache_dsa_context;
class llama_kv_cache_msa_context;
class llama_kv_cache_dsv4_raw_context;
class llama_kv_cache_dsv4_context;
class llama_kv_cache_iswa_context;
@@ -425,6 +426,26 @@ public:
const llama_kv_cache_dsa_context * mctx;
};
// standard K/V attention input against the base cache, plus destination indices for the indexer key cache
class llm_graph_input_attn_kv_msa : public llm_graph_input_attn_kv {
public:
llm_graph_input_attn_kv_msa(
const llama_hparams & hparams,
const llama_cparams & cparams,
const llama_kv_cache_msa_context * mctx);
~llm_graph_input_attn_kv_msa() = default;
void set_input(const llama_ubatch * ubatch) override;
bool can_reuse(const llm_graph_params & params) override;
ggml_tensor * get_k_idxs_idx() const { return self_k_idxs_idx; }
ggml_tensor * self_k_idxs_idx = nullptr; // I64 [n_batch]
const llama_kv_cache_msa_context * mctx_msa;
};
class llm_graph_input_attn_kv_iswa : public llm_graph_input_i {
public:
llm_graph_input_attn_kv_iswa(
@@ -1169,6 +1190,8 @@ struct llm_graph_context {
llm_graph_input_attn_k_dsa * build_attn_inp_k_dsa() const;
llm_graph_input_attn_kv_msa * build_attn_inp_kv_msa(bool msa_enabled) const;
ggml_tensor * build_attn(
llm_graph_input_attn_k_dsa * inp,
ggml_tensor * wo,
src/llama-hparams.cpp
-10
View File
@@ -180,16 +180,6 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const {
return val;
}
uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const {
if (!indexer_kv || indexer_head_size == 0) {
return 0; // arch without a MSA indexer
}
if (il < n_layer_dense_lead) {
return 0; // leading dense layers carry no indexer
}
return indexer_head_size; // 128
}
uint32_t llama_hparams::n_embd_r() const {
if (wkv_head_size != 0) {
// for RWKV models
src/llama-hparams.h
-5
View File
@@ -230,8 +230,6 @@ struct llama_hparams {
// MSA
uint32_t indexer_block_size = 0;
uint32_t indexer_local_blocks = 0;
// MSA stores its indexer keys in the main KV cache (k_idx tensors);
bool indexer_kv = false;
// Indexer is "full" (1) or "shared" (0)
// Shared indexers reuse top-k from previous full layer
@@ -356,9 +354,6 @@ struct llama_hparams {
uint32_t n_embd_k_gqa_max() const;
uint32_t n_embd_v_gqa_max() const;
// dimension of the single-head MSA indexer key stream
uint32_t n_embd_k_idx(uint32_t il = 0) const;
// dimension of the rolling state embeddings
// corresponds to Mamba's conv_states size or RWKV's token_shift states size
uint32_t n_embd_r() const;
src/llama-kv-cache-dsa.cpp
+4 -3
View File
@@ -23,7 +23,8 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_filter_cb & filter_mla,
const layer_filter_cb & filter_lid,
const layer_reuse_cb & reuse) :
hparams_lid(model.hparams), n_stream(unified ? 1 : n_seq_max) {
@@ -32,7 +33,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
kv_mla = std::make_unique<llama_kv_cache>(
model, model.hparams, type_k, type_v,
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
n_swa, swa_type, nullptr, filter, reuse, nullptr);
n_swa, swa_type, nullptr, filter_mla, reuse, nullptr);
// we use llama_kv_cache for caching indexer keys
// by hand-tweaking some hparams we fool it to create
@@ -49,7 +50,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
kv_lid = std::make_unique<llama_kv_cache>(
model, hparams_lid, type_k, type_v,
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
n_swa, swa_type, nullptr, filter, reuse, nullptr);
n_swa, swa_type, nullptr, filter_lid, reuse, nullptr);
}
void llama_kv_cache_dsa::clear(bool data) {
src/llama-kv-cache-dsa.h
+2 -1
View File
@@ -26,7 +26,8 @@ public:
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_filter_cb & filter_mla,
const layer_filter_cb & filter_lid,
const layer_reuse_cb & reuse);
~llama_kv_cache_dsa() = default;
src/llama-kv-cache-msa.cpp
+395
View File
@@ -0,0 +1,395 @@
#include "llama-kv-cache-msa.h"
#include "llama-impl.h"
#include "llama-batch.h"
#include "llama-model.h"
#include <algorithm>
#include <cassert>
#include <cmath>
// llama_kv_cache_msa
llama_kv_cache_msa::llama_kv_cache_msa(
const llama_model & model,
ggml_type type_k,
ggml_type type_v,
bool v_trans,
bool offload,
bool unified,
uint32_t kv_size,
uint32_t n_seq_max,
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_filter_cb & filter_idx,
const layer_reuse_cb & reuse) :
hparams_idx(model.hparams),
n_stream(unified ? 1 : n_seq_max), n_seq_max(n_seq_max), n_pad(n_pad),
n_swa(n_swa), swa_type(swa_type) {
LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size);
kv_base = std::make_unique<llama_kv_cache>(
model, model.hparams, type_k, type_v,
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
n_swa, swa_type, nullptr, filter, reuse, nullptr);
// the MSA indexer uses a single key head per layer
std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
// the rope parameters are kept identical to the main cache
LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
kv_idx = std::make_unique<llama_kv_cache>(
model, hparams_idx, type_k, type_v,
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
n_swa, swa_type, nullptr, filter_idx, reuse, nullptr);
}
void llama_kv_cache_msa::clear(bool data) {
kv_base->clear(data);
kv_idx ->clear(data);
}
bool llama_kv_cache_msa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
bool res = true;
res = res & kv_base->seq_rm(seq_id, p0, p1);
res = res & kv_idx ->seq_rm(seq_id, p0, p1);
return res;
}
void llama_kv_cache_msa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
kv_base->seq_cp(seq_id_src, seq_id_dst, p0, p1);
kv_idx ->seq_cp(seq_id_src, seq_id_dst, p0, p1);
}
void llama_kv_cache_msa::seq_keep(llama_seq_id seq_id) {
kv_base->seq_keep(seq_id);
kv_idx ->seq_keep(seq_id);
}
void llama_kv_cache_msa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
kv_base->seq_add(seq_id, p0, p1, shift);
kv_idx ->seq_add(seq_id, p0, p1, shift);
}
void llama_kv_cache_msa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
kv_base->seq_div(seq_id, p0, p1, d);
kv_idx ->seq_div(seq_id, p0, p1, d);
}
llama_pos llama_kv_cache_msa::seq_pos_min(llama_seq_id seq_id) const {
return kv_base->seq_pos_min(seq_id);
}
llama_pos llama_kv_cache_msa::seq_pos_max(llama_seq_id seq_id) const {
return kv_base->seq_pos_max(seq_id);
}
std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache_msa::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> mb = kv_base->memory_breakdown();
for (const auto & buft_size : kv_idx->memory_breakdown()) {
mb[buft_size.first] += buft_size.second;
}
return mb;
}
llama_memory_context_ptr llama_kv_cache_msa::init_batch(
llama_batch_allocr & balloc,
uint32_t n_ubatch,
bool embd_all) {
GGML_UNUSED(embd_all);
do {
balloc.split_reset();
std::vector<llama_ubatch> ubatches;
while (true) {
auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
if (ubatch.n_tokens == 0) {
break;
}
ubatches.push_back(std::move(ubatch));
}
if (balloc.get_n_used() < balloc.get_n_tokens()) {
// failed to find a suitable split
break;
}
auto sinfos_base = kv_base->prepare(ubatches);
if (sinfos_base.empty()) {
break;
}
auto sinfos_idx = kv_idx->prepare(ubatches);
if (sinfos_idx.empty()) {
break;
}
assert(sinfos_base.size() == sinfos_idx.size());
return std::make_unique<llama_kv_cache_msa_context>(
this, std::move(sinfos_base), std::move(sinfos_idx), std::move(ubatches));
} while (false);
return std::make_unique<llama_kv_cache_msa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
}
llama_memory_context_ptr llama_kv_cache_msa::init_full() {
return std::make_unique<llama_kv_cache_msa_context>(this);
}
llama_memory_context_ptr llama_kv_cache_msa::init_update(llama_context * lctx, bool optimize) {
return std::make_unique<llama_kv_cache_msa_context>(this, lctx, optimize);
}
bool llama_kv_cache_msa::get_can_shift() const {
return kv_base->get_can_shift() &&
kv_idx ->get_can_shift() &&
kv_base->get_size() == kv_idx->get_size();
}
void llama_kv_cache_msa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
kv_base->state_write(io, seq_id, flags);
kv_idx ->state_write(io, seq_id, flags);
}
void llama_kv_cache_msa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
kv_base->state_read(io, seq_id, flags);
kv_idx ->state_read(io, seq_id, flags);
}
llama_kv_cache * llama_kv_cache_msa::get_base() const {
return kv_base.get();
}
llama_kv_cache * llama_kv_cache_msa::get_idx() const {
return kv_idx.get();
}
// llama_kv_cache_msa_context
llama_kv_cache_msa_context::llama_kv_cache_msa_context(llama_memory_status status) :
kv(nullptr), status(status) {}
llama_kv_cache_msa_context::llama_kv_cache_msa_context(
llama_kv_cache_msa * kv) :
kv(kv),
ctx_base(kv->get_base()->init_full()),
ctx_idx (kv->get_idx ()->init_full()),
status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
}
llama_kv_cache_msa_context::llama_kv_cache_msa_context(
llama_kv_cache_msa * kv,
llama_context * lctx,
bool optimize) :
kv(kv),
ctx_base(kv->get_base()->init_update(lctx, optimize)),
ctx_idx (kv->get_idx ()->init_update(lctx, optimize)),
status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
}
llama_kv_cache_msa_context::llama_kv_cache_msa_context(
llama_kv_cache_msa * kv,
slot_info_vec_t sinfos_base,
slot_info_vec_t sinfos_idx,
std::vector<llama_ubatch> ubatches) :
kv(kv),
ubatches(std::move(ubatches)),
// here we copy the ubatches. not sure if this is ideal
ctx_base(new llama_kv_cache_context(kv->get_base(), std::move(sinfos_base), this->ubatches)),
ctx_idx (new llama_kv_cache_context(kv->get_idx (), std::move(sinfos_idx), this->ubatches)),
status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
}
llama_kv_cache_msa_context::~llama_kv_cache_msa_context() = default;
bool llama_kv_cache_msa_context::next() {
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
ctx_base->next();
ctx_idx ->next();
if (++i_next >= ubatches.size()) {
return false;
}
return true;
}
bool llama_kv_cache_msa_context::apply() {
assert(!llama_memory_status_is_fail(status));
bool res = true;
res = res & ctx_base->apply();
res = res & ctx_idx ->apply();
return res;
}
llama_memory_status llama_kv_cache_msa_context::get_status() const {
return status;
}
const llama_ubatch & llama_kv_cache_msa_context::get_ubatch() const {
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
return ubatches[i_next];
}
const llama_kv_cache_context * llama_kv_cache_msa_context::get_base() const {
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
return static_cast<const llama_kv_cache_context *>(ctx_base.get());
}
const llama_kv_cache_context * llama_kv_cache_msa_context::get_idx() const {
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
return static_cast<const llama_kv_cache_context *>(ctx_idx.get());
}
uint32_t llama_kv_cache_msa_context::get_n_pos() const {
// pad the value so that the graph remains constant across batches and can be reused
const uint32_t n_pad_cur = std::max(kv->get_n_pad(), 256u);
llama_pos pos_max = -1;
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) kv->get_n_seq_max(); ++seq_id) {
pos_max = std::max(pos_max, kv->seq_pos_max(seq_id));
}
return std::max(n_pad_cur, GGML_PAD((uint32_t) (pos_max + 1), n_pad_cur));
}
void llama_kv_cache_msa_context::set_input_cell_pos(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const {
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
GGML_ASSERT(dst->type == GGML_TYPE_I32);
GGML_ASSERT(div > 0);
const int64_t n_tokens = ubatch->n_tokens;
const int64_t n_kv = dst->ne[0];
const int64_t n_stream_ub = dst->ne[1];
GGML_ASSERT(n_tokens % n_stream_ub == 0);
const int64_t n_tps = n_tokens/n_stream_ub;
int32_t * data = (int32_t *) dst->data;
for (int64_t s = 0; s < n_stream_ub; ++s) {
const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];
const auto & cells = kv->get_base()->get_cells(seq_id);
for (int64_t j = 0; j < n_kv; ++j) {
// the value for empty or other-sequence cells is irrelevant as consumers mask them
data[s*n_kv + j] =
cells.is_empty(j) || !cells.seq_has(j, seq_id)
? 0
: (int32_t) (cells.pos_get(j)/div);
}
}
}
void llama_kv_cache_msa_context::set_input_pos_slot(ggml_tensor * dst, const llama_ubatch * ubatch) const {
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
GGML_ASSERT(dst->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_F32);
const int64_t n_tokens = ubatch->n_tokens;
const int64_t n_pos = dst->ne[0];
const int64_t n_stream_ub = dst->ne[1];
GGML_ASSERT(n_tokens % n_stream_ub == 0);
const int64_t n_tps = n_tokens/n_stream_ub;
for (int64_t s = 0; s < n_stream_ub; ++s) {
const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];
const auto & cells = kv->get_base()->get_cells(seq_id);
std::vector<int32_t> map(n_pos, 0);
for (uint32_t j = 0; j < cells.size(); ++j) {
if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
continue;
}
const llama_pos p0 = cells.pos_get(j);
if (p0 < 0 || p0 >= n_pos) {
continue;
}
map[p0] = (int32_t) j;
}
if (dst->type == GGML_TYPE_I32) {
int32_t * data = (int32_t *) dst->data + s*n_pos;
std::copy(map.begin(), map.end(), data);
} else {
float * data = (float *) dst->data + s*n_pos;
for (int64_t p = 0; p < n_pos; ++p) {
data[p] = (float) map[p];
}
}
}
}
void llama_kv_cache_msa_context::set_input_pos_mask(ggml_tensor * dst, const llama_ubatch * ubatch) const {
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
GGML_ASSERT(dst->type == GGML_TYPE_F32);
const int64_t n_tokens = ubatch->n_tokens;
const int64_t n_pos = dst->ne[0];
GGML_ASSERT(dst->ne[1] == n_tokens);
const uint32_t n_swa = kv->get_n_swa();
const llama_swa_type swa_type = kv->get_swa_type();
float * data = (float *) dst->data;
std::fill(data, data + n_pos*n_tokens, -INFINITY);
for (int64_t i = 0; i < n_tokens; ++i) {
const llama_seq_id seq_id = ubatch->seq_id[i][0];
const auto & cells = kv->get_base()->get_cells(seq_id);
const llama_pos p1 = ubatch->pos[i];
for (uint32_t j = 0; j < cells.size(); ++j) {
if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
continue;
}
const llama_pos p0 = cells.pos_get(j);
if (p0 < 0 || p0 >= n_pos) {
continue;
}
// causal mask
if (p0 > p1) {
continue;
}
// apply SWA if any
if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {
continue;
}
data[i*n_pos + p0] = 0.0f;
}
}
}
src/llama-kv-cache-msa.h
+153
View File
@@ -0,0 +1,153 @@
#pragma once
#include "llama-kv-cache.h"
#include <vector>
// llama_kv_cache_msa
// uses two instances of llama_kv_cache, one for K/V tensors, and one for the MSA indexer tensors
// both receive identical sequence operations and identical ubatches, so their cell layouts stay in synced.
// the context also exposes per-ubatch pos - cell translation maps populated from llama_kv_cells via
// llama_kv_cache::get_cells(), which the model graph uses to run MSA block selection in position space
class llama_kv_cache_msa : public llama_memory_i {
public:
llama_kv_cache_msa(
const llama_model & model,
ggml_type type_k,
ggml_type type_v,
bool v_trans,
bool offload,
bool unified,
uint32_t kv_size,
uint32_t n_seq_max,
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_filter_cb & filter_idx,
const layer_reuse_cb & reuse);
~llama_kv_cache_msa() = default;
// llama_memory_i
llama_memory_context_ptr init_batch(
llama_batch_allocr & balloc,
uint32_t n_ubatch,
bool embd_all) override;
llama_memory_context_ptr init_full() override;
llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override;
bool get_can_shift() const override;
void clear(bool data) override;
bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override;
void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override;
void seq_keep(llama_seq_id seq_id) override;
void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override;
void seq_div (llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override;
llama_pos seq_pos_min(llama_seq_id seq_id) const override;
llama_pos seq_pos_max(llama_seq_id seq_id) const override;
std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;
// state write/load
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
// llama_kv_cache_msa specific API
llama_kv_cache * get_base() const;
llama_kv_cache * get_idx () const;
uint32_t get_n_pad() const { return n_pad; }
uint32_t get_n_seq_max() const { return n_seq_max; }
uint32_t get_n_swa() const { return n_swa; }
llama_swa_type get_swa_type() const { return swa_type; }
private:
// keep the indexer KV cache hparams instance here as llama_kv_cache stores only a reference
llama_hparams hparams_idx;
const uint32_t n_stream = 1;
const uint32_t n_seq_max = 1;
const uint32_t n_pad = 1;
const uint32_t n_swa = 0;
const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
std::unique_ptr<llama_kv_cache> kv_base;
std::unique_ptr<llama_kv_cache> kv_idx;
};
class llama_kv_cache_msa_context : public llama_memory_context_i {
public:
using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
// used for errors
llama_kv_cache_msa_context(llama_memory_status status);
// used to create a full-cache context
llama_kv_cache_msa_context(
llama_kv_cache_msa * kv);
// used to create an update context
llama_kv_cache_msa_context(
llama_kv_cache_msa * kv,
llama_context * lctx,
bool optimize);
// used to create a batch processing context from a batch
llama_kv_cache_msa_context(
llama_kv_cache_msa * kv,
slot_info_vec_t sinfos_base,
slot_info_vec_t sinfos_idx,
std::vector<llama_ubatch> ubatches);
virtual ~llama_kv_cache_msa_context();
// llama_memory_context_i
bool next() override;
bool apply() override;
llama_memory_status get_status() const override;
const llama_ubatch & get_ubatch() const override;
// llama_kv_cache_msa_context specific API
const llama_kv_cache_context * get_base() const;
const llama_kv_cache_context * get_idx () const;
// max position currently present in the cache plus one, padded MSA blocks are defined over token positions
// so the block-selection tensors are sized by this value rather than by the number of cells
uint32_t get_n_pos() const;
// position <-> cell translation maps, populated from the base cache cells
// the model graph relates cache contents to token positions only through these per ubatch inputs
// value for empty or other-sequence cells is 0 so consumers must mask them
void set_input_cell_pos(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const;
// positions without a cell map to cell 0, consumers must mask them assumes one sequence per stream
void set_input_pos_slot(ggml_tensor * dst, const llama_ubatch * ubatch) const;
void set_input_pos_mask(ggml_tensor * dst, const llama_ubatch * ubatch) const;
private:
llama_kv_cache_msa * kv;
// the index of the next ubatch to process
size_t i_next = 0;
std::vector<llama_ubatch> ubatches;
const llama_memory_context_ptr ctx_base;
const llama_memory_context_ptr ctx_idx;
const llama_memory_status status;
};
src/llama-kv-cache.cpp
+20 -278
View File
@@ -112,7 +112,7 @@ llama_kv_cache::llama_kv_cache(
auto it = ctx_map.find(buft);
if (it == ctx_map.end()) {
ggml_init_params params = {
/*.mem_size =*/ size_t(3u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), //Reserve tensor metadata for up to 3 tensors per layer (K, V, and optional K_idx), plus one view per tensor per stream.
/*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
@@ -242,25 +242,9 @@ llama_kv_cache::llama_kv_cache(
v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr);
}
const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il);
ggml_tensor * k_idx = n_embd_k_idx > 0
? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream)
: nullptr;
if (k_idx) {
ggml_format_name(k_idx, "cache_k_idx_l%d", il);
msa_strict_slots = (n_stream == n_seq_max);
}
std::vector<ggml_tensor *> k_idx_stream;
for (uint32_t s = 0; s < n_stream; ++s) {
k_idx_stream.push_back(k_idx
? ggml_view_2d(ctx, k_idx, n_embd_k_idx, kv_size, k_idx->nb[1], s*k_idx->nb[2])
: nullptr);
}
map_layer_ids[il] = layers.size();
layers.push_back({ il, k, v, k_idx, k_stream, v_stream, k_idx_stream });
layers.push_back({ il, k, v, k_stream, v_stream, });
}
if (reuse) {
@@ -309,24 +293,13 @@ llama_kv_cache::llama_kv_cache(
}
{
const size_t memory_size_k = size_k_bytes();
const size_t memory_size_v = size_v_bytes();
const size_t memory_size_k_idx = size_k_idx_bytes();
const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx;
const size_t memory_size_k = size_k_bytes();
const size_t memory_size_v = size_v_bytes();
constexpr float mib = 1024.0f * 1024.0f;
const std::string k_log = format(", K (%s): %7.2f MiB", ggml_type_name(type_k), (float) memory_size_k / mib);
const std::string v_log = format(", V (%s): %7.2f MiB", ggml_type_name(type_v), (float) memory_size_v / mib);
std::string k_idx_log;
if (memory_size_k_idx > 0) {
k_idx_log = format(", K_idx (%s): %7.2f MiB", ggml_type_name(GGML_TYPE_F32), (float) memory_size_k_idx / mib);
}
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs)%s%s%s\n", __func__,
(float) memory_size_total / mib, kv_size, (int) layers.size(), n_seq_max, n_stream,
k_log.c_str(), v_log.c_str(), k_idx_log.c_str());
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__,
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,
ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),
ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));
}
// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
@@ -419,39 +392,6 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
p1 = std::numeric_limits<llama_pos>::max();
}
// empty range - nothing to remove
if (p0 >= p1) {
return true;
}
// MSA anchors block selection to absolute cache slots (slot == position). Tail trim and full removal preserve this invariant, but removing a prefix
// or middle range would free slots while later cells survive, desynchronizing the indexer cache. Reject such removals before modifying the cache.
if (msa_strict_slots) {
for (llama_seq_id sid = 0; sid < (llama_seq_id) seq_to_stream.size(); ++sid) {
if (seq_id >= 0 && sid != seq_id) {
continue;
}
const auto & cells = v_cells[seq_to_stream[sid]];
const llama_pos pmin = cells.seq_pos_min(sid);
const llama_pos pmax = cells.seq_pos_max(sid);
if (pmin < 0) {
continue; // empty sequence
}
const bool overlaps = p0 <= pmax && p1 > pmin; // the range removes something
const bool leaves_tail = p1 <= pmax; // cells beyond the range survive
if (overlaps && leaves_tail) {
LLAMA_LOG_WARN("%s: MSA: partial (non-suffix) removal [%d, %d) for seq %d is not supported "
"(block selection is anchored to cache slots) - rejected\n", __func__, p0, p1, sid);
return false;
}
}
}
if (seq_id >= 0) {
auto & cells = v_cells[seq_to_stream[seq_id]];
auto & head = v_heads[seq_to_stream[seq_id]];
@@ -906,10 +846,6 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co
if (layer.v_stream[ssrc]) {
ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);
}
if (layer.k_idx_stream[ssrc]) {
GGML_ASSERT(layer.k_idx_stream[sdst]);
ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]);
}
}
}
}
@@ -1058,44 +994,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
const auto & cells = v_cells[seq_to_stream[seq_id]];
if (n_tokens > cells.size()) {
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
return { };
}
// MSA block selection assumes slot == logical position (append-only streams).
if (msa_strict_slots) {
for (uint32_t ii = 0; ii < n_tokens; ++ii) {
const llama_pos pos = ubatch.pos[s*n_tokens + ii];
if (pos < 0 || (uint64_t) pos >= cells.size()) {
LLAMA_LOG_WARN("%s: MSA: position %d is outside the cache range [0, %u)\n",
__func__, pos, cells.size());
return { };
}
const uint32_t idx = (uint32_t) pos;
if (!cells.is_empty(idx)) {
LLAMA_LOG_WARN("%s: MSA: required slot %u is already occupied (stream %u)\n",
__func__, idx, seq_to_stream[seq_id]);
return { };
}
// strictly increasing positions, rules out duplicates and, for contiguous requests, is tightened to exact adjacency
if (!res.idxs[s].empty() && (cont ? idx != res.idxs[s].back() + 1
: idx <= res.idxs[s].back())) {
LLAMA_LOG_WARN("%s: MSA: token positions are not %s within the ubatch\n",
__func__, cont ? "contiguous" : "strictly increasing");
return { };
}
res.idxs[s].push_back(idx);
}
continue;
}
uint32_t head_cur = v_heads[seq_to_stream[seq_id]];
// if we have enough unused cells before the current head ->
@@ -1104,6 +1002,11 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
head_cur = 0;
}
if (n_tokens > cells.size()) {
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
return { };
}
uint32_t n_tested = 0;
// for continuous slots, we test that all tokens in the ubatch fit, starting from the current head
@@ -1210,15 +1113,6 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
const auto idx = sinfo.idxs[s][ii];
if (msa_strict_slots && (llama_pos) idx != ubatch.pos[i]) {
LLAMA_LOG_ERROR("%s: MSA slot/position invariant violated: "
"writing pos %d into cell %u (stream %u). The indexer cache "
"would desync and block selection would silently corrupt. "
"This is a bug, please report it with reproduction steps.\n",
__func__, ubatch.pos[i], idx, sinfo.strm[s]);
GGML_ABORT("MSA: slot != pos");
}
if (!cells.is_empty(idx)) {
assert(cells.seq_count(idx) == 1);
@@ -1262,8 +1156,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",
__func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);
// under MSA strict slots this path should be unreachable, since strict MSA placement never selects occupied cells
GGML_ASSERT(seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1));
seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);
}
}
@@ -1283,12 +1176,6 @@ bool llama_kv_cache::get_can_shift() const {
if (hparams.n_pos_per_embd() > 1) {
return false;
}
// shifting would leave k_idx stale
for (const auto & layer : layers) {
if (layer.k_idx) {
return false;
}
}
return true;
}
@@ -1337,6 +1224,12 @@ ggml_tensor * llama_kv_cache::get_k_storage(int32_t il) const {
return layers[ikv].k;
}
const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const {
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
return v_cells[seq_to_stream[seq_id]];
}
uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
uint32_t result = 0;
@@ -1405,23 +1298,6 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k
ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
}
ggml_tensor * llama_kv_cache::get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
const int32_t ikv = map_layer_ids.at(il);
auto * k_idx = layers[ikv].k_idx;
GGML_ASSERT(k_idx);
const uint64_t kv_size = get_size();
const int64_t n_idx = k_idx->ne[0]; // 128
const uint32_t ns = sinfo.s1 - sinfo.s0 + 1;
return ggml_view_4d(ctx, k_idx,
n_idx, 1, n_kv, ns,
ggml_row_size(k_idx->type, n_idx), // nb1 (single head)
ggml_row_size(k_idx->type, n_idx), // nb2 (per cell)
ggml_row_size(k_idx->type, n_idx*kv_size), // nb3 (per stream)
ggml_row_size(k_idx->type, n_idx*kv_size)*sinfo.s0);
}
ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(sinfo);
@@ -1523,28 +1399,6 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama
return k_idxs;
}
ggml_tensor * llama_kv_cache::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(sinfo);
const int32_t ikv = map_layer_ids.at(il);
ggml_tensor * k_idx = layers[ikv].k_idx;
GGML_ASSERT(k_idx && "cpy_k_idx on a layer with no indexer cache");
const int64_t n_embd_head = k_idx_cur->ne[0]; // 128
const int64_t n_head = k_idx_cur->ne[1]; // 1
const int64_t n_tokens = k_idx_cur->ne[2];
const int64_t n_embd_gqa = n_embd_head*n_head; // 128
GGML_ASSERT(ggml_row_size(k_idx_cur->type, n_embd_head) == k_idx_cur->nb[1]);
k_idx_cur = ggml_view_2d(ctx, k_idx_cur, n_embd_gqa, n_tokens, k_idx_cur->nb[2], 0);
const int64_t n_stream = k_idx->ne[2];
if (n_stream > 1) {
const int64_t kv_size = get_size();
k_idx = ggml_reshape_2d(ctx, k_idx, n_embd_gqa, kv_size*n_stream);
}
return ggml_set_rows(ctx, k_idx, k_idx_cur, k_idxs); // same k_idxs as the K store
}
ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
const uint32_t n_tokens = ubatch.n_tokens;
@@ -1979,18 +1833,6 @@ size_t llama_kv_cache::size_v_bytes() const {
return size_v_bytes;
}
size_t llama_kv_cache::size_k_idx_bytes() const {
size_t size_k_idx_bytes = 0;
for (const auto & layer : layers) {
if (layer.k_idx) {
size_k_idx_bytes += ggml_nbytes(layer.k_idx);
}
}
return size_k_idx_bytes;
}
ggml_tensor * llama_kv_cache::build_rope_shift(
const llama_cparams & cparams,
ggml_context * ctx,
@@ -2303,36 +2145,6 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
}
}
if (size_k_idx_bytes() > 0) {
const uint32_t has_k_idx_u32 = 1;
io.write(&has_k_idx_u32, sizeof(has_k_idx_u32));
for (const auto & layer : layers) {
const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0;
io.write(&layer_has_k_idx, sizeof(layer_has_k_idx));
if (!layer_has_k_idx) {
continue;
}
GGML_ASSERT(layer.k_idx_stream[cr.strm]);
const int32_t k_idx_type_i = (int32_t) layer.k_idx->type;
io.write(&k_idx_type_i, sizeof(k_idx_type_i));
const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
io.write(&k_idx_size_row, sizeof(k_idx_size_row));
for (const auto & range : cr.data) {
const size_t range_size = range.second - range.first;
const size_t buf_size = range_size * k_idx_size_row;
const size_t offset = range.first * k_idx_size_row;
io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size);
}
}
}
if (!v_trans) {
for (const auto & layer : layers) {
const uint32_t il = layer.il;
@@ -2581,68 +2393,6 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
}
}
if (size_k_idx_bytes() > 0) {
uint32_t has_k_idx_u32 = 0;
io.read(&has_k_idx_u32, sizeof(has_k_idx_u32));
if (has_k_idx_u32 != 1) {
LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__);
return false;
}
for (const auto & layer : layers) {
uint32_t layer_has_k_idx = 0;
io.read(&layer_has_k_idx, sizeof(layer_has_k_idx));
const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0;
if (layer_has_k_idx != expected_layer_has_k_idx) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx state for layer: got %u, expected %u\n",
__func__, layer_has_k_idx, expected_layer_has_k_idx);
return false;
}
if (!layer_has_k_idx) {
continue;
}
GGML_ASSERT(layer.k_idx_stream[strm]);
int32_t k_idx_type_i = -1;
io.read(&k_idx_type_i, sizeof(k_idx_type_i));
if (k_idx_type_i != (int32_t) layer.k_idx->type) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx type: got %d, expected %d\n",
__func__, k_idx_type_i, (int32_t) layer.k_idx->type);
return false;
}
uint64_t k_idx_size_row = 0;
io.read(&k_idx_size_row, sizeof(k_idx_size_row));
const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
if (k_idx_size_row != expected_k_idx_size_row) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx row size: got %zu, expected %zu\n",
__func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row);
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row);
} else {
for (uint32_t i = 0; i < cell_count; ++i) {
io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row);
}
}
}
}
}
if (!this->v_trans) {
for (const auto & layer : layers) {
const uint32_t il = layer.il;
@@ -2844,10 +2594,6 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons
return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::get_k_idx(ggml_context * ctx, int32_t il) const {
return kv->get_k_idx(ctx, il, n_kv, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
}
@@ -2856,10 +2602,6 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_
return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const {
return kv->cpy_k_idx(ctx, k_idx_cur, k_idxs, il, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
return kv->build_input_k_idxs(ctx, ubatch);
}
src/llama-kv-cache.h
+2 -10
View File
@@ -164,6 +164,8 @@ public:
std::vector<uint32_t> get_layer_ids() const;
ggml_tensor * get_k_storage(int32_t il) const;
const llama_kv_cells & get_cells(llama_seq_id seq_id) const;
//
// graph_build API
//
@@ -173,12 +175,10 @@ public:
// get views of the current state of the cache
ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
// store k_cur and v_cur in the cache based on the provided head location
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
//
// preparation API
@@ -230,11 +230,9 @@ private:
ggml_tensor * k;
ggml_tensor * v;
ggml_tensor * k_idx; // MSA single-head indexer keys, F32
std::vector<ggml_tensor *> k_stream;
std::vector<ggml_tensor *> v_stream;
std::vector<ggml_tensor *> k_idx_stream;
};
bool v_trans = true; // the value tensor is transposed
@@ -263,9 +261,6 @@ private:
// env: LLAMA_KV_CACHE_DEBUG
int debug = 0;
// set when a k_idx (indexer) cache exists and the stream layout supports MSA (single seq, or one stream per seq)
bool msa_strict_slots = false;
// this is the SWA type of the cache - not to be confused with the model SWA type
const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
@@ -298,7 +293,6 @@ private:
size_t size_k_bytes() const;
size_t size_v_bytes() const;
size_t size_k_idx_bytes() const;
ggml_tensor * build_rope_shift(
const llama_cparams & cparams,
@@ -378,7 +372,6 @@ public:
// get views of the current state of the cache
ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il) const;
// store k_cur and v_cur in the cache based on the provided head location
// note: the heads in k_cur and v_cur should be laid out contiguously in memory
@@ -388,7 +381,6 @@ public:
// - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const;
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const;
// create destination indices for each head of the current batch for where it would be written in the KV cache
// the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but
src/llama-model.cpp
+28 -3
View File
@@ -11,6 +11,7 @@
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
#include "llama-kv-cache-dsa.h"
#include "llama-kv-cache-msa.h"
#include "llama-kv-cache-dsv4.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-iswa.h"
@@ -2071,6 +2072,28 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
{
res = nullptr;
} break;
case LLM_ARCH_MINIMAX_M3:
{
// sparse (MSA) layers carry an indexer key cache, but leading dense layers do not
llama_kv_cache::layer_filter_cb filter_idx =
[&](int32_t il) { return (uint32_t) il >= hparams.n_layer_dense_lead; };
res = new llama_kv_cache_msa(
*this,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
1,
hparams.n_swa,
hparams.swa_type,
nullptr,
filter_idx,
nullptr);
} break;
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_DEEPSEEK32:
{
@@ -2101,10 +2124,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
} else {
// Main context: DSA cache for the trunk layers only - the nextn
// layer(s) are never attended by the trunk graph.
llama_kv_cache::layer_filter_cb filter = nullptr;
llama_kv_cache::layer_filter_cb filter_mla = nullptr;
if (hparams.n_layer_nextn > 0) {
filter = [&](uint32_t il) { return il < hparams.n_layer(); };
filter_mla = [&](uint32_t il) { return il < hparams.n_layer(); };
}
llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (arch != LLM_ARCH_GLM_DSA || hparams.is_indexer_full(il)); };
res = new llama_kv_cache_dsa(
*this,
@@ -2118,7 +2142,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
1,
hparams.n_swa,
hparams.swa_type,
filter,
filter_mla,
filter_lid,
nullptr);
}
} break;
src/llama-sampler.cpp
+221 -20
View File
@@ -589,6 +589,7 @@ static bool llama_sampler_backend_support(
/*.probs = */ nullptr,
/*.sampled = */ nullptr,
/*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
/*.n_vocab = */ n,
};
ggml_cgraph * gf = ggml_new_graph(ctx);
@@ -2638,7 +2639,7 @@ struct llama_sampler * llama_sampler_init_grammar_lazy_patterns(
// penalties
struct llama_sampler_penalties {
struct llama_sampler_penalties : public llama_sampler_backend {
const int32_t penalty_last_n;
const float penalty_repeat;
const float penalty_freq;
@@ -2648,10 +2649,49 @@ struct llama_sampler_penalties {
// a frequency map to count token occurrences
std::unordered_map<llama_token, int> token_count;
// backend graph inputs
ggml_tensor * inp_token_ids = nullptr;
ggml_tensor * inp_counts = nullptr;
// backend helpers
int32_t n_vocab = 0;
int32_t n_max = 0;
bool has_candidates = false;
std::vector<int32_t> host_token_ids;
std::vector<int32_t> host_counts;
static bool is_disabled(
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present) {
return penalty_last_n == 0 ||
(penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f);
}
bool is_disabled() const {
return is_disabled(penalty_last_n, penalty_repeat, penalty_freq, penalty_present);
}
llama_sampler_penalties(
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present)
: llama_sampler_backend("penalties")
, penalty_last_n (penalty_last_n)
, penalty_repeat (penalty_repeat)
, penalty_freq (penalty_freq)
, penalty_present (penalty_present)
, prev (penalty_last_n) {
}
};
static const char * llama_sampler_penalties_name(const struct llama_sampler * /*smpl*/) {
return "penalties";
static const char * llama_sampler_penalties_name(const struct llama_sampler * smpl) {
auto * ctx = (llama_sampler_penalties *) smpl->ctx;
return ctx->get_name();
}
static void llama_sampler_penalties_accept(struct llama_sampler * smpl, llama_token token) {
@@ -2688,8 +2728,7 @@ static void llama_sampler_penalties_accept(struct llama_sampler * smpl, llama_to
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)) {
if (ctx->is_disabled()) {
return;
}
@@ -2736,7 +2775,8 @@ static struct llama_sampler * llama_sampler_penalties_clone(const struct llama_s
{
auto * result_ctx = (llama_sampler_penalties *) result->ctx;
result_ctx->prev = ctx->prev;
result_ctx->prev = ctx->prev;
result_ctx->token_count = ctx->token_count;
}
return result;
@@ -2746,6 +2786,171 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) {
delete (llama_sampler_penalties *) smpl->ctx;
}
static bool llama_sampler_penalties_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
auto * sctx = (llama_sampler_penalties *) smpl->ctx;
const bool res = llama_sampler_backend_support(smpl, buft);
sctx->init(res);
return res;
}
static void llama_sampler_penalties_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_penalties *) smpl->ctx;
if (sctx->is_disabled()) {
return;
}
GGML_ASSERT(data->n_vocab > 0 && data->n_vocab <= INT32_MAX);
sctx->has_candidates = data->candidates != nullptr;
sctx->n_vocab = (int32_t) data->n_vocab;
sctx->n_max = std::min(sctx->penalty_last_n, sctx->n_vocab);
sctx->inp_token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sctx->n_max);
ggml_set_name(sctx->inp_token_ids, "penalties_token_ids");
ggml_set_input(sctx->inp_token_ids);
sctx->inp_counts = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sctx->n_max);
ggml_set_name(sctx->inp_counts, "penalties_counts");
ggml_set_input(sctx->inp_counts);
if ((int32_t) sctx->host_token_ids.size() != sctx->n_max) {
sctx->host_token_ids.assign(sctx->n_max, 0);
sctx->host_counts.assign(sctx->n_max, 0);
}
// flatten
ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
ggml_tensor * gathered = logits;
ggml_tensor * counts_f32 = ggml_cast(ctx, sctx->inp_counts, GGML_TYPE_F32);
if (sctx->has_candidates) {
ggml_tensor * candidates = ggml_reshape_1d(
ctx, data->candidates, ggml_nelements(data->candidates));
const int64_t n_candidates = candidates->ne[0];
GGML_ASSERT(n_candidates == ggml_nelements(logits));
ggml_tensor * counts_rows = ggml_fill(
ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, sctx->n_vocab), 0.0f);
ggml_tensor * scatter_rows = ggml_reshape_2d(ctx, counts_f32, 1, sctx->n_max);
counts_rows = ggml_set_rows(ctx, counts_rows, scatter_rows, sctx->inp_token_ids);
counts_f32 = ggml_get_rows(ctx, counts_rows, candidates);
counts_f32 = ggml_reshape_1d(ctx, counts_f32, n_candidates);
} else {
ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits));
gathered = ggml_get_rows(ctx, logits_rows, sctx->inp_token_ids);
gathered = ggml_reshape_1d(ctx, gathered, sctx->n_max);
}
ggml_tensor * active_mask = ggml_step(ctx, counts_f32);
ggml_tensor * inactive_mask = ggml_sub(ctx, ggml_fill(ctx, active_mask, 1.0f), active_mask);
ggml_tensor * penalized = gathered;
if (sctx->penalty_repeat != 1.0f) {
ggml_tensor * pos_mask = ggml_step(ctx, penalized);
ggml_tensor * neg_mask = ggml_sub(ctx, ggml_fill(ctx, pos_mask, 1.0f), pos_mask);
ggml_tensor * pos_scale = ggml_scale(ctx, pos_mask, 1.0f/sctx->penalty_repeat);
ggml_tensor * neg_scale = ggml_scale(ctx, neg_mask, sctx->penalty_repeat);
ggml_tensor * repeat_scale = ggml_add(ctx, pos_scale, neg_scale);
// scale inactive entries with 1 to avoid -INF * 0 = NaN for values masked by top-p
repeat_scale = ggml_mul(ctx, repeat_scale, active_mask);
repeat_scale = ggml_add(ctx, repeat_scale, inactive_mask);
penalized = ggml_mul(ctx, gathered, repeat_scale);
}
if (sctx->penalty_freq != 0.0f) {
ggml_tensor * penalty_freq = ggml_scale(ctx, counts_f32, sctx->penalty_freq);
penalized = ggml_sub(ctx, penalized, penalty_freq);
}
if (sctx->penalty_present != 0.0f) {
ggml_tensor * penalty_present = ggml_scale(ctx, active_mask, sctx->penalty_present);
penalized = ggml_sub(ctx, penalized, penalty_present);
}
if (sctx->has_candidates) {
data->logits = penalized;
} else {
ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits));
ggml_tensor * scatter_rows = ggml_reshape_2d(ctx, penalized, 1, sctx->n_max);
logits_rows = ggml_set_rows(ctx, logits_rows, scatter_rows, sctx->inp_token_ids);
data->logits = ggml_reshape_1d(ctx, logits_rows, ggml_nelements(logits));
}
}
static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_penalties *) smpl->ctx;
if (!sctx->inp_token_ids || !sctx->inp_counts || sctx->n_max <= 0 || sctx->n_vocab <= 0) {
return;
}
if (sctx->is_disabled()) {
return;
}
// fill active entries from the map
int32_t n_active = 0;
for (const auto & it : sctx->token_count) {
GGML_ASSERT(n_active < sctx->n_max);
sctx->host_token_ids[n_active] = it.first;
sctx->host_counts [n_active] = it.second;
++n_active;
}
// Sorting is required because backend_apply uses ggml_set_rows (a scatter-back operation)
std::vector<std::pair<int32_t, int32_t>> entries;
entries.reserve(n_active);
for (int32_t i = 0; i < n_active; ++i) {
entries.emplace_back(sctx->host_token_ids[i], sctx->host_counts[i]);
}
std::sort(entries.begin(), entries.end(), [](const auto & a, const auto & b) {
return a.first < b.first;
});
for (int32_t i = 0; i < n_active; ++i) {
sctx->host_token_ids[i] = entries[i].first;
sctx->host_counts [i] = entries[i].second;
}
// Padding: Finds a filler token id that is not present in token_count.
// Use it to do padding for the arrays, it avoids resizing every time.
// The arrays must always have exactly n_max entries (the GPU tensor is a fixed size).
int32_t filler = 0;
if (n_active < sctx->n_max) {
while (sctx->token_count.find(filler) != sctx->token_count.end()) {
++filler;
}
GGML_ASSERT(filler < sctx->n_vocab);
}
// Fill the rest of the arrays with the filler token id and count 0.
// Inactive slots are padded with a unique dummy token ID (count = 0).
// The uniqueness matters because ggml_set_rows with duplicate indices can produce non-deterministic or incorrect results.
// Using a filler token with count 0 that isn't in the active set is safe, because the active_mask step in backend_apply filters them out via ggml_step(counts_f32)
for (int32_t i = n_active; i < sctx->n_max; ++i) {
sctx->host_token_ids[i] = filler;
sctx->host_counts [i] = 0;
}
ggml_backend_tensor_set(sctx->inp_token_ids, sctx->host_token_ids.data(), 0, sctx->n_max * sizeof(int32_t));
ggml_backend_tensor_set(sctx->inp_counts, sctx->host_counts.data(), 0, sctx->n_max * sizeof(int32_t));
}
static struct llama_sampler_i llama_sampler_penalties_i = {
/* .name = */ llama_sampler_penalties_name,
/* .accept = */ llama_sampler_penalties_accept,
@@ -2753,10 +2958,10 @@ static struct llama_sampler_i llama_sampler_penalties_i = {
/* .reset = */ llama_sampler_penalties_reset,
/* .clone = */ llama_sampler_penalties_clone,
/* .free = */ llama_sampler_penalties_free,
/* .backend_init = */ nullptr,
/* .backend_init = */ llama_sampler_penalties_backend_init,
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_apply = */ llama_sampler_penalties_backend_apply,
/* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
};
struct llama_sampler * llama_sampler_init_penalties(
@@ -2766,22 +2971,18 @@ struct llama_sampler * llama_sampler_init_penalties(
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) {
if (llama_sampler_penalties::is_disabled(
penalty_last_n, penalty_repeat, penalty_freq, penalty_present)) {
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 = */ {},
}
/* .ctx = */ new llama_sampler_penalties(
penalty_last_n,
penalty_repeat,
penalty_freq,
penalty_present)
);
}
src/llama-vocab.cpp
+6
View File
@@ -2532,6 +2532,12 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
const std::string & key = kv(std::get<0>(it));
int32_t & id = std::get<1>(it);
if (id >= 0 && static_cast<size_t>(id) >= id_to_token.size()) {
LLAMA_LOG_WARN("%s: default special token '%s' = %d out of vocab range, disabling\n",
__func__, key.c_str(), id);
id = LLAMA_TOKEN_NULL;
}
uint32_t new_id;
if (!ml.get_key(std::get<0>(it), new_id, false)) {
continue;
src/models/minimax-m3.cpp
+157 -75
View File
@@ -1,5 +1,5 @@
#include "models.h"
#include "llama-kv-cache.h"
#include "llama-kv-cache-msa.h"
#include <cmath>
#include <vector>
#include <cstdint>
@@ -7,7 +7,8 @@
// MiniMax-M3: MiniMax-M2 style GQA (per-head QK-norm, partial rotary) with
// DeepSeek-V3 leading-dense + routed/shared experts (sigmoid gating, routed scaling),
// swigluoai activation, and MiniMax Sparse Attention (MSA). MTP is not in released model weights.
// Notes: Blocks are anchored to absolute KV cache slots.
// MSA blocks are defined over token positions. The graph translates between position space (block
// selection) and cell space (K/V/indexer storage) via per-ubatch pos<->cell maps populated from llama_kv_cells
void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
@@ -23,7 +24,6 @@ void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size);
ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks };
hparams.indexer_kv = true;
switch (hparams.n_layer()) {
case 60: type = LLM_TYPE_428B_A23B; break;
@@ -86,43 +86,83 @@ std::unique_ptr<llm_graph_context> llama_model_minimax_m3::build_arch_graph(cons
return std::make_unique<graph>(*this, params);
}
// per-query local-force bias for MSA selection
// local window always wins a slot
class llm_graph_input_msa_local : public llm_graph_input_i {
class llm_graph_input_msa : public llm_graph_input_i {
public:
llm_graph_input_msa_local(int blk, int local, int64_t nblk) : blk(blk), local(local), nblk(nblk) {}
llm_graph_input_msa(const llama_kv_cache_msa_context * mctx, int blk, int local) :
mctx(mctx), blk(blk), local(local) {}
void set_input(const llama_ubatch * ubatch) override {
if (!bias || !ubatch->pos) {
return;
}
const int64_t n_tokens = ubatch->n_tokens;
std::vector<float> data((size_t) nblk * n_tokens, 0.0f);
for (int64_t i = 0; i < n_tokens; ++i) {
const int64_t L = ubatch->pos[i] / blk;
for (int l = 0; l < local && L - l >= 0; ++l) {
if (L - l < nblk) {
data[(size_t) i * nblk + (L - l)] = 1e30f;
if (pos_slot_i) { mctx->set_input_pos_slot(pos_slot_i, ubatch); }
if (pos_slot_f) { mctx->set_input_pos_slot(pos_slot_f, ubatch); }
if (cell_blk) { mctx->set_input_cell_pos(cell_blk, ubatch, blk); }
if (pos_mask) { mctx->set_input_pos_mask(pos_mask, ubatch); }
// local-force bias over position blocks
if (bias && ubatch->pos) {
const int64_t n_tokens = ubatch->n_tokens;
const int64_t nblk = bias->ne[0];
std::vector<float> data((size_t) nblk * n_tokens, 0.0f);
for (int64_t i = 0; i < n_tokens; ++i) {
const int64_t L = ubatch->pos[i] / blk;
for (int l = 0; l < local && L - l >= 0; ++l) {
if (L - l < nblk) {
data[(size_t) i * nblk + (L - l)] = 1e30f;
}
}
}
ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float));
}
ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float));
}
// valid as long as the bias tensor dims still match the new ubatch/cache window
// valid as long as the tensor dims still match the new ubatch/cache window and the
// ubatch is in the same regime (decode graphs have pos_slot_f, batch graphs cell_blk)
bool can_reuse(const llm_graph_params & params) override {
const auto * mctx = static_cast<const llama_kv_cache_context *>(params.mctx);
const auto * mctx_new = static_cast<const llama_kv_cache_msa_context *>(params.mctx);
this->mctx = mctx_new;
const int64_t n_ps = GGML_PAD((int64_t) mctx_new->get_n_pos(), blk);
const int64_t ns = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq;
const bool decode = params.ubatch.n_tokens == ns; // one token per stream
bool res = true;
res &= bias->ne[1] == params.ubatch.n_tokens;
res &= bias->ne[0] * blk == (int64_t) mctx->get_n_kv();
res &= bias->ne[0] * blk == n_ps;
res &= bias->ne[1] == params.ubatch.n_tokens;
res &= pos_mask->ne[0] == n_ps;
res &= pos_mask->ne[1] == params.ubatch.n_tokens;
res &= pos_slot_i->ne[0] == n_ps;
res &= pos_slot_i->ne[1] == ns;
res &= decode == (pos_slot_f != nullptr);
res &= decode == (cell_blk == nullptr);
if (pos_slot_f) {
res &= pos_slot_f->ne[0] == n_ps;
res &= pos_slot_f->ne[1] == ns;
}
if (cell_blk) {
res &= cell_blk->ne[0] == (int64_t) mctx_new->get_base()->get_n_kv();
res &= cell_blk->ne[1] == ns;
}
return res;
}
ggml_tensor * bias = nullptr;
int blk;
int local;
int64_t nblk;
ggml_tensor * bias = nullptr; // F32 [nblk, n_tokens] local-force bias (position blocks)
ggml_tensor * pos_mask = nullptr; // F32 [n_ps, n_tokens] 0/-inf visibility, by position
ggml_tensor * pos_slot_i = nullptr; // I32 [n_ps, ns] pos -> cell (get_rows index)
ggml_tensor * pos_slot_f = nullptr; // F32 [n_ps, ns] pos -> cell (gatherable values, decode)
ggml_tensor * cell_blk = nullptr; // I32 [n_kv, ns] cell -> position block (batch)
const llama_kv_cache_msa_context * mctx;
int blk;
int local;
};
// One FA call for all GQA groups (and at multi-stream decode, all streams) by mapping them onto the FA sequence dim (ne[3])
@@ -173,7 +213,9 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
inpL = build_inp_embd(model.tok_embd);
ggml_tensor * inp_pos = build_inp_pos();
auto inp_attn = build_attn_inp_kv();
// ==========================================
// TODO: avoid such kind of complexity in the model graphs
// MSA calls ggml_flash_attn_ext directly and assumes the non-transposed V layout that
// llama.cpp only provides when flash attention is enabled. Block selection is anchored
@@ -185,6 +227,8 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
const bool streams_ok = cparams.n_seq_max == 1 || !cparams.kv_unified;
const bool msa_enabled = fa_on && streams_ok;
auto * inp_attn = build_attn_inp_kv_msa(msa_enabled);
static bool warned_no_fa = false;
if (!fa_on && !warned_no_fa) {
LLAMA_LOG_WARN("%s: flash attention disabled; MSA requires it -> running DENSE attention "
@@ -197,36 +241,54 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
"-> running DENSE attention. Output may be degraded. Drop --kv-unified to enable MSA.\n", __func__);
warned_unified = true;
}
// ==========================================
// hoisted per-graph MSA state (shared by every sparse layer)
llm_graph_input_msa_local * msa_loc = nullptr;
llm_graph_input_msa * msa = nullptr;
ggml_tensor * msa_kqm = nullptr;
ggml_tensor * msa_mf = nullptr;
int64_t n_kv = 0, nblk = 0, ns = 1, n_tps = 0;
ggml_tensor * msa_mf = nullptr; // F32 copy of the FA mask for the final mask add
int64_t n_kv = 0, n_ps = 0, nblk = 0, ns = 1, n_tps = 0;
bool msa_decode = false; // gather (1 token per stream) vs mask
const int blk = mm.msa_p.blk;
const int64_t Hd = hparams.indexer_n_head; // one indexer head per GQA group
if (msa_enabled) {
const auto * mctx_msa = static_cast<const llama_kv_cache_msa_context *>(mctx);
msa_kqm = inp_attn->get_kq_mask();
n_kv = msa_kqm->ne[0];
n_tps = msa_kqm->ne[1]; // tokens per stream
ns = msa_kqm->ne[3]; // streams in this ubatch
GGML_ASSERT(msa_kqm->type == GGML_TYPE_F16 && "MSA requires the FA (f16) mask");
GGML_ASSERT(n_tps*ns == n_tokens);
GGML_ASSERT(n_kv % blk == 0 &&
"MSA: KV/mask n_kv must be a multiple of indexer.block_size (128); "
"the flash-attention KV padding must be a multiple of the block size. "
"A non-multiple would silently drop the partial tail block.");
nblk = n_kv / blk;
// the position axis covers every position currently in the cache and is padded to whole blocks
n_ps = GGML_PAD((int64_t) mctx_msa->get_n_pos(), blk);
nblk = n_ps / blk;
msa_decode = n_tps == 1;
msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32);
auto inp = std::make_unique<llm_graph_input_msa>(mctx_msa, blk, mm.msa_p.local);
auto loc = std::make_unique<llm_graph_input_msa_local>(blk, mm.msa_p.local, nblk);
loc->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk, n_tokens); // stream-grouped tokens
ggml_set_input(loc->bias);
msa_loc = (llm_graph_input_msa_local *) res->add_input(std::move(loc));
inp->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk, n_tokens); // stream-grouped tokens
ggml_set_input(inp->bias);
inp->pos_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, n_tokens);
ggml_set_input(inp->pos_mask);
inp->pos_slot_i = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_ps, ns);
ggml_set_input(inp->pos_slot_i);
if (msa_decode) {
inp->pos_slot_f = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, ns);
ggml_set_input(inp->pos_slot_f);
} else {
inp->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, ns);
ggml_set_input(inp->cell_blk);
msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32);
}
msa = (llm_graph_input_msa *) res->add_input(std::move(inp));
}
ggml_tensor * inp_out_ids = build_inp_out_ids();
@@ -283,9 +345,11 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,
freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
const auto * mctx_cur = inp_attn->mctx;
ggml_build_forward_expand(gf, mctx_cur->cpy_k_idx(ctx0, ik, inp_attn->get_k_idxs(), il));
ggml_tensor * ik_kv = mctx_cur->get_k_idx(ctx0, il);
const auto * mctx_msa_l = static_cast<const llama_kv_cache_msa_context *>(mctx);
const auto * mctx_cur = mctx_msa_l->get_base();
const auto * mctx_idx = mctx_msa_l->get_idx();
ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, ik, inp_attn->get_k_idxs_idx(), il));
ggml_tensor * ik_kv = mctx_idx->get_k(ctx0, il);
if (inp_attn->self_k_rot) {
Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot);
@@ -316,42 +380,52 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
if (msa_decode) {
// decode: batched over streams top-k + gather, one grouped FA
// scores: per-stream batched matmul over the stream dim (ne[3]).
// the cache views are not contiguous across streams (stride = kv_size, not n_kv)
ggml_tensor * ikv4 = ggml_view_4d(ctx0, ik_kv, n_idx_dim, n_kv, 1, ns,
ik_kv->nb[2], ik_kv->nb[3], ik_kv->nb[3], 0);
// gather the indexer keys through the pos -> cell map
ggml_tensor * ik3 = ggml_view_3d(ctx0, ik_kv, n_idx_dim, n_kv, ns,
ik_kv->nb[2], ik_kv->nb[3], 0);
ggml_tensor * ikp = ggml_get_rows(ctx0, ik3, msa->pos_slot_i); // [n_idx_dim, n_ps, ns]
ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns);
ggml_tensor * sc = ggml_mul_mat(ctx0, ikv4, iq4);
ggml_tensor * sc = ggml_mul_mat(ctx0,
ggml_reshape_4d(ctx0, ikp, n_idx_dim, n_ps, 1, ns), iq4);
ggml_mul_mat_set_prec(sc, GGML_PREC_F32);
sc = ggml_add_inplace(ctx0, sc, msa_mf);
// unmapped positions come out -inf, so they can never rank into the top-k
sc = ggml_add_inplace(ctx0, sc,
ggml_reshape_4d(ctx0, msa->pos_mask, n_ps, 1, 1, ns));
ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);
cb(bs, "msa_bs", il);
ggml_tensor * bsf = ggml_add(ctx0, bs,
ggml_reshape_4d(ctx0, msa_loc->bias, nblk, 1, 1, ns));
ggml_tensor * idx = ggml_top_k(ctx0, bsf, K);
ggml_reshape_4d(ctx0, msa->bias, nblk, 1, 1, ns));
ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // position blocks
// token idx: tj[t,k,h,s] = blk*idx[k,h,s] + t (for the mask gather)
// row idx: tr[t,k,h,s] = tj*HKV + h (for the per-stream K/V gather)
// pos idx: tj[t,k,h,s] = blk*idx[k,h,s] + t (positions - mask gather)
// cell idx: cs[t,k,h,s] = pos_slot[tj] (pos -> cell translation)
// row idx: tr[t,k,h,s] = cs*HKV + h (per-stream K/V gather)
ggml_tensor * a = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), (float) blk);
a = ggml_reshape_4d(ctx0, a, 1, K, Hd, ns);
ggml_tensor * tj = ggml_add(ctx0,
ggml_repeat_4d(ctx0, a, blk, K, Hd, ns),
ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) blk, 1.0f), blk, 1, 1));
ggml_tensor * tr = ggml_add(ctx0,
ggml_scale(ctx0, tj, (float) HKV),
ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) HKV, 1.0f), 1, 1, Hd));
ggml_tensor * tokj = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tj, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);
ggml_tensor * cs = ggml_get_rows(ctx0,
ggml_reshape_3d(ctx0, msa->pos_slot_f, 1, n_ps, ns), tokj); // [1, blk*K*Hd, ns]
cs = ggml_reshape_4d(ctx0, cs, blk, K, Hd, ns);
ggml_tensor * tr = ggml_add(ctx0,
ggml_scale(ctx0, cs, (float) HKV),
ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) HKV, 1.0f), 1, 1, Hd));
ggml_tensor * tokr = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tr, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);
ggml_tensor * k3 = ggml_view_3d(ctx0, k, D, HKV*n_kv, ns, k->nb[1], k->nb[3], 0);
ggml_tensor * v3 = ggml_view_3d(ctx0, v, D, HKV*n_kv, ns, v->nb[1], v->nb[3], 0);
ggml_tensor * m3 = ggml_reshape_3d(ctx0, msa_kqm, 1, n_kv, ns);
ggml_tensor * mp = ggml_reshape_3d(ctx0, msa->pos_mask, 1, n_ps, ns);
ggml_tensor * kg = ggml_get_rows(ctx0, k3, tokr);
ggml_tensor * vg = ggml_get_rows(ctx0, v3, tokr);
ggml_tensor * mg = ggml_get_rows(ctx0, m3, tokj);
ggml_tensor * mg = ggml_get_rows(ctx0, mp, tokj);
// fold (group, stream) onto the FA channel dim
const ggml_type kt = ggml_is_quantized(k->type) ? GGML_TYPE_F16 : k->type;
@@ -372,12 +446,16 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
iq->nb[1], iq->nb[2], st*n_tps*iq->nb[2]);
ggml_tensor * ik_s = ggml_view_2d(ctx0, ik_kv, n_idx_dim, n_kv,
ik_kv->nb[2], st*ik_kv->nb[3]);
ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, 1, n_tps,
msa_mf->nb[1], msa_mf->nb[1], st*msa_mf->nb[3]);
ggml_tensor * km_s = ggml_view_3d(ctx0, msa_kqm, n_kv, n_tps, 1,
msa_kqm->nb[1], msa_kqm->nb[3], st*msa_kqm->nb[3]);
ggml_tensor * bias_s = ggml_view_3d(ctx0, msa_loc->bias, nblk, 1, n_tps,
msa_loc->bias->nb[1], msa_loc->bias->nb[1], st*n_tps*msa_loc->bias->nb[1]);
ggml_tensor * psl_s = ggml_view_1d(ctx0, msa->pos_slot_i, n_ps,
st*msa->pos_slot_i->nb[1]);
ggml_tensor * pm_s = ggml_view_3d(ctx0, msa->pos_mask, n_ps, 1, n_tps,
msa->pos_mask->nb[1], msa->pos_mask->nb[1], st*n_tps*msa->pos_mask->nb[1]);
ggml_tensor * cb_s = ggml_view_1d(ctx0, msa->cell_blk, n_kv,
st*msa->cell_blk->nb[1]);
ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, n_tps, 1,
msa_mf->nb[1], msa_mf->nb[3], st*msa_mf->nb[3]);
ggml_tensor * bias_s = ggml_view_3d(ctx0, msa->bias, nblk, 1, n_tps,
msa->bias->nb[1], msa->bias->nb[1], st*n_tps*msa->bias->nb[1]);
ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps,
Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]);
ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1,
@@ -385,14 +463,16 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1,
v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]);
// block scores: bs = maxpool_blk(idx_q * idx_k^T + causal mask)
// block scores: the indexer keys are gathered through the pos -> cell map first
// scores are unscaled, only the top-k ordering matters
ggml_tensor * sc = ggml_mul_mat(ctx0, ik_s,
ggml_tensor * ikp = ggml_get_rows(ctx0, ik_s, psl_s); // [n_idx_dim, n_ps]
ggml_tensor * sc = ggml_mul_mat(ctx0, ikp,
ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps));
// indexer scores run in F32
ggml_mul_mat_set_prec(sc, GGML_PREC_F32);
sc = ggml_reshape_3d(ctx0, sc, n_kv, Hd, n_tps);
sc = ggml_add_inplace(ctx0, sc, mf_s);
sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, n_tps);
// unmapped positions (holes, padding, empty cells) come out -inf
sc = ggml_add_inplace(ctx0, sc, pm_s);
ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);
cb(bs, "msa_bs", il);
@@ -416,14 +496,16 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
bm = ggml_cont(ctx0, ggml_permute(ctx0, bm, 0, 2, 1, 3)); // [nblk, n_tps, Hd]
cb(bm, "msa_block_mask", il);
// expand block -> token granularity (j = bk*blk + t),
// then combine with the causal mask in place
ggml_tensor * bmx = ggml_repeat_4d(ctx0,
ggml_reshape_3d(ctx0, bm, 1, nblk, n_tps*Hd),
blk, nblk, n_tps*Hd, 1);
// expand block -> cell granularity through the cell -> position block
// map, then combine with the causal mask. empty cells are masked by the causal mask.
ggml_tensor * bm2 = ggml_cont(ctx0, ggml_transpose(ctx0,
ggml_reshape_2d(ctx0, bm, nblk, n_tps*Hd))); // [n_tps*Hd, nblk]
ggml_tensor * bmc = ggml_get_rows(ctx0, bm2, cb_s); // [n_tps*Hd, n_kv] F32
ggml_tensor * bmx = ggml_cont(ctx0, ggml_transpose(ctx0, bmc));
bmx = ggml_reshape_3d(ctx0, bmx, n_kv, n_tps, Hd);
ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, km_s);
mask4 = ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd);
ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, mf_s);
mask4 = ggml_cast(ctx0,
ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd), GGML_TYPE_F16);
cb(mask4, "msa_mask4", il);
// cache views with groups on ne[3];
tests/test-arg-parser.cpp
+28
View File
@@ -99,6 +99,34 @@ static void test(void) {
argv = {"binary_name", "-sm", "hello"};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON));
{
common_params penalty_params;
argv = {"binary_name", "--repeat-penalty", "0"};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), penalty_params, LLAMA_EXAMPLE_COMMON));
argv = {"binary_name", "--repeat-penalty", "-1"};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), penalty_params, LLAMA_EXAMPLE_COMMON));
argv = {"binary_name", "--repeat-penalty", "nan"};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), penalty_params, LLAMA_EXAMPLE_COMMON));
argv = {"binary_name", "--repeat-penalty", "inf"};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), penalty_params, LLAMA_EXAMPLE_COMMON));
argv = {"binary_name", "--repeat-penalty", "-inf"};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), penalty_params, LLAMA_EXAMPLE_COMMON));
const char * penalty_options[] = {"--frequency-penalty", "--presence-penalty"};
const char * nonfinite_values[] = {"nan", "inf", "-inf"};
for (const char * option : penalty_options) {
for (const char * value : nonfinite_values) {
argv = {"binary_name", option, value};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), penalty_params, LLAMA_EXAMPLE_COMMON));
}
}
}
// non-existence arg in specific example (--draft cannot be used outside llama-speculative)
argv = {"binary_name", "--draft", "123"};
assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_EMBEDDING));
tests/test-backend-sampler.cpp
+561
View File
@@ -8,12 +8,15 @@
#endif
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <functional>
#include <map>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
struct test_args {
@@ -761,6 +764,563 @@ static void test_backend_logit_bias_sampling(const test_params & params) {
printf("backend logit bias sampling test PASSED\n");
}
static void accept_prompt(llama_sampler * smpl, const llama_vocab * vocab, const std::string & prompt) {
const llama_token bos = llama_vocab_bos(vocab);
if (bos != LLAMA_TOKEN_NULL) {
llama_sampler_accept(smpl, bos);
}
std::vector<llama_token> tokens(64);
int32_t n_tokens = llama_tokenize(vocab, prompt.c_str(), (int32_t) prompt.size(),
tokens.data(), (int32_t) tokens.size(), false, false);
if (n_tokens < 0) {
tokens.resize(-n_tokens);
n_tokens = llama_tokenize(vocab, prompt.c_str(), (int32_t) prompt.size(),
tokens.data(), (int32_t) tokens.size(), false, false);
}
for (int32_t i = 0; i < n_tokens; ++i) {
llama_sampler_accept(smpl, tokens[i]);
}
}
static std::vector<float> decode_raw_logits(const test_params & params, const std::string & prompt) {
const int seq_id = 0;
const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(params.model.get()));
std::vector<llama_sampler_seq_config> empty_configs;
test_context ctx(params, empty_configs);
GGML_ASSERT(ctx.decode({{ seq_id, prompt }}));
float * logits = llama_get_logits_ith(ctx.ctx.get(), ctx.idx_for_seq(seq_id));
GGML_ASSERT(logits != nullptr);
return std::vector<float>(logits, logits + n_vocab);
}
static std::vector<llama_token_data> apply_cpu_sampler(
const std::vector<float> & raw_logits,
llama_sampler * sampler) {
std::vector<llama_token_data> data;
data.reserve(raw_logits.size());
for (llama_token token = 0; token < (llama_token) raw_logits.size(); ++token) {
data.push_back({ token, raw_logits[token], 0.0f });
}
llama_token_data_array cur_p = { data.data(), data.size(), -1, false };
llama_sampler_apply(sampler, &cur_p);
data.resize(cur_p.size);
return data;
}
using sampler_setup_fn = std::function<void(llama_sampler *)>;
using sampler_init_fn = std::function<llama_sampler *()>;
enum class penalties_position {
before_filter,
after_filter,
};
static void add_filter_and_penalties(
llama_sampler * chain,
const sampler_init_fn & init_filter,
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present,
penalties_position position) {
const auto add_penalties = [&]() {
llama_sampler_chain_add(chain, llama_sampler_init_penalties(
penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
};
if (position == penalties_position::before_filter) {
add_penalties();
llama_sampler_chain_add(chain, init_filter());
} else {
llama_sampler_chain_add(chain, init_filter());
add_penalties();
}
}
static llama_sampler_ptr make_sampler_chain(
const sampler_setup_fn & add_samplers,
const sampler_setup_fn & accept_history) {
llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
add_samplers(chain.get());
accept_history(chain.get());
return chain;
}
struct backend_sampler_output {
std::vector<float> logits;
std::vector<llama_token> candidates;
};
static backend_sampler_output run_backend_sampler(
const test_params & params,
const std::string & prompt,
llama_sampler * sampler) {
const int seq_id = 0;
std::vector<llama_sampler_seq_config> configs = {{ seq_id, sampler }};
test_context ctx(params, configs);
GGML_ASSERT(ctx.decode({{ seq_id, prompt }}));
llama_synchronize(ctx.ctx.get());
const int32_t idx = ctx.idx_for_seq(seq_id);
const uint32_t n_logits = llama_get_sampled_logits_count_ith(ctx.ctx.get(), idx);
const uint32_t n_candidates = llama_get_sampled_candidates_count_ith(ctx.ctx.get(), idx);
float * logits = llama_get_sampled_logits_ith(ctx.ctx.get(), idx);
llama_token * candidates = llama_get_sampled_candidates_ith(ctx.ctx.get(), idx);
GGML_ASSERT(logits != nullptr);
backend_sampler_output result;
result.logits.assign(logits, logits + n_logits);
result.candidates.resize(n_logits);
if (n_candidates == 0) {
for (uint32_t i = 0; i < n_logits; ++i) {
result.candidates[i] = (llama_token) i;
}
} else {
GGML_ASSERT(candidates != nullptr);
GGML_ASSERT(n_candidates == n_logits);
std::memcpy(result.candidates.data(), candidates, n_candidates * sizeof(llama_token));
}
return result;
}
struct sampler_comparison_output {
std::vector<llama_token_data> expected;
backend_sampler_output actual;
};
static sampler_comparison_output run_sampler_comparison(
const test_params & params,
const std::string & prompt,
const std::vector<float> & raw_logits,
const sampler_setup_fn & add_samplers,
const sampler_setup_fn & accept_history) {
llama_sampler_ptr cpu_chain = make_sampler_chain(add_samplers, accept_history);
llama_sampler_ptr backend_chain = make_sampler_chain(add_samplers, accept_history);
return {
apply_cpu_sampler(raw_logits, cpu_chain.get()),
run_backend_sampler(params, prompt, backend_chain.get()),
};
}
static std::unordered_map<llama_token, float> map_logits(const std::vector<llama_token_data> & data) {
std::unordered_map<llama_token, float> result;
result.reserve(data.size());
for (const auto & item : data) {
result[item.id] = item.logit;
}
return result;
}
struct sampler_comparison_stats {
int n_mismatch = 0;
int n_masked = 0;
float max_diff = 0.0f;
};
static sampler_comparison_stats compare_sampler_outputs(
const char * name,
const std::unordered_map<llama_token, float> & expected,
const backend_sampler_output & actual,
bool allow_extra_candidates = false) {
GGML_ASSERT(actual.logits.size() == actual.candidates.size());
sampler_comparison_stats result;
std::unordered_set<llama_token> seen;
seen.reserve(actual.candidates.size());
for (size_t i = 0; i < actual.logits.size(); ++i) {
const llama_token token = actual.candidates[i];
const float logit = actual.logits[i];
if (!seen.insert(token).second || std::isnan(logit)) {
if (result.n_mismatch < 5) {
printf("%s token %d has invalid backend output\n", name, token);
}
++result.n_mismatch;
continue;
}
const auto it = expected.find(token);
if (it == expected.end()) {
if (std::isinf(logit) && logit < 0.0f) {
++result.n_masked;
} else if (!allow_extra_candidates) {
if (result.n_mismatch < 5) {
printf("%s token %d was not masked\n", name, token);
}
++result.n_mismatch;
}
continue;
}
const float diff = fabsf(it->second - logit);
result.max_diff = std::max(result.max_diff, diff);
if (!std::isfinite(logit) || diff > 1e-3f) {
if (result.n_mismatch < 5) {
printf("%s mismatch token %d: cpu=%.6f backend=%.6f diff=%.6f\n",
name, token, it->second, logit, diff);
}
++result.n_mismatch;
}
}
for (const auto & item : expected) {
if (seen.find(item.first) == seen.end()) {
if (result.n_mismatch < 5) {
printf("%s missing backend token %d\n", name, item.first);
}
++result.n_mismatch;
}
}
printf("%s logits: max_diff=%.6f n_masked=%d n_mismatch=%d\n",
name, result.max_diff, result.n_masked, result.n_mismatch);
return result;
}
static float find_backend_logit(const backend_sampler_output & output, llama_token token) {
for (size_t i = 0; i < output.candidates.size(); ++i) {
if (output.candidates[i] == token) {
return output.logits[i];
}
}
GGML_ABORT("backend token not found");
}
static sampler_comparison_output run_penalties_comparison(
const test_params & params,
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present,
const std::string & prompt,
const std::function<void(llama_sampler *)> & extra_accept = {}) {
const auto * vocab = llama_model_get_vocab(params.model.get());
const std::vector<float> raw_logits = decode_raw_logits(params, prompt);
const auto add_samplers = [&](llama_sampler * chain) {
llama_sampler_chain_add(chain, llama_sampler_init_penalties(
penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
};
const auto accept_history = [&](llama_sampler * chain) {
accept_prompt(chain, vocab, prompt);
if (extra_accept) {
extra_accept(chain);
}
};
return run_sampler_comparison(
params, prompt, raw_logits, add_samplers, accept_history);
}
static void compare_penalties_logits(
const test_params & params,
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present,
const std::string & prompt,
const std::function<void(llama_sampler *)> & extra_accept = {}) {
const sampler_comparison_output output = run_penalties_comparison(
params, penalty_last_n, penalty_repeat, penalty_freq, penalty_present, prompt, extra_accept);
GGML_ASSERT(output.expected.size() == output.actual.logits.size());
const sampler_comparison_stats stats = compare_sampler_outputs(
"penalties", map_logits(output.expected), output.actual);
GGML_ASSERT(stats.n_masked == 0);
GGML_ASSERT(stats.n_mismatch == 0);
}
static void test_penalty_parameter_values(const test_params & params) {
struct penalty_test_case {
const char * name;
float repeat;
float frequency;
float presence;
};
const penalty_test_case cases[] = {
{ "frequency -1", 1.0f, -1.0f, 0.0f },
{ "frequency 0", 1.0f, 0.0f, 0.0f },
{ "frequency 1", 1.0f, 1.0f, 0.0f },
{ "presence -1", 1.0f, 0.0f, -1.0f },
{ "presence 0", 1.0f, 0.0f, 0.0f },
{ "presence 1", 1.0f, 0.0f, 1.0f },
{ "repeat 1", 1.0f, 0.0f, 0.0f },
};
int n_failed = 0;
for (const auto & test : cases) {
const sampler_comparison_output output = run_penalties_comparison(
params, 64, test.repeat, test.frequency, test.presence, "Hello Hello world");
GGML_ASSERT(output.expected.size() == output.actual.logits.size());
const sampler_comparison_stats stats = compare_sampler_outputs(
test.name, map_logits(output.expected), output.actual);
n_failed += stats.n_mismatch != 0;
}
GGML_ASSERT(n_failed == 0);
}
static void compare_top_k_penalties_logits(
const test_params & params,
int32_t k,
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present,
const std::string & prompt,
penalties_position position) {
const auto * vocab = llama_model_get_vocab(params.model.get());
const std::vector<float> raw_logits = decode_raw_logits(params, prompt);
const int n_vocab = (int) raw_logits.size();
GGML_ASSERT(n_vocab > k);
const sampler_init_fn init_top_k = [k]() {
return llama_sampler_init_top_k(k);
};
llama_sampler_ptr top_k(init_top_k());
const std::vector<llama_token_data> top_k_data = apply_cpu_sampler(raw_logits, top_k.get());
GGML_ASSERT(top_k_data.size() == (size_t) k);
const llama_token retained_history_token = top_k_data[0].id;
llama_token excluded_history_token = LLAMA_TOKEN_NULL;
for (llama_token token = 0; token < n_vocab; ++token) {
const auto it = std::find_if(top_k_data.begin(), top_k_data.end(), [token](const llama_token_data & data) {
return data.id == token;
});
if (it == top_k_data.end()) {
excluded_history_token = token;
break;
}
}
GGML_ASSERT(excluded_history_token != LLAMA_TOKEN_NULL);
const auto add_samplers = [&](llama_sampler * chain) {
add_filter_and_penalties(chain, init_top_k,
penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position);
};
auto accept_history = [&](llama_sampler * smpl) {
accept_prompt(smpl, vocab, prompt);
llama_sampler_accept(smpl, excluded_history_token);
llama_sampler_accept(smpl, excluded_history_token);
llama_sampler_accept(smpl, retained_history_token);
llama_sampler_accept(smpl, retained_history_token);
};
const sampler_comparison_output output = run_sampler_comparison(
params, prompt, raw_logits, add_samplers, accept_history);
GGML_ASSERT(output.expected.size() == (size_t) k);
GGML_ASSERT(output.actual.logits.size() == (size_t) k);
const std::unordered_map<llama_token, float> expected_logits = map_logits(output.expected);
if (position == penalties_position::after_filter) {
GGML_ASSERT(expected_logits.find(retained_history_token) != expected_logits.end());
GGML_ASSERT(fabsf(expected_logits.at(retained_history_token) - raw_logits[retained_history_token]) > 1e-6f);
GGML_ASSERT(expected_logits.find(excluded_history_token) == expected_logits.end());
GGML_ASSERT(std::find(output.actual.candidates.begin(), output.actual.candidates.end(),
excluded_history_token) == output.actual.candidates.end());
} else {
const std::unordered_map<llama_token, float> unpenalized_logits = map_logits(top_k_data);
bool changed = false;
for (const auto & item : expected_logits) {
const auto it = unpenalized_logits.find(item.first);
if (it == unpenalized_logits.end() || fabsf(it->second - item.second) > 1e-6f) {
changed = true;
break;
}
}
GGML_ASSERT(changed);
}
const char * name = position == penalties_position::before_filter
? "penalties top-k"
: "top-k penalties";
const sampler_comparison_stats stats = compare_sampler_outputs(
name, expected_logits, output.actual);
GGML_ASSERT(stats.n_masked == 0);
GGML_ASSERT(stats.n_mismatch == 0);
}
static void compare_masking_penalties_logits(
const test_params & params,
const char * filter_name,
const sampler_init_fn & init_filter,
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present,
const std::string & prompt,
penalties_position position,
bool allow_extra_candidates,
bool add_history = true) {
const auto * vocab = llama_model_get_vocab(params.model.get());
const std::vector<float> raw_logits = decode_raw_logits(params, prompt);
const int n_vocab = (int) raw_logits.size();
llama_sampler_ptr filter(init_filter());
const std::vector<llama_token_data> filtered_data = apply_cpu_sampler(raw_logits, filter.get());
GGML_ASSERT(!filtered_data.empty());
GGML_ASSERT(filtered_data.size() < (size_t) n_vocab);
const llama_token penalized_token = filtered_data[0].id;
std::unordered_set<llama_token> retained_tokens;
retained_tokens.reserve(filtered_data.size());
for (const auto & data : filtered_data) {
retained_tokens.insert(data.id);
}
llama_token masked_token = LLAMA_TOKEN_NULL;
for (llama_token token = 0; token < n_vocab; ++token) {
if (retained_tokens.find(token) == retained_tokens.end()) {
masked_token = token;
break;
}
}
GGML_ASSERT(masked_token != LLAMA_TOKEN_NULL);
const auto add_samplers = [&](llama_sampler * chain) {
add_filter_and_penalties(chain, init_filter,
penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position);
};
auto accept_history = [&](llama_sampler * smpl) {
if (!add_history) {
return;
}
accept_prompt(smpl, vocab, prompt);
llama_sampler_accept(smpl, penalized_token);
llama_sampler_accept(smpl, penalized_token);
llama_sampler_accept(smpl, masked_token);
llama_sampler_accept(smpl, masked_token);
};
const sampler_comparison_output output = run_sampler_comparison(
params, prompt, raw_logits, add_samplers, accept_history);
GGML_ASSERT(output.actual.logits.size() == (size_t) n_vocab);
const std::unordered_map<llama_token, float> expected_logits = map_logits(output.expected);
GGML_ASSERT(expected_logits.find(masked_token) == expected_logits.end());
if (add_history) {
if (position == penalties_position::after_filter) {
GGML_ASSERT(expected_logits.find(penalized_token) != expected_logits.end());
GGML_ASSERT(fabsf(expected_logits.at(penalized_token) - raw_logits[penalized_token]) > 1e-6f);
} else {
llama_sampler_ptr penalties(llama_sampler_init_penalties(
penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
accept_history(penalties.get());
const std::unordered_map<llama_token, float> penalized_logits =
map_logits(apply_cpu_sampler(raw_logits, penalties.get()));
GGML_ASSERT(fabsf(penalized_logits.at(penalized_token) - raw_logits[penalized_token]) > 1e-6f);
}
}
const std::string name = position == penalties_position::before_filter
? "penalties " + std::string(filter_name)
: std::string(filter_name) + " penalties";
const sampler_comparison_stats stats = compare_sampler_outputs(
name.c_str(), expected_logits, output.actual, allow_extra_candidates);
const float masked_logit = find_backend_logit(output.actual, masked_token);
GGML_ASSERT(stats.n_masked > 0);
GGML_ASSERT(std::isinf(masked_logit) && masked_logit < 0.0f);
GGML_ASSERT(stats.n_mismatch == 0);
}
static void test_backend_penalties_sampling(const test_params & params) {
printf("Testing backend penalties (repeat + freq + presence)\n");
compare_penalties_logits(params, 64, 1.1f, 0.5f, 0.25f, "Hello Hello world");
printf("Testing backend penalties with penalty_last_n > 64\n");
const auto * vocab = llama_model_get_vocab(params.model.get());
std::vector<llama_token> tokens(8);
int32_t n_tok = llama_tokenize(vocab, "a", 1, tokens.data(), (int32_t) tokens.size(), false, false);
if (n_tok < 0) {
tokens.resize(-n_tok);
n_tok = llama_tokenize(vocab, "a", 1, tokens.data(), (int32_t) tokens.size(), false, false);
}
GGML_ASSERT(n_tok > 0);
const llama_token tok = tokens[0];
compare_penalties_logits(params, 80, 1.15f, 0.1f, 0.05f, "a", [tok](llama_sampler * smpl) {
// accept_prompt already accepted BOS + one 'a'; fill the ring to n=80
for (int i = 0; i < 78; ++i) {
llama_sampler_accept(smpl, tok);
}
});
printf("Testing backend penalties without filler entries\n");
compare_penalties_logits(params, 64, 1.1f, 0.5f, 0.25f, "Hello", [](llama_sampler * smpl) {
for (llama_token token = 0; token < 64; ++token) {
llama_sampler_accept(smpl, token);
}
});
printf("Testing backend top-k followed by penalties\n");
compare_top_k_penalties_logits(params, 8, 64, 1.1f, 0.5f, 0.25f, "Hello",
penalties_position::after_filter);
printf("Testing backend penalties followed by top-k\n");
compare_top_k_penalties_logits(params, 8, 64, 1.1f, 0.5f, 0.25f, "Hello",
penalties_position::before_filter);
printf("Testing backend top-p followed by penalties\n");
compare_masking_penalties_logits(params, "top-p", []() {
return llama_sampler_init_top_p(0.9f, 0);
}, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, true);
printf("Testing backend top-p followed by penalties with a large history window\n");
compare_masking_penalties_logits(params, "top-p large-window", []() {
return llama_sampler_init_top_p(0.9f, 0);
}, 4096, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, true);
printf("Testing backend penalties followed by top-p\n");
compare_masking_penalties_logits(params, "top-p", []() {
return llama_sampler_init_top_p(0.9f, 0);
}, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::before_filter, true);
printf("Testing backend min-p followed by penalties\n");
compare_masking_penalties_logits(params, "min-p", []() {
return llama_sampler_init_min_p(0.1f, 0);
}, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, false);
printf("Testing backend penalties followed by min-p\n");
compare_masking_penalties_logits(params, "min-p", []() {
return llama_sampler_init_min_p(0.1f, 0);
}, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::before_filter, false);
printf("Testing backend top-p followed by penalties with empty history\n");
compare_masking_penalties_logits(params, "top-p empty", []() {
return llama_sampler_init_top_p(0.9f, 0);
}, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, true, false);
printf("Testing backend top-p followed by individual penalties\n");
compare_masking_penalties_logits(params, "top-p repeat", []() {
return llama_sampler_init_top_p(0.9f, 0);
}, 64, 1.1f, 0.0f, 0.0f, "Hello", penalties_position::after_filter, true);
compare_masking_penalties_logits(params, "top-p frequency", []() {
return llama_sampler_init_top_p(0.9f, 0);
}, 64, 1.0f, 0.5f, 0.0f, "Hello", penalties_position::after_filter, true);
compare_masking_penalties_logits(params, "top-p presence", []() {
return llama_sampler_init_top_p(0.9f, 0);
}, 64, 1.0f, 0.0f, 0.25f, "Hello", penalties_position::after_filter, true);
printf("Testing backend penalty parameter values\n");
test_penalty_parameter_values(params);
printf("backend penalties sampling test PASSED\n");
}
// This test verifies that it is possible to have two different backend samplers,
// one that uses the backend dist sampler, and another that uses CPU dist sampler.
static void test_backend_mixed_sampling(const test_params & params) {
@@ -1014,6 +1574,7 @@ struct backend_test_case {
static const backend_test_case BACKEND_TESTS[] = {
{ "greedy", test_backend_greedy_sampling, true },
{ "logit_bias", test_backend_logit_bias_sampling, true },
{ "penalties", test_backend_penalties_sampling, true },
{ "temp", test_backend_temp_sampling, true },
{ "temp_ext", test_backend_temp_ext_sampling, true },
{ "top_k", test_backend_top_k_sampling, true },
tools/server/server-context.cpp
+2 -1
View File
@@ -1807,7 +1807,8 @@ private:
// initialize samplers
if (task.need_sampling()) {
try {
slot.smpl.reset(common_sampler_init(model_tgt, task.params.sampling));
slot.smpl.reset(common_sampler_init(
model_tgt, task.params.sampling, (int32_t) llama_n_ctx(ctx_tgt)));
} catch (std::exception & e) {
std::string err_msg = std::string("Failed to initialize samplers: ") + e.what();
send_error(task, err_msg, ERROR_TYPE_INVALID_REQUEST);
tools/server/server.cpp
+7
View File
@@ -489,6 +489,13 @@ int llama_server(common_params & params, int argc, char ** argv) {
SRV_INF("listening on %s\n", ctx_http.listening_address.c_str());
// TODO: remove this in the future
// check the string to also handle the .sock case
if (string_ends_with(ctx_http.listening_address, ":8080")) {
SRV_WRN("%s", "NOTICE: server default port will be changed to :9931 in a future release\n");
SRV_WRN("%s", " ref: https://github.com/ggml-org/llama.cpp/pull/26508\n");
}
if (is_router_server) {
if (!params.models_preset_hf.empty()) {
SRV_WRN( "NOTE: using preset.ini from HF repo '%s'\n", params.models_preset_hf.c_str());