[ofirzaf]

What does this PR do?

Changes the DFlash draft export (DFlash, and DeepSpec / DSpark drafts with markov_rank == 0) so the draft runs on PagedAttention in OpenVINO GenAI continuous batching. This replaces the continuous-batching PoC in #2028.

Changes to the exported draft:

• Per-row inputs: inputs_embeds, hidden_states, position_ids and the new token_type_ids input all cover the new rows of a step, [context; block]. token_type_ids is 0 for context rows, which read the target hidden states, and 1 for the seed + mask block rows, which read their embeddings.
• Attention: standard Qwen3 attention with the plain KV-cache update. The mask is built from token_type_ids: causal, bidirectional inside the block, and sliding-window for the sliding layers. SDPAToPagedAttention converts it like any decoder attention: token_type_ids reaches every PagedAttentionExtension, and each layer gets its own sliding window.
• Cache: the committed-prefix cache helpers, the target/block split inside attention, and the manual cache slicing are removed. The block rows are stored in the cache, and the next step starts at the committed length and overwrites them.
• Compute: only the block rows get queries, norms, o_proj and the MLP. Context rows only produce keys and values. last_hidden_state covers all rows and is zero in the context rows.
• DSpark: Qwen3DSparkModel drafts with markov_rank == 0 now export through the same Qwen3DFlashForCausalLM.
• Requires Support DFlash and DSpark drafts through continuous batching on PagedAttention openvinotoolkit/openvino.genai#4790 for GenAI testing

1. Export the draft model:

optimum-cli export openvino -m z-lab/Qwen3.5-9B-DFlash --task text-generation-with-past --trust-remote-code --weight-format int4 Qwen3.5-9B-DFlash

2. Export the main model:

optimum-cli export openvino -m Qwen/Qwen3.5-9B --weight-format int4 Qwen3.5-9B

3. Run speculative decoding with OpenVINO GenAI:

import openvino_genai

draft = openvino_genai.draft_model("./Qwen3.5-9B-DFlash", "GPU")
pipe = openvino_genai.VLMPipeline("./Qwen3.5-9B", "GPU", draft_model=draft)

config = openvino_genai.GenerationConfig(max_new_tokens=256, num_assistant_tokens=7)
print(pipe.generate("Explain speculative decoding in one paragraph.", generation_config=config))

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[echarlaix]

failing tests https://github.com/huggingface/optimum-intel/actions/runs/41068096136/job/125328119818?pr=2096, please take a look

[ofirzaf]

Hi @echarlaix pushed a fix. Thanks!

[echarlaix]

@ofirzaf this PR has merge conflicts with main, could you rebase/merge and resolve them

[ofirzaf]

Updated the branch. Please review @echarlaix

[MaximProshin]

please don't merge it before #2062

[ofirzaf]

Preview models support validation test failed because nightly GenAI doesn't have the new DFlash pipeline merged yet openvinotoolkit/openvino.genai#4790

[rkazants]

Preview models support validation test failed because nightly GenAI doesn't have the new DFlash pipeline merged yet openvinotoolkit/openvino.genai#4790

please skip temporarily this test case with a comment.

[ofirzaf]

@rkazants pushed a commit to skip the test with a comment that we can enable it once the genai PR lands in nightly

[rkazants]

can you also update test_cb_compatible.py to check that dflash exported model is continous batching compatible and SDPAToPA transformation is applied?