Export DSpark block drafters for Qwen 3.5 and 3.8 - #2490
Open
Tianlei Wu (tianleiwu) wants to merge 3 commits into
Open
Export DSpark block drafters for Qwen 3.5 and 3.8#2490Tianlei Wu (tianleiwu) wants to merge 3 commits into
Tianlei Wu (tianleiwu) wants to merge 3 commits into
Conversation
5 tasks
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/20260828/builder-dflash2
branch
2 times, most recently
from
August 29, 2026 00:15
dbd5a08 to
b6b2e57
Compare
An EAGLE3/DFlash-style block drafter does not re-run the target's layers. It reads the target's intermediate residual streams instead, so the target graph has to expose them. Add an `aux_hidden_state_layers` extra option naming the decoder layers to tap. Each tap is concatenated on the last axis into one `aux_hidden_states` graph output, cast to the model's io_dtype when the layer norm ran in a different precision. Absent or blank, nothing is emitted and the graph is byte-for-byte what it was before. Entry `i` is the residual stream *entering* decoder layer i. That value only exists as a tensor inside layer i's input SkipLayerNorm, so the tap reads that node's output_3 rather than the layer boundary. The boundary is an off-by-one that still trains and still runs, but produces a drafter whose measured acceptance rate saturates at exactly 1.000 -- it looks like a perfect drafter while reading the wrong tensor. Layer 0 is rejected because its incoming residual stream is the raw embedding, which no skip layer norm produces. The tap records the pre-cast output_3 together with new_io_dtype; the post-cast tensor of the same name carries old_dtype, so pairing that name with new_io_dtype would mislabel the tap whenever the layer norm cast fires.
A DFlash 2 drafter never re-runs the target's layers. It reads the target's auxiliary hidden states, predicts a whole block of tokens in one pass, and returns a candidate lattice the Engine walks greedily. Add the drafter graph builder and wire it into the composite Qwen3.5 builder beside the existing MTP head: `dflash2_path` points at the draft checkpoint and produces an auxiliary `dflash2.onnx` plus a `dflash2` section in genai_config.json. Without the option nothing is built and the export is unchanged. The drafter has no embedding and no LM head of its own; both come from the target and are shared on disk through the same initializer-sharing path the MTP head uses, so the section and the decoder both record the shared entries. The drafter indexes the target's residual streams by position, so aux_hidden_state_layers must equal the checkpoint's target_layer_ids exactly, including order. A mismatch is rejected at init rather than silently feeding the drafter the wrong tensors. DFlash2Builder owns its own small IR helpers and does not subclass Model, so it carries no dependency on the shared builder base.
DSpark is a parallel block drafter with a low-rank Markov head. It exposes the same candidate lattice to the runtime as DFlash 2, so it reuses that proposal and verification path and only needs its own graph export. Add the drafter graph builder and wire it into the composite Qwen3.5 builder next to DFlash 2: `dspark_path` produces an auxiliary `dspark.onnx` plus a `dspark` section in genai_config.json, sharing the target's embedding and LM head on disk. The two block drafters are mutually exclusive, since the runtime drives one lattice per model. SpecForge names the layers whose *output* it taps, while the builder's aux_hidden_state_layers names the residual stream *entering* a layer, so the required tap set is each `target_layer_ids` entry plus one. Passing SpecForge's ids through unchanged is an off-by-one that leaves acceptance pinned at exactly 1.0, so a mismatch is rejected at init; the tests cover that case explicitly. `dspark_top_k` controls how many candidates the lattice keeps per block slot.
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/20260828/builder-dspark
branch
from
August 29, 2026 00:23
1b59851 to
23babc1
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add builder support for exporting a DSpark block drafter alongside Qwen 3.5 and 3.8 target models. The exporter uses selected target-model residual streams and emits the DSpark graph and configuration as an alternative block-drafting path to DFlash 2.
Stack
Depends on #2489.
Changes
dspark_path,dspark_num_draft_tokens, anddspark_top_kbuilder options.dspark.onnx, export its configuration, and share compatible target-model initializers.Testing
15 passed.