From 715993bfcfa2ab58c3631470607b6b60b769f136 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Fri, 24 Jul 2026 11:38:05 +0000 Subject: [PATCH 1/3] support dflash --- docs/en/advance/spec_decoding.md | 41 + docs/zh_cn/advance/spec_decoding.md | 39 + lmdeploy/cli/cli.py | 7 +- lmdeploy/cli/utils.py | 14 +- lmdeploy/messages.py | 15 + lmdeploy/pytorch/config.py | 49 +- lmdeploy/pytorch/configurations/qwen3_5.py | 15 +- lmdeploy/pytorch/engine/config_builder.py | 8 + lmdeploy/pytorch/engine/model_agent/agent.py | 18 +- lmdeploy/pytorch/model_inputs.py | 3 + lmdeploy/pytorch/models/module_map.py | 3 + lmdeploy/pytorch/models/qwen3.py | 17 +- lmdeploy/pytorch/models/qwen3_5.py | 49 +- lmdeploy/pytorch/models/qwen3_5_moe.py | 5 +- lmdeploy/pytorch/models/qwen3_dflash.py | 578 +++++++++++ lmdeploy/pytorch/models/utils/cudagraph.py | 34 +- lmdeploy/pytorch/spec_decode/base.py | 9 + lmdeploy/pytorch/spec_decode/dflash_debug.py | 48 + lmdeploy/pytorch/spec_decode/dflash_utils.py | 192 ++++ .../pytorch/spec_decode/guided_spec_helper.py | 2 +- .../pytorch/spec_decode/proposers/__init__.py | 1 + .../pytorch/spec_decode/proposers/base.py | 60 +- .../pytorch/spec_decode/proposers/dflash.py | 354 +++++++ lmdeploy/pytorch/spec_decode/spec_agent.py | 124 ++- tests/pytorch/engine/test_model_agent.py | 70 ++ .../spec_decode/test_cudagraph_strategy.py | 19 +- .../pytorch/spec_decode/test_dflash_utils.py | 906 ++++++++++++++++++ tests/pytorch/spec_decode/test_spec_agent.py | 235 ++++- 28 files changed, 2841 insertions(+), 74 deletions(-) create mode 100644 lmdeploy/pytorch/models/qwen3_dflash.py create mode 100644 lmdeploy/pytorch/spec_decode/dflash_debug.py create mode 100644 lmdeploy/pytorch/spec_decode/dflash_utils.py create mode 100644 lmdeploy/pytorch/spec_decode/proposers/dflash.py create mode 100644 tests/pytorch/spec_decode/test_dflash_utils.py diff --git a/docs/en/advance/spec_decoding.md b/docs/en/advance/spec_decoding.md index f804532da1..2fbd4d60f4 100644 --- a/docs/en/advance/spec_decoding.md +++ b/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. diff --git a/docs/zh_cn/advance/spec_decoding.md b/docs/zh_cn/advance/spec_decoding.md index f68ea27d4e..090326a3d0 100644 --- a/docs/zh_cn/advance/spec_decoding.md +++ b/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、正则表达式），从而显著提高接受率。 diff --git a/lmdeploy/cli/cli.py b/lmdeploy/cli/cli.py index 1276ea859d..93c3232a6f 100644 --- a/lmdeploy/cli/cli.py +++ b/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 diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 8d5f7d31d0..2bf04bb018 100644 --- a/lmdeploy/cli/utils.py +++ b/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 diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index e07a46ae5e..4ef1edd76c 100644 --- a/lmdeploy/messages.py +++ b/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 diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index f730c985fd..de79079796 100644 --- a/lmdeploy/pytorch/config.py +++ b/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 diff --git a/lmdeploy/pytorch/configurations/qwen3_5.py b/lmdeploy/pytorch/configurations/qwen3_5.py index d6e585b217..f4e12c1217 100644 --- a/lmdeploy/pytorch/configurations/qwen3_5.py +++ b/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 diff --git a/lmdeploy/pytorch/engine/config_builder.py b/lmdeploy/pytorch/engine/config_builder.py index 450ab58a60..46fb5898be 100644 --- a/lmdeploy/pytorch/engine/config_builder.py +++ b/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 diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py index 09f0a78237..92d782b3eb 100644 --- a/lmdeploy/pytorch/engine/model_agent/agent.py +++ b/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: diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index 39d68dd854..5fb9df097a 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/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: diff --git a/lmdeploy/pytorch/models/module_map.py b/lmdeploy/pytorch/models/module_map.py index 27e74e8e67..6c7932334a 100644 --- a/lmdeploy/pytorch/models/module_map.py +++ b/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'}) diff --git a/lmdeploy/pytorch/models/qwen3.py b/lmdeploy/pytorch/models/qwen3.py index 11326d3225..8b517d8b09 100644 --- a/lmdeploy/pytorch/models/qwen3.py +++ b/lmdeploy/pytorch/models/qwen3.py @@ -12,7 +12,7 @@ from lmdeploy.pytorch.nn.linear import build_down_linear, build_gateup_linear, build_o_proj, build_qkv_proj from lmdeploy.pytorch.weight_loader.model_weight_loader import load_weight -from .patch import add_prefix +from .patch import add_prefix, get_build_model_context from .utils.cudagraph import CudaGraphMixin from .utils.model import DeployModelMixinV1, build_embedding @@ -265,6 +265,8 @@ def __init__(self, prefix=add_prefix(f'layers.{layer_idx}', prefix)) for layer_idx in range(config.num_hidden_layers) ]) + self.aux_hidden_state_layers: tuple[int, ...] = get_build_model_context().target_aux_hidden_state_layers + self._aux_hidden_state_layers_set: frozenset[int] = frozenset(self.aux_hidden_state_layers) # build norm self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, dtype=dtype, device=device) @@ -294,6 +296,7 @@ def forward( rotary_pos_emb = (cos, sin) # decoding + aux_hidden_states = [] residual = None for idx, decoder_layer in enumerate(self.layers): past_key_value = past_key_values[idx] @@ -304,10 +307,14 @@ def forward( residual=residual, attn_metadata=attn_metadata, ) + if idx in self._aux_hidden_state_layers_set: + aux_hidden_states.append(hidden_states if residual is None else hidden_states + residual) # norm hidden_states, _ = self.norm(hidden_states, residual) + if len(aux_hidden_states) > 0: + return dict(hidden_states=hidden_states, aux_hidden_states=torch.cat(aux_hidden_states, dim=-1)) return hidden_states def get_input_embeddings(self): @@ -367,6 +374,14 @@ def get_input_embeddings(self): """Get input embeddings.""" return self.model.get_input_embeddings() + def get_outputs_cudagraph(self, output_buffers: dict[str, torch.Tensor], input_ids: torch.Tensor, **kwargs): + """Return Qwen3 target outputs captured by a decode graph.""" + outputs = super().get_outputs_cudagraph(output_buffers, input_ids, **kwargs) + aux_hidden_states = output_buffers.get('aux_hidden_states') + if aux_hidden_states is not None: + outputs['aux_hidden_states'] = aux_hidden_states[:, :input_ids.size(-1)] + return outputs + def prepare_inputs_for_generation( self, past_key_values: list[list[torch.Tensor]], diff --git a/lmdeploy/pytorch/models/qwen3_5.py b/lmdeploy/pytorch/models/qwen3_5.py index 2544e38ef7..1e1b8a1201 100644 --- a/lmdeploy/pytorch/models/qwen3_5.py +++ b/lmdeploy/pytorch/models/qwen3_5.py @@ -836,6 +836,8 @@ def __init__(self, prefix=add_prefix(f'layers.{layer_idx}', prefix)) for layer_idx in range(self.config.num_hidden_layers) ]) + self.aux_hidden_state_layers: tuple[int, ...] = get_build_model_context().target_aux_hidden_state_layers + self._aux_hidden_state_layers_set: frozenset[int] = frozenset(self.aux_hidden_state_layers) # build norm self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, dtype=dtype, device=device) @@ -877,6 +879,7 @@ def forward( attn_metadata) # decoding + aux_hidden_states = [] residual = None for idx, decoder_layer in enumerate(self.layers): hidden_states, residual = decoder_layer( @@ -888,10 +891,14 @@ def forward( gated_delta_meta=gated_delta_meta, all_routed_experts=all_routed_experts, ) + if idx in self._aux_hidden_state_layers_set: + aux_hidden_states.append(hidden_states if residual is None else hidden_states + residual) # norm hidden_states, _ = self.norm(hidden_states, residual) + if len(aux_hidden_states) > 0: + return dict(hidden_states=hidden_states, aux_hidden_states=torch.cat(aux_hidden_states, dim=-1)) return hidden_states def get_input_embeddings(self): @@ -972,7 +979,7 @@ def forward( output_inputs_embeds = inputs_embeds if return_input_embeds else None - hidden_states = self.language_model( + language_outputs = self.language_model( input_ids=input_ids, position_ids=position_ids, past_key_values=past_key_values, @@ -982,7 +989,13 @@ def forward( mrope_position_ids=mrope_position_ids, all_routed_experts=all_routed_experts, ) - return hidden_states, output_inputs_embeds + aux_hidden_states = None + if isinstance(language_outputs, dict): + hidden_states = language_outputs['hidden_states'] + aux_hidden_states = language_outputs.get('aux_hidden_states') + else: + hidden_states = language_outputs + return hidden_states, output_inputs_embeds, aux_hidden_states def get_input_embeddings(self): """Get input embeddings.""" @@ -1027,8 +1040,9 @@ def __init__(self, device=device) # dense model self.enable_return_routed_experts = False - self.is_spec_decoding = get_build_model_context().num_spec_tokens > 0 - + bm_ctx = get_build_model_context() + self.is_spec_decoding = bm_ctx.num_spec_tokens > 0 + self.requires_target_inputs_embeds = bm_ctx.requires_target_inputs_embeds def forward( self, @@ -1060,7 +1074,7 @@ def forward( all_routed_experts = position_ids.new_empty( (num_tokens, config.num_hidden_layers, config.num_experts_per_tok), dtype=torch.uint16) - hidden_states, target_inputs_embeds = self.model( + hidden_states, target_inputs_embeds, aux_hidden_states = self.model( input_ids=input_ids, position_ids=position_ids, past_key_values=past_key_values, @@ -1081,14 +1095,31 @@ def forward( ts_lens=ts_lens, ts_sr=ts_sr, ) - return dict(hidden_states=hidden_states, - all_routed_experts=all_routed_experts, - target_inputs_embeds=target_inputs_embeds) + outputs = dict(hidden_states=hidden_states, + all_routed_experts=all_routed_experts, + target_inputs_embeds=target_inputs_embeds) + if aux_hidden_states is not None: + outputs['aux_hidden_states'] = aux_hidden_states + return outputs def get_input_embeddings(self): """Get input embeddings.""" return self.model.get_input_embeddings() + def get_outputs_cudagraph(self, output_buffers: dict[str, torch.Tensor], input_ids: torch.Tensor, **kwargs): + """Return Qwen3.5 target outputs captured by a decode graph.""" + outputs = super().get_outputs_cudagraph(output_buffers, input_ids, **kwargs) + aux_hidden_states = output_buffers.get('aux_hidden_states') + if aux_hidden_states is not None: + outputs['aux_hidden_states'] = aux_hidden_states[:, :input_ids.size(-1)] + return outputs + + def _should_return_target_inputs_embeds(self, pixel_values: torch.Tensor | None, + context: StepContext) -> bool: + """Whether a shifted-input proposer needs multimodal embeddings.""" + return (self.is_spec_decoding and self.requires_target_inputs_embeds + and (pixel_values is not None or context.is_chunk_multimodal)) + def prepare_inputs_for_generation( self, past_key_values: list[list[torch.Tensor]], @@ -1158,7 +1189,7 @@ def prepare_inputs_for_generation( inputs_embeds[:, vision_embedding_indexing, :] = vision_embeddings.to(inputs_embeds) # return input embeds for spec decoding - return_input_embeds = self.is_spec_decoding and (pixel_values is not None or context.is_chunk_multimodal) + return_input_embeds = self._should_return_target_inputs_embeds(pixel_values, context) # inputs of forward return dict( diff --git a/lmdeploy/pytorch/models/qwen3_5_moe.py b/lmdeploy/pytorch/models/qwen3_5_moe.py index 616a4d01d0..944a2eeb54 100644 --- a/lmdeploy/pytorch/models/qwen3_5_moe.py +++ b/lmdeploy/pytorch/models/qwen3_5_moe.py @@ -206,6 +206,8 @@ def __init__(self, prefix=add_prefix(f'layers.{layer_idx}', prefix)) for layer_idx in range(self.config.num_hidden_layers) ]) + self.aux_hidden_state_layers: tuple[int, ...] = get_build_model_context().target_aux_hidden_state_layers + self._aux_hidden_state_layers_set: frozenset[int] = frozenset(self.aux_hidden_state_layers) # build norm self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, dtype=dtype, device=device) @@ -275,7 +277,8 @@ def __init__(self, # for router replay bm_ctx = get_build_model_context() self.enable_return_routed_experts = bm_ctx.enable_return_routed_experts - self.is_spec_decoding = get_build_model_context().num_spec_tokens > 0 + self.is_spec_decoding = bm_ctx.num_spec_tokens > 0 + self.requires_target_inputs_embeds = bm_ctx.requires_target_inputs_embeds def _load_weight_experts(self, name: str, loaded_weight: torch.Tensor, params_dict: dict[str, nn.Parameter]): """Load weight experts.""" diff --git a/lmdeploy/pytorch/models/qwen3_dflash.py b/lmdeploy/pytorch/models/qwen3_dflash.py new file mode 100644 index 0000000000..de6680d4ac --- /dev/null +++ b/lmdeploy/pytorch/models/qwen3_dflash.py @@ -0,0 +1,578 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import replace +from typing import Any + +import torch +from torch import nn +from transformers.configuration_utils import PretrainedConfig + +from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager +from lmdeploy.pytorch.nn import ApplyRotaryEmb, Attention, RMSNorm, build_rotary_embedding_from_config +from lmdeploy.pytorch.nn.linear import ( + build_colwise_linear, + build_o_proj, + build_qkv_proj, +) +from lmdeploy.pytorch.weight_loader.model_weight_loader import load_weight + +from .patch import add_prefix, get_build_model_context +from .qwen3 import Qwen3MLP +from .utils.cudagraph import CudaGraphMeta, CudaGraphMixin + +_SLIDING_ATTENTION = 'sliding_attention' +_DFLASH_FULL_SCHEDULER_METADATA = 'dflash_full_scheduler_metadata' + + +def _normalize_dflash_weight_name(name: str) -> str | None: + """Normalize common DFlash checkpoint weight names to this module + layout.""" + if 'rotary_emb.inv_freq' in name: + return None + if 'rotary_emb.cos_cached' in name or 'rotary_emb.sin_cached' in name: + return None + if name.startswith('model.'): + name = name[len('model.'):] + if name.startswith('midlayer.'): + name = name.replace('midlayer.', 'layers.0.', 1) + return name + + +def _resolve_dflash_layer_attention(config: PretrainedConfig, layer_idx: int) -> tuple[int | None, bool]: + """Consume the attention policy already validated by the outer parser.""" + layer_type = config.layer_types[layer_idx] + sliding_window = config.sliding_window if layer_type == _SLIDING_ATTENTION else None + default_causal = layer_type == _SLIDING_ATTENTION + causal_override = config.dflash_config.get('causal', getattr(config, 'causal', None)) + causal = default_causal if causal_override is None else bool(causal_override) + return sliding_window, causal + + +class DFlashQwen3Attention(nn.Module): + """Qwen3 attention used by DFlash draft layers.""" + + def __init__(self, + config: PretrainedConfig, + layer_idx: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + prefix: str = ''): + super().__init__() + quantization_config = getattr(config, 'quantization_config', None) + num_heads = config.num_attention_heads + num_key_value_heads = config.num_key_value_heads + hidden_size = config.hidden_size + head_dim = getattr(config, 'head_dim', hidden_size // num_heads) + num_replicate_kv_heads = getattr(config, 'num_replicate_key_value_heads', 1) + sliding_window, causal = _resolve_dflash_layer_attention(config, layer_idx) + self.layer_idx = layer_idx + self.head_dim = head_dim + self.causal = causal + self.sliding_window = sliding_window + + self.qkv_proj = build_qkv_proj( + hidden_size, + num_q_heads=num_heads, + num_kv_heads=num_key_value_heads, + head_size=head_dim, + bias=getattr(config, 'attention_bias', False), + quant_config=quantization_config, + dtype=dtype, + device=device, + num_replicate_kv_heads=num_replicate_kv_heads, + prefix=add_prefix('qkv_proj', prefix), + ) + self.apply_rotary_pos_emb = ApplyRotaryEmb() + self.attn_fwd = Attention( + num_heads, + head_dim, + num_kv_heads=num_key_value_heads, + v_head_size=head_dim, + sliding_window=sliding_window, + causal=causal, + ) + self.o_proj = build_o_proj( + num_heads * head_dim, + hidden_size, + bias=getattr(config, 'attention_bias', False), + quant_config=quantization_config, + dtype=dtype, + device=device, + is_tp=True, + prefix=add_prefix('o_proj', prefix), + ) + self.q_norm = RMSNorm( + head_dim, + config.rms_norm_eps, + quant_config=quantization_config, + dtype=dtype, + device=device, + prefix=add_prefix('q_norm', prefix), + ) + self.k_norm = RMSNorm( + head_dim, + config.rms_norm_eps, + quant_config=quantization_config, + dtype=dtype, + device=device, + prefix=add_prefix('k_norm', prefix), + ) + + def _split_qkv(self, hidden_states: torch.Tensor): + """Project and split Q/K/V states.""" + qkv_states = self.qkv_proj(hidden_states) + qkv_states = qkv_states.flatten(0, -2) + return self.qkv_proj.split_qkv(qkv_states) + + def kv_proj_only(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Project hidden states to K/V for DFlash context materialization.""" + _, key_states, value_states = self._split_qkv(hidden_states) + return key_states, value_states + + def apply_k_norm(self, key_states: torch.Tensor) -> torch.Tensor: + """Apply per-head K RMSNorm.""" + return self.k_norm(key_states) + + def apply_k_rope(self, key_states: torch.Tensor, + rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor]) -> torch.Tensor: + """Apply RoPE to K states using a disposable Q tensor.""" + cos, sin = rotary_pos_emb + dummy_query = torch.empty_like(key_states) + _, key_states = self.apply_rotary_pos_emb(dummy_query, key_states, cos, sin, inplace=True) + return key_states + + def project_context_kv( + self, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return projected, normalized, RoPE'd context K/V states.""" + key_states, value_states = self.kv_proj_only(hidden_states) + key_states = self.apply_k_norm(key_states) + key_states = self.apply_k_rope(key_states, rotary_pos_emb) + return key_states, value_states + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor], + past_key_value: tuple[torch.Tensor], + attn_metadata: Any, + ): + """Forward a DFlash query block.""" + query_states, key_states, value_states = self._split_qkv(hidden_states) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + cos, sin = rotary_pos_emb + query_states, key_states = self.apply_rotary_pos_emb( + query_states, + key_states, + cos, + sin, + inplace=True, + ) + + attn_output = self.attn_fwd( + query_states, + key_states, + value_states, + past_key_value[0], + past_key_value[1], + attn_metadata, + k_scales_zeros=None if len(past_key_value) == 2 else past_key_value[2], + v_scales_zeros=None if len(past_key_value) == 2 else past_key_value[3], + inplace=True, + ) + attn_output = attn_output.reshape(*hidden_states.shape[:-1], -1) + return self.o_proj(attn_output) + + +class DFlashQwen3DecoderLayer(nn.Module): + """Decoder layer for standard Qwen-family DFlash draft checkpoints.""" + + def __init__(self, + config: PretrainedConfig, + layer_idx: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + prefix: str = ''): + super().__init__() + self.layer_idx = layer_idx + quantization_config = getattr(config, 'quantization_config', None) + self.self_attn = DFlashQwen3Attention( + config, + layer_idx, + dtype=dtype, + device=device, + prefix=add_prefix('self_attn', prefix), + ) + self.mlp = Qwen3MLP(config, dtype=dtype, device=device, prefix=add_prefix('mlp', prefix)) + self.input_layernorm = RMSNorm( + config.hidden_size, + config.rms_norm_eps, + quant_config=quantization_config, + dtype=dtype, + device=device, + prefix=add_prefix('input_layernorm', prefix), + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, + config.rms_norm_eps, + quant_config=quantization_config, + dtype=dtype, + device=device, + prefix=add_prefix('post_attention_layernorm', prefix), + ) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor], + past_key_value: list[torch.FloatTensor] | None, + attn_metadata: Any, + residual: torch.Tensor | None = None, + dflash_full_scheduler_metadata: torch.Tensor | None = None, + ): + """Forward one decoder layer.""" + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + layer_attn_metadata = self._get_layer_attn_metadata(attn_metadata, dflash_full_scheduler_metadata) + hidden_states = self.self_attn( + hidden_states=hidden_states, + rotary_pos_emb=rotary_pos_emb, + past_key_value=past_key_value, + attn_metadata=layer_attn_metadata, + ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + def _get_layer_attn_metadata(self, attn_metadata: Any, + dflash_full_scheduler_metadata: torch.Tensor | None): + """Return a layer-local view only for non-causal full decode.""" + is_full_noncausal = self.self_attn.sliding_window is None and not self.self_attn.causal + if not attn_metadata.is_decoding or not is_full_noncausal: + return attn_metadata + return replace(attn_metadata, scheduler_metadata=dflash_full_scheduler_metadata) + + +class DFlashDraftModel(nn.Module, CudaGraphMixin): + """Qwen-family DFlash draft model. + + z-lab DFlash draft checkpoints carry only the draft transformer and feature-fusion weights. Embeddings and logits + are shared from the target model by the proposer. + """ + + packed_modules_mapping = { + 'qkv_proj': [ + 'q_proj', + 'k_proj', + 'v_proj', + ], + 'gate_up_proj': [ + 'gate_proj', + 'up_proj', + ], + } + + def __init__(self, + config: PretrainedConfig, + ctx_mgr: StepContextManager, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + prefix: str = ''): + super().__init__() + self.config = config + self.ctx_mgr = ctx_mgr + self.dtype = dtype + build_ctx = get_build_model_context() + self.target_layer_ids = build_ctx.target_aux_hidden_state_layers + if not self.target_layer_ids: + raise ValueError('DFlash draft construction requires resolved target_aux_hidden_state_layers metadata.') + self.mask_token_id = build_ctx.speculative_mask_token_id + if self.mask_token_id is None: + raise ValueError('DFlash draft construction requires resolved speculative_mask_token_id metadata.') + self.num_context_features = len(self.target_layer_ids) + target_hidden_size = int(getattr(config, 'target_hidden_size', config.hidden_size)) + fc_input_size = target_hidden_size * self.num_context_features + quantization_config = getattr(config, 'quantization_config', None) + + self.embed_tokens = None + self.mask_embedding = nn.Parameter( + torch.zeros(config.hidden_size, dtype=dtype or torch.float16, device=device), + requires_grad=False, + ) + self.has_separate_mask_embedding = False + + self.layers = nn.ModuleList([ + DFlashQwen3DecoderLayer(config, + layer_idx, + dtype=dtype, + device=device, + prefix=add_prefix(f'layers.{layer_idx}', prefix)) + for layer_idx in range(config.num_hidden_layers) + ]) + self.fc = build_colwise_linear( + fc_input_size, + config.hidden_size, + bias=False, + dtype=dtype, + device=device, + is_tp=False, + quant_config=quantization_config, + check_dist=False, + prefix=add_prefix('fc', prefix), + ) + self.hidden_norm = RMSNorm( + config.hidden_size, + config.rms_norm_eps, + quant_config=quantization_config, + dtype=dtype, + device=device, + prefix=add_prefix('hidden_norm', prefix), + ) + self.norm = RMSNorm( + config.hidden_size, + config.rms_norm_eps, + quant_config=quantization_config, + dtype=dtype, + device=device, + prefix=add_prefix('norm', prefix), + ) + self.rotary_emb = build_rotary_embedding_from_config(config, device=device) + + def set_input_embeddings(self, embed_tokens: nn.Module): + """Set target-shared token embeddings.""" + self.embed_tokens = embed_tokens + + def get_input_embeddings(self): + """Get target-shared token embeddings.""" + return self.embed_tokens + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + """Embed query ids, replacing mask slots when a separate mask vector is + loaded.""" + if self.embed_tokens is None: + raise RuntimeError('DFlash draft model requires target-shared input embeddings.') + embeds = self.embed_tokens(input_ids) + if self.has_separate_mask_embedding and self.mask_token_id is not None: + mask = (input_ids == int(self.mask_token_id)).unsqueeze(-1) + embeds = torch.where(mask, self.mask_embedding.to(dtype=embeds.dtype), embeds) + return embeds + + def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: + """Project concatenated target-layer hidden states into draft hidden + size.""" + expected = int(self.fc.in_features) + if target_hidden.ndim != 2 or int(target_hidden.shape[-1]) != expected: + raise ValueError('DFlash target hidden feature dim mismatch. ' + f'Expected shape [N, {expected}] from {self.num_context_features} target layers, ' + f'got {tuple(target_hidden.shape)}.') + return self.hidden_norm(self.fc(target_hidden)) + + def _rotary_pos_emb_for_context( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build rotary embeddings for flattened context states.""" + if position_ids is None: + raise ValueError('DFlash context KV materialization requires position_ids.') + if position_ids.dim() == 1: + rope_position_ids = position_ids.unsqueeze(0) + else: + rope_position_ids = position_ids + cos, sin = self.rotary_emb(hidden_states.unsqueeze(0), rope_position_ids) + return cos[0], sin[0] + + @torch.inference_mode() + def precompute_context_kv( + self, + target_hidden: torch.Tensor, + position_ids: torch.Tensor, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Project target features into per-layer draft context K/V tensors.""" + context_states = self.project_target_hidden(target_hidden) + rotary_pos_emb = self._rotary_pos_emb_for_context(context_states, position_ids) + return [layer.self_attn.project_context_kv(context_states, rotary_pos_emb) for layer in self.layers] + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + target_hidden: torch.Tensor, + position_ids: torch.Tensor, + past_key_values: list[list[torch.Tensor]] | None = None, + attn_metadata: Any | None = None, + max_q_seqlen: int | None = None, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Project context K/V and optionally write them into the draft KV + cache. + + When cache arguments are omitted, this returns the materialized K/V tensors for tests, alignment dumps, or + proposer-side staging. + """ + per_layer_kv = self.precompute_context_kv(target_hidden, position_ids) + if past_key_values is None or attn_metadata is None: + return per_layer_kv + if max_q_seqlen is None: + raise ValueError('DFlash draft KV materialization requires CPU max_q_seqlen metadata.') + + from lmdeploy.pytorch.kernels.cuda import fill_kv_cache + + for (key_states, value_states), past_key_value in zip(per_layer_kv, past_key_values, strict=True): + fill_kv_cache( + key_states, + value_states, + past_key_value[0], + past_key_value[1], + attn_metadata.q_start_loc, + attn_metadata.q_seqlens, + kv_seq_length=attn_metadata.kv_seqlens, + max_q_seq_length=int(max_q_seqlen), + block_offsets=attn_metadata.block_offsets, + k_scales_zeros=None if len(past_key_value) == 2 else past_key_value[2], + v_scales_zeros=None if len(past_key_value) == 2 else past_key_value[3], + quant_policy=attn_metadata.quant_policy, + ) + return per_layer_kv + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + past_key_values: list[list[torch.Tensor]], + attn_metadata: Any, + inputs_embeds: torch.Tensor | None = None, + dflash_full_scheduler_metadata: torch.Tensor | None = None, + **kwargs, + ): + """Forward a DFlash query block.""" + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + + hidden_states = inputs_embeds + cos, sin = self.rotary_emb(hidden_states, position_ids) + cos, sin = cos[0], sin[0] + rotary_pos_emb = (cos, sin) + + residual = None + for idx, decoder_layer in enumerate(self.layers): + hidden_states, residual = decoder_layer( + hidden_states, + rotary_pos_emb=rotary_pos_emb, + past_key_value=past_key_values[idx], + residual=residual, + attn_metadata=attn_metadata, + dflash_full_scheduler_metadata=dflash_full_scheduler_metadata, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + def prepare_inputs_for_generation( + self, + past_key_values: list[list[torch.Tensor]], + inputs_embeds: torch.Tensor | None = None, + context: StepContext | None = None, + ): + """Prepare DFlash draft model inputs.""" + return dict( + input_ids=context.input_ids, + position_ids=context.position_ids, + past_key_values=past_key_values, + attn_metadata=context.attn_metadata, + inputs_embeds=inputs_embeds, + dflash_full_scheduler_metadata=None, + ) + + def make_buffers_cudagraph(self, graph_meta: CudaGraphMeta, **kwargs): + """Add the DFlash full/non-causal FA3 scheduler buffer.""" + input_buffers = super().make_buffers_cudagraph(graph_meta=graph_meta, **kwargs) + if graph_meta.use_fa3_decoding: + max_seqlen_k = graph_meta.num_blocks * graph_meta.block_size + input_buffers[_DFLASH_FULL_SCHEDULER_METADATA] = self.build_fa3_scheduler_metadata( + graph_meta.max_batchs, + graph_meta.decode_query_len, + block_size=graph_meta.block_size, + max_seqlen_k=max_seqlen_k, + cache_seqlens=input_buffers['kv_seqlens'], + sliding_window=None, + causal=False, + ) + return input_buffers + + def fill_buffers_cudagraph(self, graph_meta: CudaGraphMeta, **kwargs): + """Refresh and expose the graph bucket's full/non-causal buffer.""" + new_inputs = super().fill_buffers_cudagraph(graph_meta=graph_meta, **kwargs) + if graph_meta.use_fa3_decoding: + input_buffers = graph_meta.input_buffers + max_seqlen_k = graph_meta.num_blocks * graph_meta.block_size + scheduler_metadata = self.build_fa3_scheduler_metadata( + graph_meta.max_batchs, + graph_meta.decode_query_len, + block_size=graph_meta.block_size, + max_seqlen_k=max_seqlen_k, + cache_seqlens=input_buffers['kv_seqlens'], + sliding_window=None, + causal=False, + ) + num_meta = scheduler_metadata.size(0) + full_buffer = input_buffers[_DFLASH_FULL_SCHEDULER_METADATA] + full_buffer[:num_meta].copy_(scheduler_metadata) + full_buffer[num_meta:].zero_() + new_inputs[_DFLASH_FULL_SCHEDULER_METADATA] = full_buffer[:num_meta] + else: + new_inputs[_DFLASH_FULL_SCHEDULER_METADATA] = None + return new_inputs + + def update_model_metas( + self, + past_key_values: list[list[torch.Tensor]], + inputs_embeds: torch.Tensor | None = None, + context: StepContext | None = None, + ): + """Return draft model metadata.""" + return None + + @classmethod + def rename_weight(cls, name: str) -> str: + """Rename loaded checkpoint weights.""" + return _normalize_dflash_weight_name(name) or name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + """Load DFlash draft checkpoint weights.""" + stacked_params_mapping = [ + ('.qkv_proj', '.q_proj', 'q'), + ('.qkv_proj', '.k_proj', 'k'), + ('.qkv_proj', '.v_proj', 'v'), + ('.gate_up_proj', '.gate_proj', 0), + ('.gate_up_proj', '.up_proj', 1), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + name = _normalize_dflash_weight_name(name) + if name is None: + continue + if name in ('embed_tokens.weight', 'lm_head.weight'): + continue + if name == 'mask_embedding.weight': + self.mask_embedding.data.copy_(loaded_weight.reshape_as(self.mask_embedding)) + self.has_separate_mask_embedding = True + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + load_weight(param, loaded_weight, shard_id=shard_id) + break + else: + param = params_dict[name] + load_weight(param, loaded_weight) diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index 55b530f63e..c2db94e162 100644 --- a/lmdeploy/pytorch/models/utils/cudagraph.py +++ b/lmdeploy/pytorch/models/utils/cudagraph.py @@ -110,13 +110,19 @@ def make_output_buffers(self, output): output_buffers = output return output_buffers - def update_meta_flashattn(self, batch_size: int, max_seqlen_q: int, block_size: int, max_seqlen_k: int, - cache_seqlens: torch.Tensor): - """Update meta flashattn.""" + def build_fa3_scheduler_metadata(self, + batch_size: int, + max_seqlen_q: int, + block_size: int, + max_seqlen_k: int, + cache_seqlens: torch.Tensor, + *, + sliding_window: Any, + causal: bool): + """Build FA3 metadata for an explicit attention pattern.""" ctx_mgr = get_step_ctx_manager() step_ctx = ctx_mgr.current_context() model_config = step_ctx.model_config - sliding_window = model_config.sliding_window num_attention_heads, num_key_value_heads = model_config.get_num_qkv_head_by_tp() headdim = model_config.head_dim torch_dtype = model_config.dtype @@ -138,6 +144,7 @@ def update_meta_flashattn(self, batch_size: int, max_seqlen_q: int, block_size: cache_seqlens=cache_seqlens, qkv_dtype=torch_dtype, page_size=block_size, + causal=causal, window_size=window_size, ) return scheduler_metadata @@ -193,11 +200,15 @@ def make_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p # use fa3 decode kernel for spec decode elif graph_meta.use_fa3_decoding is True: max_seqlen_k = graph_meta.num_blocks * graph_meta.block_size - input_buffers['scheduler_metadata'] = self.update_meta_flashattn(graph_meta.max_batchs, - graph_meta.decode_query_len, - block_size=graph_meta.block_size, - max_seqlen_k=max_seqlen_k, - cache_seqlens=input_buffers['kv_seqlens']) + model_config = get_step_ctx_manager().current_context().model_config + input_buffers['scheduler_metadata'] = self.build_fa3_scheduler_metadata( + graph_meta.max_batchs, + graph_meta.decode_query_len, + block_size=graph_meta.block_size, + max_seqlen_k=max_seqlen_k, + cache_seqlens=input_buffers['kv_seqlens'], + sliding_window=model_config.sliding_window, + causal=True) # mrope if graph_meta.use_mrope: @@ -285,12 +296,15 @@ def fill_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p # graph_meta.num_blocks * graph_meta.block_size to match the buffer shape # allocated in make_buffers_cudagraph, not the runtime attn_metadata value. max_seqlen_k = graph_meta.num_blocks * graph_meta.block_size - scheduler_metadata = self.update_meta_flashattn( + model_config = get_step_ctx_manager().current_context().model_config + scheduler_metadata = self.build_fa3_scheduler_metadata( new_batch_size, decode_query_len, block_size=graph_meta.block_size, max_seqlen_k=max_seqlen_k, cache_seqlens=input_buffers['kv_seqlens'], + sliding_window=model_config.sliding_window, + causal=True, ) num_meta = scheduler_metadata.size(0) input_buffers['scheduler_metadata'][:num_meta] = scheduler_metadata diff --git a/lmdeploy/pytorch/spec_decode/base.py b/lmdeploy/pytorch/spec_decode/base.py index 80a122cf29..13d71f6fbe 100644 --- a/lmdeploy/pytorch/spec_decode/base.py +++ b/lmdeploy/pytorch/spec_decode/base.py @@ -59,6 +59,15 @@ def __init__( def is_enabled(self): return self._enabled + def requires_target_inputs_embeds(self) -> bool: + """Whether the proposer reuses target multimodal input embeddings.""" + proposer = self.proposer + if proposer is not None: + return bool(getattr(proposer, 'requires_target_inputs_embeds', True)) + # Follower ranks do not own a proposer. Keep their target-build + # capability consistent with the proposer-owning rank. + return self.method != 'dflash' + def set_cache_config(self, cache_config: CacheConfig): """Set all cache config.""" pass diff --git a/lmdeploy/pytorch/spec_decode/dflash_debug.py b/lmdeploy/pytorch/spec_decode/dflash_debug.py new file mode 100644 index 0000000000..dc09f09731 --- /dev/null +++ b/lmdeploy/pytorch/spec_decode/dflash_debug.py @@ -0,0 +1,48 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +from __future__ import annotations + +import json +import os +from collections.abc import Callable +from pathlib import Path + +import torch + + +def dflash_debug_dir() -> Path | None: + """Return the optional DFlash debug output directory.""" + path = os.getenv('LMDEPLOY_DFLASH_DEBUG_DIR') + return None if not path else Path(path) + + +def dflash_debug_enabled() -> bool: + """Return whether opt-in DFlash tracing is enabled.""" + return dflash_debug_dir() is not None + + +def debug_tensor(value: torch.Tensor | None, limit: int = 64): + """Serialize a small tensor, truncating large tensors.""" + if value is None: + return None + tensor = value.detach().cpu() + if tensor.numel() <= limit: + return tensor.tolist() + return { + 'shape': list(tensor.shape), + 'head': tensor.flatten()[:limit].tolist(), + } + + +def write_dflash_debug(rank: int, event: str, payload: dict | Callable[[], dict]): + """Append one DFlash debug event when opt-in tracing is enabled.""" + debug_dir = dflash_debug_dir() + if debug_dir is None: + return + if callable(payload): + payload = payload() + debug_dir.mkdir(parents=True, exist_ok=True) + path = debug_dir / f'rank{rank}.jsonl' + record = {'event': event, **payload} + with path.open('a', encoding='utf-8') as f: + f.write(json.dumps(record, sort_keys=True) + '\n') diff --git a/lmdeploy/pytorch/spec_decode/dflash_utils.py b/lmdeploy/pytorch/spec_decode/dflash_utils.py new file mode 100644 index 0000000000..d5506d148a --- /dev/null +++ b/lmdeploy/pytorch/spec_decode/dflash_utils.py @@ -0,0 +1,192 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from __future__ import annotations + +from typing import Any + + +def _parse_layer_ids(value: Any) -> tuple[int, ...] | None: + """Parse target layer ids.""" + if value is None: + return None + if not isinstance(value, (list, tuple)): + raise ValueError('DFlash dflash_config.target_layer_ids must be a list of integers.') + if len(value) == 0: + raise ValueError('DFlash dflash_config.target_layer_ids must not be empty.') + if any(isinstance(layer_id, bool) or not isinstance(layer_id, int) for layer_id in value): + raise ValueError(f'DFlash target_layer_ids must contain only integers, got {value!r}.') + layer_ids = tuple(value) + _validate_target_layer_id_order(layer_ids, 'DFlash dflash_config.target_layer_ids') + return layer_ids + + +def _validate_target_layer_id_order(layer_ids: tuple[int, ...], field_name: str) -> None: + """Validate DFlash target feature ids are in checkpoint feature order.""" + seen: set[int] = set() + prev_layer_id: int | None = None + for pos, layer_id in enumerate(layer_ids): + if layer_id in seen: + raise ValueError(f'{field_name} must be duplicate-free; duplicate layer id {layer_id} at position {pos}.') + if prev_layer_id is not None and layer_id <= prev_layer_id: + raise ValueError(f'{field_name} must be strictly increasing and duplicate-free; got {layer_ids}.') + seen.add(layer_id) + prev_layer_id = layer_id + + +def build_target_layer_ids(num_target_layers: int, num_draft_layers: int) -> tuple[int, ...]: + """Select evenly spaced DFlash target layer ids. + + DFlash consumes hidden features sampled from the target model. This fallback matches the SGLang/vLLM convention used + when checkpoint metadata does not explicitly list target layers. + """ + if num_target_layers < 1: + raise ValueError(f'DFlash num_target_layers must be positive, got {num_target_layers!r}.') + if num_draft_layers < 1: + raise ValueError(f'DFlash num_hidden_layers must be positive, got {num_draft_layers!r}.') + if num_draft_layers == 1: + return (num_target_layers // 2,) + + start = 1 + end = num_target_layers - 3 + if end < start: + raise ValueError(f'DFlash target layer fallback requires at least 4 target layers, got {num_target_layers}.') + span = end - start + layer_ids = tuple(int(round(start + i * span / (num_draft_layers - 1))) for i in range(num_draft_layers)) + _validate_target_layer_id_order(layer_ids, 'DFlash fallback target_layer_ids') + return layer_ids + + +def _normalize_sliding_window(sliding_window: Any) -> int | None: + """Normalize the no-window values used by HF and ``ModelConfig``.""" + if sliding_window in (None, 0, -1): + return None + if isinstance(sliding_window, bool) or not isinstance(sliding_window, int): + raise ValueError(f'Invalid DFlash sliding_window: {sliding_window!r}.') + if sliding_window < 1: + raise ValueError(f'DFlash sliding_window must be positive, got {sliding_window!r}.') + return sliding_window + + +def _validate_dflash_v1_supported(draft_hf_config: Any, dflash_config: Any) -> None: + """Validate the complete DFlash V1 checkpoint contract once.""" + if not isinstance(dflash_config, dict): + raise ValueError('DFlash checkpoint requires a dflash_config dictionary.') + + num_hidden_layers = draft_hf_config.num_hidden_layers + layer_types = getattr(draft_hf_config, 'layer_types', None) + if not isinstance(layer_types, (list, tuple)): + raise ValueError('DFlash V1 requires an explicit layer_types list.') + if len(layer_types) != num_hidden_layers: + raise ValueError('DFlash layer_types length must equal num_hidden_layers. ' + f'Got len(layer_types)={len(layer_types)}, num_hidden_layers={num_hidden_layers}.') + if any(not isinstance(layer_type, str) for layer_type in layer_types): + raise ValueError('DFlash layer_types must contain only strings.') + unsupported_layer_types = sorted( + {layer_type for layer_type in layer_types if layer_type not in ('full_attention', 'sliding_attention')}) + if unsupported_layer_types: + raise ValueError('DFlash supports only full_attention and sliding_attention draft layers. ' + f'Got unsupported layer_types={unsupported_layer_types}.') + + use_sliding_window = getattr(draft_hf_config, 'use_sliding_window', True) + default_sliding_window = None + if use_sliding_window: + default_sliding_window = _normalize_sliding_window(getattr(draft_hf_config, 'sliding_window', None)) + + has_sliding_layer = 'sliding_attention' in layer_types + if has_sliding_layer and default_sliding_window is None: + raise ValueError('DFlash sliding_attention layers require the model-default sliding_window.') + + explicit_sliding_window = dflash_config.get('swa_window_size', dflash_config.get('sliding_window')) + if explicit_sliding_window is not None: + explicit_sliding_window = _normalize_sliding_window(explicit_sliding_window) + if explicit_sliding_window != default_sliding_window: + raise ValueError('DFlash V1 requires the SWA window to equal ModelConfig.sliding_window. ' + f'Got SWA window={explicit_sliding_window}, ' + f'model default={default_sliding_window}.') + + use_swa = bool(dflash_config.get('use_swa', getattr(draft_hf_config, 'use_swa', False))) + if use_swa and not has_sliding_layer: + raise ValueError('DFlash V1 does not support use_swa forcing non-causal sliding attention; ' + 'checkpoint layer_types must describe the supported attention pattern explicitly.') + + causal_override = dflash_config.get('causal', getattr(draft_hf_config, 'causal', None)) + if causal_override is not None and not isinstance(causal_override, bool): + raise ValueError(f'DFlash causal override must be boolean, got {causal_override!r}.') + allowed_patterns = {(default_sliding_window, True), (None, False)} + for layer_idx, layer_type in enumerate(layer_types): + sliding_window = default_sliding_window if layer_type == 'sliding_attention' else None + default_causal = layer_type == 'sliding_attention' + causal = default_causal if causal_override is None else bool(causal_override) + pattern = (sliding_window, causal) + if pattern not in allowed_patterns: + raise ValueError('DFlash V1 supports only the model-default causal attention pattern and ' + 'non-causal full attention. ' + f'Layer {layer_idx} resolves to unsupported pattern={pattern}, ' + f'allowed={sorted(allowed_patterns, key=str)}.') + + if dflash_config.get('attention_sink_bias', getattr(draft_hf_config, 'add_swa_attention_sink_bias', False)): + raise ValueError('DFlash attention-sink bias checkpoints are not supported yet.') + + +def parse_dflash_config(draft_hf_config: Any, num_speculative_tokens: int, + target_num_layers: int) -> tuple[tuple[int, ...], int]: + """Return resolved ``(target_layer_ids, mask_token_id)`` metadata.""" + num_hidden_layers = draft_hf_config.num_hidden_layers + num_target_layers = draft_hf_config.num_target_layers + dflash_config = draft_hf_config.dflash_config + _validate_dflash_v1_supported(draft_hf_config, dflash_config) + if num_target_layers != target_num_layers: + raise ValueError('DFlash draft/target depth mismatch: ' + f'draft declares num_target_layers={num_target_layers}, ' + f'but target ModelConfig has num_layers={target_num_layers}.') + + max_query_length = dflash_config['block_size'] + query_length = num_speculative_tokens + 1 + if query_length > max_query_length: + raise ValueError('DFlash query length (1 + speculative_num_draft_tokens) must not exceed checkpoint ' + 'dflash_config.block_size. ' + f'Got block_size={max_query_length}, query_length={query_length}.') + + mask_token_id = dflash_config.get('mask_token_id') + if mask_token_id is None: + raise ValueError('DFlash checkpoint requires dflash_config.mask_token_id.') + + target_layer_ids = _parse_layer_ids(dflash_config.get('target_layer_ids')) + if target_layer_ids is None: + target_layer_ids = build_target_layer_ids(num_target_layers, num_hidden_layers) + for pos, layer_id in enumerate(target_layer_ids): + if layer_id < 0 or layer_id >= num_target_layers: + raise ValueError('DFlash target_layer_ids contains an out-of-range value: ' + f'target_layer_ids[{pos}]={layer_id}, num_target_layers={num_target_layers}.') + + return target_layer_ids, mask_token_id + + +def validate_dflash_dist_config(dist_config: Any): + """Validate the initially supported DFlash distribution shape.""" + if dist_config is None: + return + if getattr(dist_config, 'dp', 1) != 1: + raise ValueError('DFlash V1 does not support data parallelism. Set dp=1.') + if getattr(dist_config, 'ep', 1) != 1: + raise ValueError('DFlash V1 does not support expert parallel draft execution. Set ep=1.') + + +def validate_dflash_cache_config(cache_config: Any): + """Validate cache options that affect DFlash draft KV correctness.""" + if cache_config is None: + return + if getattr(cache_config, 'enable_prefix_caching', False): + raise ValueError('DFlash V1 does not support prefix-cache reuse yet. Disable enable_prefix_caching.') + quant_policy = getattr(cache_config, 'quant_policy', 0) + if int(quant_policy) != 0: + raise ValueError('DFlash V1 does not support KV-cache quantization yet. Set quant_policy=0.') + + +def validate_dflash_runtime_config(cache_config: Any = None, backend_config: Any = None): + """Validate the currently supported DFlash runtime envelope.""" + device_type = getattr(backend_config, 'device_type', None) + if device_type is None: + device_type = getattr(cache_config, 'device_type', 'cuda') + if device_type != 'cuda': + raise ValueError('DFlash V1 requires CUDA because draft context K/V materialization uses CUDA kernels. ' + f'Got device_type={device_type!r}.') diff --git a/lmdeploy/pytorch/spec_decode/guided_spec_helper.py b/lmdeploy/pytorch/spec_decode/guided_spec_helper.py index d3e5e6cf25..55a8f11c93 100644 --- a/lmdeploy/pytorch/spec_decode/guided_spec_helper.py +++ b/lmdeploy/pytorch/spec_decode/guided_spec_helper.py @@ -102,7 +102,7 @@ async def accept_draft_tokens(self, draft_token_ids: torch.Tensor, """Accept draft tokens on the provided (forked) grammar matchers. In speculative decoding the matchers are typically forked from the - originals (created in :meth:`SpecModelAgent._async_model_forward`), + originals (created in :meth:`SpecModelAgent.async_model_forward`), so this method accepts on whichever matchers are passed in. """ if not processors or self._mgr is None: diff --git a/lmdeploy/pytorch/spec_decode/proposers/__init__.py b/lmdeploy/pytorch/spec_decode/proposers/__init__.py index a425779233..ce78c6c2aa 100644 --- a/lmdeploy/pytorch/spec_decode/proposers/__init__.py +++ b/lmdeploy/pytorch/spec_decode/proposers/__init__.py @@ -2,6 +2,7 @@ # Copyright (c) OpenMMLab. All rights reserved. from .deepseek_mtp import DeepseekMTP # noqa F401 +from .dflash import DFlash # noqa F401 from .eagle import Eagle # noqa F401 from .eagle3 import Eagle3 # noqa F401 from .qwen3_5_mtp import Qwen3_5MTP # noqa F401 diff --git a/lmdeploy/pytorch/spec_decode/proposers/base.py b/lmdeploy/pytorch/spec_decode/proposers/base.py index 2ba1e36b78..4264175fa6 100644 --- a/lmdeploy/pytorch/spec_decode/proposers/base.py +++ b/lmdeploy/pytorch/spec_decode/proposers/base.py @@ -1,6 +1,8 @@ # Copyright (c) OpenMMLab. All rights reserved. from __future__ import annotations +from dataclasses import dataclass +from enum import Enum from typing import Any import torch @@ -9,7 +11,7 @@ from lmdeploy.utils import get_logger -from ...config import ModelConfig, SpecDecodeConfig +from ...config import CacheConfig, ModelConfig, SpecDecodeConfig from ...engine.cache_engine import CacheEngine from ...model_inputs import ModelInputs, step_ctx_manager from ...models.patch import build_patched_model, update_custom_module_map @@ -22,6 +24,39 @@ logger = get_logger('lmdeploy') +class ProposalMethod(str, Enum): + """How the agent should prepare and execute draft proposal.""" + + AUTOREGRESSIVE = 'autoregressive' + DIFFUSION = 'diffusion' + + +@dataclass(frozen=True) +class ProposalContext: + """Explicit runtime dependencies for a non-autoregressive proposer.""" + + cache_engine: CacheEngine | None + rank: int + debug_step: int + + +@dataclass(frozen=True) +class ProposalWarmupCase: + """One declarative draft warmup input shape.""" + + batch_size: int + is_decoding: bool + max_q_seqlen: int + target_hidden_size: int + + +@dataclass(frozen=True) +class ProposalWarmupPlan: + """Ordered proposer-specific cases executed by the agent.""" + + cases: tuple[ProposalWarmupCase, ...] + + @torch.inference_mode() def draft_model_forward( model: torch.nn.Module, @@ -60,6 +95,9 @@ def draft_model_forward( class BaseSpecProposer: + proposal_method = ProposalMethod.AUTOREGRESSIVE + requires_target_inputs_embeds = True + def __init__(self, specdecode_config: SpecDecodeConfig, device: torch.device = None): self.specdecode_config = specdecode_config self.model = None @@ -98,6 +136,26 @@ async def get_outputs(self, """Get outputs.""" raise NotImplementedError() + async def propose(self, + model_inputs: ModelInputs, + extra_inputs: ExtraInputs, + sampling_inputs, + proposal_ctx: ProposalContext | None = None): + """Run a non-autoregressive proposal method.""" + raise NotImplementedError(f'{type(self).__name__} does not implement its proposal method.') + + def get_warmup_plan(self, + max_batches: int, + target_model_config: ModelConfig, + capture_batch_sizes: list[int], + cache_config: CacheConfig) -> ProposalWarmupPlan | None: + """Return custom warmup shapes, or ``None`` for generic AR warmup.""" + return None + + def prepare_warmup_forward(self, inputs: ModelInputs, cache_engine: CacheEngine) -> ModelInputs | None: + """Prepare one declarative case for forwarding by the agent.""" + return inputs + @record_function('draft_model_forward') def _forward(self, model_inputs: ModelInputs, cache_engine: CacheEngine = None): """Forward.""" diff --git a/lmdeploy/pytorch/spec_decode/proposers/dflash.py b/lmdeploy/pytorch/spec_decode/proposers/dflash.py new file mode 100644 index 0000000000..c53ced335a --- /dev/null +++ b/lmdeploy/pytorch/spec_decode/proposers/dflash.py @@ -0,0 +1,354 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +import torch + +from lmdeploy.utils import get_logger + +from ...config import CacheConfig, ModelConfig +from ...engine.cache_engine import CacheEngine +from ...model_inputs import ModelInputs, step_ctx_manager +from ...strategies.ar_spec.model_agent import ARSpecExtraInputs +from ..dflash_debug import debug_tensor, write_dflash_debug +from .base import ( + SPEC_PROPOSERS, + BaseSpecProposer, + ProposalContext, + ProposalMethod, + ProposalWarmupCase, + ProposalWarmupPlan, +) + +logger = get_logger('lmdeploy') + + +@SPEC_PROPOSERS.register_module(name='dflash') +class DFlash(BaseSpecProposer): + """DFlash proposer with one-shot block drafting.""" + + proposal_method = ProposalMethod.DIFFUSION + requires_target_inputs_embeds = False + + def build_model(self, empty_init: bool, target_model: torch.nn.Module = None, build_model_ctx=None): + if target_model is None: + raise RuntimeError('DFlash requires the target model for shared input embeddings.') + super().build_model(empty_init, target_model=target_model, build_model_ctx=build_model_ctx) + if not hasattr(self.model, 'set_input_embeddings'): + raise RuntimeError('DFlash draft model must implement set_input_embeddings().') + logger.info('Using embed_tokens from target model for DFlash draft.') + self.model.set_input_embeddings(target_model.get_input_embeddings()) + if self.model.get_input_embeddings() is None: + raise RuntimeError('DFlash target input embeddings are not available.') + + def get_target_hidden_size(self, model_config): + """Get concatenated DFlash target hidden size.""" + return model_config.hidden_size * len(self.specdecode_config.target_layer_ids) + + def get_warmup_plan(self, + max_batches: int, + target_model_config: ModelConfig, + capture_batch_sizes: list[int], + cache_config: CacheConfig) -> ProposalWarmupPlan: + """Declare materialization shapes and fixed block-query graph cases.""" + if cache_config is None: + raise RuntimeError('DFlash warmup requires a draft cache configuration.') + target_hidden_size = self.get_target_hidden_size(target_model_config) + cache_block_size = max(1, int(cache_config.block_size)) + prefill_token_budget = max(1, int(cache_config.max_prefill_token_num)) + prefill_q_len = min(cache_block_size, prefill_token_budget) + prefill_batch_size = min(max_batches, max(1, prefill_token_budget // prefill_q_len)) + decode_q_len = self.num_speculative_tokens + 1 + cases = [ + ProposalWarmupCase(prefill_batch_size, + is_decoding=False, + max_q_seqlen=prefill_q_len, + target_hidden_size=target_hidden_size) + ] + cases.extend( + ProposalWarmupCase(batch_size, + is_decoding=True, + max_q_seqlen=decode_q_len, + target_hidden_size=target_hidden_size) for batch_size in capture_batch_sizes) + return ProposalWarmupPlan(cases=tuple(cases)) + + def prepare_warmup_forward(self, inputs: ModelInputs, cache_engine: CacheEngine) -> ModelInputs | None: + """Materialize target features and return only production query + inputs.""" + if cache_engine is None: + raise RuntimeError('DFlash warmup requires a draft cache engine.') + cache_block_size = max(1, int(cache_engine.cache_config.block_size)) + required_blocks = (inputs.max_q_seqlen + cache_block_size - 1) // cache_block_size + if inputs.block_offsets.size(1) < required_blocks: + inputs.block_offsets = torch.nn.functional.pad(inputs.block_offsets, + (0, required_blocks - inputs.block_offsets.size(1)), + value=0) + + batch_size = int(inputs.seq_length.numel()) + extra_inputs = ARSpecExtraInputs( + next_token_ids=inputs.input_ids.new_zeros(batch_size), + target_hidden_states=inputs.target_hidden_states, + ) + context_inputs, target_hidden, context_lengths, context_position_ids = \ + self._prepare_context_materialization(inputs, extra_inputs) + self._materialize_context(context_inputs, target_hidden, cache_engine) + if not inputs.is_decoding: + return None + return self._build_query_inputs(inputs, context_lengths, extra_inputs.next_token_ids, + context_position_ids) + + def _draft_model(self): + """Return the underlying draft nn.Module when graph runner wraps it.""" + if hasattr(self.model, 'get_model'): + return self.model.get_model() + return self.model + + def _flatten_target_hidden(self, extra_inputs: ARSpecExtraInputs) -> torch.Tensor: + """Return flattened target aux hidden states.""" + target_hidden = extra_inputs.target_hidden_states + if target_hidden is None: + raise RuntimeError('DFlash requires target aux hidden states from the main model.') + if target_hidden.dim() == 3: + target_hidden = target_hidden[0] + if target_hidden.dim() != 2: + raise RuntimeError(f'DFlash expected target hidden states with shape [N, H], got {target_hidden.shape}.') + return target_hidden + + def _context_lengths(self, model_inputs: ModelInputs, extra_inputs: ARSpecExtraInputs): + """Resolve committed context lengths visible to DFlash.""" + context_lengths = model_inputs.seq_length + if extra_inputs.num_rejected_tokens is not None: + context_lengths = context_lengths - extra_inputs.num_rejected_tokens.to(context_lengths) + return context_lengths + + @staticmethod + def _slice_by_lengths(tensor: torch.Tensor, + seq_lengths: torch.Tensor, + keep_lengths: torch.Tensor, + max_seq_length: int, + preserve_features: bool = False): + """Slice a flattened per-token tensor by per-request valid lengths.""" + if preserve_features: + if tensor.dim() == 3 and tensor.size(0) == 1: + flat_tensor = tensor[0] + elif tensor.dim() >= 2: + flat_tensor = tensor.flatten(0, -2) + else: + flat_tensor = tensor.reshape(-1) + elif tensor.dim() == 2 and tensor.size(0) == 1: + flat_tensor = tensor[0] + elif tensor.dim() >= 2: + flat_tensor = tensor.flatten(0, -2) + else: + flat_tensor = tensor.reshape(-1) + starts = seq_lengths.cumsum(0) - seq_lengths + offsets = torch.arange(max_seq_length, device=seq_lengths.device) + valid = offsets[None, :] < keep_lengths[:, None] + indices = starts[:, None] + offsets[None, :] + indices = indices[valid] + return flat_tensor.index_select(0, indices) + + def _slice_target_position_ids(self, model_inputs: ModelInputs, context_lengths: torch.Tensor): + """Slice explicit target positions when the target path provides + them.""" + target_position_ids = model_inputs.target_position_ids + if target_position_ids is None: + return None + if target_position_ids.dim() == 2 and target_position_ids.size(0) == 1: + target_position_ids = target_position_ids[0] + if target_position_ids.dim() != 1: + raise RuntimeError('DFlash supports only 1D target_position_ids for draft context materialization, ' + f'got shape={tuple(target_position_ids.shape)}.') + return self._slice_by_lengths(target_position_ids, + model_inputs.seq_length, + context_lengths, + max_seq_length=model_inputs.max_q_seqlen) + + def _build_context_inputs(self, + model_inputs: ModelInputs, + context_lengths: torch.Tensor, + context_position_ids: torch.Tensor | None = None): + """Build draft-cache materialization inputs for committed context + tokens.""" + context_ids = self._slice_by_lengths(model_inputs.input_ids, + model_inputs.seq_length, + context_lengths, + max_seq_length=model_inputs.max_q_seqlen) + target_position_ids = None if context_position_ids is None else context_position_ids.unsqueeze(0) + return model_inputs.clone( + input_ids=context_ids.unsqueeze(0), + seq_length=context_lengths, + max_q_seqlen=model_inputs.max_q_seqlen, + max_kv_seqlen=model_inputs.max_kv_seqlen, + sum_kv_seqlen=model_inputs.sum_kv_seqlen, + is_decoding=False, + target_hidden_states=None, + target_position_ids=target_position_ids, + target_inputs_embeds=None, + ) + + def _build_query_inputs(self, model_inputs: ModelInputs, context_lengths: torch.Tensor, + next_token_ids: torch.Tensor, context_position_ids: torch.Tensor | None = None): + """Build one DFlash query block per request: [next, mask, ...].""" + batch_size = int(model_inputs.seq_length.numel()) + query_len = self.num_speculative_tokens + 1 + query_ids = model_inputs.input_ids.new_full((batch_size, query_len), int(self.specdecode_config.mask_token_id)) + query_ids[:, 0] = next_token_ids + query_history = model_inputs.history_lengths + context_lengths + if context_position_ids is None: + query_start_positions = query_history + else: + starts = context_lengths.cumsum(0) - context_lengths + query_start_positions = context_position_ids[starts + context_lengths - 1] + 1 + query_positions = query_start_positions[:, None] + torch.arange(query_len, device=query_history.device)[None, :] + return model_inputs.clone( + input_ids=query_ids.reshape(1, -1), + seq_length=model_inputs.seq_length.new_full((batch_size, ), query_len), + history_lengths=query_history, + max_q_seqlen=query_len, + max_kv_seqlen=model_inputs.max_kv_seqlen + query_len, + sum_kv_seqlen=model_inputs.sum_kv_seqlen + batch_size * query_len, + # DFlash queries are fixed-size speculative blocks over paged + # context. Mark them as decoding so FA3 reads the block table + # directly instead of flattening the full KV history as prefill. + is_decoding=True, + target_hidden_states=None, + target_position_ids=query_positions.reshape(1, -1), + target_inputs_embeds=None, + ) + + def _materialize_context( + self, + context_inputs: ModelInputs, + target_hidden: torch.Tensor, + cache_engine: CacheEngine, + ): + """Project target aux hidden states into the draft KV cache.""" + if target_hidden.numel() == 0: + return + kv_caches = cache_engine.gpu_cache + ctx_mgr = self.model.ctx_mgr + with step_ctx_manager(ctx_mgr): + context = ctx_mgr.build_context( + inputs=context_inputs, + model_config=self.specdecode_config.model_config, + cache_config=cache_engine.cache_config, + kv_caches=kv_caches, + ) + with ctx_mgr.context(context): + self._draft_model().precompute_and_store_context_kv( + target_hidden=target_hidden, + position_ids=context.position_ids.flatten(), + past_key_values=kv_caches, + attn_metadata=context.attn_metadata, + max_q_seqlen=context_inputs.max_q_seqlen, + ) + + def _prepare_context_materialization(self, model_inputs: ModelInputs, extra_inputs: ARSpecExtraInputs): + """Build inputs and target states for DFlash context K/V + materialization.""" + target_hidden = self._flatten_target_hidden(extra_inputs) + context_lengths = self._context_lengths(model_inputs, extra_inputs) + context_position_ids = self._slice_target_position_ids(model_inputs, context_lengths) + context_inputs = self._build_context_inputs(model_inputs, context_lengths, context_position_ids) + target_hidden = self._slice_by_lengths(target_hidden, + model_inputs.seq_length, + context_lengths, + max_seq_length=model_inputs.max_q_seqlen, + preserve_features=True) + return context_inputs, target_hidden, context_lengths, context_position_ids + + def materialize_context( + self, + model_inputs: ModelInputs, + extra_inputs: ARSpecExtraInputs, + cache_engine: CacheEngine, + ): + """Materialize DFlash context K/V without running the mask-query + draft.""" + if cache_engine is None: + raise RuntimeError('DFlash requires a draft cache engine for context K/V materialization.') + + context_inputs, target_hidden, _, _ = self._prepare_context_materialization(model_inputs, extra_inputs) + self._materialize_context(context_inputs, target_hidden, cache_engine) + + async def propose_block( + self, + model_inputs: ModelInputs, + extra_inputs: ARSpecExtraInputs, + cache_engine: CacheEngine, + guided_processors: dict | None = None, + ): + """Run a one-shot DFlash draft proposal.""" + if guided_processors: + raise NotImplementedError('DFlash guided decoding is not implemented yet.') + if cache_engine is None: + raise RuntimeError('DFlash requires a draft cache engine for context K/V materialization.') + if extra_inputs.next_token_ids is None: + raise RuntimeError('DFlash requires sampled next_token_ids from the target model.') + + context_inputs, target_hidden, context_lengths, context_position_ids = self._prepare_context_materialization( + model_inputs, extra_inputs) + query_inputs = self._build_query_inputs(model_inputs, context_lengths, extra_inputs.next_token_ids, + context_position_ids) + + self._materialize_context(context_inputs, target_hidden, cache_engine) + outputs = self._forward(query_inputs, cache_engine=cache_engine) + hidden_states = outputs['hidden_states'] + if hidden_states.dim() == 3: + hidden_states = hidden_states[0] + + batch_size = int(query_inputs.seq_length.numel()) + query_len = self.num_speculative_tokens + 1 + mask_indices = torch.arange(batch_size * query_len, device=hidden_states.device).view(batch_size, query_len) + mask_indices = mask_indices[:, 1:].reshape(-1) + logits = self.get_logits(hidden_states[mask_indices][None])[0] + draft_token_ids = logits.argmax(dim=-1).view(batch_size, self.num_speculative_tokens) + return draft_token_ids + + async def propose(self, + model_inputs: ModelInputs, + extra_inputs: ARSpecExtraInputs, + sampling_inputs, + proposal_ctx: ProposalContext | None = None): + """Produce all DFlash draft tokens in one block proposal.""" + if proposal_ctx is None: + raise RuntimeError('DFlash requires ProposalContext with a draft cache engine.') + orig_processors = self.guided_helper.get_processors( + sampling_inputs.session_ctx if sampling_inputs else None, + sampling_inputs.response_formats if sampling_inputs else None) + if orig_processors: + raise NotImplementedError('DFlash guided decoding is not implemented yet.') + + if model_inputs.is_chunk and not model_inputs.is_last_chunk: + self.materialize_context(model_inputs, extra_inputs, proposal_ctx.cache_engine) + output_draft_ids = model_inputs.input_ids.new_zeros(model_inputs.seq_length.size(0), + self.num_speculative_tokens) + else: + output_draft_ids = await self.propose_block(model_inputs, extra_inputs, proposal_ctx.cache_engine) + + write_dflash_debug(proposal_ctx.rank, 'proposal', lambda: { + 'step': proposal_ctx.debug_step, + 'is_decoding': bool(model_inputs.is_decoding), + 'seq_length': debug_tensor(model_inputs.seq_length), + 'history_lengths': debug_tensor(model_inputs.history_lengths), + 'input_ids': debug_tensor(model_inputs.input_ids), + 'next_token_ids': debug_tensor(extra_inputs.next_token_ids), + 'prev_num_rejected_tokens': debug_tensor(extra_inputs.num_rejected_tokens), + 'draft_token_ids': debug_tensor(output_draft_ids), + }) + + return ARSpecExtraInputs( + output_draft_token_ids=output_draft_ids, + next_token_ids=extra_inputs.next_token_ids, + num_rejected_tokens=extra_inputs.num_rejected_tokens, + output_token_ids=extra_inputs.output_token_ids, + logprobs=extra_inputs.logprobs, + ) + + async def get_outputs(self, + model_outputs: dict[str, torch.Tensor], + model_inputs: ModelInputs, + extra_inputs: ARSpecExtraInputs = None, + guided_processors: dict | None = None): + """Get DFlash draft outputs.""" + raise NotImplementedError('DFlash requires the block-proposal hook; the autoregressive draft loop is invalid ' + 'because context K/V must be materialized before a one-shot mask query forward.') diff --git a/lmdeploy/pytorch/spec_decode/spec_agent.py b/lmdeploy/pytorch/spec_decode/spec_agent.py index f73bcdf2c2..15bdb1608a 100644 --- a/lmdeploy/pytorch/spec_decode/spec_agent.py +++ b/lmdeploy/pytorch/spec_decode/spec_agent.py @@ -22,8 +22,13 @@ from ..strategies.ar_spec.model_agent import ARSpecExtraInputs from ..strategies.base.model_agent import ExtraInputs from .base import BaseSpecModelAgent +from .dflash_debug import debug_tensor, write_dflash_debug from .guided_spec_helper import GuidedSpecHelper -from .proposers.base import build_specdecode_proposer +from .proposers.base import ( + ProposalContext, + ProposalMethod, + build_specdecode_proposer, +) if TYPE_CHECKING: pass @@ -180,11 +185,13 @@ def __init__( def _init_runtime_state(self): """Initialize request-local draft carry state.""" self._prev_chunk_last = {} + self._dflash_debug_step = 0 def reset_runtime_state(self): """Discard request-local draft carry state after sleep cancels sessions.""" self._prev_chunk_last.clear() + self._dflash_debug_step = 0 @contextmanager def draft_context(self): @@ -489,6 +496,20 @@ def __compute_logprobs(raw_logprobs: torch.Tensor, token_ids: torch.LongTensor, next_token_ids, sampling_inputs=draft_sampling_inputs, ) + if self.method == 'dflash': + write_dflash_debug(self.rank, 'rejection', lambda: { + 'step': self._dflash_debug_step, + 'is_decoding': bool(model_inputs.is_decoding), + 'seq_length': debug_tensor(model_inputs.seq_length), + 'history_lengths': debug_tensor(model_inputs.history_lengths), + 'input_ids': debug_tensor(model_inputs.input_ids), + 'draft_token_ids': debug_tensor(extra_inputs.output_draft_token_ids), + 'target_argmax': debug_tensor(target_draft_logits.argmax(dim=-1)), + 'bonus_token_ids': debug_tensor(next_token_ids), + 'output_token_ids': debug_tensor(output_token_ids), + 'num_rejected_tokens': debug_tensor(num_rejected_tokens), + }) + self._dflash_debug_step += 1 last_token_indices = last_token_indices - num_rejected_tokens # Guided: accept final tokens on original matchers. @@ -591,9 +612,9 @@ async def async_sampling_logits(self, model_inputs: ModelInputs, extra_inputs: A with record_function('spec_rejection_sampling'): return await self._rejection_sampling(model_inputs, extra_inputs, sampling_inputs) - async def _async_model_forward(self, inputs: ModelInputs, extra_inputs: ARSpecExtraInputs, - sampling_inputs: SamplingInputs): - """Model forward. + async def _async_autoregressive_model_forward(self, inputs: ModelInputs, extra_inputs: ARSpecExtraInputs, + sampling_inputs: SamplingInputs): + """Default autoregressive draft model forward. Args: inputs (dict): The input data comes from _make_inputs. @@ -688,8 +709,23 @@ async def async_model_forward( sampling_inputs: SamplingInputs, ): """Draft model forward.""" - draft_model_inputs, draft_extra_inputs = self._prepare_inputs_from_main(model_inputs, extra_inputs) - return await self._async_model_forward(draft_model_inputs, draft_extra_inputs, sampling_inputs) + proposal_method = getattr(self.proposer, 'proposal_method', ProposalMethod.AUTOREGRESSIVE) + if proposal_method == ProposalMethod.AUTOREGRESSIVE: + draft_model_inputs, draft_extra_inputs = self._prepare_inputs_from_main(model_inputs, extra_inputs) + return await self._async_autoregressive_model_forward(draft_model_inputs, draft_extra_inputs, + sampling_inputs) + if proposal_method == ProposalMethod.DIFFUSION: + proposal_ctx = ProposalContext( + cache_engine=self.cache_engine, + rank=self.rank, + debug_step=self._dflash_debug_step, + ) + with self.draft_context(): + return await self.proposer.propose(model_inputs, + extra_inputs, + sampling_inputs, + proposal_ctx=proposal_ctx) + raise RuntimeError(f'Unsupported speculative proposal method: {proposal_method!r}.') def _build_warmup_dp_meta(self, inputs: ModelInputs): """Build dp_meta for warmup dummy inputs. @@ -714,48 +750,70 @@ def warmup(self, max_batches: int, target_model_config: ModelConfig): if dist_config.dp > 1: dist.barrier(group=self.draft_dist_ctx.cpu_group) + capture_batch_sizes = self.proposer.model.get_capture_batch_sizes() + capture_batch_sizes = sorted(capture_batch_sizes, reverse=True) + + def _make_dummy(batch_size: int, + is_decoding: bool, + max_q_seqlen: int, + target_hidden_size: int): + return self.inputs_strategy.make_dummy(batch_size, + is_decoding=is_decoding, + device='cuda', + vocab_size=self.model_config.vocab_size, + max_q_seqlen=max_q_seqlen, + target_hidden_size=target_hidden_size, + target_dtype=self.model_config.dtype, + meta=self.make_dummy_meta) + + cache_engine = getattr(self, 'cache_engine', None) + get_warmup_plan = getattr(self.proposer, 'get_warmup_plan', None) + warmup_plan = None + if get_warmup_plan is not None: + warmup_plan = get_warmup_plan(max_batches, + target_model_config, + capture_batch_sizes, + cache_engine.cache_config if cache_engine is not None else None) + if warmup_plan is not None: + for case in warmup_plan.cases: + inputs = _make_dummy(case.batch_size, + is_decoding=case.is_decoding, + max_q_seqlen=case.max_q_seqlen, + target_hidden_size=case.target_hidden_size) + self._build_warmup_dp_meta(inputs) + forward_inputs = self.proposer.prepare_warmup_forward(inputs, cache_engine) + if forward_inputs is not None: + self._forward_impl(forward_inputs) + torch.cuda.synchronize() + return + target_hidden_size = self.proposer.get_target_hidden_size(target_model_config) - # warmup prefill - inputs = self.inputs_strategy.make_dummy(max_batches, - is_decoding=False, - device='cuda', - vocab_size=self.model_config.vocab_size, - target_hidden_size=target_hidden_size, - target_dtype=self.model_config.dtype, - meta=self.make_dummy_meta) + inputs = _make_dummy(max_batches, + is_decoding=False, + max_q_seqlen=1, + target_hidden_size=target_hidden_size) self._build_warmup_dp_meta(inputs) # warmup prefill self._forward_impl(inputs) torch.cuda.synchronize() - capture_batch_sizes = self.proposer.model.get_capture_batch_sizes() - capture_batch_sizes = sorted(capture_batch_sizes, reverse=True) - # warmup decode for batch_size in capture_batch_sizes: # decode with num_spec_tokens + 1 per seq - inputs = self.inputs_strategy.make_dummy(batch_size, - is_decoding=True, - device='cuda', - vocab_size=self.model_config.vocab_size, - max_q_seqlen=self.num_spec_tokens + 1, - target_hidden_size=target_hidden_size, - target_dtype=self.model_config.dtype, - meta=self.make_dummy_meta) + inputs = _make_dummy(batch_size, + is_decoding=True, + max_q_seqlen=self.num_spec_tokens + 1, + target_hidden_size=target_hidden_size) self._build_warmup_dp_meta(inputs) self._forward_impl(inputs) torch.cuda.synchronize() # decode 1 tokens per sequence - inputs = self.inputs_strategy.make_dummy(batch_size, - is_decoding=True, - device='cuda', - vocab_size=self.model_config.vocab_size, - max_q_seqlen=1, - target_hidden_size=self.model_config.hidden_size, - target_dtype=self.model_config.dtype, - meta=self.make_dummy_meta) + inputs = _make_dummy(batch_size, + is_decoding=True, + max_q_seqlen=1, + target_hidden_size=self.model_config.hidden_size) self._build_warmup_dp_meta(inputs) self._forward_impl(inputs) torch.cuda.synchronize() diff --git a/tests/pytorch/engine/test_model_agent.py b/tests/pytorch/engine/test_model_agent.py index c074a47b5d..e241c87583 100644 --- a/tests/pytorch/engine/test_model_agent.py +++ b/tests/pytorch/engine/test_model_agent.py @@ -174,6 +174,76 @@ def test_build_spec_agent_shares_guided_helper_with_proposer(monkeypatch): assert spec_agent.guided_helper.manager is guided_manager assert proposer.guided_helper is spec_agent.guided_helper +def _make_minimal_build_agent(target_layer_ids, mask_token_id=99): + from lmdeploy.pytorch.engine.model_agent.agent import BaseModelAgent + + spec_agent = SimpleNamespace( + is_enabled=lambda: True, + method='dflash', + specdecode_config=SimpleNamespace(target_layer_ids=target_layer_ids, mask_token_id=mask_token_id), + num_spec_tokens=15, + requires_target_inputs_embeds=lambda: False, + ) + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.model_path = 'target-model' + agent.adapters = None + agent.device = torch.device('cpu') + agent.rank = 0 + agent.model_config = SimpleNamespace( + custom_module_map=None, + quant_config=None, + fp32_lm_head=False, + tie_word_embeddings=False, + ) + agent.misc_config = SimpleNamespace( + enable_return_routed_experts=False, + language_model_only=False, + dllm_config=None, + empty_init=True, + ) + agent.need_output = False + agent.strategy_factory = None + agent.cache_config = SimpleNamespace(max_batches=4) + agent.spec_agent = spec_agent + return agent + + +def test_model_agent_dflash_build_requires_parsed_metadata(): + agent = _make_minimal_build_agent(target_layer_ids=None) + + with pytest.raises(ValueError, match='parsed target_layer_ids metadata'): + agent._build_model() + + agent = _make_minimal_build_agent(target_layer_ids=(1, 3, 5), mask_token_id=None) + with pytest.raises(ValueError, match='parsed mask_token_id metadata'): + agent._build_model() + + +def test_model_agent_dflash_layers_flow_through_build_context(monkeypatch): + import lmdeploy.pytorch.engine.model_agent.agent as agent_mod + + agent = _make_minimal_build_agent(target_layer_ids=(1, 3, 5)) + captured = {} + patched_model = object() + + def build_patched_model(model_config, device=None, build_model_ctx=None): + captured['model_config'] = model_config + captured['device'] = device + captured['build_model_ctx'] = build_model_ctx + return patched_model + + monkeypatch.setattr(agent_mod, 'build_patched_model', build_patched_model) + + agent._build_model() + + assert captured['model_config'] is agent.model_config + assert captured['device'] == torch.device('cpu') + assert captured['build_model_ctx'].target_aux_hidden_state_layers == (1, 3, 5) + assert captured['build_model_ctx'].speculative_mask_token_id == 99 + assert captured['build_model_ctx'].requires_target_inputs_embeds is False + assert agent.patched_model is patched_model + assert agent.build_model_ctx is captured['build_model_ctx'] + def test_spec_agent_reset_runtime_state_discards_chunk_carry(): from lmdeploy.pytorch.spec_decode.spec_agent import SpecModelAgent diff --git a/tests/pytorch/spec_decode/test_cudagraph_strategy.py b/tests/pytorch/spec_decode/test_cudagraph_strategy.py index c8d03dc46b..aff76a93bc 100644 --- a/tests/pytorch/spec_decode/test_cudagraph_strategy.py +++ b/tests/pytorch/spec_decode/test_cudagraph_strategy.py @@ -20,20 +20,35 @@ def test_arspec_cudagraph_keeps_full_spec_capture_for_eagle3(): assert strategy.get_max_tokens(batch_size=8, origin_batch_size=8, num_tokens=40) == 40 -def test_cudagraph_fa3_metadata_uses_single_query_len_for_single_token_capture(): +def test_cudagraph_fa3_metadata_uses_explicit_swa_policy_for_single_token_capture(monkeypatch): from types import SimpleNamespace import torch + from lmdeploy.pytorch.models.utils import cudagraph as cudagraph_mod from lmdeploy.pytorch.models.utils.cudagraph import CudaGraphMeta, CudaGraphMixin + step_context = SimpleNamespace(model_config=SimpleNamespace(sliding_window=4096)) + monkeypatch.setattr(cudagraph_mod, 'get_step_ctx_manager', + lambda: SimpleNamespace(current_context=lambda: step_context)) + class DummyCudaGraphModel(CudaGraphMixin): def __init__(self): self.max_seqlen_q_calls = [] - def update_meta_flashattn(self, batch_size, max_seqlen_q, block_size, max_seqlen_k, cache_seqlens): + def build_fa3_scheduler_metadata(self, + batch_size, + max_seqlen_q, + block_size, + max_seqlen_k, + cache_seqlens, + *, + sliding_window, + causal): self.max_seqlen_q_calls.append(max_seqlen_q) + assert sliding_window == 4096 + assert causal is True return torch.zeros(4, dtype=torch.int32) model = DummyCudaGraphModel() diff --git a/tests/pytorch/spec_decode/test_dflash_utils.py b/tests/pytorch/spec_decode/test_dflash_utils.py new file mode 100644 index 0000000000..0cadbc11a9 --- /dev/null +++ b/tests/pytorch/spec_decode/test_dflash_utils.py @@ -0,0 +1,906 @@ +import argparse +import inspect +from types import SimpleNamespace + +import pytest +import torch + +from lmdeploy.cli.utils import ArgumentHelper, get_speculative_config +from lmdeploy.messages import PytorchEngineConfig, SpeculativeConfig +from lmdeploy.pytorch.backends.cuda.attention.default import TritonAttentionMetadata +from lmdeploy.pytorch.config import CacheConfig, DistConfig, ModelConfig, SpecDecodeConfig +from lmdeploy.pytorch.engine.config_builder import ConfigBuilder +from lmdeploy.pytorch.model_inputs import BuildModelContext, ModelInputs +from lmdeploy.pytorch.models.module_map import MODULE_MAP +from lmdeploy.pytorch.models.patch import build_model_context +from lmdeploy.pytorch.models.qwen3 import Qwen3ForCausalLM +from lmdeploy.pytorch.models.qwen3_5 import Qwen3_5ForConditionalGeneration, Qwen3_5Model +from lmdeploy.pytorch.models.qwen3_dflash import ( + DFlashDraftModel, + DFlashQwen3Attention, + DFlashQwen3DecoderLayer, + _normalize_dflash_weight_name, + _resolve_dflash_layer_attention, +) +from lmdeploy.pytorch.models.utils.cudagraph import CudaGraphMixin +from lmdeploy.pytorch.spec_decode.dflash_utils import ( + build_target_layer_ids, + parse_dflash_config, + validate_dflash_cache_config, + validate_dflash_dist_config, + validate_dflash_runtime_config, +) +from lmdeploy.pytorch.spec_decode.proposers.dflash import DFlash + + +def _draft_config(**kwargs): + values = dict( + architectures=['DFlashDraftModel'], + bos_token_id=1, + eos_token_id=2, + hidden_size=128, + model_type='qwen3', + num_attention_heads=4, + num_hidden_layers=4, + num_target_layers=16, + num_key_value_heads=2, + vocab_size=32000, + dflash_config=dict( + block_size=4, + mask_token_id=32001, + target_layer_ids=[1, 5, 9, 13], + ), + layer_types=['full_attention'] * 4, + use_sliding_window=False, + ) + values.update(kwargs) + return SimpleNamespace(**values) + + +def _parse_dflash(draft_config, num_speculative_tokens, target_num_layers=None): + if target_num_layers is None: + target_num_layers = draft_config.num_target_layers + return parse_dflash_config(draft_config, + num_speculative_tokens=num_speculative_tokens, + target_num_layers=target_num_layers) + + +def test_parse_dflash_config_valid(): + target_layer_ids, mask_token_id = _parse_dflash(_draft_config(), num_speculative_tokens=3) + + assert target_layer_ids == (1, 5, 9, 13) + assert mask_token_id == 32001 + + +def test_specdecode_config_stores_resolved_dflash_fields_directly(): + target_layer_ids, mask_token_id = _parse_dflash(_draft_config(), num_speculative_tokens=3) + cfg = SpecDecodeConfig(model='draft-model', + method='dflash', + num_speculative_tokens=3, + target_layer_ids=target_layer_ids, + mask_token_id=mask_token_id) + + assert cfg.target_layer_ids == (1, 5, 9, 13) + assert cfg.mask_token_id == 32001 + assert not hasattr(cfg, 'dflash_config') + + +def test_dflash_block_size_overrides_draft_token_count(): + cfg = SpeculativeConfig(method='dflash', num_speculative_tokens=15, dflash_block_size=8) + + assert cfg.dflash_block_size == 8 + assert cfg.num_speculative_tokens == 7 + + +@pytest.mark.parametrize('block_size', [True, 1, 0, -1]) +def test_dflash_block_size_rejects_invalid_values(block_size): + with pytest.raises(ValueError, match='integer greater than or equal to 2'): + SpeculativeConfig(method='dflash', dflash_block_size=block_size) + + +def test_dflash_block_size_rejects_invalid_method(): + with pytest.raises(ValueError, match='only when method="dflash"'): + SpeculativeConfig(method='eagle3', dflash_block_size=8) + + +def test_dflash_block_size_cli_override_and_algorithm_validation(): + parser = argparse.ArgumentParser() + ArgumentHelper.add_spec_group(parser) + args = parser.parse_args([ + '--speculative-algorithm', + 'dflash', + '--speculative-num-draft-tokens', + '15', + '--speculative-dflash-block-size', + '8', + ]) + cfg = get_speculative_config(args) + assert cfg.dflash_block_size == 8 + assert cfg.num_speculative_tokens == 7 + + args = parser.parse_args(['--speculative-dflash-block-size', '8']) + with pytest.raises(ValueError, match='requires --speculative-algorithm dflash'): + get_speculative_config(args) + + +def test_specdecode_config_checks_matching_dflash_target_depth(monkeypatch): + draft_model_config = SimpleNamespace(hf_config=_draft_config(num_target_layers=16)) + target_model_config = SimpleNamespace(num_layers=4, llm_config=SimpleNamespace(num_hidden_layers=16)) + calls = [] + + def from_pretrained(model, **kwargs): + calls.append((model, kwargs.get('is_draft_model'))) + return draft_model_config if model == 'draft-model' else target_model_config + + monkeypatch.setattr(ModelConfig, 'from_pretrained', from_pretrained) + cfg = SpecDecodeConfig.from_config(method='dflash', + num_speculative_tokens=3, + model='draft-model', + target_model='target-model', + target_cache_cfg=CacheConfig(max_batches=1, + block_size=64, + num_cpu_blocks=0, + num_gpu_blocks=1)) + + assert calls == [('draft-model', True), ('target-model', False)] + assert cfg.target_layer_ids == (1, 5, 9, 13) + + +def test_specdecode_config_rejects_mismatched_dflash_target_depth(monkeypatch): + draft_model_config = SimpleNamespace(hf_config=_draft_config(num_target_layers=16)) + target_model_config = SimpleNamespace(num_layers=4, llm_config=SimpleNamespace(num_hidden_layers=15)) + monkeypatch.setattr(ModelConfig, 'from_pretrained', + lambda model, **kwargs: draft_model_config if model == 'draft-model' else target_model_config) + + with pytest.raises(ValueError, match='draft/target depth mismatch.*16.*15'): + SpecDecodeConfig.from_config(method='dflash', + num_speculative_tokens=3, + model='draft-model', + target_model='target-model', + target_cache_cfg=CacheConfig(max_batches=1, + block_size=64, + num_cpu_blocks=0, + num_gpu_blocks=1)) + + +def test_parse_dflash_config_requires_mask_token_id(): + draft_config = _draft_config(dflash_config=dict(block_size=4, target_layer_ids=[1, 5])) + + with pytest.raises(ValueError, match='mask_token_id'): + _parse_dflash(draft_config, num_speculative_tokens=3) + + +@pytest.mark.parametrize('num_speculative_tokens', [3, 7, 15]) +def test_parse_dflash_config_allows_runtime_query_up_to_checkpoint_block_size(num_speculative_tokens): + draft_config = _draft_config(dflash_config=dict(block_size=16, mask_token_id=32001)) + + target_layer_ids, mask_token_id = _parse_dflash(draft_config, + num_speculative_tokens=num_speculative_tokens) + + assert target_layer_ids == build_target_layer_ids(16, 4) + assert mask_token_id == 32001 + + +def test_parse_dflash_config_rejects_query_above_checkpoint_block_size(): + draft_config = _draft_config(dflash_config=dict(block_size=16, mask_token_id=32001)) + + with pytest.raises(ValueError, match='must not exceed.*block_size'): + _parse_dflash(draft_config, num_speculative_tokens=16) + + +def test_parse_dflash_config_resolves_target_layers_from_num_target_layers(): + draft_config = _draft_config(dflash_config=dict(block_size=4, mask_token_id=32001)) + + target_layer_ids, mask_token_id = _parse_dflash(draft_config, num_speculative_tokens=3) + + assert target_layer_ids == build_target_layer_ids(16, 4) + assert mask_token_id == 32001 + + +def test_parse_dflash_config_rejects_non_increasing_target_layers(): + draft_config = _draft_config( + dflash_config=dict( + block_size=4, + mask_token_id=32001, + target_layer_ids=[1, 9, 5, 13], + )) + + with pytest.raises(ValueError, match='strictly increasing'): + _parse_dflash(draft_config, num_speculative_tokens=3) + + +def test_parse_dflash_config_rejects_duplicate_target_layers(): + draft_config = _draft_config( + dflash_config=dict( + block_size=4, + mask_token_id=32001, + target_layer_ids=[1, 5, 5, 13], + )) + + with pytest.raises(ValueError, match='duplicate-free'): + _parse_dflash(draft_config, num_speculative_tokens=3) + + +def test_parse_dflash_config_rejects_out_of_range_target_layers(): + draft_config = _draft_config( + num_target_layers=8, + dflash_config=dict( + block_size=4, + mask_token_id=32001, + target_layer_ids=[1, 5, 8], + )) + + with pytest.raises(ValueError, match='out-of-range'): + _parse_dflash(draft_config, num_speculative_tokens=3) + + +def test_parse_dflash_config_rejects_mismatched_target_depth(): + draft_config = _draft_config(num_target_layers=40) + + with pytest.raises(ValueError, match='draft/target depth mismatch.*40.*39'): + _parse_dflash(draft_config, num_speculative_tokens=3, target_num_layers=39) + + +def test_parse_dflash_config_matches_real_checkpoint_schema(): + draft_config = _draft_config( + num_hidden_layers=6, + num_target_layers=40, + dflash_config={ + 'block_size': 16, + 'mask_token_id': 248077, + 'target_layer_ids': [1, 6, 11, 16, 22, 27, 32, 37], + }, + layer_types=['sliding_attention'] * 5 + ['full_attention'], + use_sliding_window=True, + sliding_window=4096, + ) + + target_layer_ids, mask_token_id = _parse_dflash(draft_config, num_speculative_tokens=15) + + assert target_layer_ids == (1, 6, 11, 16, 22, 27, 32, 37) + assert mask_token_id == 248077 + + +def test_validate_dflash_allows_hybrid_sliding_attention(): + draft_config = _draft_config( + layer_types=['sliding_attention', 'sliding_attention', 'full_attention', 'full_attention'], + sliding_window=4096, + use_sliding_window=True, + ) + target_layer_ids, mask_token_id = _parse_dflash(draft_config, num_speculative_tokens=3) + + assert target_layer_ids == (1, 5, 9, 13) + assert mask_token_id == 32001 + + +def test_validate_dflash_rejects_unknown_layer_type(): + draft_config = _draft_config(layer_types=['linear_attention'] * 4) + + with pytest.raises(ValueError, match='full_attention and sliding_attention'): + _parse_dflash(draft_config, num_speculative_tokens=3) + + +def test_parse_dflash_config_rejects_invalid_schema_before_model_construction(): + with pytest.raises(ValueError, match='dflash_config dictionary'): + _parse_dflash(_draft_config(dflash_config=SimpleNamespace()), num_speculative_tokens=3) + + with pytest.raises(ValueError, match='layer_types length'): + _parse_dflash(_draft_config(layer_types=['full_attention']), num_speculative_tokens=3) + + with pytest.raises(ValueError, match='layer_types must contain only strings'): + _parse_dflash(_draft_config(layer_types=['full_attention', 1, 'full_attention', 'full_attention']), + num_speculative_tokens=3) + + sink_config = _draft_config( + dflash_config=dict(block_size=4, mask_token_id=32001, attention_sink_bias=True)) + with pytest.raises(ValueError, match='attention-sink bias'): + _parse_dflash(sink_config, num_speculative_tokens=3) + + +def test_parse_dflash_config_rejects_non_default_scheduler_patterns(): + forced_swa = _draft_config( + dflash_config=dict(block_size=4, + mask_token_id=32001, + use_swa=True, + swa_window_size=2048), + sliding_window=2048, + use_sliding_window=True, + ) + with pytest.raises(ValueError, match='does not support use_swa forcing'): + _parse_dflash(forced_swa, num_speculative_tokens=3) + + alternate_window = _draft_config( + dflash_config=dict(block_size=4, mask_token_id=32001, swa_window_size=2048), + layer_types=['sliding_attention'] * 4, + sliding_window=4096, + use_sliding_window=True, + ) + with pytest.raises(ValueError, match='SWA window.*ModelConfig.sliding_window'): + _parse_dflash(alternate_window, num_speculative_tokens=3) + + noncausal_swa = _draft_config( + dflash_config=dict(block_size=4, mask_token_id=32001, causal=False), + layer_types=['sliding_attention'] * 4, + sliding_window=4096, + use_sliding_window=True, + ) + with pytest.raises(ValueError, match='unsupported pattern=.*4096, False'): + _parse_dflash(noncausal_swa, num_speculative_tokens=3) + + +def test_validate_dflash_dist_rejects_dp_and_ep_allows_tp(): + validate_dflash_dist_config(DistConfig(tp=2)) + + with pytest.raises(ValueError, match='dp=1'): + validate_dflash_dist_config(DistConfig(dp=2, tp=2)) + + with pytest.raises(ValueError, match='ep=1'): + validate_dflash_dist_config(DistConfig(ep=2)) + + +def test_validate_dflash_cache_rejects_prefix_cache_and_kv_quant(): + validate_dflash_cache_config(SimpleNamespace(enable_prefix_caching=False, quant_policy=0)) + + with pytest.raises(ValueError, match='prefix-cache'): + validate_dflash_cache_config(SimpleNamespace(enable_prefix_caching=True, quant_policy=0)) + + with pytest.raises(ValueError, match='KV-cache quantization'): + validate_dflash_cache_config(SimpleNamespace(enable_prefix_caching=False, quant_policy=8)) + + +def test_validate_dflash_runtime_rejects_non_cuda_and_allows_cudagraph(): + validate_dflash_runtime_config(cache_config=SimpleNamespace(device_type='cuda'), + backend_config=SimpleNamespace(device_type='cuda', eager_mode=True)) + validate_dflash_runtime_config(cache_config=SimpleNamespace(device_type='cuda'), + backend_config=SimpleNamespace(device_type='cuda', eager_mode=False)) + + with pytest.raises(ValueError, match='requires CUDA'): + validate_dflash_runtime_config(cache_config=SimpleNamespace(device_type='ascend')) + + +@pytest.mark.parametrize('model_cls', [Qwen3ForCausalLM, Qwen3_5ForConditionalGeneration]) +def test_qwen_cudagraph_outputs_preserve_only_target_aux_hidden_states(model_cls): + output_buffers = { + 'hidden_states': torch.arange(24).view(1, 6, 4), + 'aux_hidden_states': torch.arange(60).view(1, 6, 10), + 'target_inputs_embeds': torch.arange(30).view(1, 6, 5), + 'all_routed_experts': torch.arange(36).view(6, 2, 3), + } + input_ids = torch.zeros((1, 3), dtype=torch.long) + + model = object.__new__(model_cls) + outputs = model.get_outputs_cudagraph(output_buffers, input_ids) + + assert outputs['hidden_states'].shape == (1, 3, 4) + assert outputs['aux_hidden_states'].shape == (1, 3, 10) + assert 'target_inputs_embeds' not in outputs + assert outputs['aux_hidden_states'].data_ptr() == output_buffers['aux_hidden_states'].data_ptr() + torch.testing.assert_close(outputs['aux_hidden_states'], output_buffers['aux_hidden_states'][:, :3]) + torch.testing.assert_close(outputs['all_routed_experts'], output_buffers['all_routed_experts'][:3]) + + +def test_dflash_draft_graph_inputs_exclude_target_hidden_and_target_embeds(): + assert 'target_hidden_states' not in inspect.getsource(DFlashDraftModel.make_buffers_cudagraph) + assert 'target_hidden_states' not in inspect.getsource(DFlashDraftModel.fill_buffers_cudagraph) + + model = object.__new__(DFlashDraftModel) + explicit_embeds = torch.randn(1, 2, 4) + context = SimpleNamespace( + input_ids=torch.tensor([[1, 2]]), + position_ids=torch.tensor([[0, 1]]), + attn_metadata=object(), + target_inputs_embeds=torch.full((1, 2, 4), 42.0), + ) + outputs = model.prepare_inputs_for_generation([], inputs_embeds=explicit_embeds, context=context) + + assert outputs['inputs_embeds'] is explicit_embeds + assert 'target_hidden_states' not in outputs + + +def test_dflash_graph_scheduler_metadata_reuses_generic_swa_and_separates_full(monkeypatch): + model = object.__new__(DFlashDraftModel) + torch.nn.Module.__init__(model) + + generic = torch.tensor([1, 2, 3], dtype=torch.int32) + monkeypatch.setattr(CudaGraphMixin, 'make_buffers_cudagraph', + lambda self, graph_meta, **kwargs: { + 'kv_seqlens': torch.zeros(2, dtype=torch.int32), + 'scheduler_metadata': generic.clone(), + }) + patterns = [] + lengths = iter((4, 4, 2, 3)) + + def build_fa3_scheduler_metadata(*args, sliding_window, causal, **kwargs): + patterns.append((sliding_window, causal)) + return torch.arange(next(lengths), dtype=torch.int32) + + monkeypatch.setattr(model, 'build_fa3_scheduler_metadata', build_fa3_scheduler_metadata) + graph_meta1 = SimpleNamespace(use_fa3_decoding=True, + num_blocks=2, + block_size=64, + max_batchs=2, + decode_query_len=16) + graph_meta2 = SimpleNamespace(**vars(graph_meta1)) + + buffers1 = model.make_buffers_cudagraph(graph_meta1) + buffers2 = model.make_buffers_cudagraph(graph_meta2) + graph_meta1.input_buffers = buffers1 + graph_meta2.input_buffers = buffers2 + assert buffers1['scheduler_metadata'].data_ptr() != buffers1['dflash_full_scheduler_metadata'].data_ptr() + assert buffers1['dflash_full_scheduler_metadata'].data_ptr() != buffers2[ + 'dflash_full_scheduler_metadata'].data_ptr() + + attn_metadata = TritonAttentionMetadata(is_decoding=True, + block_offsets=torch.zeros(2, 1, dtype=torch.int32), + scheduler_metadata=buffers1['scheduler_metadata']) + monkeypatch.setattr(CudaGraphMixin, 'fill_buffers_cudagraph', + lambda self, graph_meta, **kwargs: {'attn_metadata': attn_metadata}) + inputs1 = model.fill_buffers_cudagraph(graph_meta1) + inputs2 = model.fill_buffers_cudagraph(graph_meta2) + + assert patterns == [(None, False)] * 4 + full1 = inputs1['dflash_full_scheduler_metadata'] + full2 = inputs2['dflash_full_scheduler_metadata'] + assert full1.shape == (2, ) + assert full2.shape == (3, ) + assert full1.data_ptr() == buffers1['dflash_full_scheduler_metadata'].data_ptr() + assert full2.data_ptr() == buffers2['dflash_full_scheduler_metadata'].data_ptr() + assert not hasattr(attn_metadata, 'scheduler_metadata_overrides') + assert attn_metadata.scheduler_metadata is buffers1['scheduler_metadata'] + + +def test_dflash_layer_scheduler_metadata_is_local_and_never_mutates_shared_input(): + draft_config = _draft_config( + layer_types=['sliding_attention'] * 3 + ['full_attention'], + sliding_window=4096, + use_sliding_window=True, + ) + _parse_dflash(draft_config, num_speculative_tokens=3) + layer_patterns = [ + _resolve_dflash_layer_attention(draft_config, layer_idx) + for layer_idx in range(draft_config.num_hidden_layers) + ] + assert layer_patterns == [(4096, True)] * 3 + [(None, False)] + + generic = torch.tensor([1, 2, 3], dtype=torch.int32) + attn_metadata = TritonAttentionMetadata(is_decoding=True, + block_offsets=torch.zeros(1, 1, dtype=torch.int32), + scheduler_metadata=generic) + swa_layer = SimpleNamespace(self_attn=SimpleNamespace(sliding_window=4096, causal=True)) + full_layer = SimpleNamespace(self_attn=SimpleNamespace(sliding_window=None, causal=False)) + + swa_metadata = DFlashQwen3DecoderLayer._get_layer_attn_metadata(swa_layer, attn_metadata, None) + eager_full_metadata = DFlashQwen3DecoderLayer._get_layer_attn_metadata(full_layer, attn_metadata, None) + full_buffer = torch.tensor([4, 5, 6], dtype=torch.int32) + graph_full_metadata = DFlashQwen3DecoderLayer._get_layer_attn_metadata(full_layer, attn_metadata, full_buffer) + prefill_metadata = TritonAttentionMetadata(is_decoding=False, + block_offsets=attn_metadata.block_offsets, + scheduler_metadata=generic) + + assert swa_metadata is attn_metadata + assert eager_full_metadata is not attn_metadata + assert eager_full_metadata.scheduler_metadata is None + assert graph_full_metadata is not attn_metadata + assert graph_full_metadata.scheduler_metadata is full_buffer + assert DFlashQwen3DecoderLayer._get_layer_attn_metadata(full_layer, prefill_metadata, + full_buffer) is prefill_metadata + assert attn_metadata.scheduler_metadata is generic + + scheduler_policy_by_ptr = { + generic.data_ptr(): (4096, True), + full_buffer.data_ptr(): (None, False), + } + for sliding_window, causal in layer_patterns: + layer = SimpleNamespace(self_attn=SimpleNamespace(sliding_window=sliding_window, causal=causal)) + layer_metadata = DFlashQwen3DecoderLayer._get_layer_attn_metadata(layer, attn_metadata, full_buffer) + assert scheduler_policy_by_ptr[layer_metadata.scheduler_metadata.data_ptr()] == (sliding_window, causal) + + +def test_dflash_production_forward_requires_attention_metadata(): + assert inspect.signature(DFlashDraftModel.forward).parameters['attn_metadata'].default is inspect.Parameter.empty + assert inspect.signature(DFlashQwen3DecoderLayer.forward).parameters[ + 'attn_metadata'].default is inspect.Parameter.empty + assert inspect.signature(DFlashQwen3Attention.forward).parameters[ + 'attn_metadata'].default is inspect.Parameter.empty + + +def test_qwen35_dflash_target_embed_policy_covers_image_and_chunked_multimodal(): + model = object.__new__(Qwen3_5ForConditionalGeneration) + model.is_spec_decoding = True + model.requires_target_inputs_embeds = False + + image_context = SimpleNamespace(is_chunk_multimodal=False) + chunk_context = SimpleNamespace(is_chunk_multimodal=True) + assert model._should_return_target_inputs_embeds(torch.empty(1), image_context) is False + assert model._should_return_target_inputs_embeds(None, chunk_context) is False + + model.requires_target_inputs_embeds = True + assert model._should_return_target_inputs_embeds(torch.empty(1), image_context) is True + assert model._should_return_target_inputs_embeds(None, chunk_context) is True + + +def test_qwen35_visual_processing_and_aux_output_do_not_require_returned_input_embeds(): + + class FakeEmbeddings(torch.nn.Module): + + def forward(self, input_ids): + return torch.zeros((*input_ids.shape, 2), dtype=torch.float32) + + class FakeLanguageModel(torch.nn.Module): + + def __init__(self): + super().__init__() + self.embeddings = FakeEmbeddings() + + def get_input_embeddings(self): + return self.embeddings + + def forward(self, **kwargs): + inputs_embeds = kwargs['inputs_embeds'] + return dict(hidden_states=inputs_embeds, aux_hidden_states=inputs_embeds + 1) + + class FakeVisual(torch.nn.Module): + spatial_merge_size = 1 + + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, pixel_values, **kwargs): + self.calls += 1 + return pixel_values + + model = object.__new__(Qwen3_5Model) + torch.nn.Module.__init__(model) + model.language_model = FakeLanguageModel() + model.visual = FakeVisual() + image_embeds = torch.tensor([[9.0, 10.0], [11.0, 12.0]]) + + hidden_states, target_inputs_embeds, aux_hidden_states = model( + input_ids=torch.tensor([[1, 2, 3]]), + position_ids=torch.tensor([[0, 1, 2]]), + past_key_values=[], + attn_metadata=None, + state_ids=torch.tensor([0]), + pixel_values=image_embeds, + vis_cu_seqlens=torch.tensor([2], dtype=torch.int32), + vis_pos_emb=(torch.ones(2, 1), torch.ones(2, 1)), + multimodal_mask=torch.tensor([[True, True, False]]), + grid_thw=torch.tensor([[1, 1, 2]]), + return_input_embeds=False, + ) + + assert model.visual.calls == 1 + assert target_inputs_embeds is None + torch.testing.assert_close(hidden_states[0, :2], image_embeds) + torch.testing.assert_close(aux_hidden_states[0, :2], image_embeds + 1) + + +def test_dflash_specdecode_builder_preserves_tp(monkeypatch): + captured = {} + + def _fake_from_config(**kwargs): + captured.update(kwargs) + return SimpleNamespace(dist_config=kwargs['dist_config']) + + monkeypatch.setattr('lmdeploy.pytorch.engine.config_builder.SpecDecodeConfig.from_config', _fake_from_config) + + cache_config = CacheConfig( + max_batches=1, + block_size=64, + num_cpu_blocks=0, + num_gpu_blocks=1, + device_type='cuda', + ) + dist_config = DistConfig(tp=2) + speculative_config = SpeculativeConfig(method='dflash', + model=None, + num_speculative_tokens=15, + dflash_block_size=8) + specdecode_config = ConfigBuilder.build_specdecode_config( + target_model='target-model', + speculative_config=speculative_config, + engine_config=PytorchEngineConfig(tp=2, eager_mode=True), + cache_config=cache_config, + dist_config=dist_config, + trust_remote_code=True, + ) + + assert specdecode_config.dist_config.tp == 2 + assert specdecode_config.dist_config.attn_tp == 2 + assert captured['dist_config'] is not dist_config + assert captured['num_speculative_tokens'] == 7 + + +def test_dflash_target_model_config_uses_ar_spec_not_dllm(): + hf_config = _draft_config(architectures=['Qwen3ForCausalLM']) + + model_config = ModelConfig.from_hf_config(hf_config, spec_method='dflash') + + assert model_config.model_paradigm == 'ar_spec' + assert not hasattr(model_config, 'dflash_config') + + +def test_qwen3_dflash_uses_only_resolved_build_context_metadata(monkeypatch): + import lmdeploy.pytorch.models.qwen3_dflash as dflash_model_mod + + class FakeAttention(torch.nn.Module): + + def __init__(self): + super().__init__() + self.causal = True + self.sliding_window = None + + class FakeLayer(torch.nn.Module): + + def __init__(self, *args, **kwargs): + super().__init__() + self.self_attn = FakeAttention() + + monkeypatch.setattr(dflash_model_mod, 'DFlashQwen3DecoderLayer', FakeLayer) + monkeypatch.setattr(dflash_model_mod, 'build_colwise_linear', + lambda in_features, out_features, **kwargs: torch.nn.Linear(in_features, out_features, + bias=False)) + monkeypatch.setattr(dflash_model_mod, 'RMSNorm', lambda *args, **kwargs: torch.nn.Identity()) + monkeypatch.setattr(dflash_model_mod, 'build_rotary_embedding_from_config', + lambda *args, **kwargs: torch.nn.Identity()) + + # Raw checkpoint values deliberately disagree with the resolved metadata. + hf_config = _draft_config( + target_hidden_size=8, + rms_norm_eps=1e-6, + dflash_config=dict(mask_token_id=32001, target_layer_ids=[0]), + ) + resolved_ctx = BuildModelContext(target_aux_hidden_state_layers=(1, 6, 11), + speculative_mask_token_id=99) + with build_model_context(resolved_ctx): + model = DFlashDraftModel(hf_config, ctx_mgr=SimpleNamespace(), device=torch.device('cpu')) + + assert model.target_layer_ids == (1, 6, 11) + assert model.mask_token_id == 99 + assert model.num_context_features == 3 + assert model.fc.in_features == 24 + + +def test_dflash_module_map_keeps_only_supported_architecture_name(): + assert MODULE_MAP['DFlashDraftModel'].endswith('.qwen3_dflash.DFlashDraftModel') + assert 'DFlashQwen3ForCausalLM' not in MODULE_MAP + assert 'Qwen3DFlashForCausalLM' not in MODULE_MAP + + +def test_qwen3_dflash_weight_name_normalization(): + assert _normalize_dflash_weight_name('model.fc.weight') == 'fc.weight' + assert _normalize_dflash_weight_name('midlayer.self_attn.q_proj.weight') == 'layers.0.self_attn.q_proj.weight' + assert _normalize_dflash_weight_name('layers.0.self_attn.rotary_emb.inv_freq') is None + + +def test_qwen3_dflash_attention_resolver_consumes_trusted_checkpoint_fields(): + sliding_config = _draft_config( + layer_types=['sliding_attention', 'full_attention', 'full_attention', 'full_attention'], + sliding_window=4096, + use_sliding_window=True, + ) + _parse_dflash(sliding_config, num_speculative_tokens=3) + assert _resolve_dflash_layer_attention(sliding_config, 0) == (4096, True) + assert _resolve_dflash_layer_attention(sliding_config, 1) == (None, False) + + resolver_source = inspect.getsource(_resolve_dflash_layer_attention) + assert '_get_dflash_cfg' not in resolver_source + assert 'raise ValueError' not in resolver_source + + +def test_qwen3_dflash_materialization_requires_cpu_max_q_seqlen(): + model = object.__new__(DFlashDraftModel) + model.precompute_context_kv = lambda target_hidden, position_ids: [] + + with pytest.raises(ValueError, match='CPU max_q_seqlen'): + model.precompute_and_store_context_kv( + torch.empty(0, 1), + torch.empty(0, dtype=torch.long), + past_key_values=[], + attn_metadata=SimpleNamespace(), + ) + + +def test_dflash_block_proposer_builds_context_and_query_inputs(): + proposer = DFlash( + SimpleNamespace( + mask_token_id=99, + target_layer_ids=(1, 5), + num_speculative_tokens=3, + model_config=None, + ), + device='cpu', + ) + model_inputs = ModelInputs( + input_ids=torch.tensor([[10, 11, 12, 20, 21]]), + seq_length=torch.tensor([3, 2]), + history_lengths=torch.tensor([0, 5]), + block_offsets=torch.zeros((2, 4), dtype=torch.int32), + is_decoding=False, + num_ignored_history=torch.zeros(2, dtype=torch.long), + max_q_seqlen=3, + max_kv_seqlen=7, + sum_kv_seqlen=10, + target_inputs_embeds=torch.randn(1, 5, 4), + ) + context_lengths = torch.tensor([2, 2]) + hidden = torch.arange(20).reshape(5, 4) + + sliced_hidden = proposer._slice_by_lengths(hidden, + model_inputs.seq_length, + torch.tensor([2, 1]), + max_seq_length=model_inputs.max_q_seqlen) + context_inputs = proposer._build_context_inputs(model_inputs, context_lengths) + query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8])) + + assert sliced_hidden.tolist() == hidden[[0, 1, 3]].tolist() + assert context_inputs.input_ids.tolist() == [[10, 11, 20, 21]] + assert context_inputs.seq_length.tolist() == [2, 2] + assert context_inputs.max_q_seqlen == 3 + assert context_inputs.max_kv_seqlen == 7 + assert context_inputs.sum_kv_seqlen == 10 + assert context_inputs.target_inputs_embeds is None + assert query_inputs.input_ids.tolist() == [[7, 99, 99, 99, 8, 99, 99, 99]] + assert query_inputs.history_lengths.tolist() == [2, 7] + assert query_inputs.target_position_ids.tolist() == [[2, 3, 4, 5, 7, 8, 9, 10]] + assert query_inputs.is_decoding is True + assert query_inputs.max_q_seqlen == 4 + assert query_inputs.max_kv_seqlen == 11 + assert query_inputs.sum_kv_seqlen == 18 + assert query_inputs.target_inputs_embeds is None + + +def test_dflash_block_proposer_metadata_helpers_do_not_call_tensor_item(monkeypatch): + proposer = DFlash( + SimpleNamespace( + mask_token_id=99, + target_layer_ids=(1, 5), + num_speculative_tokens=3, + model_config=None, + ), + device='cpu', + ) + model_inputs = ModelInputs( + input_ids=torch.tensor([[10, 11, 12, 20, 21]]), + seq_length=torch.tensor([3, 2]), + history_lengths=torch.tensor([0, 5]), + block_offsets=torch.zeros((2, 4), dtype=torch.int32), + is_decoding=False, + num_ignored_history=torch.zeros(2, dtype=torch.long), + max_q_seqlen=3, + max_kv_seqlen=7, + sum_kv_seqlen=10, + target_position_ids=torch.tensor([[100, 101, 102, 200, 201]]), + ) + extra_inputs = SimpleNamespace( + num_rejected_tokens=torch.tensor([1, 0]), + target_hidden_states=torch.arange(5 * 4, dtype=torch.float32).view(5, 4), + ) + + def _fail_item(self): + raise AssertionError('DFlash proposer helper path should not call Tensor.item().') + + monkeypatch.setattr(torch.Tensor, 'item', _fail_item) + + context_inputs, target_hidden, context_lengths, context_position_ids = proposer._prepare_context_materialization( + model_inputs, extra_inputs) + query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8]), + context_position_ids) + + assert context_inputs.max_q_seqlen == 3 + assert context_inputs.max_kv_seqlen == 7 + assert context_inputs.sum_kv_seqlen == 10 + assert target_hidden.shape == (4, 4) + assert query_inputs.max_q_seqlen == 4 + assert query_inputs.max_kv_seqlen == 11 + assert query_inputs.sum_kv_seqlen == 18 + + +def test_dflash_slice_by_lengths_preserves_single_feature_row(): + proposer = DFlash( + SimpleNamespace( + mask_token_id=99, + target_layer_ids=(1, 5), + num_speculative_tokens=3, + model_config=None, + ), + device='cpu', + ) + + feature = torch.tensor([[1, 2, 3, 4]]) + token_ids = torch.tensor([[7]]) + + sliced_feature = proposer._slice_by_lengths( + feature, + seq_lengths=torch.tensor([1]), + keep_lengths=torch.tensor([1]), + max_seq_length=1, + preserve_features=True, + ) + sliced_tokens = proposer._slice_by_lengths( + token_ids, + seq_lengths=torch.tensor([1]), + keep_lengths=torch.tensor([1]), + max_seq_length=1, + ) + + assert sliced_feature.tolist() == [[1, 2, 3, 4]] + assert sliced_tokens.tolist() == [7] + + +def test_dflash_block_proposer_uses_explicit_target_positions(): + proposer = DFlash( + SimpleNamespace( + mask_token_id=99, + target_layer_ids=(1, 5), + num_speculative_tokens=3, + model_config=None, + ), + device='cpu', + ) + model_inputs = ModelInputs( + input_ids=torch.tensor([[10, 11, 12, 20, 21]]), + seq_length=torch.tensor([3, 2]), + history_lengths=torch.tensor([0, 5]), + block_offsets=torch.zeros((2, 4), dtype=torch.int32), + is_decoding=False, + num_ignored_history=torch.zeros(2, dtype=torch.long), + max_q_seqlen=3, + max_kv_seqlen=7, + sum_kv_seqlen=10, + target_position_ids=torch.tensor([[100, 101, 102, 200, 201]]), + ) + context_lengths = torch.tensor([2, 1]) + context_position_ids = proposer._slice_target_position_ids(model_inputs, context_lengths) + context_inputs = proposer._build_context_inputs(model_inputs, context_lengths, context_position_ids) + query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8]), + context_position_ids) + + assert context_position_ids.tolist() == [100, 101, 200] + assert context_inputs.target_position_ids.tolist() == [[100, 101, 200]] + assert query_inputs.target_position_ids.tolist() == [[102, 103, 104, 105, 201, 202, 203, 204]] + + +def test_dflash_materialize_context_only_slices_committed_hidden(monkeypatch): + proposer = DFlash( + SimpleNamespace( + cache_config=SimpleNamespace(block_size=8), + mask_token_id=99, + target_layer_ids=(1, 5), + num_speculative_tokens=3, + model_config=None, + ), + device='cpu', + ) + model_inputs = ModelInputs( + input_ids=torch.tensor([[10, 11, 12, 20, 21]]), + seq_length=torch.tensor([3, 2]), + history_lengths=torch.tensor([0, 5]), + block_offsets=torch.zeros((2, 4), dtype=torch.int32), + is_decoding=False, + num_ignored_history=torch.zeros(2, dtype=torch.long), + max_q_seqlen=3, + max_kv_seqlen=7, + sum_kv_seqlen=10, + ) + extra_inputs = SimpleNamespace( + num_rejected_tokens=torch.tensor([1, 0]), + target_hidden_states=torch.arange(5 * 4, dtype=torch.float32).view(5, 4), + ) + captured = {} + + def _materialize_context(context_inputs, target_hidden, cache_engine): + captured['context_inputs'] = context_inputs + captured['target_hidden'] = target_hidden + captured['cache_engine'] = cache_engine + + monkeypatch.setattr(proposer, '_materialize_context', _materialize_context) + cache_engine = object() + + proposer.materialize_context(model_inputs, extra_inputs, cache_engine) + + assert captured['cache_engine'] is cache_engine + assert captured['context_inputs'].input_ids.tolist() == [[10, 11, 20, 21]] + assert captured['context_inputs'].seq_length.tolist() == [2, 2] + assert captured['target_hidden'].tolist() == extra_inputs.target_hidden_states[[0, 1, 3, 4]].tolist() diff --git a/tests/pytorch/spec_decode/test_spec_agent.py b/tests/pytorch/spec_decode/test_spec_agent.py index 2dcef80bf4..30dc3e11a9 100644 --- a/tests/pytorch/spec_decode/test_spec_agent.py +++ b/tests/pytorch/spec_decode/test_spec_agent.py @@ -1,10 +1,21 @@ import asyncio +import inspect +from contextlib import nullcontext +from dataclasses import FrozenInstanceError from types import SimpleNamespace +import pytest import torch from lmdeploy.pytorch.model_inputs import DPMeta, ModelInputs from lmdeploy.pytorch.spec_decode.guided_spec_helper import GuidedSpecHelper +from lmdeploy.pytorch.spec_decode.proposers.base import ( + ProposalContext, + ProposalMethod, + ProposalWarmupCase, + ProposalWarmupPlan, +) +from lmdeploy.pytorch.spec_decode.proposers.dflash import DFlash from lmdeploy.pytorch.spec_decode.spec_agent import SpecModelAgent, _expand_sampling_inputs from lmdeploy.pytorch.strategies.ar_spec.model_agent import ARSpecExtraInputs @@ -62,6 +73,8 @@ def update_inputs(self, inputs): class _DummyProposer: + proposal_method = ProposalMethod.AUTOREGRESSIVE + def __init__(self): self.get_outputs_calls = 0 self.update_inputs_decoding_calls = 0 @@ -147,6 +160,11 @@ def accept_token(self, processor, token): assert manager.seen_inference_mode == [True, True, True] assert manager.accepted == [10, 20, 11, 21] +class _NoForkGuidedHelper: + + def get_processors(self, session_ctx, response_formats): + return {0: object()} + def test_prepare_inputs_from_main_dp_non_last_first_chunk_shifts_last_token_indices(): """DP non-last first chunks run draft forwards, so indices must match @@ -343,7 +361,7 @@ def _forward_impl(_inputs): agent._forward_impl = _forward_impl - output = asyncio.run(agent._async_model_forward(inputs, extra_inputs, sampling_inputs=None)) + output = asyncio.run(agent._async_autoregressive_model_forward(inputs, extra_inputs, sampling_inputs=None)) expected = torch.zeros((2, 3), dtype=torch.long) torch.testing.assert_close(output.output_draft_token_ids, expected) @@ -378,7 +396,7 @@ def _forward_impl(_inputs): agent._forward_impl = _forward_impl - output = asyncio.run(agent._async_model_forward(inputs, extra_inputs, sampling_inputs=None)) + output = asyncio.run(agent._async_autoregressive_model_forward(inputs, extra_inputs, sampling_inputs=None)) expected = torch.tensor([[0, 1, 2], [0, 1, 2]], dtype=torch.long) torch.testing.assert_close(output.output_draft_token_ids, expected) @@ -418,7 +436,7 @@ def _forward_impl(_inputs): agent._forward_impl = _forward_impl - asyncio.run(agent._async_model_forward(inputs, extra_inputs, sampling_inputs=None)) + asyncio.run(agent._async_autoregressive_model_forward(inputs, extra_inputs, sampling_inputs=None)) assert agent.proposer.model.update_inputs_dp_is_decoding == [True, True] @@ -549,6 +567,217 @@ def forward_impl(inputs): ] +def test_dflash_diffusion_warmup_materializes_context_and_captures_only_block_queries(monkeypatch): + """DFlash warmup must not capture its dead prefill or q_len=1 graphs.""" + + class DummyInputsStrategy: + + def __init__(self): + self.calls = [] + + def make_dummy(self, + batch_size, + is_decoding, + device='cpu', + vocab_size=1, + max_q_seqlen=1, + target_hidden_size=None, + target_dtype=torch.float32, + meta=None): + self.calls.append((batch_size, is_decoding, max_q_seqlen, target_hidden_size)) + num_tokens = batch_size * max_q_seqlen + return ModelInputs( + input_ids=torch.zeros((1, num_tokens), dtype=torch.long), + seq_length=torch.full((batch_size, ), max_q_seqlen, dtype=torch.long), + history_lengths=torch.zeros(batch_size, dtype=torch.long), + block_offsets=torch.zeros((batch_size, 1), dtype=torch.long), + is_decoding=is_decoding, + num_ignored_history=torch.zeros(batch_size, dtype=torch.long), + max_q_seqlen=max_q_seqlen, + max_kv_seqlen=max_q_seqlen, + sum_kv_seqlen=num_tokens, + target_hidden_states=torch.zeros((1, num_tokens, target_hidden_size), dtype=target_dtype), + target_position_ids=torch.zeros((1, num_tokens), dtype=torch.long), + ) + + class DummyGraphModel: + + def get_capture_batch_sizes(self): + return [2, 4] + + proposer = _make_dflash_proposer() + proposer.model = DummyGraphModel() + materialized = [] + monkeypatch.setattr(proposer, '_materialize_context', + lambda inputs, hidden, cache: materialized.append( + (inputs.is_decoding, inputs.seq_length.numel(), inputs.max_q_seqlen, + tuple(hidden.shape), cache))) + + inputs_strategy = DummyInputsStrategy() + forward_graph_keys = [] + sync_calls = [] + monkeypatch.setattr(torch.cuda, 'synchronize', lambda: sync_calls.append(True)) + + agent = object.__new__(SpecModelAgent) + agent.draft_dist_ctx = SimpleNamespace(dist_config=SimpleNamespace(dp=1)) + agent.draft_context = lambda: nullcontext() + agent.inputs_strategy = inputs_strategy + agent.proposer = proposer + agent.cache_engine = SimpleNamespace( + cache_config=SimpleNamespace(block_size=8, max_prefill_token_num=10)) + agent.model_config = SimpleNamespace(vocab_size=11, dtype=torch.float32, hidden_size=4) + agent.num_spec_tokens = 3 + agent.make_dummy_meta = None + dp_markers = [] + + def build_dp_meta(inputs): + marker = object() + inputs.dp_meta = marker + dp_markers.append(marker) + + agent._build_warmup_dp_meta = build_dp_meta + agent._forward_impl = lambda inputs: forward_graph_keys.append( + (inputs.seq_length.numel(), inputs.max_q_seqlen, inputs.dp_meta)) or {} + + agent.warmup(max_batches=4, target_model_config=SimpleNamespace(hidden_size=4)) + + assert inputs_strategy.calls == [(1, False, 8, 8), (4, True, 4, 8), (2, True, 4, 8)] + assert [(is_decoding, batches, q_len, shape) for is_decoding, batches, q_len, shape, _ in materialized] == [ + (False, 1, 8, (8, 8)), + (False, 4, 4, (16, 8)), + (False, 2, 4, (8, 8)), + ] + assert all(cache is agent.cache_engine for *_, cache in materialized) + assert forward_graph_keys == [(4, 4, dp_markers[1]), (2, 4, dp_markers[2])] + assert len(sync_calls) == 3 + assert 'synchronize' not in inspect.getsource(DFlash.prepare_warmup_forward) + + +def test_proposal_warmup_plan_is_immutable_and_shape_only(): + case = ProposalWarmupCase(batch_size=4, is_decoding=True, max_q_seqlen=16, target_hidden_size=1024) + plan = ProposalWarmupPlan(cases=(case, )) + + assert tuple(ProposalWarmupCase.__dataclass_fields__) == ( + 'batch_size', 'is_decoding', 'max_q_seqlen', 'target_hidden_size') + assert tuple(ProposalWarmupPlan.__dataclass_fields__) == ('cases', ) + with pytest.raises(FrozenInstanceError): + case.batch_size = 2 + with pytest.raises(FrozenInstanceError): + plan.cases = () + + +def _make_dflash_proposer(): + return DFlash( + SimpleNamespace( + mask_token_id=99, + target_layer_ids=(1, 5), + num_speculative_tokens=3, + model_config=None, + ), + device='cpu', + ) + + +def test_dflash_proposer_hook_rejects_guided_before_fork(): + agent = object.__new__(SpecModelAgent) + agent.guided_helper = _NoForkGuidedHelper() + agent.proposer = _make_dflash_proposer() + agent.proposer.guided_helper = agent.guided_helper + agent.cache_engine = None + agent.rank = 0 + agent._dflash_debug_step = 0 + agent.draft_context = lambda: nullcontext() + + model_inputs = ModelInputs( + input_ids=torch.tensor([[1]]), + seq_length=torch.ones(1, dtype=torch.long), + history_lengths=torch.zeros(1, dtype=torch.long), + block_offsets=torch.ones((1, 1), dtype=torch.int32), + is_decoding=False, + num_ignored_history=torch.zeros(1, dtype=torch.long), + max_q_seqlen=1, + max_kv_seqlen=1, + sum_kv_seqlen=1, + ) + extra_inputs = ARSpecExtraInputs(next_token_ids=torch.tensor([2])) + sampling_inputs = SimpleNamespace(session_ctx=object(), response_formats=object()) + + with pytest.raises(NotImplementedError, match='DFlash guided decoding'): + asyncio.run(agent.async_model_forward(model_inputs, extra_inputs, sampling_inputs)) + + +def test_dflash_proposer_hook_materializes_non_last_chunk_without_drafting(monkeypatch): + import lmdeploy.pytorch.spec_decode.proposers.dflash as dflash_mod + + inputs, extra_inputs = _make_non_last_chunk_inputs() + + agent = object.__new__(SpecModelAgent) + agent.num_spec_tokens = 3 + agent.rank = 0 + agent._dflash_debug_step = 0 + agent.guided_helper = GuidedSpecHelper() + agent.proposer = _make_dflash_proposer() + agent.proposer.guided_helper = agent.guided_helper + agent.cache_engine = object() + agent.draft_context = lambda: nullcontext() + captured = { + 'materialize_context_calls': 0, + 'propose_block_calls': 0, + } + + def materialize_context(model_inputs, extra_inputs, cache_engine): + captured['materialize_context_calls'] += 1 + assert cache_engine is agent.cache_engine + + async def propose_block(model_inputs, extra_inputs, cache_engine, guided_processors=None): + captured['propose_block_calls'] += 1 + return model_inputs.input_ids.new_full((model_inputs.seq_length.size(0), 3), 9) + + monkeypatch.setattr(agent.proposer, 'materialize_context', materialize_context) + monkeypatch.setattr(agent.proposer, 'propose_block', propose_block) + monkeypatch.delenv('LMDEPLOY_DFLASH_DEBUG_DIR', raising=False) + + def fail_debug_tensor(*args, **kwargs): + raise AssertionError('Disabled DFlash tracing must not serialize tensors.') + + monkeypatch.setattr(dflash_mod, 'debug_tensor', fail_debug_tensor) + + output = asyncio.run(agent.async_model_forward(inputs, extra_inputs, sampling_inputs=None)) + + torch.testing.assert_close(output.output_draft_token_ids, torch.zeros((2, 3), dtype=torch.long)) + assert captured['materialize_context_calls'] == 1 + assert captured['propose_block_calls'] == 0 + assert output.next_token_ids is extra_inputs.next_token_ids + + +def test_dflash_proposer_api_uses_explicit_context_not_spec_agent(): + proposer = _make_dflash_proposer() + + assert proposer.proposal_method == ProposalMethod.DIFFUSION + assert proposer.requires_target_inputs_embeds is False + assert not hasattr(proposer, 'input_mode') + assert not hasattr(proposer, 'proposal_mode') + assert tuple(ProposalContext.__dataclass_fields__) == ('cache_engine', 'rank', 'debug_step') + + +def test_dflash_proposer_requires_explicit_proposal_context(): + proposer = _make_dflash_proposer() + + with pytest.raises(RuntimeError, match='requires ProposalContext'): + asyncio.run(proposer.propose(None, None, None)) + + +def test_dflash_spec_agent_reset_runtime_state_discards_chunk_carry_and_debug_step(): + agent = SpecModelAgent.__new__(SpecModelAgent) + agent._prev_chunk_last = {'hidden_states': object()} + agent._dflash_debug_step = 7 + + agent.reset_runtime_state() + + assert agent._prev_chunk_last == {} + assert agent._dflash_debug_step == 0 + + def test_slice_sampling_inputs_decode(): """Test _slice_sampling_inputs with decoding (num_tokens_per_batch > 1).""" from lmdeploy.pytorch.engine.logits_process import SamplingInputs From 5eb971152e5879ec4451534936b0fb5fa7401059 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Tue, 28 Jul 2026 08:15:03 +0000 Subject: [PATCH 2/3] fix(requirements): require FLA 0.4.2 --- requirements/runtime_cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/runtime_cuda.txt b/requirements/runtime_cuda.txt index eb091e0f66..7b1a7b90d1 100644 --- a/requirements/runtime_cuda.txt +++ b/requirements/runtime_cuda.txt @@ -3,7 +3,7 @@ accelerate>=0.29.3 aiohttp apache-tvm-ffi==0.1.11; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine -flash-linear-attention +flash-linear-attention>=0.4.2 opencv-python-headless peft<=0.14.0 prometheus_client From ee462c7cffae84189f41c75dc5193d78ba755ce2 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Thu, 30 Jul 2026 01:07:56 +0000 Subject: [PATCH 3/3] remove sync --- lmdeploy/cli/serve.py | 2 + lmdeploy/cli/utils.py | 10 ++ .../pytorch/spec_decode/proposers/dflash.py | 155 ++++++++---------- .../pytorch/spec_decode/test_dflash_utils.py | 137 +++++++++------- 4 files changed, 166 insertions(+), 138 deletions(-) diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 9f7af90a32..1f0f07b58c 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -127,6 +127,7 @@ def add_parser_api_server(): session_len_act = ArgumentHelper.session_len(pt_group) max_batch_size_act = ArgumentHelper.max_batch_size(pt_group) cache_max_entry_act = ArgumentHelper.cache_max_entry_count(pt_group) + ArgumentHelper.num_gpu_blocks(pt_group) cache_block_seq_len_act = ArgumentHelper.cache_block_seq_len(pt_group) prefix_caching_act = ArgumentHelper.enable_prefix_caching(pt_group) max_prefill_token_num_act = ArgumentHelper.max_prefill_token_num(pt_group) @@ -242,6 +243,7 @@ def api_server(args): ep=args.ep, max_batch_size=max_batch_size, cache_max_entry_count=args.cache_max_entry_count, + num_gpu_blocks=args.num_gpu_blocks, block_size=args.cache_block_seq_len, kernel_block_size=args.kernel_block_size, session_len=args.session_len, diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 2bf04bb018..06068ce557 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -531,6 +531,16 @@ def cache_max_entry_count(parser): help='The percentage of free gpu memory occupied by the k/v ' 'cache, excluding weights ') + @staticmethod + def num_gpu_blocks(parser): + """Add argument num_gpu_blocks to parser.""" + + return parser.add_argument('--num-gpu-blocks', + type=int, + default=0, + help='Explicit number of GPU KV cache blocks for PyTorch engine. ' + 'Use 0 to auto-size from cache-max-entry-count.') + @staticmethod def adapters(parser): """Add argument adapters to parser.""" diff --git a/lmdeploy/pytorch/spec_decode/proposers/dflash.py b/lmdeploy/pytorch/spec_decode/proposers/dflash.py index c53ced335a..391011fd76 100644 --- a/lmdeploy/pytorch/spec_decode/proposers/dflash.py +++ b/lmdeploy/pytorch/spec_decode/proposers/dflash.py @@ -82,18 +82,18 @@ def prepare_warmup_forward(self, inputs: ModelInputs, cache_engine: CacheEngine) (0, required_blocks - inputs.block_offsets.size(1)), value=0) - batch_size = int(inputs.seq_length.numel()) + batch_size = inputs.seq_length.numel() extra_inputs = ARSpecExtraInputs( next_token_ids=inputs.input_ids.new_zeros(batch_size), target_hidden_states=inputs.target_hidden_states, ) - context_inputs, target_hidden, context_lengths, context_position_ids = \ + context_inputs, target_hidden, context_lengths, query_start_positions = \ self._prepare_context_materialization(inputs, extra_inputs) self._materialize_context(context_inputs, target_hidden, cache_engine) if not inputs.is_decoding: return None return self._build_query_inputs(inputs, context_lengths, extra_inputs.next_token_ids, - context_position_ids) + query_start_positions=query_start_positions) def _draft_model(self): """Return the underlying draft nn.Module when graph runner wraps it.""" @@ -119,85 +119,52 @@ def _context_lengths(self, model_inputs: ModelInputs, extra_inputs: ARSpecExtraI context_lengths = context_lengths - extra_inputs.num_rejected_tokens.to(context_lengths) return context_lengths - @staticmethod - def _slice_by_lengths(tensor: torch.Tensor, - seq_lengths: torch.Tensor, - keep_lengths: torch.Tensor, - max_seq_length: int, - preserve_features: bool = False): - """Slice a flattened per-token tensor by per-request valid lengths.""" - if preserve_features: - if tensor.dim() == 3 and tensor.size(0) == 1: - flat_tensor = tensor[0] - elif tensor.dim() >= 2: - flat_tensor = tensor.flatten(0, -2) - else: - flat_tensor = tensor.reshape(-1) - elif tensor.dim() == 2 and tensor.size(0) == 1: - flat_tensor = tensor[0] - elif tensor.dim() >= 2: - flat_tensor = tensor.flatten(0, -2) - else: - flat_tensor = tensor.reshape(-1) - starts = seq_lengths.cumsum(0) - seq_lengths - offsets = torch.arange(max_seq_length, device=seq_lengths.device) - valid = offsets[None, :] < keep_lengths[:, None] - indices = starts[:, None] + offsets[None, :] - indices = indices[valid] - return flat_tensor.index_select(0, indices) - - def _slice_target_position_ids(self, model_inputs: ModelInputs, context_lengths: torch.Tensor): - """Slice explicit target positions when the target path provides - them.""" + def _query_start_positions(self, model_inputs: ModelInputs, context_lengths: torch.Tensor): + """Resolve DFlash query start positions without ragged compaction. + + Decode verifier inputs are fixed-width ``[batch, max_q_seqlen]`` + blocks, while prefill inputs are packed by their real per-request + lengths. In both cases the draft context materialization may write + rejected tail positions into KV cache, but subsequent draft queries + expose only the accepted prefix through ``history_lengths`` and these + query positions. Use fixed-size gather/index-select arithmetic instead + of compacting position ids with boolean indexing/nonzero. + """ + query_history = model_inputs.history_lengths + context_lengths target_position_ids = model_inputs.target_position_ids if target_position_ids is None: - return None + return query_history if target_position_ids.dim() == 2 and target_position_ids.size(0) == 1: target_position_ids = target_position_ids[0] if target_position_ids.dim() != 1: - raise RuntimeError('DFlash supports only 1D target_position_ids for draft context materialization, ' + raise RuntimeError('DFlash supports only 1D target_position_ids for draft query positioning, ' f'got shape={tuple(target_position_ids.shape)}.') - return self._slice_by_lengths(target_position_ids, - model_inputs.seq_length, - context_lengths, - max_seq_length=model_inputs.max_q_seqlen) - - def _build_context_inputs(self, - model_inputs: ModelInputs, - context_lengths: torch.Tensor, - context_position_ids: torch.Tensor | None = None): - """Build draft-cache materialization inputs for committed context - tokens.""" - context_ids = self._slice_by_lengths(model_inputs.input_ids, - model_inputs.seq_length, - context_lengths, - max_seq_length=model_inputs.max_q_seqlen) - target_position_ids = None if context_position_ids is None else context_position_ids.unsqueeze(0) - return model_inputs.clone( - input_ids=context_ids.unsqueeze(0), - seq_length=context_lengths, - max_q_seqlen=model_inputs.max_q_seqlen, - max_kv_seqlen=model_inputs.max_kv_seqlen, - sum_kv_seqlen=model_inputs.sum_kv_seqlen, - is_decoding=False, - target_hidden_states=None, - target_position_ids=target_position_ids, - target_inputs_embeds=None, - ) - def _build_query_inputs(self, model_inputs: ModelInputs, context_lengths: torch.Tensor, - next_token_ids: torch.Tensor, context_position_ids: torch.Tensor | None = None): + last_offsets = context_lengths.clamp_min(1) - 1 + if model_inputs.is_decoding: + batch_size = model_inputs.seq_length.numel() + query_len = model_inputs.max_q_seqlen + target_position_ids = target_position_ids.reshape(batch_size, query_len) + last_positions = target_position_ids.gather(1, last_offsets[:, None]).squeeze(1) + else: + starts = model_inputs.seq_length.cumsum(0) - model_inputs.seq_length + last_positions = target_position_ids.index_select(0, starts + last_offsets) + return last_positions + 1 + + def _build_query_inputs(self, + model_inputs: ModelInputs, + context_lengths: torch.Tensor, + next_token_ids: torch.Tensor, + *, + query_start_positions: torch.Tensor | None = None): """Build one DFlash query block per request: [next, mask, ...].""" - batch_size = int(model_inputs.seq_length.numel()) + batch_size = model_inputs.seq_length.numel() query_len = self.num_speculative_tokens + 1 query_ids = model_inputs.input_ids.new_full((batch_size, query_len), int(self.specdecode_config.mask_token_id)) query_ids[:, 0] = next_token_ids query_history = model_inputs.history_lengths + context_lengths - if context_position_ids is None: + if query_start_positions is None: query_start_positions = query_history - else: - starts = context_lengths.cumsum(0) - context_lengths - query_start_positions = context_position_ids[starts + context_lengths - 1] + 1 query_positions = query_start_positions[:, None] + torch.arange(query_len, device=query_history.device)[None, :] return model_inputs.clone( input_ids=query_ids.reshape(1, -1), @@ -247,14 +214,35 @@ def _prepare_context_materialization(self, model_inputs: ModelInputs, extra_inpu materialization.""" target_hidden = self._flatten_target_hidden(extra_inputs) context_lengths = self._context_lengths(model_inputs, extra_inputs) - context_position_ids = self._slice_target_position_ids(model_inputs, context_lengths) - context_inputs = self._build_context_inputs(model_inputs, context_lengths, context_position_ids) - target_hidden = self._slice_by_lengths(target_hidden, - model_inputs.seq_length, - context_lengths, - max_seq_length=model_inputs.max_q_seqlen, - preserve_features=True) - return context_inputs, target_hidden, context_lengths, context_position_ids + if model_inputs.is_decoding: + # The decode verifier always provides a fixed speculative block. + # Avoid compacting per-row accepted prefixes with boolean indexing: + # that lowers to aten::nonzero and forces cudaStreamSynchronize. + # Instead write the full verifier block into the draft KV cache and + # expose only accepted prefixes through the next query's + # history_lengths/query positions. + query_start_positions = self._query_start_positions(model_inputs, context_lengths) + context_inputs = model_inputs.clone( + is_decoding=False, + target_hidden_states=None, + target_inputs_embeds=None, + ) + return context_inputs, target_hidden, context_lengths, query_start_positions + + # Prefill inputs are already packed by ``model_inputs.seq_length``. + # Avoid ragged accepted-prefix compaction here as well: boolean indexing + # lowers through ``aten::nonzero`` and can synchronize the CUDA stream. + # It is safe to materialize the full prefill block; if a caller carries + # shorter ``context_lengths`` (for example from rejected speculative + # tails), subsequent query history/positions expose only the committed + # prefix. + query_start_positions = self._query_start_positions(model_inputs, context_lengths) + context_inputs = model_inputs.clone( + is_decoding=False, + target_hidden_states=None, + target_inputs_embeds=None, + ) + return context_inputs, target_hidden, context_lengths, query_start_positions def materialize_context( self, @@ -285,10 +273,10 @@ async def propose_block( if extra_inputs.next_token_ids is None: raise RuntimeError('DFlash requires sampled next_token_ids from the target model.') - context_inputs, target_hidden, context_lengths, context_position_ids = self._prepare_context_materialization( + context_inputs, target_hidden, context_lengths, query_start_positions = self._prepare_context_materialization( model_inputs, extra_inputs) query_inputs = self._build_query_inputs(model_inputs, context_lengths, extra_inputs.next_token_ids, - context_position_ids) + query_start_positions=query_start_positions) self._materialize_context(context_inputs, target_hidden, cache_engine) outputs = self._forward(query_inputs, cache_engine=cache_engine) @@ -296,11 +284,12 @@ async def propose_block( if hidden_states.dim() == 3: hidden_states = hidden_states[0] - batch_size = int(query_inputs.seq_length.numel()) + batch_size = query_inputs.seq_length.numel() query_len = self.num_speculative_tokens + 1 - mask_indices = torch.arange(batch_size * query_len, device=hidden_states.device).view(batch_size, query_len) - mask_indices = mask_indices[:, 1:].reshape(-1) - logits = self.get_logits(hidden_states[mask_indices][None])[0] + hidden_size = hidden_states.size(-1) + mask_hidden_states = hidden_states.reshape(batch_size, query_len, hidden_size)[:, 1:].reshape( + 1, batch_size * self.num_speculative_tokens, hidden_size) + logits = self.get_logits(mask_hidden_states)[0] draft_token_ids = logits.argmax(dim=-1).view(batch_size, self.num_speculative_tokens) return draft_token_ids diff --git a/tests/pytorch/spec_decode/test_dflash_utils.py b/tests/pytorch/spec_decode/test_dflash_utils.py index 0cadbc11a9..056fb706a5 100644 --- a/tests/pytorch/spec_decode/test_dflash_utils.py +++ b/tests/pytorch/spec_decode/test_dflash_utils.py @@ -702,7 +702,7 @@ def test_qwen3_dflash_materialization_requires_cpu_max_q_seqlen(): ) -def test_dflash_block_proposer_builds_context_and_query_inputs(): +def test_dflash_block_proposer_builds_full_context_and_query_inputs(): proposer = DFlash( SimpleNamespace( mask_token_id=99, @@ -724,23 +724,26 @@ def test_dflash_block_proposer_builds_context_and_query_inputs(): sum_kv_seqlen=10, target_inputs_embeds=torch.randn(1, 5, 4), ) - context_lengths = torch.tensor([2, 2]) - hidden = torch.arange(20).reshape(5, 4) - - sliced_hidden = proposer._slice_by_lengths(hidden, - model_inputs.seq_length, - torch.tensor([2, 1]), - max_seq_length=model_inputs.max_q_seqlen) - context_inputs = proposer._build_context_inputs(model_inputs, context_lengths) - query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8])) - - assert sliced_hidden.tolist() == hidden[[0, 1, 3]].tolist() - assert context_inputs.input_ids.tolist() == [[10, 11, 20, 21]] - assert context_inputs.seq_length.tolist() == [2, 2] + extra_inputs = SimpleNamespace( + num_rejected_tokens=torch.tensor([1, 0]), + target_hidden_states=torch.arange(5 * 4, dtype=torch.float32).view(5, 4), + ) + + context_inputs, target_hidden, context_lengths, query_start_positions = proposer._prepare_context_materialization( + model_inputs, extra_inputs) + query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8]), + query_start_positions=query_start_positions) + + assert not hasattr(DFlash, '_slice_by_lengths') + assert context_inputs.input_ids.tolist() == [[10, 11, 12, 20, 21]] + assert context_inputs.seq_length.tolist() == [3, 2] assert context_inputs.max_q_seqlen == 3 assert context_inputs.max_kv_seqlen == 7 assert context_inputs.sum_kv_seqlen == 10 assert context_inputs.target_inputs_embeds is None + assert target_hidden.tolist() == extra_inputs.target_hidden_states.tolist() + assert context_lengths.tolist() == [2, 2] + assert query_start_positions.tolist() == [2, 7] assert query_inputs.input_ids.tolist() == [[7, 99, 99, 99, 8, 99, 99, 99]] assert query_inputs.history_lengths.tolist() == [2, 7] assert query_inputs.target_position_ids.tolist() == [[2, 3, 4, 5, 7, 8, 9, 10]] @@ -783,21 +786,23 @@ def _fail_item(self): monkeypatch.setattr(torch.Tensor, 'item', _fail_item) - context_inputs, target_hidden, context_lengths, context_position_ids = proposer._prepare_context_materialization( + context_inputs, target_hidden, context_lengths, query_start_positions = proposer._prepare_context_materialization( model_inputs, extra_inputs) query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8]), - context_position_ids) + query_start_positions=query_start_positions) assert context_inputs.max_q_seqlen == 3 assert context_inputs.max_kv_seqlen == 7 assert context_inputs.sum_kv_seqlen == 10 - assert target_hidden.shape == (4, 4) + assert context_inputs.seq_length.tolist() == [3, 2] + assert target_hidden.shape == (5, 4) + assert query_start_positions.tolist() == [102, 202] assert query_inputs.max_q_seqlen == 4 assert query_inputs.max_kv_seqlen == 11 assert query_inputs.sum_kv_seqlen == 18 -def test_dflash_slice_by_lengths_preserves_single_feature_row(): +def test_dflash_block_proposer_uses_explicit_target_positions(): proposer = DFlash( SimpleNamespace( mask_token_id=99, @@ -807,29 +812,37 @@ def test_dflash_slice_by_lengths_preserves_single_feature_row(): ), device='cpu', ) - - feature = torch.tensor([[1, 2, 3, 4]]) - token_ids = torch.tensor([[7]]) - - sliced_feature = proposer._slice_by_lengths( - feature, - seq_lengths=torch.tensor([1]), - keep_lengths=torch.tensor([1]), - max_seq_length=1, - preserve_features=True, + model_inputs = ModelInputs( + input_ids=torch.tensor([[10, 11, 12, 20, 21]]), + seq_length=torch.tensor([3, 2]), + history_lengths=torch.tensor([0, 5]), + block_offsets=torch.zeros((2, 4), dtype=torch.int32), + is_decoding=False, + num_ignored_history=torch.zeros(2, dtype=torch.long), + max_q_seqlen=3, + max_kv_seqlen=7, + sum_kv_seqlen=10, + target_position_ids=torch.tensor([[100, 101, 102, 200, 201]]), ) - sliced_tokens = proposer._slice_by_lengths( - token_ids, - seq_lengths=torch.tensor([1]), - keep_lengths=torch.tensor([1]), - max_seq_length=1, + extra_inputs = SimpleNamespace( + num_rejected_tokens=torch.tensor([1, 1]), + target_hidden_states=torch.arange(5 * 4, dtype=torch.float32).view(5, 4), ) - assert sliced_feature.tolist() == [[1, 2, 3, 4]] - assert sliced_tokens.tolist() == [7] + context_inputs, target_hidden, context_lengths, query_start_positions = proposer._prepare_context_materialization( + model_inputs, extra_inputs) + query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8]), + query_start_positions=query_start_positions) + assert context_inputs.input_ids.tolist() == [[10, 11, 12, 20, 21]] + assert context_inputs.target_position_ids.tolist() == [[100, 101, 102, 200, 201]] + assert target_hidden.tolist() == extra_inputs.target_hidden_states.tolist() + assert context_lengths.tolist() == [2, 1] + assert query_start_positions.tolist() == [102, 201] + assert query_inputs.target_position_ids.tolist() == [[102, 103, 104, 105, 201, 202, 203, 204]] -def test_dflash_block_proposer_uses_explicit_target_positions(): + +def test_dflash_decode_materialization_uses_full_block_without_ragged_slice(monkeypatch): proposer = DFlash( SimpleNamespace( mask_token_id=99, @@ -840,29 +853,41 @@ def test_dflash_block_proposer_uses_explicit_target_positions(): device='cpu', ) model_inputs = ModelInputs( - input_ids=torch.tensor([[10, 11, 12, 20, 21]]), - seq_length=torch.tensor([3, 2]), - history_lengths=torch.tensor([0, 5]), + input_ids=torch.tensor([[10, 11, 12, 13, 20, 21, 22, 23]]), + seq_length=torch.tensor([4, 4]), + history_lengths=torch.tensor([100, 200]), block_offsets=torch.zeros((2, 4), dtype=torch.int32), - is_decoding=False, + is_decoding=True, num_ignored_history=torch.zeros(2, dtype=torch.long), - max_q_seqlen=3, - max_kv_seqlen=7, - sum_kv_seqlen=10, - target_position_ids=torch.tensor([[100, 101, 102, 200, 201]]), + max_q_seqlen=4, + max_kv_seqlen=204, + sum_kv_seqlen=308, + target_position_ids=torch.tensor([[100, 101, 102, 103, 200, 201, 202, 203]]), + ) + extra_inputs = SimpleNamespace( + num_rejected_tokens=torch.tensor([2, 0]), + target_hidden_states=torch.arange(8 * 4, dtype=torch.float32).view(8, 4), ) - context_lengths = torch.tensor([2, 1]) - context_position_ids = proposer._slice_target_position_ids(model_inputs, context_lengths) - context_inputs = proposer._build_context_inputs(model_inputs, context_lengths, context_position_ids) + + assert not hasattr(DFlash, '_slice_by_lengths') + + context_inputs, target_hidden, context_lengths, query_start_positions = \ + proposer._prepare_context_materialization(model_inputs, extra_inputs) query_inputs = proposer._build_query_inputs(model_inputs, context_lengths, torch.tensor([7, 8]), - context_position_ids) + query_start_positions=query_start_positions) - assert context_position_ids.tolist() == [100, 101, 200] - assert context_inputs.target_position_ids.tolist() == [[100, 101, 200]] - assert query_inputs.target_position_ids.tolist() == [[102, 103, 104, 105, 201, 202, 203, 204]] + assert context_inputs.is_decoding is False + assert context_inputs.input_ids.tolist() == model_inputs.input_ids.tolist() + assert context_inputs.seq_length.tolist() == [4, 4] + assert context_inputs.target_position_ids.tolist() == model_inputs.target_position_ids.tolist() + assert target_hidden.tolist() == extra_inputs.target_hidden_states.tolist() + assert context_lengths.tolist() == [2, 4] + assert query_start_positions.tolist() == [102, 204] + assert query_inputs.history_lengths.tolist() == [102, 204] + assert query_inputs.target_position_ids.tolist() == [[102, 103, 104, 105, 204, 205, 206, 207]] -def test_dflash_materialize_context_only_slices_committed_hidden(monkeypatch): +def test_dflash_prefill_materialize_context_uses_full_block_without_ragged_slice(monkeypatch): proposer = DFlash( SimpleNamespace( cache_config=SimpleNamespace(block_size=8), @@ -895,12 +920,14 @@ def _materialize_context(context_inputs, target_hidden, cache_engine): captured['target_hidden'] = target_hidden captured['cache_engine'] = cache_engine + assert not hasattr(DFlash, '_slice_by_lengths') + monkeypatch.setattr(proposer, '_materialize_context', _materialize_context) cache_engine = object() proposer.materialize_context(model_inputs, extra_inputs, cache_engine) assert captured['cache_engine'] is cache_engine - assert captured['context_inputs'].input_ids.tolist() == [[10, 11, 20, 21]] - assert captured['context_inputs'].seq_length.tolist() == [2, 2] - assert captured['target_hidden'].tolist() == extra_inputs.target_hidden_states[[0, 1, 3, 4]].tolist() + assert captured['context_inputs'].input_ids.tolist() == [[10, 11, 12, 20, 21]] + assert captured['context_inputs'].seq_length.tolist() == [3, 2] + assert captured['target_hidden'].tolist() == extra_inputs.target_hidden_states.tolist()