Skip to content

Support dflash for qwen3.5 - #4789

Open
RunningLeon wants to merge 4 commits into
InternLM:mainfrom
RunningLeon:lmdeploy-dflash
Open

Support dflash for qwen3.5 #4789
RunningLeon wants to merge 4 commits into
InternLM:mainfrom
RunningLeon:lmdeploy-dflash

Conversation

@RunningLeon

@RunningLeon RunningLeon commented Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Motivation

Support dflash in pt

Use cases (Optional)

pipeline

from lmdeploy import PytorchEngineConfig, pipeline
from lmdeploy.messages import SpeculativeConfig

spec_cfg = SpeculativeConfig(
    method='dflash',
    model='z-lab/Qwen3.5-35B-A3B-DFlash',
    dflash_block_size=8,
)
pipe = pipeline(
    'Qwen/Qwen3.5-35B-A3B',
    backend_config=PytorchEngineConfig(tp=2),
    speculative_config=spec_cfg,
)

serving

lmdeploy serve api_server \
Qwen/Qwen3.5-35B-A3B \
--backend pytorch \
--tp 2 \
--speculative-algorithm dflash \
--speculative-draft-model z-lab/Qwen3.5-35B-A3B-DFlash \
--speculative-dflash-block-size 8

When a DFlash block size is provided, it overrides
--speculative-num-draft-tokens by setting the number of newly proposed
tokens to block_size - 1.

Checklist

  1. Pre-commit or other linting tools are used to fix the potential lint issues.
  2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
  3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
  4. The documentation has been modified accordingly, like docstring or example tutorials.

@RunningLeon RunningLeon changed the title [WIP]: Support dflash for qwen3.5 Support dflash for qwen3.5 Jul 31, 2026
@RunningLeon
RunningLeon marked this pull request as ready for review July 31, 2026 10:46
Copilot AI review requested due to automatic review settings July 31, 2026 10:46
Copilot AI reviewed Jul 31, 2026
View reviewed changes

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.

🟡 Human review recommended

It introduces a new speculative decoding method spanning core runtime, model implementations, CLI/config surface area, and benchmarking, which warrants final human validation despite strong unit test coverage.

Pull request overview

This PR adds end-to-end PyTorch speculative decoding support for the DFlash draft method (target-side aux hidden-state extraction + draft KV materialization + one-shot block proposal), including CLI/config plumbing, model implementations for Qwen3/Qwen3.5, and comprehensive tests/docs/benchmark updates.

Changes:

  • Introduce a non-autoregressive “diffusion” proposal path via ProposalMethod/ProposalContext, and implement the DFlash proposer with warmup planning + context KV materialization.
  • Add Qwen-family DFlash draft model (DFlashDraftModel) and extend Qwen3/Qwen3.5 targets to optionally emit concatenated aux_hidden_states for DFlash.
  • Add CLI/config support for --speculative-algorithm dflash and --speculative-dflash-block-size, plus docs, benchmarks, and extensive new unit tests.
File summaries
File Description
tests/pytorch/spec_decode/test_spec_agent.py Extends spec-agent tests to cover DFlash proposal/warmup/reset behaviors and new proposal abstractions.
tests/pytorch/spec_decode/test_dflash_utils.py Adds a broad DFlash-focused test suite (config parsing/validation, model behaviors, proposer utilities).
tests/pytorch/spec_decode/test_cudagraph_strategy.py Updates cudagraph FA3 metadata tests to validate explicit SWA policy handling.
tests/pytorch/engine/test_model_agent.py Adds tests ensuring DFlash build context requires/propagates parsed metadata.
requirements/runtime_cuda.txt Pins flash-linear-attention to >=0.4.2 for CUDA runtime.
lmdeploy/pytorch/spec_decode/spec_agent.py Adds diffusion proposal dispatch, DFlash debug tracing hooks, and proposer-driven warmup plans.
lmdeploy/pytorch/spec_decode/proposers/dflash.py Implements the DFlash proposer (context materialization + one-shot masked block proposal).
lmdeploy/pytorch/spec_decode/proposers/base.py Introduces ProposalMethod, ProposalContext, and proposer warmup plan abstractions.
lmdeploy/pytorch/spec_decode/proposers/init.py Registers DFlash proposer import for module discovery.
lmdeploy/pytorch/spec_decode/guided_spec_helper.py Updates doc reference to the new forward entrypoint.
lmdeploy/pytorch/spec_decode/dflash_utils.py Adds DFlash checkpoint/runtime validation and layer-id resolution logic.
lmdeploy/pytorch/spec_decode/dflash_debug.py Adds opt-in JSONL debug tracing utilities for DFlash.
lmdeploy/pytorch/spec_decode/base.py Adds requires_target_inputs_embeds() capability plumbing for spec agents (notably DFlash).
lmdeploy/pytorch/models/utils/cudagraph.py Refactors FA3 metadata creation to an explicit build_fa3_scheduler_metadata(...) API.
lmdeploy/pytorch/models/qwen3.py Adds aux hidden-state capture and cudagraph output support for DFlash target features.
lmdeploy/pytorch/models/qwen3_dflash.py Adds the DFlash draft model implementation for Qwen-family checkpoints.
lmdeploy/pytorch/models/qwen3_5.py Adds aux hidden-state capture and target-embed policy gating for spec decoding/DFlash.
lmdeploy/pytorch/models/qwen3_5_moe.py Propagates build-context flags needed for spec decoding embed policy decisions.
lmdeploy/pytorch/models/module_map.py Registers DFlashDraftModel in the PyTorch module map.
lmdeploy/pytorch/model_inputs.py Extends BuildModelContext with DFlash-related metadata fields.
lmdeploy/pytorch/engine/model_agent/agent.py Injects DFlash metadata into BuildModelContext and validates required parsed fields.
lmdeploy/pytorch/engine/config_builder.py Validates DFlash runtime envelope and dist constraints when building specdecode config.
lmdeploy/pytorch/configurations/qwen3_5.py Allows spec_method='dflash' for Qwen3.5 configurations.
lmdeploy/pytorch/config.py Adds DFlash spec config fields and parses/validates DFlash checkpoint metadata in SpecDecodeConfig.
lmdeploy/messages.py Adds dflash_block_size and resolves it into num_speculative_tokens in SpeculativeConfig.
lmdeploy/cli/utils.py Adds CLI flags for DFlash and --num-gpu-blocks; wires DFlash block-size override behavior.
lmdeploy/cli/serve.py Exposes --num-gpu-blocks for the PyTorch serving entrypoint.
lmdeploy/cli/cli.py Removes speculative_dflash_block_size from forwarded CLI kwargs (handled via SpeculativeConfig).
docs/zh_cn/advance/spec_decoding.md Documents DFlash usage and the block-size override behavior (Chinese).
docs/en/advance/spec_decoding.md Documents DFlash usage and the block-size override behavior (English).
benchmark/profile_restful_api.py Improves prompt token counting for chat-completion backends via apply_chat_template.
benchmark/benchmark_serving.py Adds Prometheus scraping and CSV summary reporting for speculative decoding acceptance metrics.
Review details
  • Files reviewed: 32/32 changed files
  • Comments generated: 0
  • Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

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

Reviewers

Copilot code review Copilot Copilot left review comments

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@RunningLeon