upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-10-01 10:27:30 -05:00
Code Issues Packages Projects Releases Wiki Activity

llama : preserve original batch order for speculative decoding layer inputs (#29019)

Browse Source 
* llama: preserve original batch order for layer inputs

Assisted-by: Codex

* tests: cover layer-input order across KV layouts

Assisted-by: Codex

* tests: exercise layer-input ordering on CUDA devices

Assisted-by: Codex

* llama: make layer input reordering compatible with tensor split

Copy each microbatch tensor from offset zero and restore original row order after synchronization. Extend the layer-input regression to cover tensor split and repeated reads and decodes.

Assisted-by: Codex

* llama: restore token order for unmasked NextN embeddings

Use the original-token mapping for unmasked NextN rows, including when
layer-input capture is disabled. Keep masked NextN rows on the logits
output mapping and preserve offset-zero tensor copies.

Extend the existing regression to cover NextN alone, combined layer
capture, and masked outputs with repeated decodes and getters.

Validation: all 256 CPU/CUDA/tensor configurations pass. Qwen3.8-27B
Q4_K_M MTP completes MT-Bench at concurrency 16 before and after.

Assisted-by: Codex

* ggml: fix WebGPU reservation and OpenVINO hidden-state capture

Reserve WebGPU vector attention scratch across batch sizes and refresh reservations when NextN capture settings change. Preserve requested OpenVINO outputs, dynamic shapes, sequence counts, and current graph bindings.

Extend existing WebGPU regression coverage and enable strict allocation checks.

Assisted-by: Codex

* llama: defer regression test and backend fixes to follow-ups

Keep this PR focused on restoring token order for layer inputs and unmasked NextN embeddings. Remove the added regression test, OpenVINO and WebGPU changes, and the separate NextN reservation change.

Assisted-by: Codex

* llama: keep n_embd declaration in its original position

Assisted-by: Codex

* llama : pass token count to layer input extraction

Assisted-by: Codex

* llama : name original batch indices batch_idxs

Assisted-by: Codex

* llama : name extracted embedding indices embd_batch_idxs

Assisted-by: Codex

* llama : tag target embedding reordering

Assisted-by: Codex

* llama : tag extraction and name the index capture flag

Assisted-by: Codex
This commit is contained in:
Hrishith Thadicherla
2026-09-30 21:17:40 +03:00
committed by GitHub
parent 4f31296a90
commit 4453b535fd
4 changed files with 45 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-batch.cpp
+1
View File
@@ -800,6 +800,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
udata->seq_idx .resize(LLAMA_MAX_SEQ, -1);
udata->output .resize(n_tokens);
udata->batch_idxs = idxs;
udata->seq_id_data.reserve(n_tokens);
seq_set_t seq_set_unq;
src/llama-batch.h
+1
View File
@@ -62,6 +62,7 @@ struct llama_ubatch {
std::vector<llama_seq_id> seq_id_unq;
std::vector<int32_t> seq_idx;
std::vector<int8_t> output;
std::vector<int32_t> batch_idxs; // original batch index for each token
std::vector<llama_seq_id> seq_id_data;
};
src/llama-context.cpp
+41 -13
View File
@@ -1799,6 +1799,7 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
n_queued_tokens += n_tokens_all;
output_swaps.clear();
embd_batch_idxs.clear();
sched_reserve();
@@ -2008,7 +2009,8 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
}
}
extract_layer_inputs(res, n_tokens_prev, ubatch.n_tokens);
// [TAG_EXTRACT_TARGET_EMBEDDINGS]
bool extract_all_idxs = extract_layer_inputs(res, n_tokens_prev, ubatch.n_tokens);
// extract nextn embeddings before
// only meaningful in LLAMA_POOLING_TYPE_NONE (per-token); other pooling modes are ignored.
@@ -2026,9 +2028,17 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
GGML_ASSERT((offset + n_rows)*n_embd <= (int64_t) embd_nextn.size);
ggml_backend_tensor_get_async(backend_h, t_h_nextn, embd_nextn_out, 0, n_rows*n_embd*sizeof(float));
extract_all_idxs = extract_all_idxs || !masked;
}
}
if (extract_all_idxs) {
GGML_ASSERT(ubatch.data && ubatch.data->batch_idxs.size() == ubatch.n_tokens);
GGML_ASSERT(embd_batch_idxs.size() == (size_t) n_tokens_prev);
const auto & batch_idxs = ubatch.data->batch_idxs;
embd_batch_idxs.insert(embd_batch_idxs.end(), batch_idxs.begin(), batch_idxs.end());
}
if (has_samplers) {
const auto stride = n_vocab;
@@ -2265,7 +2275,8 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) {
return n_outputs_max;
}
void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens) {
bool llama_context::extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens) {
bool extracted = false;
for (uint32_t il = 0; il < cparams.embeddings_layer_inp.size(); ++il) {
if (!cparams.embeddings_layer_inp[il]) {
continue;
@@ -2284,13 +2295,17 @@ void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t to
GGML_ASSERT(nfloats % n_tokens == 0);
const size_t row_floats = nfloats / n_tokens;
GGML_ASSERT(row_floats == model.hparams.n_embd);
const size_t dst_offset = token_offset * row_floats;
GGML_ASSERT(dst_offset + nfloats <= embd_layer_inp[il].size);
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched.get(), t);
GGML_ASSERT(backend != nullptr);
// Tensor-split backends require a zero source offset.
ggml_backend_tensor_get_async(backend, t, embd_layer_inp[il].data + dst_offset, 0, nbytes);
extracted = true;
}
return extracted;
}
void llama_context::output_reorder() {
@@ -2314,22 +2329,12 @@ void llama_context::output_reorder() {
}
}
if (embd_nextn.size > 0) {
if (embd_nextn.size > 0 && cparams.embeddings_nextn_masked) {
for (uint64_t k = 0; k < n_embd_out; k++) {
std::swap(embd_nextn.data[i0*n_embd_out + k], embd_nextn.data[i1*n_embd_out + k]);
}
}
if (embd_layer_inp.size() > 0) {
for (int lid = 0; lid < (int) embd_layer_inp.size(); ++lid) {
if (embd_layer_inp[lid].size > 0) {
for (uint64_t k = 0; k < n_embd; ++k) {
std::swap(embd_layer_inp[lid].data[i0*n_embd + k], embd_layer_inp[lid].data[i1*n_embd + k]);
}
}
}
}
if (!sampling.samplers.empty()) {
assert(sampling.logits.size > 0);
assert(sampling.probs.size > 0);
@@ -2359,6 +2364,29 @@ void llama_context::output_reorder() {
}
output_swaps.clear();
// [TAG_EXTRACT_TARGET_EMBEDDINGS]
// Layer inputs and unmasked NextN embeddings contain all token rows, independent of logits selection.
for (size_t i = 0; i < embd_batch_idxs.size(); ++i) {
while (embd_batch_idxs[i] != (int32_t) i) {
const int32_t j = embd_batch_idxs[i];
GGML_ASSERT(j >= 0 && (size_t) j < embd_batch_idxs.size());
if (embd_nextn.has_data() && !cparams.embeddings_nextn_masked) {
for (size_t k = 0; k < n_embd_out; ++k) {
std::swap(embd_nextn.data[i*n_embd_out + k], embd_nextn.data[j*n_embd_out + k]);
}
}
for (auto & layer : embd_layer_inp) {
if (layer.has_data()) {
for (size_t k = 0; k < n_embd; ++k) {
std::swap(layer.data[i*n_embd + k], layer.data[j*n_embd + k]);
}
}
}
std::swap(embd_batch_idxs[i], embd_batch_idxs[j]);
}
}
embd_batch_idxs.clear();
}
//
src/llama-context.h
+2 -1
View File
@@ -239,7 +239,7 @@ private:
// async-copy enabled layer-input tensors (per cparams.output_layer_inp)
// from backend into host-side embd_layer_inp buffers
void extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens);
bool extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens);
//
// graph
@@ -312,6 +312,7 @@ private:
// host buffers for output layer input embeddings, per layer
// populated when cparams.output_layer_inp[il] is true
std::vector<buffer_view<float>> embd_layer_inp;
std::vector<int32_t> embd_batch_idxs; // extracted index -> original batch index
struct sampling_info {
// !samplers.empty() to check if any samplers are active