Commit 4453b535f for llama.cpp
commit 4453b535fd15cd5b9d5ccb956ebd38dc325c98fc
Author: Hrishith Thadicherla <99313418+hthadicherla@users.noreply.github.com>
Date: Wed Sep 30 11:17:40 2026 -0700
llama : preserve original batch order for speculative decoding layer inputs (#29019)
* 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
diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 89a1f3f37..7a17b3a4e 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -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;
diff --git a/src/llama-batch.h b/src/llama-batch.h
index 201d48cce..18e8144c8 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -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;
};
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index ba89ca794..4f04db0ee 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -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();
}
//
diff --git a/src/llama-context.h b/src/llama-context.h
index 35a7071ed..95b722628 100644
--- a/src/llama-context.h
+++ b/src/llama-context.h
@@ -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