Skip to content

speculative: take the draft block layout from the draft model, not the CLI type - #122

Open
bri-prism wants to merge 2 commits into
prism-v7from
dspark-layout-autodetect-v7
Open

speculative: take the draft block layout from the draft model, not the CLI type#122
bri-prism wants to merge 2 commits into
prism-v7from
dspark-layout-autodetect-v7

Conversation

@bri-prism

Copy link
Copy Markdown

Lands the draft block-layout fix on prism-v7. Replaces #114, which was closed: it was broken as pushed and based on prism-v6, now marked do-not-build.

The bug

A DSpark-lineage drafter reads its drafted block anchor-first, predicting from row 0. A DFlash drafter puts the anchor token in row 0 and predicts from row 1. Which convention gets used was decided solely by --spec-type, so running a DSpark drafter as draft-dflash took the DFlash branch and read every drafted token one row late.

Nothing errored. The drafter loaded, drafted at full speed, and produced tokens the target rejected almost every time. Because rejected drafts are simply discarded, the generated text stayed correct and only throughput moved, so this reads as "speculative decoding isn't helping much" rather than as a bug.

Measured on one RTX 4090, same binary, same weights, same prompts:

block layout acceptance mean accepted len
dflash (as requested by the flag) 0.272% 1.01
dspark (as the model requires) 51.385% 3.04

The fix

Derive the layout from the draft model instead of the requested type.

These files declare general.architecture = dflash whichever lineage they come from, so the type can't identify them. The Markov head can: src/models/dflash.cpp defines DSpark as "DFlash + a semi-autoregressive Markov head and Confidence head" and creates that head as required whenever the tensor is present. It is already loaded by the time the speculator is constructed, so dspark_markov_w1 != nullptr is exactly the loader's own definition of the lineage, available for free.

  • llama_model_has_dspark_markov_head() in src/llama-ext.h / src/llama-model.cpp
  • common/speculative.cpp sets is_dspark from the model, and warns when --spec-type disagrees rather than silently honoring a request that cannot work
  • the resolved layout is now in the init log next to sample_from_anchor, so the trap is visible without reading source

Validation: with the fix applied and the server launched with the wrong flag (--spec-type draft-dflash, the exact invocation that produced 0.272%), acceptance is 51.385% and the warning appears in the log. The failure mode is now unreachable through flag misuse rather than merely avoidable.

Acceptance-floor test

test_draft_acceptance_floor in tools/server/tests/unit/test_speculative.py.

The existing tests assert that drafting happens (draft_n > 0), never that anything is accepted. Those are different failures, and a 185x acceptance collapse passed every test in the file. The fixture drafter is deliberately mismatched, so its acceptance is only 15.0% and 17.4% on the two prompts (deterministic at temperature 0); the floor sits at 5%, well under that and well over a collapse. I red-checked it by raising the floor above the measured rate and confirming it fails with a message naming the rate and counts.

It covers the draft-simple path only. No tiny fixture model exercises the block-layout-sensitive draft types, so this guards the general "acceptance quietly went to zero" class rather than this specific bug.

Deliberately not included

Shifting the target layer tap by one (input-of-layer-k vs output-of-layer-k) is a real discrepancy against the private implementation and worth fixing on its own merits, but it is not this bug: measured, it moves acceptance by less than a point (50.456% vs 51.385%), inside run-to-run noise. I left it out rather than pass an unvalidated change alongside a proven one.

Note on formatting

The repo's .clang-format disagrees with the committed style of common/speculative.cpp, including the constructor signature I only touched to remove one initializer. Running it would bury a four-line fix in unrelated reformatting, so the added lines follow the file's existing style instead.

…e CLI type

A DSpark-lineage drafter reads its drafted block anchor-first; a DFlash drafter
puts the anchor token in row 0 and predicts from row 1. That choice was driven
solely by --spec-type, so running a DSpark drafter as draft-dflash took the
DFlash branch and read every drafted token one row late. Nothing errored: the
drafter loaded, drafted at full speed, and produced tokens the target rejected
almost every time. Output stayed correct, so only throughput showed it.

Measured on one RTX 4090, same binary, same weights, same prompts:

  dflash layout (as requested): 0.272% acceptance, mean len 1.01
  dspark layout (as required):  51.385% acceptance, mean len 3.04

These files declare general.architecture = dflash whichever lineage they come
from, so the requested type cannot be trusted here. The Markov head is the
on-disk marker of the DSpark lineage and is already loaded, so derive the
layout from the model and warn when the requested type disagrees. The resolved
layout now appears in the init log next to sample_from_anchor.

Adds an acceptance-floor test. The existing coverage asserts that drafting
happens, not that anything is accepted, which is why a 185x collapse passed
every test in the file. The fixture drafter is deliberately mismatched so its
acceptance is only 15-17%; the floor sits at 5%, well under that and well over
a collapse.
@bri-prism

Copy link
Copy Markdown
Author

Correcting the scope of the validation note above, and flagging something a reviewer should know before trusting the detection.

The 51.385% number was measured on a build of the v6-lineage tree with this patch applied, not on this branch. I tried to reproduce it here and could not: the DSpark drafter checkpoint I have does not load on prism-v7 at all.

llama_model_load: error loading model: done_getting_tensors: wrong number of tensors; expected 79, got 75

The file carries 79 tensors and the loader creates 75. The four it does not create are exactly markov_w1.weight, markov_w2.weight, conf_proj.weight and conf_proj.bias, and the "DFlash with DSpark markov head" line never appears. So load_arch_tensors is not taking the if (markov_meta) branch even though the file contains a tensor named exactly markov_w1.weight. This reproduces on a clean build of prism-v7 with this patch, and this patch adds no tensors, so it is pre-existing on the branch rather than something the change introduces.

That matters here beyond being an unrelated bug. The detection in this PR reads model->dspark_markov_w1, so anything that stops the loader from populating that pointer also stops the detection from firing, and a genuine DSpark drafter would be classified as DFlash. That is the same wrong-layout outcome this PR exists to prevent, reached by a different route. On the v6-lineage tree the head does load and the detection works, which is what the 51.385% result demonstrates.

Where that leaves the change:

  • compiles clean on this branch, CUDA SM89 and Metal
  • all seven tests in tools/server/tests/unit/test_speculative.py pass on this branch, including the new acceptance floor
  • the auto-detect has not been exercised on this branch against a real DSpark drafter, because no such checkpoint currently loads on it

I would rather land the layout fix and the acceptance floor now, since both are correct on their own terms and the floor is what makes a future collapse visible at all, and treat the tensor-count mismatch as its own issue. But if you would rather hold this until a DSpark checkpoint loads on prism-v7 and the auto-detect can be confirmed end to end here, that is reasonable and I am happy to wait.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Derives speculative draft layout from the loaded model rather than the requested CLI type.

Changes:

  • Adds DSpark Markov-head detection.
  • Resolves and logs DFlash/DSpark decoding behavior.
  • Adds a speculative acceptance-rate regression test.

Reviewed changes

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

File Description
common/speculative.cpp Selects draft semantics from model tensors and reports mismatches.
src/llama-model.cpp Implements Markov-head detection.
src/llama-ext.h Exposes the detection API.
tools/server/tests/unit/test_speculative.py Adds an acceptance floor test.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread src/llama-ext.h Outdated
Comment on lines +127 to +129
// returns true if the draft model carries a DSpark Markov head, which implies the
// anchor-first block layout -- the file says general.architecture = dflash either way,
// so this is the only reliable way to tell the two lineages apart

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct, fixed in c4c4667. The comment now only claims what the head establishes, which is the lineage, and says both lineages declare general.architecture = dflash so this is how to tell them apart. It no longer mentions layout.

Comment thread common/speculative.cpp Outdated
Comment on lines +926 to +930
// draft-dspark: the draft carries a Markov head and uses an anchor-first block layout.
// Derived from the draft model itself, not from the requested type: reading a DSpark
// draft with the DFlash layout is off by one row and silently collapses acceptance
// (measured 0.27% against 51% on the same binary and weights).
bool is_dspark = false;

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, fixed in c4c4667. The comment now says the DSpark path truncates on confidence and starts one row earlier only for sample_from_anchor models, instead of claiming every DSpark model is anchor-first.

Comment on lines +57 to +69
# test_with_and_without_draft asserts that drafting happens (draft_n > 0), not that
# any draft is accepted. Those are different failures: a drafter can run at full
# speed, produce well-formed blocks, and have essentially every token rejected. That
# is what a wrong block layout or a mis-tapped feature looks like, and it is a pure
# slowdown with identical output, so every existing assertion here still passes.
#
# This drafter/target pair is deliberately mismatched (q4_0 stories15M drafting for
# the F16 MoE), so acceptance is low by nature: 15.0% and 17.4% measured on the two
# prompts below, deterministic at temperature 0. The floor sits well under that and
# well over a collapse -- the regression this guards against measured 0.27%.
#
# Scope: this covers the draft-simple path only. The block-layout-sensitive draft
# types need their own drafters and are not exercised by any tiny fixture model.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in c4c4667, cut from 13 lines to 2. Kept only the stable reason the threshold is 5 percent, which is that the fixture drafter and target are mismatched on purpose so accept is 15-17 percent. The measurement details moved to the PR description where they do not age into the source.

Comment thread common/speculative.cpp Outdated
Comment on lines +990 to +992
LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s, block_layout=%s\n", __func__,
block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false",
is_dspark ? "dspark (anchor-first)" : "dflash (anchor at row 0)");

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, the log was wrong for that case. Fixed in c4c4667 by renaming the field rather than expanding it: it now prints lineage=dspark|dflash next to the existing sample_from_anchor=true|false. Those are two independent facts and the previous string implied one could be read off the other. The start row is derivable from sample_from_anchor, which was already in the same line.

Comment thread common/speculative.cpp Outdated
Comment on lines +963 to +965
LOG_WRN("%s: draft model carries %s, but --spec-type requested %s. Using the %s block "
"layout that the model requires; the other layout reads drafts one row off and "
"acceptance collapses to near zero.\n", __func__,

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair, and this was the most misleading of the five. The warning now says the wrong path drops confidence truncation, and only adds that it reads drafts one row late for sample_from_anchor models. The measured collapse I quoted, 0.272 percent against 51.385 percent, was on an anchor-first model, so it was never evidence for the bonus-anchor case. Fixed in c4c4667.

…shes

Review caught that the Markov head does not imply an anchor-first block layout.
A DSpark export can set dflash.sample_from_anchor=false (bonus anchor, slot 0 is
a bonus token), and that path starts drafts at row 1, same as DFlash. For those
models the DSpark branch differs by confidence truncation, not a row offset.

The detection itself is unchanged and still correct. What was wrong was the
surrounding documentation and diagnostics:

- the public accessor comment claimed the head implies the layout
- the is_dspark comment claimed every DSpark model is anchor-first
- the init log always printed "anchor-first", even when the code started at row 1
- the mismatch warning promised acceptance collapse for a case where both paths
  read the same rows

The log field is now lineage=dspark|dflash next to the existing
sample_from_anchor, so the two facts are reported separately instead of one being
inferred from the other. The warning names the real consequences.

Also shortens the acceptance-floor test comment to two lines per AGENTS.md.
@bri-prism

Copy link
Copy Markdown
Author

Following up on my earlier note about this being unvalidated on prism-v7: it is validated now, and the blocker turned out to be unrelated to this change.

The drafter would not load on prism-v7 because the branch has no support for DSpark log-SNR conditioning. These drafters carry log_snr_fc1 and log_snr_fc2, four tensors, and set log_snr_conditioning in metadata, so the loader created four fewer tensors than the file holds and refused. That is what "expected 79, got 75" was. My earlier guess in this thread, that the Markov head was not loading, was wrong: instrumenting the loader showed the head loads fine and the four unclaimed tensors are the log-SNR ones. #123 ports that feature.

With #123 applied, measured on Metal against a low-bit target:

run acceptance mean len
correct flag, --spec-type draft-dspark 65.891% 3.58
deliberately wrong flag, --spec-type draft-dflash 65.891% 3.58

Identical, and the wrong-flag run logs the mismatch warning and lineage=dspark. That is the property this PR is meant to guarantee: the requested type can no longer change the result, only the model can. The earlier 0.272% against 51.385% pair was measured on the v6-lineage tree on a 4090, so the absolute numbers are not comparable across hardware, but the collapse is gone on both.

So this PR now has a hardware-validated auto-detect, and it needs #123 to be exercisable at all on this branch. #123 does not depend on this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants