diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0c90ad3021fc..5f43818278a6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -996,8 +996,16 @@ float * llama_context::get_embeddings_capture_ith(int32_t i) { const uint32_t n_embd = model.hparams.n_embd; const uint32_t row = n_cap * n_embd; // width of one concatenated row - // capture rows always follow the masked (output-row) layout, mirroring the - // pre-norm masked path: the buffer holds one row per output position. + if (!cparams.embeddings_capture_masked) { + // unmasked: capture rows are stored densely, indexed by raw token + // position, mirroring get_embeddings_nextn_ith's unmasked path. + if (i < 0 || (size_t) (i + 1) * row > embd_capture.size) { + throw std::runtime_error(format("out of range [0, %zu)", embd_capture.size / row)); + } + return embd_capture.data + (size_t) i * row; + } + + // masked (default): the buffer holds one row per output position. const int64_t j = output_resolve_row(i); if (j < 0 || (size_t)(j + 1) * row > embd_capture.size) { throw std::runtime_error(format("out of range [0, %zu)", embd_capture.size / row)); @@ -1196,11 +1204,12 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) { cparams.embeddings_nextn_masked = masked; } -void llama_context::set_capture_layers(const std::vector & layer_ids) { +void llama_context::set_capture_layers(const std::vector & layer_ids, bool masked) { // reset - cparams.embeddings_capture = false; - cparams.n_capture_layers = 0; - cparams.capture_layer_idx = {}; + cparams.embeddings_capture = false; + cparams.n_capture_layers = 0; + cparams.capture_layer_idx = {}; + cparams.embeddings_capture_masked = masked; // enabling/disabling capture adds/removes the t_h_capture node from the // graph (see llm_graph_result::set_outputs()), so the scheduler's @@ -1236,8 +1245,6 @@ void llama_context::set_capture_layers(const std::vector & layer_ids) { cparams.n_capture_layers = n; cparams.embeddings_capture = n > 0; - // capture rows reuse the masked output-row layout; force masked extraction on. - cparams.embeddings_nextn_masked = true; } void llama_context::set_dspark_ctx( @@ -2103,33 +2110,57 @@ int llama_context::decode(const llama_batch & batch_inp) { } } - // extract multi-layer capture embeddings, concatenated per output position. - // capture always uses the masked (output-row) layout, so t_h_capture is - // [n_capture * n_embd, n_outputs]; copy in one shot per ubatch. - if (embd_capture.data && cparams.n_capture_layers > 0 && n_outputs > 0 && + // extract multi-layer capture embeddings, concatenated per position. + // masked (default): t_h_capture is [n_capture * n_embd, n_outputs], one row + // per output position. unmasked: t_h_capture is dense, one row per raw + // ubatch token regardless of batch.logits -- mirrors the t_h_nextn + // masked/unmasked split above. + { + const bool cap_masked = cparams.embeddings_capture_masked; + const int64_t n_rows_cap = cap_masked ? n_outputs : (int64_t) ubatch.n_tokens; + const int64_t offset_cap = cap_masked ? n_outputs_prev : n_tokens_prev; + + // Dense (unmasked) rows are stored and later indexed in raw ubatch-token + // order (get_embeddings_capture_ith(i) reads row i directly), and + // output_reorder()'s swap list is built to fix up *output-row* order -- + // neither accounts for split_equal()'s per-sequence interleaving of a + // multi-sequence ubatch (llama-batch.cpp), so a dense capture row for + // token i could come from the wrong sequence, or get scrambled by an + // output-row swap meant for a different token. Every current dense + // capture consumer (dspark) is single-sequence; fail loudly rather than + // silently return another sequence's capture if that ever changes. + GGML_ASSERT((cap_masked || ubatch.n_seqs_unq <= 1) && + "dense (unmasked) capture is only validated for single-sequence ubatches; " + "multi-sequence interleaving is not accounted for in its row ordering"); + + if (embd_capture.data && cparams.n_capture_layers > 0 && n_rows_cap > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { - ggml_tensor * t_cap = res->get_h_capture(); - const size_t row = (size_t) cparams.n_capture_layers * hparams.n_embd; - float * embd_capture_out = embd_capture.data + (size_t) n_outputs_prev * row; - GGML_ASSERT((n_outputs_prev + n_outputs)*(int64_t) row <= (int64_t) embd_capture.size); - if (t_cap) { - ggml_backend_t backend_c = ggml_backend_sched_get_tensor_backend(sched.get(), t_cap); - GGML_ASSERT(backend_c != nullptr); - ggml_backend_tensor_get_async(backend_c, t_cap, embd_capture_out, 0, n_outputs*row*sizeof(float)); - } else { - // capture was requested (n_capture_layers > 0) but this model's - // graph never produced a capture tensor -- only qwen35 builds it. - // output_reserve() already allocated embd_capture, so zero the - // rows for this ubatch rather than leave uninitialized memory that - // llama_get_embeddings_capture*() would hand back. Warn once so the - // misconfiguration (capture on an unsupported arch) is visible. - static bool warned_no_capture = false; - if (!warned_no_capture) { - LLAMA_LOG_WARN("%s: capture layers were requested but this architecture does not " - "produce capture embeddings; returning zeros\n", __func__); - warned_no_capture = true; + ggml_tensor * t_cap = res->get_h_capture(); + const size_t row = (size_t) cparams.n_capture_layers * hparams.n_embd; + float * embd_capture_out = embd_capture.data + (size_t) offset_cap * row; + GGML_ASSERT((offset_cap + n_rows_cap) * (int64_t) row <= (int64_t) embd_capture.size); + if (t_cap) { + ggml_backend_t backend_c = ggml_backend_sched_get_tensor_backend(sched.get(), t_cap); + GGML_ASSERT(backend_c != nullptr); + ggml_backend_tensor_get_async(backend_c, t_cap, embd_capture_out, 0, + n_rows_cap * row * sizeof(float)); + } else { + // capture was requested (n_capture_layers > 0) but this model's + // graph never produced a capture tensor -- only qwen35 builds it. + // output_reserve() already allocated embd_capture, so zero the + // rows for this ubatch rather than leave uninitialized memory that + // llama_get_embeddings_capture*() would hand back. Warn once so the + // misconfiguration (capture on an unsupported arch) is visible. + static bool warned_no_capture = false; + if (!warned_no_capture) { + LLAMA_LOG_WARN( + "%s: capture layers were requested but this architecture does not " + "produce capture embeddings; returning zeros\n", + __func__); + warned_no_capture = true; + } + memset(embd_capture_out, 0, n_rows_cap * row * sizeof(float)); } - memset(embd_capture_out, 0, n_outputs*row*sizeof(float)); } } @@ -2247,6 +2278,12 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) { embd_nextn.size = (size_t) n_embd_out * n_batch; } + if (has_embd_capture && !cparams.embeddings_capture_masked) { + // unmasked: same as embeddings_nextn above -- a capture row exists for + // every token in the batch, not just output rows, so size by token count. + embd_capture.size = (size_t) cparams.n_capture_layers * model.hparams.n_embd * n_batch; + } + // Allocate backend sampling output buffers if there are backend samplers configured. const bool has_sampling = !sampling.samplers.empty(); if (has_sampling) { @@ -3989,13 +4026,13 @@ float * llama_get_embeddings_nextn_ith(llama_context * ctx, int32_t i) { // multi-layer hidden-state tap C API (staging) ------------------------------- -void llama_set_capture_layers(llama_context * ctx, const int32_t * layer_ids, size_t n_layers) { +void llama_set_capture_layers(llama_context * ctx, const int32_t * layer_ids, size_t n_layers, bool masked) { std::vector ids; ids.reserve(n_layers); for (size_t i = 0; i < n_layers; ++i) { ids.push_back(layer_ids[i]); } - ctx->set_capture_layers(ids); + ctx->set_capture_layers(ids, masked); } uint32_t llama_get_n_capture(llama_context * ctx) { diff --git a/src/llama-context.h b/src/llama-context.h index 656d7dcf652b..de93d0c8eaf2 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -122,7 +122,7 @@ struct llama_context { // register the ordered set of intermediate layers to capture. pass an empty // list to disable. the concatenation order follows the order of layer_ids. - void set_capture_layers(const std::vector & layer_ids); + void set_capture_layers(const std::vector & layer_ids, bool masked = true); // dspark drafter: stage the target-tap context window consumed by the next // decode() call. feat is [n_ctx_rows * n_embd_cap] row-major (row i is diff --git a/src/llama-cparams.h b/src/llama-cparams.h index cb42fdd9c0d9..1a9024025a18 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -42,6 +42,17 @@ struct llama_cparams { bool embeddings_capture = false; uint32_t n_capture_layers = 0; std::array capture_layer_idx = {}; + // if true (default), the capture tap is narrowed to output rows (batch.logits + // != 0) at the tap point itself, same as embeddings_nextn_masked -- cheap when + // few rows need a capture row, but forces every captured row to also be an + // output row, so requesting capture on every prompt position (e.g. to condition + // a speculative drafter) also forces the final norm + lm_head to run on every + // one of those rows. If false, the tap stays full-width through the rest of the + // layer stack and the output-row narrowing is deferred until just before lm_head + // (mirrors embeddings_nextn's own masked=false path), so a caller can request a + // dense per-position capture while still keeping batch.logits (and therefore the + // lm_head projection) narrow. + bool embeddings_capture_masked = true; bool causal_attn; bool offload_kqv; diff --git a/src/llama-ext.h b/src/llama-ext.h index e5de37dbbf60..59cf21fbb26f 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -114,7 +114,18 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); // This is the shared primitive both EAGLE3-proper and dspark consume: where the // pre-norm path above exposes one final-layer hidden vector, this exposes an // arbitrary set of intermediate layers in one concatenated row. -LLAMA_API void llama_set_capture_layers(struct llama_context * ctx, const int32_t * layer_ids, size_t n_layers); +// If masked == true (default), capture is narrowed to output rows (batch.logits +// != 0) at the tap point -- requesting a capture row for every position therefore +// also forces every one of those rows through the final norm + lm_head. If +// masked == false, capture stays dense (every position, regardless of +// batch.logits) and the output-row narrowing is deferred to just before lm_head, +// so batch.logits can stay narrow (e.g. only the sampled row) while still getting +// a full per-position capture buffer -- avoids the wasted full-vocab projection +// on rows that are only needed for their capture features, not their logits. +LLAMA_API void llama_set_capture_layers(struct llama_context * ctx, + const int32_t * layer_ids, + size_t n_layers, + bool masked = true); LLAMA_API uint32_t llama_get_n_capture(struct llama_context * ctx); // mirrors llama_get_embeddings_nextn / _ith LLAMA_API float * llama_get_embeddings_capture (struct llama_context * ctx); diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 37dfc51bc8d5..195062d47b3b 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -158,6 +158,43 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para // capture order, then concatenate them along dim0 after the layer loop. std::vector h_capture(cparams.n_capture_layers, nullptr); + // The last-layer residual stream and the final norm+lm_head can only be + // narrowed to output rows (inp_out_ids) *before* the last layer runs if every + // active consumer of the un-narrowed rows agrees to that -- i.e. embeddings_nextn + // wants the masked (narrow-early) layout, AND capture is either inactive or + // also wants it narrowed at the tap point. Capture only cares about this at all + // if the LAST layer itself is one of the requested capture layers -- taps at any + // earlier layer have already branched off `cur` before this point in the loop + // (see the per-layer tap below), so narrowing the last layer's own compute doesn't + // touch them. If a dense (embeddings_capture_masked == false) tap of the last + // layer specifically is requested while embeddings_nextn_masked is on, narrowing + // early would clip that dense row too (they share inp_out_ids at the same point), + // so defer to the post-loop narrowing instead -- capture then sees the full, + // unnarrowed last-layer output and only the final projection (result_norm + + // lm_head) is limited to inp_out_ids. + bool capture_taps_last_layer = false; + for (uint32_t c = 0; c < cparams.n_capture_layers; ++c) { + if (cparams.capture_layer_idx[c] == n_layer - 1) { + capture_taps_last_layer = true; + break; + } + } + const bool capture_wants_dense = capture_taps_last_layer && !cparams.embeddings_capture_masked; + + // t_h_nextn's readback (llama-context.cpp) trusts embeddings_nextn_masked to know + // whether t_h_nextn is narrow (n_outputs rows) or full-width (ubatch.n_tokens rows); + // t_h_nextn is assigned from `cur` right after this loop, so it inherits whatever + // narrow_before_last_layer decided. If nextn is simultaneously active and asked for + // the narrow layout, letting capture's dense request silently widen `cur` here would + // widen t_h_nextn too without nextn's own readback knowing -- wrong offsets, not just + // wrong rows. DSpark capture and MTP nextn are never engaged together in practice; fail + // loudly instead of silently corrupting nextn's output if that assumption is ever broken. + GGML_ASSERT(!(capture_wants_dense && cparams.embeddings_nextn && cparams.embeddings_nextn_masked) && + "dspark dense capture (embeddings_capture_masked=false) is incompatible with simultaneous " + "masked MTP nextn extraction -- they share the same narrow-timing decision"); + + const bool narrow_before_last_layer = cparams.embeddings_nextn_masked && !capture_wants_dense; + // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass. for (int il = 0; il < n_layer; ++il) { ggml_tensor * inpSA = inpL; @@ -176,7 +213,7 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il); } - if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { + if (il == n_layer - 1 && inp_out_ids && narrow_before_last_layer) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -214,7 +251,7 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para for (uint32_t c = 0; c < cparams.n_capture_layers; ++c) { if (cparams.capture_layer_idx[c] == il) { ggml_tensor * cap = cur; - if (cparams.embeddings_nextn_masked && inp_out_ids) { + if (cparams.embeddings_capture_masked && inp_out_ids) { cap = ggml_get_rows(ctx0, cap, inp_out_ids); } cb(cap, "h_capture", il); @@ -249,7 +286,7 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para ggml_build_forward_expand(gf, cap); } - if (!cparams.embeddings_nextn_masked && inp_out_ids) { + if (!narrow_before_last_layer && inp_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); } diff --git a/tests/test-dspark-real-eval.cpp b/tests/test-dspark-real-eval.cpp index 596d70970e65..349738aacc1c 100644 --- a/tests/test-dspark-real-eval.cpp +++ b/tests/test-dspark-real-eval.cpp @@ -286,8 +286,12 @@ int main(int argc, char ** argv) { // engage the target-layer tap capture ONCE, permanently, on the target // context (see file header comment -- this is the missing piece no - // existing CLI/server path wires up for dspark). - llama_set_capture_layers(ctx_tgt, target_layers.data(), target_layers.size()); + // existing CLI/server path wires up for dspark). masked=false: capture stays + // dense (every prompt position) independent of batch.logits, so the prefill + // below can request logits=false on context rows like the AR baseline does, + // instead of paying a full-vocab lm_head projection on every prompt position + // just to get a capture row for it (see PrismML-Eng/llama.cpp-private#33). + llama_set_capture_layers(ctx_tgt, target_layers.data(), target_layers.size(), /* masked = */ false); common_params_speculative sparams; sparams.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK }; @@ -368,7 +372,8 @@ int main(int argc, char ** argv) { // overhead, i.e. it is genuinely comparable to plain decoding, not // "dspark plumbing with drafting turned off". === llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, 0, -1); - llama_set_capture_layers(ctx_tgt, nullptr, 0); + llama_set_capture_layers(ctx_tgt, nullptr, 0, + /* masked = */ true); // disabling (n_layers=0); masked value unused const auto t_ar0 = std::chrono::steady_clock::now(); @@ -406,7 +411,7 @@ int main(int argc, char ** argv) { // === DSpark pass (existing logic below, now timed) === llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, 0, -1); - llama_set_capture_layers(ctx_tgt, target_layers.data(), target_layers.size()); + llama_set_capture_layers(ctx_tgt, target_layers.data(), target_layers.size(), /* masked = */ false); const auto t_sp0 = std::chrono::steady_clock::now(); @@ -414,12 +419,17 @@ int main(int argc, char ** argv) { common_sampler_ptr smpl(common_sampler_init(model_tgt, sparams_smpl)); - // manual prefill with per-row logits requested (llama_batch_get_one() - // only requests the last row -- dspark's process() needs a capture - // row for EVERY prompt position, see common/speculative.cpp). + // manual prefill, logits=false on every context row (same as the AR + // baseline above) -- none of these rows are sampled from here (id_last is + // staged separately via dp.id_last below and verified in its own batch), + // and dense capture (masked=false, set above) no longer needs logits=true + // to populate a capture row for every position. Previously this requested + // logits=true on every row solely to get a capture row for it, which forced + // the full-vocab lm_head projection ~n_prompt_tokens times instead of the + // AR baseline's 1 -- see PrismML-Eng/llama.cpp-private#33. common_batch_clear(batch_tgt); for (size_t i = 0; i < prompt_tgt.size(); ++i) { - common_batch_add(batch_tgt, prompt_tgt[i], (llama_pos) i, { seq_id }, /* logits = */ true); + common_batch_add(batch_tgt, prompt_tgt[i], (llama_pos) i, { seq_id }, /* logits = */ false); } if (llama_decode(ctx_tgt, batch_tgt) != 0) fail("prefill decode failed for prompt " + std::to_string(pi)); if (!common_speculative_process(spec, batch_tgt)) fail("common_speculative_process (prefill) failed for prompt " + std::to_string(pi));