diff --git a/common/arg.cpp b/common/arg.cpp index 6afac8bd9c..071b5c9f36 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -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( diff --git a/common/common.cpp b/common/common.cpp index ff27d392fb..c941fd505a 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -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()) }; } diff --git a/common/sampling.cpp b/common/sampling.cpp index 256ac161e2..5698c0263b 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -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; diff --git a/common/sampling.h b/common/sampling.h index 4191988bb8..91e2cea787 100644 --- a/common/sampling.h +++ b/common/sampling.h @@ -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); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 7f4e252dca..f6fb91798c 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -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); diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 33be16dc5c..d27d8acb1d 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -627,7 +627,8 @@ template struct block_reduce_policy { }; template -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(val); const unsigned int block_size = block_size_template == 0 ? blockDim.x : block_size_template; if (block_size > WARP_SIZE) { diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 09d9f3a7d6..c3758cd50c 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -64,7 +64,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr tmp += xi * xi; } - tmp = block_reduce(tmp, s_sum); + tmp = block_reduce(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<<>>(x, dst, group_size, ne_elements, eps); } else { const dim3 block_dims(1024, 1, 1); - group_norm_f32<1024><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>(x, dst, group_size, ne_elements, eps); + group_norm_f32<1024><< WARP_SIZE ? 2 * 32 * sizeof(float): 0, stream>>>(x, dst, group_size, ne_elements, eps); } } diff --git a/ggml/src/ggml-cuda/softmax.cu b/ggml/src/ggml-cuda/softmax.cu index 285c0e9543..f320c6f004 100644 --- a/ggml/src/ggml-cuda/softmax.cu +++ b/ggml/src/ggml-cuda/softmax.cu @@ -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(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(local_max, shared_vals); + local_max = block_reduce(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(local_max, shared_vals); + local_max = block_reduce(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(tmp_expf, shared_vals); + tmp_expf = block_reduce(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(tmp_expf, shared_vals); + tmp_expf = block_reduce(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); } } diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 915c8e7b90..fc0fce0d78 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -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; } diff --git a/include/llama.h b/include/llama.h index 6e53e22972..f2d7e38858 100644 --- a/include/llama.h +++ b/include/llama.h @@ -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( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 320784c3a8..24f05cc916 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -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 diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index e12a8cdc2a..9dde345df4 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -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(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(mctx); + + auto inp = std::make_unique(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); diff --git a/src/llama-graph.h b/src/llama-graph.h index 160e294135..32d8d395aa 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -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, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 50af97f358..846d4c69a6 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -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 diff --git a/src/llama-hparams.h b/src/llama-hparams.h index fc770bf003..6e8336c987 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -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; diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 241c50365a..96cb045d2e 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -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( 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( 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) { diff --git a/src/llama-kv-cache-dsa.h b/src/llama-kv-cache-dsa.h index e2b330993b..e74fc4d910 100644 --- a/src/llama-kv-cache-dsa.h +++ b/src/llama-kv-cache-dsa.h @@ -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; diff --git a/src/llama-kv-cache-msa.cpp b/src/llama-kv-cache-msa.cpp new file mode 100644 index 0000000000..55ef286caf --- /dev/null +++ b/src/llama-kv-cache-msa.cpp @@ -0,0 +1,395 @@ +#include "llama-kv-cache-msa.h" + +#include "llama-impl.h" +#include "llama-batch.h" +#include "llama-model.h" + +#include +#include +#include + +// 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( + 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( + 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 llama_kv_cache_msa::memory_breakdown() const { + std::map 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 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( + this, std::move(sinfos_base), std::move(sinfos_idx), std::move(ubatches)); + } while (false); + + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); +} + +llama_memory_context_ptr llama_kv_cache_msa::init_full() { + return std::make_unique(this); +} + +llama_memory_context_ptr llama_kv_cache_msa::init_update(llama_context * lctx, bool optimize) { + return std::make_unique(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 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(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(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 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; + } + } +} diff --git a/src/llama-kv-cache-msa.h b/src/llama-kv-cache-msa.h new file mode 100644 index 0000000000..f09b6d32b0 --- /dev/null +++ b/src/llama-kv-cache-msa.h @@ -0,0 +1,153 @@ +#pragma once + +#include "llama-kv-cache.h" + +#include + +// 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 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 kv_base; + std::unique_ptr 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 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 ubatches; + + const llama_memory_context_ptr ctx_base; + const llama_memory_context_ptr ctx_idx; + + const llama_memory_status status; +}; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 44cb1668da..8678a326d9 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -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 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::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); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index d5a92f4405..6cb6dbd2f9 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -164,6 +164,8 @@ public: std::vector 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 k_stream; std::vector v_stream; - std::vector 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 diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4b4fe4712c..8fff1a4326 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -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; diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index a9cb6bee5f..b2f1abe737 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -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 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 host_token_ids; + std::vector 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> 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(penalty_last_n), - /* .token_count = */ {}, - } + /* .ctx = */ new llama_sampler_penalties( + penalty_last_n, + penalty_repeat, + penalty_freq, + penalty_present) ); } diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 443cd46408..10032a8c64 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -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(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; diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index 0773ad5435..854d5aed0f 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -1,5 +1,5 @@ #include "models.h" -#include "llama-kv-cache.h" +#include "llama-kv-cache-msa.h" #include #include #include @@ -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 llama_model_minimax_m3::build_arch_graph(cons return std::make_unique(*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 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 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(params.mctx); + const auto * mctx_new = static_cast(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(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(mctx_msa, blk, mm.msa_p.local); - auto loc = std::make_unique(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(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]; diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 1d3584f903..fd5adb740e 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -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)); diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index c24076e313..1a46468ba2 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -8,12 +8,15 @@ #endif #include +#include #include #include #include +#include #include #include #include +#include #include 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 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 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 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(logits, logits + n_vocab); +} + +static std::vector apply_cpu_sampler( + const std::vector & raw_logits, + llama_sampler * sampler) { + std::vector 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; +using sampler_init_fn = std::function; + +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 logits; + std::vector 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 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 expected; + backend_sampler_output actual; +}; + +static sampler_comparison_output run_sampler_comparison( + const test_params & params, + const std::string & prompt, + const std::vector & 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 map_logits(const std::vector & data) { + std::unordered_map 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 & 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 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 & extra_accept = {}) { + const auto * vocab = llama_model_get_vocab(params.model.get()); + const std::vector 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 & 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 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 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 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 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 raw_logits = decode_raw_logits(params, prompt); + const int n_vocab = (int) raw_logits.size(); + llama_sampler_ptr filter(init_filter()); + const std::vector 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 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 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 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 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 }, diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 4655b518e2..5d2798cc14 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -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); diff --git a/tools/server/server.cpp b/tools/server/server.cpp index bf403d27f8..1b2e6edb4e 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -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());