Skip to content

Support DFlash and DSpark drafts through continuous batching on PagedAttention - #4790

Open
ofirzaf wants to merge 30 commits into
openvinotoolkit:masterfrom
ofirzaf:dflash-cb-draft
Open

ofirzaf wants to merge 30 commits into
openvinotoolkit:masterfrom
ofirzaf:dflash-cb-draft

Conversation

@ofirzaf

@ofirzaf ofirzaf commented Nov 1, 2026 •

Copy link
Copy Markdown
Contributor

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 f09a56ac with transformers 5.15.1, which that version needs to export Qwen3.5. On GPU, the draft needs token_type_ids in 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

Speedup over auto-regressive (throughput tokens/s)

Decoding Arc B390 iGPU¹ Arc Pro B70
Auto-regressive (no draft) 1.00x (20.84) 1.00x (81.83)
DFlash, SDPA draft (before this PR) 3.12x (65.10) 2.47x (202.02)
DFlash, PagedAttention draft (this PR) 3.18x (66.18) 2.50x (204.50)

¹ 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:

Speedup over auto-regressive (throughput tokens/s)

Decoding Arc B390 iGPU¹ Arc Pro B70
Auto-regressive (no draft) 1.00x (41.70) 1.00x (89.68)
DSpark, PagedAttention draft (this PR) 1.55x (64.63) 1.82x (163.22)

E2E instructions:

  • Main model export

    optimum-cli export openvino \
      --model Qwen/Qwen3.5-9B \
      --weight-format int4 \
      models/qwen3.5-9b-int4-ov
  • Draft model export

    optimum-cli export openvino \
      --model z-lab/Qwen3.5-9B-DFlash \
      --weight-format int4 \
      models/qwen3.5-9b-dflash-int4-ov
  • Run concurrent requests with DFlash

    import openvino_genai as ov_genai
    
    draft = ov_genai.draft_model("models/qwen3.5-9b-dflash-int4-ov", "GPU")
    pipe = ov_genai.ContinuousBatchingPipeline(
        "models/qwen3.5-9b-int4-ov", ov_genai.SchedulerConfig(), "GPU", {"draft_model": draft}
    )
    
    config = ov_genai.GenerationConfig(max_new_tokens=256, num_assistant_tokens=7)
    prompts = ["Explain speculative decoding in one paragraph.", "Write a Python function that reverses a list."]
    for result in pipe.generate(prompts, generation_config=[config] * len(prompts)):
        print(result.m_generation_ids[0])

CVS-191402, CVS-194482

Checklist:

  • This PR follows GenAI Contributing guidelines.
  • Tests have been updated or added to cover the new code.
    • C++ 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 near max_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.
    • C++ 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.
    • C++ tests/cpp/sampler.cpp (SamplerBlockSampling): the candidates of several requests in one sampler call, DeepSpec candidates from the seed row, the cap near max_new_tokens, and partly-run windows.
    • C++ 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.
    • C++ tests/cpp/model_runner.cpp (TestModelRunnerDFlashDraft): token types, hidden rows and candidate rows of batched windows, the DeepSpec gather, and a context-only window.
    • C++ tests/cpp/scheduler.cpp (dflash_draft_windows_are_scheduled_whole).
    • Python tests/python_tests/test_continuous_batching.py: greedy output equals the target alone for DFlash and DeepSpec drafts; requests without drafting (num_assistant_tokens=0 through generate(), and an echo-only request); batched prompts with chunked prefill; staggered requests with mixed num_assistant_tokens, with and without dynamic split fuse; prefix caching; a prefix-caching mismatch; and a draft cache smaller than the target's.
    • Python 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 the xfail case above.
    • Python 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; LLMPipeline with a draft without properties; and the Markov head errors.
  • This PR fully addresses the ticket.
  • I have made corresponding changes to the documentation. Preview: https://ofirzaf.github.io/openvino.genai/docs/concepts/optimization-techniques/speculative-decoding#dflash

ofirzaf and others added 6 commits October 20, 2026 13:53
…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>
Copilot AI balanced review requested due to automatic review settings November 1, 2026 20:30
@github-actions github-actions Bot added category: visual language Visual language pipeline category: continuous batching Continuous batching category: speculative decoding Speculative decoding category: CPP API Changes in GenAI C++ public headers no-match-files category: GGUF GGUF file reader labels Nov 1, 2026
@ofirzaf ofirzaf added this to the 2026.5 milestone Nov 1, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 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
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

Comment thread tests/python_tests/test_vlm_pipeline.py Resolved
Comment thread tests/python_tests/test_vlm_pipeline.py Resolved
ofirzaf and others added 8 commits November 2, 2026 09:01
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>
Copilot AI balanced review requested due to automatic review settings November 2, 2026 17:10
@github-actions github-actions Bot added category: sampling Sampling / Decoding algorithms category: GHA CI based on Github actions labels Nov 2, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 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
3 resolved since last review

Review effort: Balanced

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🔵 Needs a closer look

Human assessment is needed because partial batch-admission failures can leave target/draft requests queued and block subsequent generation calls.

No open issues 🎉

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>
Copilot AI balanced review requested due to automatic review settings November 8, 2026 23:55
@ofirzaf
ofirzaf requested a review from pavel-esir as a code owner November 8, 2026 23:55
@github-actions github-actions Bot added the category: LLM LLM pipeline (stateful, static) label Nov 8, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 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

Comment thread src/cpp/src/sampling/sampler.cpp Resolved
ofirzaf and others added 2 commits November 9, 2026 03:38
… 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>
Copilot AI balanced review requested due to automatic review settings November 9, 2026 11:04

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🔵 Needs a closer look

Human assessment is needed to evaluate test reproducibility because the downloaded configuration and tokenizer depend on a mutable upstream revision.

No open issues 🎉

4 resolved since last review

Review effort: Balanced

@imargulis imargulis changed the title Run the DFlash draft through continuous batching on PagedAttention Support DFlash and DSpark drafts through continuous batching on PagedAttention Nov 9, 2026
Comment thread src/cpp/src/sampling/sampler.cpp Resolved
Comment thread src/cpp/src/sampling/sampler.cpp Resolved
Comment thread src/cpp/src/sampling/sampler.hpp Resolved
ofirzaf and others added 3 commits November 9, 2026 06:34
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>
Copilot AI balanced review requested due to automatic review settings November 9, 2026 16:56

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🔵 Needs a closer look

Human review is needed to assess a critical empty-batch out-of-bounds regression and DFlash test dependency mismatches between CI and the documented local setup.

No open issues 🎉

Review effort: Balanced

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>
Comment on lines +252 to +270
// 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);
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

looks unrelated change. Please remove it

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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.

Comment on lines +3729 to +3739
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",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

the same comment in test_continious batching.

This branch has not been deployed

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

Labels

category: continuous batching Continuous batching category: CPP API Changes in GenAI C++ public headers category: GGUF GGUF file reader category: GH Pages Docs Github Pages documentation category: GHA CI based on Github actions category: LLM LLM pipeline (stateful, static) category: sampling Sampling / Decoding algorithms category: speculative decoding Speculative decoding category: visual language Visual language pipeline Code Freeze no-match-files priority: high High piority

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants