[WIP]: Support dflash for qwen3.5

docs/en/advance/spec_decoding.md

@@ -106,6 +106,47 @@ deepseek-ai/DeepSeek-V3 \
 --enable-metrics
 ```
 
+### DFlash
+
+For DFlash, the block size is the complete draft query/target verification
+window. It includes one current target token, so a block size of 8 proposes
+seven new draft tokens. The block size must not exceed the maximum declared by
+the draft checkpoint.
+
+#### pipeline
+
+```python
+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
+
+```shell
+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`.
+
 ## Guided Decoding with Speculative Decoding
 
 Speculative decoding (MTP) can be combined with [structured output](./structed_output.md) so that the draft tokens proposed by the spec model also respect the grammar constraints (e.g. JSON schema, regex). This significantly improves the acceptance rate compared to running spec decoding without grammar masks.

docs/zh_cn/advance/spec_decoding.md

@@ -105,6 +105,45 @@ deepseek-ai/DeepSeek-V3 \
 --enable-metrics
 ```
 
+### DFlash
+
+对于 DFlash，block size 表示完整的草稿查询/目标验证窗口，其中包含一个当前
+目标 token。因此 block size 为 8 时会提出 7 个新的草稿 token。运行时 block
+size 不能超过草稿模型 checkpoint 中声明的最大值。
+
+#### pipeline
+
+```python
+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
+
+```shell
+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
+```
+
+设置 DFlash block size 后，它会覆盖 `--speculative-num-draft-tokens`，
+并将新提出的 token 数设置为 `block_size - 1`。
+
 ## 投机解码与结构化输出
 
 投机解码（MTP）可以与[结构化输出](./structed_output.md)结合使用，使草稿模型提出的 token 也遵循语法约束（如 JSON Schema、正则表达式），从而显著提高接受率。

lmdeploy/cli/cli.py

@@ -167,7 +167,12 @@ def chat(args):
 
         kwargs = convert_args(args)
         speculative_config = get_speculative_config(args)
-        to_remove = ['speculative_algorithm', 'speculative_draft_model', 'speculative_num_draft_tokens']
+        to_remove = [
+            'speculative_algorithm',
+            'speculative_draft_model',
+            'speculative_num_draft_tokens',
+            'speculative_dflash_block_size',
+        ]
         for key in to_remove:
             kwargs.pop(key)
         kwargs['speculative_config'] = speculative_config

lmdeploy/cli/utils.py

@@ -92,12 +92,16 @@ def get_speculative_config(args):
     """Get speculative config from args."""
     from lmdeploy.messages import SpeculativeConfig
     speculative_config = None
+    dflash_block_size = getattr(args, 'speculative_dflash_block_size', None)
     if args.speculative_algorithm is not None:
         speculative_config = SpeculativeConfig(
             method=args.speculative_algorithm,
             model=args.speculative_draft_model,
             num_speculative_tokens=args.speculative_num_draft_tokens,
+            dflash_block_size=dflash_block_size,
         )
+    elif dflash_block_size is not None:
+        raise ValueError('--speculative-dflash-block-size requires --speculative-algorithm dflash.')
     return speculative_config
 
 
@@ -792,7 +796,7 @@ def add_spec_group(parser):
         spec_group.add_argument('--speculative-algorithm',
                                 type=str,
                                 default=None,
-                                choices=['eagle', 'eagle3', 'deepseek_mtp', 'qwen3_5_mtp'],
+                                choices=['eagle', 'eagle3', 'deepseek_mtp', 'qwen3_5_mtp', 'dflash'],
                                 help='The speculative algorithm to use. `None` means speculative decoding is disabled')
 
         spec_group.add_argument('--speculative-draft-model',
@@ -805,6 +809,14 @@ def add_spec_group(parser):
                                 default=1,
                                 help='The number of speculative tokens to generate per step')
 
+        spec_group.add_argument(
+            '--speculative-dflash-block-size',
+            type=int,
+            default=None,
+            help='DFlash only. Runtime draft block length, including the current target token. '
+            'Overrides --speculative-num-draft-tokens with block_size - 1 proposed tokens and must not exceed '
+            'the DFlash checkpoint block_size.')
+
         return spec_group
 
     @staticmethod

lmdeploy/messages.py

@@ -766,7 +766,22 @@ class SpeculativeConfig:
         method: the speculative decoding method.
         model: the path of speculative model.
         num_speculative_tokens: number of generated token of draft model per step
+        dflash_block_size: DFlash query/verify window length. When set, this
+            DFlash-specific value overrides ``num_speculative_tokens`` using
+            ``num_speculative_tokens = dflash_block_size - 1``.
     """
     method: str
     model: str = ''
     num_speculative_tokens: int = 1
+    dflash_block_size: int | None = None
+
+    def __post_init__(self):
+        """Resolve the DFlash block-size override."""
+        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.')
+            # DFlash's complete query contains the current/next target token in
+            # slot zero followed by the newly proposed draft tokens.
+            self.num_speculative_tokens = self.dflash_block_size - 1

lmdeploy/pytorch/config.py

@@ -555,6 +555,9 @@ def from_hf_config(
         # should after setting `hf_config` and `model_arch` attributes
         model_config = _update_torch_dtype(model_config, dtype, device_type=device_type)
 
+        if spec_method == 'dflash':
+            model_config.model_paradigm = 'ar_spec'
+
         # update eos_token_id to list
         if isinstance(model_config.eos_token_id, int):
             model_config.eos_token_id = [model_config.eos_token_id]
@@ -638,6 +641,8 @@ class SpecDecodeConfig:
     num_speculative_tokens: int = 1
     model_config: ModelConfig = None
     dist_config: DistConfig = field(default_factory=DistConfig)
+    target_layer_ids: tuple[int, ...] | None = None
+    mask_token_id: int | None = None
 
     @classmethod
     def from_config(
@@ -653,9 +658,9 @@ def from_config(
         hf_overrides: dict[str, Any] = None,
         dist_config: DistConfig = None,
     ):
-        model = model or target_model
+        draft_model = model or target_model
         dist_config = dist_config or DistConfig()
-        model_config = ModelConfig.from_pretrained(model,
+        model_config = ModelConfig.from_pretrained(draft_model,
                                                    trust_remote_code=trust_remote_code,
                                                    dtype=dtype,
                                                    dist_config=dist_config,
@@ -666,6 +671,42 @@ def from_config(
                                                    hf_overrides=hf_overrides,
                                                    device_type=target_cache_cfg.device_type,
                                                    )
+        target_layer_ids = None
+        mask_token_id = None
+        if method == 'dflash':
+            from lmdeploy.pytorch.spec_decode.dflash_utils import (
+                parse_dflash_config,
+                validate_dflash_cache_config,
+                validate_dflash_runtime_config,
+            )
+            validate_dflash_cache_config(target_cache_cfg)
+            validate_dflash_runtime_config(cache_config=target_cache_cfg)
+            if target_model is None:
+                raise ValueError('DFlash requires an explicit target_model for checkpoint compatibility checks.')
+            target_model_config = ModelConfig.from_pretrained(
+                target_model,
+                trust_remote_code=trust_remote_code,
+                dtype=dtype,
+                dist_config=dist_config,
+                is_draft_model=False,
+                spec_method=method,
+                num_spec_tokens=num_speculative_tokens,
+                model_format=model_format,
+                hf_overrides=hf_overrides,
+                device_type=target_cache_cfg.device_type,
+                block_size=target_cache_cfg.block_size,
+            )
+            # Hybrid target configs use ``ModelConfig.num_layers`` for the
+            # number of KV-cache attention layers, while DFlash layer ids
+            # address every transformer layer. Prefer the underlying text
+            # config depth and fall back to the generic ModelConfig field.
+            target_num_layers = getattr(target_model_config.llm_config, 'num_hidden_layers',
+                                        target_model_config.num_layers)
+            target_layer_ids, mask_token_id = parse_dflash_config(
+                model_config.hf_config,
+                num_speculative_tokens,
+                target_num_layers=target_num_layers,
+            )
         cache_config = None
         # include medusa
         no_caches = ['medusa']
@@ -682,12 +723,14 @@ def from_config(
                                        quant_policy=target_cache_cfg.quant_policy,
                                        migration_backend=target_cache_cfg.migration_backend)
         obj = cls(
-            model=model,
+            model=draft_model,
             method=method,
             cache_config=cache_config,
             model_config=model_config,
             dist_config=dist_config,
             num_speculative_tokens=num_speculative_tokens,
+            target_layer_ids=target_layer_ids,
+            mask_token_id=mask_token_id,
         )
         return obj
 

lmdeploy/pytorch/configurations/qwen3_5.py

@@ -85,18 +85,19 @@ def build(cls,
 
         # for spec
         if spec_method is not None:
-            assert spec_method == 'qwen3_5_mtp'
+            assert spec_method in ['qwen3_5_mtp', 'dflash']
             cfg.model_paradigm = 'ar_spec'
 
         # draft model cfg
         if is_draft_model:
-            hf_config.architectures[0] = 'Qwen3_5MTPModel'
-            # remove for correct mapping when building the patched model
-            if hasattr(hf_config, 'auto_map'):
-                del hf_config.auto_map
-
             cfg.model_paradigm = 'ar_spec'
-            cfg.num_layers = text_config.mtp_num_hidden_layers
             cfg.states_shapes = []
 
+            if spec_method == 'qwen3_5_mtp':
+                hf_config.architectures[0] = 'Qwen3_5MTPModel'
+                # remove for correct mapping when building the patched model
+                if hasattr(hf_config, 'auto_map'):
+                    del hf_config.auto_map
+                cfg.num_layers = text_config.mtp_num_hidden_layers
+
         return cfg

lmdeploy/pytorch/engine/config_builder.py

@@ -119,6 +119,14 @@ def _build_draft_dist_ctx(dist_config):
             # TODO support tp > 1, ep > 1 for other methods
             if speculative_config.method == 'qwen3_5_mtp':
                 draft_dist_config = dist_config
+            elif speculative_config.method == 'dflash':
+                from lmdeploy.pytorch.spec_decode.dflash_utils import (
+                    validate_dflash_dist_config,
+                    validate_dflash_runtime_config,
+                )
+                validate_dflash_dist_config(dist_config)
+                validate_dflash_runtime_config(cache_config=cache_config, backend_config=engine_config)
+                draft_dist_config = copy.deepcopy(dist_config)
             else:
                 draft_dist_config = DistConfig()
             return draft_dist_config

lmdeploy/pytorch/engine/model_agent/agent.py

@@ -1117,6 +1117,19 @@ def _build_model(self):
         # for router replay
         enable_return_routed_experts = self.misc_config.enable_return_routed_experts and self.need_output
 
+        target_aux_hidden_state_layers = ()
+        speculative_mask_token_id = None
+        if self.spec_agent.is_enabled() and self.spec_agent.method == 'dflash':
+            target_layer_ids = self.spec_agent.specdecode_config.target_layer_ids
+            if target_layer_ids is None:
+                raise ValueError('DFlash target model construction requires parsed target_layer_ids metadata.')
+            target_aux_hidden_state_layers = tuple(int(layer_id) for layer_id in target_layer_ids)
+            mask_token_id = self.spec_agent.specdecode_config.mask_token_id
+            if mask_token_id is None:
+                raise ValueError('DFlash model construction requires parsed mask_token_id metadata.')
+            speculative_mask_token_id = int(mask_token_id)
+        requires_target_inputs_embeds = self.spec_agent.requires_target_inputs_embeds()
+
         build_model_ctx = BuildModelContext(language_model_only=self.misc_config.language_model_only,
                                             dllm_config=self.misc_config.dllm_config,
                                             strategy_factory=self.strategy_factory,
@@ -1125,7 +1138,10 @@ def _build_model(self):
                                             fp32_lm_head=self.model_config.fp32_lm_head,
                                             tie_word_embeddings=self.model_config.tie_word_embeddings,
                                             num_spec_tokens=self.spec_agent.num_spec_tokens,
-                                            max_batch_size=self.cache_config.max_batches)
+                                            max_batch_size=self.cache_config.max_batches,
+                                            target_aux_hidden_state_layers=target_aux_hidden_state_layers,
+                                            speculative_mask_token_id=speculative_mask_token_id,
+                                            requires_target_inputs_embeds=requires_target_inputs_embeds)
         patched_model = build_patched_model(self.model_config, device=device, build_model_ctx=build_model_ctx)
         logger.debug(msg_with_rank(rank, 'loading weights.'))
         if not self.misc_config.empty_init:

lmdeploy/pytorch/model_inputs.py

@@ -476,6 +476,9 @@ class BuildModelContext:
     tie_word_embeddings: bool = False
     num_spec_tokens: int = 0
     max_batch_size: int = 0
+    target_aux_hidden_state_layers: tuple[int, ...] = ()
+    speculative_mask_token_id: int | None = None
+    requires_target_inputs_embeds: bool = True
 
     @property
     def deep_ep_max_tokens_per_rank(self) -> int:

lmdeploy/pytorch/models/module_map.py

@@ -299,3 +299,6 @@
 
 # deepseek mtp
 MODULE_MAP.update({'DeepseekMTPModel': f'{LMDEPLOY_PYTORCH_MODEL_PATH}.deepseek_mtp.DeepseekMTPModel'})
+
+# dflash qwen
+MODULE_MAP.update({'DFlashDraftModel': f'{LMDEPLOY_PYTORCH_MODEL_PATH}.qwen3_dflash.DFlashDraftModel'})