speculative: take the draft block layout from the draft model, not the CLI type - #122
speculative: take the draft block layout from the draft model, not the CLI type#122bri-prism wants to merge 2 commits into
Conversation
…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.
|
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. The file carries 79 tensors and the loader creates 75. The four it does not create are exactly That matters here beyond being an unrelated bug. The detection in this PR reads Where that leaves the change:
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. |
There was a problem hiding this comment.
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.
| // 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 |
There was a problem hiding this comment.
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.
| // 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; |
There was a problem hiding this comment.
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.
| # 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. |
There was a problem hiding this comment.
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.
| 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)"); |
There was a problem hiding this comment.
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.
| 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__, |
There was a problem hiding this comment.
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.
|
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:
Identical, and the wrong-flag run logs the mismatch warning and 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. |
Lands the draft block-layout fix on
prism-v7. Replaces #114, which was closed: it was broken as pushed and based onprism-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 asdraft-dflashtook 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:
The fix
Derive the layout from the draft model instead of the requested type.
These files declare
general.architecture = dflashwhichever lineage they come from, so the type can't identify them. The Markov head can:src/models/dflash.cppdefines 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, sodspark_markov_w1 != nullptris exactly the loader's own definition of the lineage, available for free.llama_model_has_dspark_markov_head()insrc/llama-ext.h/src/llama-model.cppcommon/speculative.cppsetsis_dsparkfrom the model, and warns when--spec-typedisagrees rather than silently honoring a request that cannot worksample_from_anchor, so the trap is visible without reading sourceValidation: 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_floorintools/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-simplepath 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-formatdisagrees with the committed style ofcommon/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.