diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 86a1a4ba18..386fab04ac 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -695,6 +695,8 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u udata->seq_idx .resize(LLAMA_MAX_SEQ, -1); udata->output .resize(n_tokens); + udata->seq_id_data.reserve(n_tokens); + seq_set_t seq_set_unq; for (size_t i = 0; i < idxs.size(); ++i) { @@ -716,11 +718,13 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u } udata->n_seq_id[i] = batch.n_seq_id[idxs[i]]; - udata->seq_id[i] = batch.seq_id[idxs[i]]; udata->output[i] = batch.logits[idxs[i]]; for (int s = 0; s < udata->n_seq_id[i]; ++s) { - seq_set_unq.set(udata->seq_id[i][s]); + const llama_seq_id seq_id = batch.seq_id[idxs[i]][s]; + + udata->seq_id_data.push_back(seq_id); + seq_set_unq.set(seq_id); } if (udata->output[i]) { @@ -728,6 +732,12 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u } } + llama_seq_id * seq_id_ptr = udata->seq_id_data.data(); + for (size_t i = 0; i < idxs.size(); ++i) { + udata->seq_id[i] = seq_id_ptr; + seq_id_ptr += udata->n_seq_id[i]; + } + for (uint32_t s = 0; s < n_seq_max; ++s) { if (seq_set_unq.test(s)) { udata->seq_idx[s] = udata->seq_id_unq.size(); diff --git a/src/llama-batch.h b/src/llama-batch.h index 209cf3699d..8e6fac0efa 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -56,13 +56,15 @@ struct llama_ubatch { std::vector embd; std::vector pos; std::vector n_seq_id; - std::vector seq_id; + std::vector seq_id; // these point into the seq_id_data below std::vector seq_id_unq; std::vector seq_idx; std::vector output; + + std::vector seq_id_data; }; - // the llama_ubatch pointers above point to this data if set. otherwise - points to non-owning data + // the llama_ubatch pointers above point to this data if set. otherwise - point to external non-owning data std::shared_ptr data; };