Support original z-lab Qwen3 DFlash checkpoints; require checkpoint target_layer_ids by harshal-96 · Pull Request #4991 · InternLM/lmdeploy

harshal-96 · 2026-09-19T09:16:57Z

Motivation

Follow-up to #4789, as discussed with @RunningLeon in #4789 (comment): enable dense Qwen3 DFlash on the merged infrastructure. The original z-lab checkpoints (z-lab/Qwen3-4B-DFlash-b16, Qwen3-8B-DFlash-b16) do not load on main for two reasons, and main also carries the incorrect target-layer fallback that the reference repo removed.

Original z-lab checkpoints store block_size (and mask_token_id on some) at the top level of config.json rather than inside dflash_config, so parse_dflash_config rejects them.
A dense fp16 target (e.g. an AWQ Qwen3) with a bf16 draft crashes on the dtype boundary. Running everything in fp16 is not a workaround: Qwen3 activation outliers overflow fp16 in the target-feature fusion and acceptance collapses to 0.06 of 15 drafts (measured on both this implementation and the reference one), so the draft must stay bf16 with casts at the boundary.
build_target_layer_ids guesses evenly spaced tap layers when dflash_config.target_layer_ids is missing. The tap layers are a training-time choice baked into the checkpoint (the fc input width and which target layers the draft was trained against), so a guess reads the wrong features; for 14 of the 20 published DFlash checkpoints the guess differs from the trained layers, and when only the counts match the load succeeds silently with wrong layers. See bug: build_target_layer_ids never runs, and its guess would be wrong for 14 of 20 checkpoints z-lab/dflash#156; the reference repo removed the fallback in bug: require dflash_config.target_layer_ids instead of guessing it z-lab/dflash#157 and sglang did the same in spec: require dflash_config.target_layer_ids instead of guessing tap layers sgl-project/sglang#37476. Every published DFlash checkpoint sets the key.

Modification

parse_dflash_config: accept block_size / mask_token_id at the top level of the draft config as a fallback, mirroring the reference implementation's resolution order (dflash_config first).
qwen3_dflash.py: cast to the draft dtype at the two boundaries where target-side tensors enter the draft (input embeddings and projected target hidden states), so a bf16 draft runs against an fp16 target.
dflash_utils.py: remove build_target_layer_ids; a missing dflash_config.target_layer_ids now raises with a message naming the key.
Tests updated accordingly in tests/pytorch/spec_decode/test_dflash_utils.py (top-level layout regression test, explicit target_layer_ids in the block-size tests, missing-key raise test).

BC-breaking (Optional)

Checkpoints relying on the guessed target layers would now fail at load with a clear error instead of silently producing near-zero acceptance. No published DFlash checkpoint omits dflash_config.target_layer_ids.

Use cases (Optional)

Verified end to end on the merged main (RTX 4070 Laptop 8GB, WSL2, eager mode, greedy, batch 1, target thewimo/Qwen3-4B-AWQ, draft z-lab/Qwen3-4B-DFlash-b16; one math and one code prompt, 256 new tokens each):

DFlash: 35.5 to 57.6 tok/s across repeated runs; baseline without speculative decoding, measured back to back on an idle GPU: 12.2 to 15.5 tok/s. Roughly 3x wall-clock on average; the absolute numbers swing with laptop clocks, so treat them as indicative.
Mean accepted drafts per step: 4.19 of 15, byte-identical across every run (greedy decoding is deterministic) and matching the value measured when these fixes were validated on the pre-merge Support dflash for qwen3.5 #4789 branch, so the port to the refactored code is behavior-preserving.
Outputs are coherent on both prompts.

Checklist

Pre-commit or other linting tools are used to fix the potential lint issues.
The modification is covered by complete unit tests: pytest tests/pytorch/spec_decode/test_dflash_utils.py passes 48 tests.
No new dependencies.
Docstrings updated where behavior changed.

Validated dense Qwen3 end-to-end on the DFlash stack (RTX 4070 8 GB, target thewimo/Qwen3-4B-AWQ float16, draft z-lab/Qwen3-4B-DFlash-b16 bfloat16, greedy, 256 new tokens): 33.6-36.8 tok/s vs 25.1 baseline, mean 4.19 of 15 drafts accepted per block, coherent lossless output. Two gaps had to be fixed to get there: - parse_dflash_config required dflash_config['block_size'], but the original z-lab Qwen3 checkpoints (Qwen3-4B/8B-DFlash-b16) store block_size at the top level of config.json. Fall back to the top-level attribute (same resolution order as the reference implementation), for mask_token_id as well, with a regression test. - The draft ran the shared target embedding output and the target aux states in the target dtype through bfloat16 draft weights, which fails with a dtype mismatch for float16 targets (e.g. AWQ). Casting the whole engine to float16 instead is not an option: Qwen3 hidden state outliers overflow float16 during feature fusion and acceptance collapses to zero (measured 0.06 of 15). Cast to the draft dtype at the two ingestion boundaries; logits already cast via get_logits. Signed-off-by: harshal-96 <harshal.dhandrut@gmail.com>

…er_ids The evenly spaced fallback cannot be correct: the tap layers are a training-time choice baked into the checkpoint (the fc input width and which target layers the draft was trained against), so any load-time guess reads the wrong features. For 14 of the 20 published DFlash checkpoints the guess differs from the trained layers, and when the counts collide the load succeeds silently with wrong layers (z-lab/dflash#156). Every published DFlash checkpoint sets dflash_config.target_layer_ids, so a missing key now fails loudly. Matches the same removal in the reference repo (z-lab/dflash#157) and sglang (sgl-project/sglang#37476). Signed-off-by: harshal-96 <harshal.dhandrut@gmail.com>

Copilot

Copilot review overview

🟡 Changes recommended

Add mixed-dtype regression coverage and address the two identified test/diagnostic gaps.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1 · 2

Open (3)

Add mixed-dtype regression tests for embedding and projected-hidden paths · New
Update block_size error to reference both supported config locations · New
Test top-level mask_token_id fallback in regression coverage · New

What changed in this PR

This PR adds support for original Qwen3 DFlash checkpoints and enforces checkpoint-defined target layers.

Changes:

Supports top-level checkpoint metadata.
Adds mixed-dtype boundary casts.
Requires explicit target_layer_ids and updates tests.

File	Summary
`tests/pytorch/spec_decode/test_dflash_utils.py`	Adds checkpoint-layout and required-layer tests; top-level `mask_token_id` coverage remains missing.
`lmdeploy/pytorch/spec_decode/dflash_utils.py`	Resolves metadata and removes inferred target layers; one diagnostic still references the nested block-size key.
`lmdeploy/pytorch/models/qwen3_dflash.py`	Adds target/draft dtype casts; mixed-dtype paths lack regression tests.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+        # the shared target embedding may run in a different dtype than the
+        # draft (e.g. float16 AWQ target with a bfloat16 draft checkpoint)
+        return embeds.to(self.dtype)


        raise ValueError('DFlash query length (1 + speculative_num_draft_tokens) must not exceed checkpoint '
                         'dflash_config.block_size. '


+def test_parse_dflash_config_top_level_checkpoint_layout():
+    """Original z-lab DFlash checkpoints (e.g. Qwen3-4B-DFlash-b16) keep
+    block_size at the top level of config.json and nest only mask_token_id and
+    target_layer_ids inside dflash_config."""
+    config = _draft_config(block_size=4,
+                           dflash_config=dict(
+                               mask_token_id=32001,
+                               target_layer_ids=[1, 5, 9, 13],
+                           ))


…rer block_size error - Unit-test the two draft-dtype ingestion boundaries (shared target embeddings and projected target hidden states) with a float16 source and a bfloat16 draft; both fail on the uncast code. - Cover the top-level mask_token_id fallback, matching the block_size fallback coverage. - Mention both supported block_size locations in the query-length error. Signed-off-by: harshal-96 <harshal.dhandrut@gmail.com>

harshal-96 added 2 commits September 19, 2026 13:20

RunningLeon requested review from RunningLeon and a lite review from Copilot September 20, 2026 03:06

RunningLeon added the improvement label Sep 20, 2026

Copilot started reviewing on behalf of RunningLeon September 20, 2026 03:07 View session

Copilot AI reviewed Sep 20, 2026

View reviewed changes

harshal-96 force-pushed the feat/dflash-qwen3-dense-rebased branch from 63c2073 to 57883c2 Compare September 20, 2026 18:41

Navigation Menu

Support original z-lab Qwen3 DFlash checkpoints; require checkpoint target_layer_ids - #4991

Support original z-lab Qwen3 DFlash checkpoints; require checkpoint target_layer_ids#4991
harshal-96 wants to merge 3 commits into
InternLM:mainfrom
harshal-96:feat/dflash-qwen3-dense-rebased

harshal-96 commented Sep 19, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		raise ValueError('DFlash query length (1 + speculative_num_draft_tokens) must not exceed checkpoint '
		'dflash_config.block_size. '

Conversation

harshal-96 commented Sep 19, 2026

Motivation

Modification

BC-breaking (Optional)

Use cases (Optional)

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants