happyz
/
llama.cpp
mirror of https://github.com/ggerganov/llama.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

mtp-batch(fix): avoid logits for mtp kv cache operations

This commit is contained in:
samuel 2025-10-16 13:42:31 -03:00
parent 0127c6beeb
commit cae85fe531
1 changed files with 17 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
src/llama-context.cpp
View File

@ -1155,16 +1155,25 @@ int llama_context::decode(const llama_batch & batch_inp) {
// extract logits
if (t_logits && n_outputs > 0) {
ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits);
GGML_ASSERT(backend_res != nullptr);
GGML_ASSERT(logits != nullptr);
// MTP operations that are purely for updating the KV cache
// (MTP_OP_WARMUP and MTP_OP_UPDATE_ACCEPTED) also produce a logit tensor
// as a side effect of running the graph. If these logits are copied
// back to the main context buffer, they will overwrite the valid logits
// produced by the main model's pass, leading to incorrect sampling.
// This condition explicitly prevents that copy for cache-only operations.
if (batch_inp.mtp_params.op_type != MTP_OP_WARMUP &&
batch_inp.mtp_params.op_type != MTP_OP_UPDATE_ACCEPTED) {
ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits);
GGML_ASSERT(backend_res != nullptr);
GGML_ASSERT(logits != nullptr);
float * logits_out = logits + n_outputs_prev*n_vocab;
float * logits_out = logits + n_outputs_prev*n_vocab;
if (n_outputs) {
GGML_ASSERT( n_outputs_prev + n_outputs <= n_outputs_all);
GGML_ASSERT((n_outputs_prev + n_outputs)*n_vocab <= (int64_t) logits_size);
ggml_backend_tensor_get_async(backend_res, t_logits, logits_out, 0, n_outputs*n_vocab*sizeof(float));
if (n_outputs) {
GGML_ASSERT( n_outputs_prev + n_outputs <= n_outputs_all);
GGML_ASSERT((n_outputs_prev + n_outputs)*n_vocab <= (int64_t) logits_size);
ggml_backend_tensor_get_async(backend_res, t_logits, logits_out, 0, n_outputs*n_vocab*sizeof(float));
}
}
}