dspark: give tap capture its own unmasked path, avoid full-vocab lm_head on every prompt row

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<int32_t> & layer_ids) {
+void llama_context::set_capture_layers(const std::vector<int32_t> & 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<int32_t> & 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<int32_t> 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) {

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<int32_t> & layer_ids);
+    void set_capture_layers(const std::vector<int32_t> & 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

src/llama-cparams.h

@@ -42,6 +42,17 @@ struct llama_cparams {
     bool     embeddings_capture = false;
     uint32_t n_capture_layers   = 0;
     std::array<int32_t, LLAMA_MAX_LAYERS> 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;

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);

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<ggml_tensor *> 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);
     }