Support dflash for qwen3.5

[RunningLeon]

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.

[Copilot]

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