This commit is contained in:
Chris Muzyn 2026-02-06 09:44:40 -06:00 committed by GitHub
commit 2dcd2b64ff
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 4 deletions

View File

@ -1281,6 +1281,8 @@ int llama_context::encode(const llama_batch & batch_inp) {
case LLAMA_POOLING_TYPE_CLS:
case LLAMA_POOLING_TYPE_LAST:
{
const int64_t n_embd_tensor = t_embd ? t_embd->ne[0] : hparams.n_embd_inp();
// extract sequence embeddings
auto & embd_seq_out = embd_seq;
@ -1288,8 +1290,8 @@ int llama_context::encode(const llama_batch & batch_inp) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
const int32_t seq_idx = ubatch.seq_idx[seq_id];
embd_seq_out[seq_id].resize(n_embd);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd*seq_idx)*sizeof(float), n_embd*sizeof(float));
embd_seq_out[seq_id].resize(n_embd_tensor);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd_tensor*seq_idx)*sizeof(float), n_embd_tensor*sizeof(float));
}
} break;
case LLAMA_POOLING_TYPE_RANK:
@ -1708,6 +1710,8 @@ int llama_context::decode(const llama_batch & batch_inp) {
case LLAMA_POOLING_TYPE_CLS:
case LLAMA_POOLING_TYPE_LAST:
{
const int64_t n_embd_tensor = t_embd ? t_embd->ne[0] : hparams.n_embd_inp();
// extract sequence embeddings (cleared before processing each batch)
auto & embd_seq_out = embd_seq;
@ -1715,8 +1719,8 @@ int llama_context::decode(const llama_batch & batch_inp) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
const int32_t seq_idx = ubatch.seq_idx[seq_id];
embd_seq_out[seq_id].resize(n_embd);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd*seq_idx)*sizeof(float), n_embd*sizeof(float));
embd_seq_out[seq_id].resize(n_embd_tensor);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd_tensor*seq_idx)*sizeof(float), n_embd_tensor*sizeof(float));
}
} break;
case LLAMA_POOLING_TYPE_RANK: