dspark: give tap capture its own unmasked path, avoid full-vocab lm_head on every prompt row by bri-prism · Pull Request #63 · PrismML-Eng/llama.cpp

bri-prism · 2026-07-14T01:27:14Z

What

Adds an independent masked flag to DSpark's multi-layer hidden-state tap capture (llama_set_capture_layers), separate from the existing embeddings_nextn_masked it was previously reusing. When a caller opts into masked=false, the capture tap stays dense (populated for every prompt position, matching embeddings_nextn's existing unmasked mode) independent of which rows have batch.logits set — so the harness/caller can go back to requesting logits=false on context rows (as the plain autoregressive path already does) while still getting a full per-position capture buffer for the drafter.

tests/test-dspark-real-eval.cpp is updated to use the new unmasked mode: the speculative-path prefill now sets logits=false on every context row, exactly mirroring the AR baseline.

Why

Filed and fully diagnosed in #33: DSpark speculative prefill measured up to ~2x slower than plain target prefill at realistic prompt lengths. The isolation there found the slowdown splits into two parts:

~1/3 was this harness artifact: capture previously reused embeddings_nextn_masked's hardcoded narrow-early behavior, which ties row selection at the capture tap to inp_out_ids. The only way to get a capture row for every prompt position was to mark every position as an output row (logits=true), which also forces the full-vocab lm_head projection to run on every position instead of just the sampled one.
~2/3 is the drafter's own genuine prefill forward pass (a second transformer-body pass over the same context length) — structural, not addressed by this PR.

This PR removes the harness-side third: capture no longer needs logits=true to populate every row, so the wasted per-position lm_head projection goes away without touching the drafter itself.

Hardware-dependent nuance (found independently by two follow-up investigations after #33, both confirming the same two-contributor split): the artifact's cost is bandwidth-dependent. On L40S (lower HBM bandwidth), the every-row lm_head GEMM is a real bottleneck, matching #33's ~1/3 share. On A100 (much higher bandwidth), an isolation probe found the same every-row-logits prefill came back net-neutral to positive versus the 1-row baseline — turning a GEMV into a wide GEMM is cheap or even favorable there. So this fix is a clear, real win on bandwidth-limited cards and likely close to a wash on high-bandwidth ones; it's still correct to land regardless since it strictly reduces wasted work and never costs anything (the flag defaults to current behavior for every existing caller). The second contributor (~2/3 of the original slowdown, the drafter's own genuine prefill forward pass building its own KV/GDN state) is unaffected by this PR and remains structural.

How

src/llama-cparams.h: new embeddings_capture_masked (default true, preserves current behavior for every existing caller).
src/llama-context.h / .cpp, src/llama-ext.h: llama_set_capture_layers() takes a masked parameter. Readback sizing (output_reserve), the per-decode capture copy (mirrors the existing embd_nextn masked/unmasked split), and get_embeddings_capture_ith's row-index resolution (mirrors get_embeddings_nextn_ith's existing unmasked path) all now branch on the new flag.
src/models/qwen35.cpp: the capture-tap narrowing and the last-layer/final-projection narrowing now key off the new flag (narrow_before_last_layer) instead of unconditionally reusing embeddings_nextn_masked. Added a GGML_ASSERT guard: if DSpark's dense capture and MTP's masked nextn were ever requested simultaneously in the same context (they aren't today — DSpark and MTP are alternative, mutually exclusive speculative mechanisms), this fails loudly instead of silently widening t_h_nextn and corrupting its readback offsets, since both currently share the same narrow-timing decision point.
tests/test-dspark-real-eval.cpp: engages capture with masked=false; drops the speculative-path prefill's logits=true back to false on context rows.

Scope is deliberately narrow — no change to MoE (qwen35moe.cpp never builds a capture tensor at all, so it's unaffected), no change to the drafter's own forward pass, no change to any existing caller's masked=true default behavior.

Validation

Built and ran on a fresh L40S 48GB pod, btl6l1 (binary cont6k) drafter against the binary cont6k Q1_0 target, 12 chat-templated prompts, k=4:

=== OVERALL: prompts=12 n_predicted=1172 rounds=299 drafted=1196 accepted=873 accept=0.7299 tau=3.9197 ar_tok_s=96.80 sp_tok_s=129.95 speedup=1.343 ===

accept=0.7299, tau=3.9197 matches the same drafter/target pair's reference numbers measured earlier this session on two other cards (accept=0.730, tau=3.92-3.93) within normal run-to-run noise — decode output is unaffected, as required for a change that should only touch which rows the harness marks as output, not what gets sampled. test-dspark-forward --tier1 also passes cleanly (finite logits, consistent argmax across positions).

Performance: this validation run used the existing ~25-token prompts, where the harness-artifact slowdown is only ~10-15% of total PP time (per #33's own note that "short prompts hide this"). The isolation that shows the ~1/3 recovery clearly needs #33's longer (~512-token) pp512_prompt.jsonl repro — @khosravipasha, if you still have that harness handy it'd be the fastest way to confirm the PP-side number; happy to also run it myself if useful.

Cc

@khosravipasha — following up on #33.

Copilot

Pull request overview

Adds dense DSpark capture to avoid unnecessary prompt-row lm_head projections.

Changes:

Adds independent masked/unmasked capture configuration.
Updates Qwen35 graph narrowing and capture readback.
Uses dense capture in the DSpark evaluation harness.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
`src/llama-cparams.h`	Adds capture masking state.
`src/llama-context.h`	Extends the internal capture API.
`src/llama-context.cpp`	Implements dense capture allocation and readback.
`src/llama-ext.h`	Extends the public staging API.
`src/models/qwen35.cpp`	Adjusts graph narrowing for dense capture.
`tests/test-dspark-real-eval.cpp`	Avoids context-row logits during DSpark prefill.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

bri-prism · 2026-07-14T01:41:14Z

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


Confirmed and fixed in 2f98c1a3d. That assignment predated this PR's independent capture flag — it was left over from when capture reused embeddings_nextn_masked directly, and stayed as dead weight after the flags were split. Removed it entirely; set_capture_layers() no longer touches nextn's own masked setting at all, so this assert now only fires in the genuine conflict case it's meant to catch.

bri-prism · 2026-07-14T01:41:16Z

+    // at the same point), so defer to the post-loop narrowing instead -- capture
+    // then sees the full, unnarrowed layer stream and only the final projection
+    // (result_norm + lm_head) is limited to inp_out_ids.
+    const bool capture_wants_dense = cparams.n_capture_layers > 0 && !cparams.embeddings_capture_masked;


Good catch, fixed in 2f98c1a3d. narrow_before_last_layer's capture-side deferral now only kicks in when the last layer itself is one of the requested capture_layer_idx entries. Confirmed this recovers a small real win on today's checkpoints (taps at 1/16/31/46/61, last layer is 63) — re-validated on the L40S pod, sp_tok_s went from 130.3 to 133.1 tok/s with accept/tau unchanged (0.7299/3.9197, bit-identical).

bri-prism · 2026-07-14T01:41:17Z

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


You're right that the general multi-sequence/interleaved case isn't handled — I don't have a way to validate a fix for it (no multi-sequence dspark consumer exists to test against), so rather than ship an unverified reorder fix I added a guard in 2f98c1a3d: dense (unmasked) capture now asserts ubatch.n_seqs_unq <= 1, failing loudly instead of silently returning another sequence's capture if a multi-sequence caller ever shows up. Scoping the general fix as follow-up work rather than blocking on it here, since every current consumer (dspark) is single-sequence.

bri-prism · 2026-07-14T01:41:19Z

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


Same issue as the sibling comment on get_embeddings_capture_ith — addressed by the same guard in 2f98c1a3d (GGML_ASSERT(cap_masked || ubatch.n_seqs_unq <= 1) right at this readback site). Multi-sequence dense capture now fails loudly rather than silently scrambling rows; general support is follow-up work.

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

bri-prism · 2026-07-14T01:46:05Z

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


This is guarded, not unhandled: 2f98c1a3d added GGML_ASSERT((cap_masked || ubatch.n_seqs_unq <= 1) && ...) right at this readback site (src/llama-context.cpp:2132), so a multi-sequence dense-capture ubatch aborts instead of silently packing rows in execution order. Every current consumer (dspark) is single-sequence (n_seq_max=1, one prompt per decode), where there's nothing to interleave and no reordering can occur — so the guard covers the real usage without requiring an original-index mapping I have no way to validate (no multi-sequence dspark test case exists). Treating the general multi-sequence case as follow-up work once there's an actual caller to test it against, rather than shipping an unverified fix.

bri-prism · 2026-07-14T01:46:06Z

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


output_reorder()'s swaps are driven by output_swaps, which is only populated to fix up cross-ubatch output-row order — within the single-sequence ubatches dense capture is guarded to (see the n_seqs_unq<=1 assert added in 2f98c1a3d, same reply as the sibling comment on line 2123), there's no reordering happening for dense capture to begin with, so there's nothing for output_reorder to corrupt in the actual supported case. Agreed this would need real handling if dense capture is ever extended to multi-sequence batches — tracking that as the same follow-up as the sibling comment, not blocking this PR on it.

bri-prism · 2026-07-14T01:46:02Z

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


This is already fixed, in the same commit (2f98c1a3d) this second review pass seems to have missed — grep -n "embeddings_nextn_masked = true" src/llama-context.cpp now returns nothing; set_capture_layers() no longer touches that field at all. Can you re-check against HEAD (2f98c1a3d)?

bri-prism · 2026-07-14T01:46:08Z

@@ -2103,33 +2112,44 @@ int llama_context::decode(const llama_batch & batch_inp) {
            }


Good catch, but this is pre-existing MTP nextn code this PR never touched (I only added the analogous capture path, mirroring the existing unmasked-nextn pattern including, apparently, this latent gap) — my diff doesn't add or modify anything at line 2112's nextn block. Worth its own issue since it's a separate, already-shipped feature with its own consumers to reason about; happy to file it if useful, but I'd rather not fold an unrelated pre-existing MTP fix into this PR's scope (the DSpark harness PP-slowdown fix).

…n every prompt row Gives DSpark's multi-layer hidden-state tap capture its own masked flag, separate from embeddings_nextn_masked which it was previously reusing. Opting into masked=false keeps capture dense (every prompt position) regardless of batch.logits, so callers can request logits=false on context rows (as the plain AR path already does) while still getting a full per-position capture buffer for the drafter. Mirrors the existing embeddings_nextn unmasked path (llama_context.cpp) at every layer: cparams flag, output_reserve sizing, per-decode readback offset/size, and get_embeddings_capture_ith row resolution. test-dspark-real-eval.cpp now engages capture with masked=false and drops the speculative-path prefill's logits back to false on context rows, matching the AR baseline. Fixes the harness-side third of the PP slowdown reported in #33: capture previously needed logits=true on every row just to populate a capture row for it, which forced the full-vocab lm_head projection to run on every prompt position instead of one.

- restore the default masked=true on llama_set_capture_layers's public declaration -- it was mandatory there, breaking source compat for any existing 3-arg caller. - set_capture_layers() no longer stomps embeddings_nextn_masked=true as a side effect; that assignment predated the independent capture flag and made the assert below unreachable in the exact case it exists to catch (dense capture silently overriding a caller's masked=false nextn config instead of tripping the guard). - narrow_before_last_layer's capture-side deferral now only applies when the LAST layer is actually one of the requested capture layers; taps at any earlier layer already branched off cur before this point in the loop, so deferring the last layer's own narrowing for them was an unnecessary regression (recovers a little more speed on today's real checkpoints, whose taps never include the last layer). - guard dense (unmasked) capture to single-sequence ubatches: its rows are indexed/reordered assuming raw-token order, which split_equal()'s per-sequence interleaving on a multi-sequence ubatch would violate. No current consumer is multi-sequence; fail loudly instead of silently returning another sequence's capture if that changes.

bri-prism requested a review from khosravipasha July 14, 2026 01:27

github-actions Bot added testing model labels Jul 14, 2026

khosravipasha requested a review from Copilot July 14, 2026 01:29

Copilot started reviewing on behalf of khosravipasha July 14, 2026 01:29 View session

Copilot AI reviewed Jul 14, 2026

View reviewed changes

khosravipasha requested a review from Copilot July 14, 2026 01:38

Copilot started reviewing on behalf of khosravipasha July 14, 2026 01:38 View session

Copilot AI reviewed Jul 14, 2026

View reviewed changes

bri-prism added 2 commits July 14, 2026 02:49

bri-prism force-pushed the fix/dspark-unmasked-capture branch from 2f98c1a to da9c580 Compare July 14, 2026 01:51

bri-prism merged commit 5d2aa86 into prism Jul 14, 2026
14 of 28 checks passed

khosravipasha mentioned this pull request Jul 14, 2026

superseded #66

Closed

khosravipasha deleted the fix/dspark-unmasked-capture branch July 14, 2026 05:30

Navigation Menu

Uh oh!

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

dspark: give tap capture its own unmasked path, avoid full-vocab lm_head on every prompt row#63
bri-prism merged 2 commits into
prismfrom
fix/dspark-unmasked-capture

bri-prism commented Jul 14, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

Uh oh!

bri-prism Jul 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		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;

		@@ -2103,33 +2112,44 @@ int llama_context::decode(const llama_batch & batch_inp) {
		}

Uh oh!

Conversation

bri-prism commented Jul 14, 2026

What

Why

How

Validation

Cc

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants