Repository navigation
Conversation
…BatchingForSpeculativeDecodingImpl is_prefix_caching_enabled, discard_awaiting_request and restore_awaiting_prefix only use the shared scheduler and awaiting-request state, so any speculative pipeline can use them. No behavior change. Co-authored-by: Cursor <cursoragent@cursor.com>
A DFlash draft forwards, per request, [target hidden-state rows ; seed + mask block]. Behind the enable_dflash_draft() flag, forward() fills these rows next to the other per-sequence inputs: the target rows passed through set_initial_hidden_state() go to hidden_states with token type 0, the remaining scheduled rows are the block with token type 1 and zero hidden states, and the block rows after the candidate position offset are added to the sampled_tokens_indices gather. The hidden_states input is sized from the model, since it is wider than the embeddings. Co-authored-by: Cursor <cursoragent@cursor.com>
… stateful draft runner The draft is a second CB pipeline on PagedAttention. Each speculative step forwards, per request, one window [new target hidden-state rows ; seed + mask block]: token_type_ids mark the block rows, which attend to each other bidirectionally, and the candidates are read from the block's mask rows. The update sets m_max_content_len to the window end and keeps the processed length at the committed rows, so the unchanged scheduler offers exactly the window and the step does not sample it. Like Eagle3 and MTP, the draft logic sits inline behind dflash_mode_enabled in update_request and multistep, and detail::make_dflash_draft_window() plans each window. The window's target rows reach the model runner through set_initial_hidden_state() for one step only; after the forward, each candidate row is sampled through the sampler, so logit processors apply. The strategy reuses the shared speculative step() and generate(), so multiple concurrent requests work as for the other speculative pipelines. The draft follows the target's dynamic_split_fuse setting and has no batch-token budget, so the scheduler never splits a block. When the target recomputes preempted tokens, the draft waits and reuses its committed rows. When the scheduler does not run a draft window whole and the target's rows no longer reach the committed positions, the draft frees that request's KV and the target finishes the request alone. DFlashCBDraftRunner, the single-request stateful SDPA draft, its helpers, and dflash_strategy_utils.hpp are removed. A draft without a token_type_ids input was exported for that runner and is rejected with a re-export hint. Co-authored-by: Cursor <cursoragent@cursor.com>
C++: the draft update after prefill chunks and after full, partial and full rejection, a deferred target, a target that recomputes below the committed rows, a missed window that leaves the request to the target, the block near max_new_tokens, request-id reuse, request admission, the draft generation config, and draft windows through the unchanged scheduler. Python: greedy output equals the target alone for DFlash and DeepSpec drafts, batched prompts with chunked prefill or without dynamic split fuse, staggered requests with mixed num_assistant_tokens with and without dynamic split fuse, a draft cache smaller than the target's, and a VLM target, including echo and prompt logprobs; visual token pruning fails the same way with and without the draft. Target preemption is marked xfail: with candidates pending validation, the replay is scheduled one token short of the validation window, which affects every speculative pipeline (related: CVS-198275). Parity tests use an uncompressed KV cache, since a compressed cache lets speculative and autoregressive decoding diverge. Co-authored-by: Cursor <cursoragent@cursor.com>
Both pipelines publish prompt blocks only, as MTP does: draft KV past the prompt holds mask and rewound block rows, so those blocks are never hashed. After both requests are added, the strategy restores target and draft to the same cached prefix; the target then recomputes and exports hidden states from that length, and the draft request starts with the restored rows committed. Main and draft must use the same enable_prefix_caching setting, and embedding inputs need one prompt token ID per embedding. Hybrid linear-attention targets restore through the existing checkpoint path. The awaiting-request lookup used by the prefix alignment moves to ContinuousBatchingForSpeculativeDecodingImpl::find_awaiting_request and replaces the Eagle3 and MTP lambdas. Co-authored-by: Cursor <cursoragent@cursor.com>
With prefix caching on, a VLM DFlash request whose num_assistant_tokens is unset or 0 must match the target alone through add_request with images and generate with a chat history. The test is marked xfail: those entry points keep the prompt token IDs only when the request config already assists generation, and DFlash needs them to hash embedding inputs. A separate change that always keeps the VLM prompt token IDs removes the xfail. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved DFlash request-handling defects can trigger unwanted drafting or reject valid EOS and echo-only requests, while the new VLM tests are excluded from compatible CI jobs.
6 open issues
Retain explicit EOS in the draft stop-token set · New Keep draft configuration valid for echo-only requests · New Preserve explicit zero assistant tokens in DFlash generation · New Avoid rebuilding logit processors for existing request contexts · New Add VLM DFlash tests to compatible CI matrix jobs · New Pass the cached draft snapshot path to the exporter · New
What changed in this PR
Replaces DFlash’s stateful draft runner with a second continuous-batching pipeline using PagedAttention, enabling concurrent requests and prefix caching within the shared speculative-decoding infrastructure.
Changes:
- Adds per-request draft windows, hidden-state transfer, candidate sampling, and paired prefix restoration.
- Removes legacy runner utilities and rejects older draft exports with a re-export message.
- Expands C++ and Python coverage for batching, rejection handling, caching, preemption, and VLM targets.
| File | Description |
|---|---|
| tests/python_tests/test_vlm_pipeline.py | Adds DFlash VLM parity, batching, and prefix-cache tests. |
| tests/python_tests/test_continuous_batching.py | Adds DFlash and DeepSpec batching and cache tests. |
| tests/cpp/test_dflash_cb_strategy.cpp | Replaces legacy utility tests with draft-window and admission tests. |
| tests/cpp/scheduler.cpp | Tests scheduling draft windows as recompute rows. |
| tests/cpp/model_runner.cpp | Tests draft input packing and candidate-row selection. |
| src/cpp/src/speculative_decoding/dflash_model_transforms.hpp | Removes the legacy hidden-state reshape declaration. |
| src/cpp/src/speculative_decoding/dflash_model_transforms.cpp | Removes legacy reshaping and validates mask-token metadata. |
| src/cpp/src/speculative_decoding/continuous_batching/pipeline_impl.hpp | Defines draft windows and shared prefix-management helpers. |
| src/cpp/src/speculative_decoding/continuous_batching/pipeline_impl.cpp | Implements DFlash window updates and batched candidate sampling. |
| src/cpp/src/speculative_decoding/continuous_batching/mtp_strategy.cpp | Reuses shared request and prefix helpers. |
| src/cpp/src/speculative_decoding/continuous_batching/eagle3_strategy.cpp | Reuses shared awaiting-request lookup. |
| src/cpp/src/speculative_decoding/continuous_batching/dflash_strategy.hpp | Switches DFlash to shared speculative-pipeline orchestration. |
| src/cpp/src/speculative_decoding/continuous_batching/dflash_strategy.cpp | Builds paired pipelines and manages admission and prefix alignment. |
| src/cpp/src/speculative_decoding/continuous_batching/dflash_strategy_utils.hpp | Deletes obsolete stateful-runner utilities. |
| src/cpp/src/continuous_batching/model_runner.hpp | Packs DFlash hidden states, token types, and candidate indices. |
| src/cpp/include/openvino/genai/continuous_batching_pipeline.hpp | Declares and grants access to the internal DFlash pipeline class. |
Review effort: Balanced
The draft update no longer marks its window as rows to recompute. It keeps the processed length at the committed rows and asks for the rows past what the request offers by itself through set_num_validated_tokens, as the Eagle3 and fast-draft updates do; a prefill chunk keeps the expected-row cap. A window that starts inside the prompt carries only prompt rows, and the first block follows the prompt: without dynamic split fuse a step runs either prompt or generation windows, so every block window must be a generation window. The model runner gathers the candidate rows through the generic gather, with the count from the DFlash row fill. Sampler::set_block_sampling selects a per-request block sampler: inside the regular sample() call, which samples requests in parallel, it replaces the mask placeholders with one greedy token per gathered row, caps the candidates near max_new_tokens, and takes the candidates of a window the scheduler ran only partly. The per-candidate sampler calls in the draft multistep are removed. When a draft request ends while its main request runs, the speculative step finishes it and drops its handle, and the main request continues alone. A draft pipeline out of KV cache (IGNORED) is logged there; a DFlash window that could not run stops the draft request (STOP) and is logged where it stops. Co-authored-by: Cursor <cursoragent@cursor.com>
The draft config cleared stop_token_ids but kept eos_token_id. The draft pipeline then refilled the empty set from its defaults, and validate() rejected a request whose explicit EOS differs from the default. ignore_eos already keeps the draft from stopping on any stop token, so the set stays. Co-authored-by: Cursor <cursoragent@cursor.com>
An echo-only request (max_new_tokens=0, echo=True) is valid for the target, but the draft config turns echo off, so the draft add_request rejected it and the whole request failed. Such a request generates nothing to draft, so it gets no draft request, and the awaiting-queue check counts it the same way. Co-authored-by: Cursor <cursoragent@cursor.com>
The inherited SpeculativeDecodingImpl::generate() replaced an explicit num_assistant_tokens=0 with the default before add_request, so a target-only request still drafted. DFlash now passes the configs through unchanged, as Eagle3 does: add_request already fills an unset value with the default and keeps an explicit 0. Co-authored-by: Cursor <cursoragent@cursor.com>
The DFlash draft calls create_logit_processor on every update. Each call built a LogitProcessor, which scans the whole prompt, and then discarded it because the request context already existed. Return early instead; the existing context was already kept, so no caller sees a change. Co-authored-by: Cursor <cursoragent@cursor.com>
The DFlash drafts need the export from huggingface/optimum-intel#2096, which the pinned optimum-intel does not have, and the VLM DFlash tests need Transformers 5.2. As for the Qwen3.5 MTP tests, a separate job on Linux, manylinux and Windows installs both and runs the continuous-batching and VLM DFlash tests; the API and VLM jobs skip them. Co-authored-by: Cursor <cursoragent@cursor.com>
Wrap the lines that darker reformats in CI. Co-authored-by: Cursor <cursoragent@cursor.com>
The pipeline tests run with the DFlash draft; a DeepSpec draft differs only in where its candidates start, which the C++ tests cover, so only the parity test runs both drafts. The requests that get no draft request share one test, the VLM parity test keeps the echo and logprobs case, and the prefix caching xfail keeps the unset case. The single-window model runner test repeats the batched one, and the logit processor test passed before its fix too. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
🟡 Changes recommended
The unbounded draft budget can exceed memory limits or break cache-rotation setup, while pending-window replay after preemption can produce incorrect tokens that a conditional xfail permits.
6 open issues
Preserve the configured draft token budget · New Fix pending-window replay instead of xfail · New Export from a locally resolved Hub snapshot · New Pass the cached draft snapshot path to the exporter Add VLM DFlash tests to compatible CI matrix jobs Avoid rebuilding logit processors for existing request contexts
3 resolved since last review
Review effort: Balanced
A DSpark draft is a DFlash backbone with a Markov head that scores every candidate of a block from the token drafted before it. optimum-intel exports the head next to the draft as openvino_markov_head_model.xml and tags the backbone with dflash/markov_rank and dflash/markov_head_type. The DFlash draft ignored the head and drew every candidate from the backbone logits alone, so a DSpark draft ran as a plain DFlash draft and accepted fewer tokens. A draft whose rt_info has a Markov rank above 0 now loads the head from the draft folder and requires it there. The head holds a dense Markov step, which maps the token before to a bias over the vocabulary, and a Top-K candidate scorer, which scores the Top-K candidates of every row after the candidates of the row before. The scorer is used by default; the markov_head_top_k draft property selects the dense step with 0 and replaces the exported K with any other value. Only the vanilla head is supported. The sampler takes an optional BlockRowSource with the block sampling start offset. It collects the blocks of a step, prepares the source once, and draws every block from the rows the source fills for the token before. A source that runs per position is stepped over the blocks that still draw, and a source can draw a whole block, or per position a token, itself when the request takes its tokens without the logit processors. A position the source has no row for ends the block. The DFlash candidate selector is that source for a draft with a Markov head, and a draft without one keeps the backbone rows. The scorer is run once over the packed block rows; a greedy request whose logit processors keep the argmax walks the lattice greedily, taking the best candidate after the token before, and any other request samples the lattice row after the token before, ending the block at a token outside the candidates. The dense step is run per position and biases the backbone rows by the token drafted before; a greedy request whose logit processors keep the argmax takes the best biased token without filling the row. The selector compiles the head with the draft properties. Since a block can now end before its last candidate row, the draft step counts the candidates its sequence generated instead of the rows it forwarded. The DFlash CI job also runs test_dflash_markov_head_cb.py, which checks tiny DFlash v1 and DSpark drafts exported with optimum-intel against the target alone. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
🟡 Changes recommended
Draft inference metrics omit scorer and Markov execution time, the greedy path unnecessarily allocates vocabulary-sized buffers, and documentation omits the new draft property semantics while incorrectly describing candidate limits.
4 open issues
Review effort: Balanced
… language model properties The Markov head runs during sampling, after the draft's inference timer stops, so its time was missing from the draft inference duration. The candidate selector now times its inferences, and the draft step adds them to its inference duration. A draft without properties takes the target's, which can carry a MODEL_PROPERTIES overlay. The backbone resolves the overlay for the language model, but the head passed it to compile_model as is, and compilation failed. The head now compiles with the language model's properties. The dense Markov step allocated a vocabulary-sized row for every block on every draft step, even when it drew every token without one. The row is allocated on the first fill_row() of a block. Co-authored-by: Cursor <cursoragent@cursor.com>
The guide said that GenAI does not limit num_assistant_tokens, but GenerationConfig::validate() rejects values above 100. GenAI does not check the value against the draft's block size. Co-authored-by: Cursor <cursoragent@cursor.com>
The block token sampler applies the logit processors to the logits in place, so it takes them by reference, and the block sampler of one request takes its sequence group and logits by const reference. Co-authored-by: Cursor <cursoragent@cursor.com>
The job also runs the DSpark tests in test_dflash_markov_head_cb.py. The drafts need the export from optimum-intel main. Co-authored-by: Cursor <cursoragent@cursor.com>
DSpark gets its own row in the strategies table and its own section after DFlash: why the Markov head helps a parallel draft, the Markov head model exported next to the draft, which --weight-format does not compress, and the Top-K candidate scorer and dense Markov walk selected by markov_head_top_k, with the default K and how the draft selects its candidates. The DFlash section gains the runtime limits, which DSpark shares. Based on c3dfe1b in openvinotoolkit#4847, limited to the drafts and decoding modes that this PR supports. Co-authored-by: Igor Margulis <igor.margulis@intel.com> Co-authored-by: Cursor <cursoragent@cursor.com>
The sampler draws the blocks of several requests in parallel, and each draw read the shape of the candidate scorer's output tensors. A CPU plugin tensor rebuilds the shape that get_shape() returns on every call and returns a reference to it, so a concurrent call could free the shape another draw was reading. The draw then computed a wrong K and read the lattice out of bounds, which crashed the pipeline now and then with several concurrent DSpark requests. The selector now keeps views over the scorer outputs, whose shape is fixed. Co-authored-by: Cursor <cursoragent@cursor.com>
015cc72 to
aa566ab
Compare
| // The main request continues alone when its draft request ended: the draft pipeline ran out of KV cache | ||
| // (IGNORED, logged here) or stopped the request itself (STOP, logged where it stopped; only a DFlash window | ||
| // that could not run). A handle is dropped only after its request ended, since dropping it stops the request. | ||
| for (auto draft_it = m_draft_generations.begin(); draft_it != m_draft_generations.end();) { | ||
| const uint64_t request_id = draft_it->first; | ||
| const GenerationStatus draft_status = draft_it->second->get_status(); | ||
| if ((draft_status != GenerationStatus::IGNORED && draft_status != GenerationStatus::STOP) || | ||
| !main_generated_requests.count(request_id)) { | ||
| ++draft_it; | ||
| continue; | ||
| } | ||
| if (draft_status == GenerationStatus::IGNORED) { | ||
| GENAI_WARN("Request %llu: the draft model ran out of KV cache; continuing without speculative decoding.", | ||
| static_cast<unsigned long long>(request_id)); | ||
| } | ||
| m_draft_pipeline->finish_request(request_id); | ||
| draft_it = m_draft_generations.erase(draft_it); | ||
| } | ||
|
|
There was a problem hiding this comment.
looks unrelated change. Please remove it
There was a problem hiding this comment.
It's in the shared SpeculativeDecodingImpl::step(), which the DFlash pipeline uses. When a draft request ends while its main request runs (the draft ran out of KV cache, or a DFlash window could not run), it finishes the draft request so the target continues alone. test_dflash_small_draft_cache_leaves_requests_to_the_target covers it.
| VLM_DFLASH_MAIN_MODEL_ID = "optimum-intel-internal-testing/tiny-random-qwen3.5" | ||
| VLM_DFLASH_DRAFT_MODEL_ID = "optimum-intel-internal-testing/tiny-random-qwen3.5-dflash" | ||
|
|
||
|
|
||
| def _get_vlm_dflash_model_paths() -> tuple[Path, Path]: | ||
| _maybe_skip_unsupported_model_export(VLM_DFLASH_MAIN_MODEL_ID) | ||
| draft_model_path = get_ov_cache_converted_models_dir() / sanitize_model_id(VLM_DFLASH_DRAFT_MODEL_ID) | ||
| AtomicDownloadManager(draft_model_path).execute( | ||
| lambda temp_path: export_with_optimum_cli( | ||
| retry_request(lambda: snapshot_download(VLM_DFLASH_DRAFT_MODEL_ID)), | ||
| "text-generation-with-past", |
There was a problem hiding this comment.
avoid all of this dflash stuff in common test suite. Use already existing test suites for testing speculative decoding methods. Avoid a multiple new dflast tests
There was a problem hiding this comment.
the same comment in test_continious batching.



Description
This PR runs the DFlash/DSpark draft as a second continuous-batching pipeline on PagedAttention, replacing the stateful SDPA draft, so DFlash/DSpark serves concurrent requests and supports prefix caching like the other speculative pipelines. It replaces the continuous-batching PoC in #4535.
Limits: greedy decoding only, one sequence per request, no adapters, no cache eviction.
Depends on huggingface/optimum-intel#2096 (merged) for the draft export. The default optimum-intel version in GenAI does not include it yet, so the DFlash/DSpark CI job installs optimum-intel at
f09a56acwith transformers 5.15.1, which that version needs to export Qwen3.5. On GPU, the draft needstoken_type_idsin the PagedAttention MIXED stage (openvinotoolkit/openvino#38897).Breaking change: DFlash/DSpark drafts exported with optimum-intel before huggingface/optimum-intel#2096, including the drafts used with GenAI 2026.4, no longer load. Re-export them with an optimum-intel version that includes that PR.
Documentation: the speculative decoding guide (
site/docs/concepts/optimization-techniques/speculative-decoding.mdx) gets a DFlash/DSpark section.Performance
Decode throughput of the target alone (auto-regressive), DFlash with the stateful SDPA draft (before this PR), and DFlash with the continuous-batching PagedAttention draft (this PR).
Setup
--weight-format int4VLMPipeline, one request at a time, greedy decoding,num_assistant_tokens=7, prefix caching off, other settings defaultmaster, draft exported with optimum-intelmainSpeedup over auto-regressive (throughput tokens/s)
¹ Intel Core Ultra X7 358H
With one request at a time, the PagedAttention draft runs as fast as the SDPA draft or slightly faster
DSpark
Decode throughput of the target alone (auto-regressive) and DSpark with the PagedAttention draft (this PR). The setup matches the DFlash run above, except:
--weight-format int4--weight-format int4with optimum-intelmain, defaultmarkov_head_top_k(Top-K candidate scorer)Speedup over auto-regressive (throughput tokens/s)
E2E instructions:
Main model export
optimum-cli export openvino \ --model Qwen/Qwen3.5-9B \ --weight-format int4 \ models/qwen3.5-9b-int4-ovDraft model export
optimum-cli export openvino \ --model z-lab/Qwen3.5-9B-DFlash \ --weight-format int4 \ models/qwen3.5-9b-dflash-int4-ovRun concurrent requests with DFlash
CVS-191402, CVS-194482
Checklist:
tests/cpp/test_dflash_cb_strategy.cpp: the draft update after prefill chunks, the prompt rows before the first block, after full, partial and full rejection, a deferred target, a recomputing target, a missed window that ends the draft request, a window cut at the seed that keeps drafting, stale extra rows, the block nearmax_new_tokens, a restored prefix, request-id reuse, request admission, the draft generation config, and old drafts. The tests of the removed runner are deleted.tests/cpp/test_dflash_cb_strategy.cpp(DSparkModelTransforms,DFlashCBCandidateSelection): reading and validating the Markov head next to a DSpark draft, selecting the Top-K candidate scorer or the dense Markov walk and replacing K, the greedy lattice walk of each block, the logit processors on lattice rows, ties, a drafted token outside the candidates that ends the block, and the dense Markov walk.tests/cpp/sampler.cpp(SamplerBlockSampling): the candidates of several requests in one sampler call, DeepSpec candidates from the seed row, the cap nearmax_new_tokens, and partly-run windows.tests/cpp/sampler.cpp(SamplerBlockSampling): a block row source that draws every token from the row after the token before, steps the blocks that still draw position by position, ends a block at a position without a row, or draws a whole block.tests/cpp/model_runner.cpp(TestModelRunnerDFlashDraft): token types, hidden rows and candidate rows of batched windows, the DeepSpec gather, and a context-only window.tests/cpp/scheduler.cpp(dflash_draft_windows_are_scheduled_whole).tests/python_tests/test_continuous_batching.py: greedy output equals the target alone for DFlash and DeepSpec drafts; requests without drafting (num_assistant_tokens=0throughgenerate(), and an echo-only request); batched prompts with chunked prefill; staggered requests with mixednum_assistant_tokens, with and without dynamic split fuse; prefix caching; a prefix-caching mismatch; and a draft cache smaller than the target's.tests/python_tests/test_vlm_pipeline.py: a VLM Qwen3.5 target with batched image and text requests on its linear-attention layers, prefix caching, echo with logprobs, and thexfailcase above.tests/python_tests/test_dflash_markov_head_cb.py: tiny DFlash and DSpark drafts exported with optimum-intel, DSpark with the Top-K candidate scorer and the dense Markov walk; greedy output equals the target alone for requests drafted together, staggered requests, chunked prefill, no dynamic split fuse, a draft KV cache too small for all requests, a fixed target cache, prefix caching, and penalties; target and draft preemption completes;LLMPipelinewith a draft without properties; and the Markov head errors.