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.

[Copilot]

🟡 Not ready to approve

There are confirmed correctness/robustness issues in the new benchmarking metadata handling and SpeculativeConfig validation that should be fixed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (5)

benchmark/benchmark_serving.py:22

• Speculative decoding detection doesn’t consider the new speculative_dflash_block_size knob. If a benchmark config sets only speculative_dflash_block_size (and leaves the other speculative keys unset), the script will skip scraping/spec summary even though spec decoding is enabled.

def _is_spec_decoding(server_config: dict) -> bool:
    """Return whether this engine config enables speculative decoding."""
    return (server_config.get('speculative_algorithm') is not None
            or server_config.get('speculative_num_draft_tokens') is not None
            or server_config.get('speculative_draft_model') is not None)

lmdeploy/pytorch/spec_decode/proposers/base.py:143

• The new propose API signature is misindented (proposal_ctx is not aligned), which reduces readability and can cause style/lint failures. Reformat the signature to a standard multi-line style.

    async def propose(self,
                      model_inputs: ModelInputs,
                      extra_inputs: ExtraInputs,
                      sampling_inputs,
        proposal_ctx: ProposalContext | None = None):

lmdeploy/pytorch/spec_decode/proposers/dflash.py:340

• get_outputs signature indentation is inconsistent with typical Python formatting, which can break linting/formatting and makes the interface harder to scan. Please reformat the signature to match the base proposer style.

    async def get_outputs(self,
                    model_outputs: dict[str, torch.Tensor],
                    model_inputs: ModelInputs,
                    extra_inputs: ARSpecExtraInputs = None,
                    guided_processors: dict | None = None):

lmdeploy/messages.py:787

• SpeculativeConfig.__post_init__ assumes dflash_block_size is an int (and not a bool). Passing a non-int (e.g. a string) will currently raise a TypeError at the < 2 comparison instead of a consistent ValueError, and True is accepted as an int subclass. Add an explicit type check so invalid values reliably raise a helpful ValueError.

        if self.dflash_block_size is not None:
            if self.method != 'dflash':
                raise ValueError('dflash_block_size is supported only when method="dflash".')
            if self.dflash_block_size < 2:
                raise ValueError('dflash_block_size must be an integer greater than or equal to 2.')

benchmark/benchmark_serving.py:135

• Run metadata uses speculative_num_draft_tokens but ignores speculative_dflash_block_size. For DFlash, users typically specify block size (which overrides the draft token count), so the exported num_spec can be wrong/misleading.

    spec_method = server_config.get('speculative_algorithm') or 'none'
    num_spec = server_config.get('speculative_num_draft_tokens')
    if num_spec is None:
        num_spec = 0
    return {

• Files reviewed: 32/32 changed files
• Comments generated: 0 new
• Review effort level: Lite

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