From 0ba94c8925e974ccfdbf77ff1f870bcc8991d5b3 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Sat, 5 Sep 2026 16:41:18 +0000 Subject: [PATCH 1/2] feat: support Qwen3.8-27B DFlash2 speculative decoding --- .../common/basemodel/attention/base_att.py | 2 +- lightllm/common/basemodel/batch_objs.py | 25 ++- lightllm/common/basemodel/hidden_collector.py | 12 ++ lightllm/common/basemodel/mtp_manager.py | 8 +- .../common/req_manager/req_sampling_params.py | 19 ++ lightllm/models/__init__.py | 1 + lightllm/models/qwen3_dflash2/__init__.py | 3 + .../qwen3_dflash2/layer_infer/__init__.py | 7 + .../layer_infer/post_layer_infer.py | 171 ++++++++++++++++ .../layer_infer/transformer_layer_infer.py | 103 ++++++++++ .../qwen3_dflash2/layer_weights/__init__.py | 11 + .../pre_and_post_layer_weight.py | 65 ++++++ .../layer_weights/transformer_layer_weight.py | 46 +++++ lightllm/models/qwen3_dflash2/model.py | 81 ++++++++ .../qwen3_dflash2/triton_kernel/__init__.py | 3 + .../triton_kernel/grouped_dynamic_conv.py | 88 ++++++++ .../triton_kernel/selector_walk.py | 93 +++++++++ lightllm/server/api_cli.py | 7 +- lightllm/server/api_start.py | 6 + lightllm/server/core/objs/start_args_type.py | 1 + .../mode_backend/chunked_prefill/impl.py | 23 +-- .../mode_backend/dp_backend/impl.py | 5 +- .../mode_backend/generic_post_process.py | 60 +++++- .../model_infer/mtp_speculative/dflash2.py | 192 ++++++++++++++++++ .../model_infer/mtp_speculative/engine.py | 86 +++++++- .../mtp_speculative/proposers/__init__.py | 4 + .../mtp_speculative/proposers/dflash.py | 16 ++ .../mtp_speculative/proposers/dflash2.py | 41 ++++ .../proposers/proposal_type.py | 8 + .../model_infer/mtp_speculative/utils.py | 37 +++- lightllm/utils/envs_utils.py | 2 + .../common/basemodel/test_mtp_manager.py | 3 + .../mtp_speculative/test_dflash2.py | 92 +++++++++ .../mtp_speculative/test_planner.py | 16 +- unit_tests/utils/test_speculative_utils.py | 11 +- 35 files changed, 1307 insertions(+), 41 deletions(-) create mode 100644 lightllm/models/qwen3_dflash2/__init__.py create mode 100644 lightllm/models/qwen3_dflash2/layer_infer/__init__.py create mode 100644 lightllm/models/qwen3_dflash2/layer_infer/post_layer_infer.py create mode 100644 lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py create mode 100644 lightllm/models/qwen3_dflash2/layer_weights/__init__.py create mode 100644 lightllm/models/qwen3_dflash2/layer_weights/pre_and_post_layer_weight.py create mode 100644 lightllm/models/qwen3_dflash2/layer_weights/transformer_layer_weight.py create mode 100644 lightllm/models/qwen3_dflash2/model.py create mode 100644 lightllm/models/qwen3_dflash2/triton_kernel/__init__.py create mode 100644 lightllm/models/qwen3_dflash2/triton_kernel/grouped_dynamic_conv.py create mode 100644 lightllm/models/qwen3_dflash2/triton_kernel/selector_walk.py create mode 100644 lightllm/server/router/model_infer/mtp_speculative/dflash2.py create mode 100644 lightllm/server/router/model_infer/mtp_speculative/proposers/dflash2.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index a7e2d8122a..5cf8958d75 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -71,7 +71,7 @@ def uses_dynamic_spec_verify_layout(self) -> bool: def uses_causal_attention(self) -> bool: args = get_env_start_args() - is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash") + is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash", "dflash2") return not is_parallel_block_draft def _find_layer_index( diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index ae645d4b7b..2cfe71a1d5 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -147,6 +147,11 @@ class ModelMtpOutputCollector: # - Vanilla MTP、EAGLE、EAGLE3、DFlash 以及未启用 MTP 的模型均不使用该字段。 draft_token_ids: Optional[torch.Tensor] = None + # DFlash2 selector 每个 proposal 位置的 top-k candidate id 与完整条件分布 q。 + # 两者形状均为 [request_count, draft_step, selector_top_k]。 + draft_candidate_ids: Optional[torch.Tensor] = None + draft_candidate_probs: Optional[torch.Tensor] = None + # DSpark confidence head 输出的原始置信度 logits，形状通常为 # [request_count, block_size]，供动态 MTP verify 计算各 draft 位置的调度分数。 # - 仅 DSpark checkpoint 启用 confidence head 时返回；动态 verify 模式要求该字段存在。 @@ -159,6 +164,10 @@ def to_no_ref_tensor(self) -> None: self.spec_hidden = tensor_to_no_ref_tensor(self.spec_hidden) if self.draft_token_ids is not None: self.draft_token_ids = tensor_to_no_ref_tensor(self.draft_token_ids) + if self.draft_candidate_ids is not None: + self.draft_candidate_ids = tensor_to_no_ref_tensor(self.draft_candidate_ids) + if self.draft_candidate_probs is not None: + self.draft_candidate_probs = tensor_to_no_ref_tensor(self.draft_candidate_probs) if self.confidence_logits is not None: self.confidence_logits = tensor_to_no_ref_tensor(self.confidence_logits) @@ -166,8 +175,22 @@ def unpad_decode(self, padded_batch_size: int, origin_batch_size: int) -> "Model collector = copy.copy(self) if collector.spec_hidden is not None: collector.spec_hidden = collector.spec_hidden[:origin_batch_size] + + def unpad_head_rows(value: torch.Tensor) -> torch.Tensor: + row_count = value.shape[0] + if row_count == padded_batch_size: + return value[:origin_batch_size] + assert row_count > 0 and padded_batch_size % row_count == 0 + physical_rows_per_output = padded_batch_size // row_count + assert origin_batch_size % physical_rows_per_output == 0 + return value[: origin_batch_size // physical_rows_per_output] + if collector.draft_token_ids is not None: - collector.draft_token_ids = collector.draft_token_ids[:origin_batch_size] + collector.draft_token_ids = unpad_head_rows(collector.draft_token_ids) + if collector.draft_candidate_ids is not None: + collector.draft_candidate_ids = unpad_head_rows(collector.draft_candidate_ids) + if collector.draft_candidate_probs is not None: + collector.draft_candidate_probs = unpad_head_rows(collector.draft_candidate_probs) if collector.confidence_logits is not None: confidence_row_count = collector.confidence_logits.shape[0] assert confidence_row_count > 0 and padded_batch_size % confidence_row_count == 0 diff --git a/lightllm/common/basemodel/hidden_collector.py b/lightllm/common/basemodel/hidden_collector.py index 3eb946fe82..62fe87550e 100644 --- a/lightllm/common/basemodel/hidden_collector.py +++ b/lightllm/common/basemodel/hidden_collector.py @@ -91,6 +91,8 @@ def add_mtp_outputs( self, draft_token_ids: Optional[torch.Tensor], confidence_logits: Optional[torch.Tensor], + draft_candidate_ids: Optional[torch.Tensor] = None, + draft_candidate_probs: Optional[torch.Tensor] = None, ) -> None: """Collect optional token/confidence outputs produced by an MTP head. @@ -133,6 +135,8 @@ class MtpHeadOutputCollector(NoopHiddenCollector): def __init__(self) -> None: self.draft_token_ids: Optional[torch.Tensor] = None + self.draft_candidate_ids: Optional[torch.Tensor] = None + self.draft_candidate_probs: Optional[torch.Tensor] = None self.confidence_logits: Optional[torch.Tensor] = None def new_instance(self) -> HiddenCollector: @@ -142,16 +146,24 @@ def add_mtp_outputs( self, draft_token_ids: Optional[torch.Tensor], confidence_logits: Optional[torch.Tensor], + draft_candidate_ids: Optional[torch.Tensor] = None, + draft_candidate_probs: Optional[torch.Tensor] = None, ) -> None: self.draft_token_ids = draft_token_ids + self.draft_candidate_ids = draft_candidate_ids + self.draft_candidate_probs = draft_candidate_probs self.confidence_logits = confidence_logits def finish_output(self, infer_state) -> ModelMtpOutputCollector: output = ModelMtpOutputCollector( draft_token_ids=self.draft_token_ids, + draft_candidate_ids=self.draft_candidate_ids, + draft_candidate_probs=self.draft_candidate_probs, confidence_logits=self.confidence_logits, ) self.draft_token_ids = None + self.draft_candidate_ids = None + self.draft_candidate_probs = None self.confidence_logits = None return output diff --git a/lightllm/common/basemodel/mtp_manager.py b/lightllm/common/basemodel/mtp_manager.py index be6c477b99..b7ffda49b4 100644 --- a/lightllm/common/basemodel/mtp_manager.py +++ b/lightllm/common/basemodel/mtp_manager.py @@ -16,7 +16,7 @@ class MtpManager: _instance: ClassVar[Optional["MtpManager"]] = None _CHAINED_DRAFT_MODES = ("vanilla_with_att", "vanilla_no_att") _RECURRENT_DRAFT_MODES = ("eagle_with_att", "eagle_no_att", "eagle3") - _BLOCK_DRAFT_MODES = ("dspark", "dflash") + _BLOCK_DRAFT_MODES = ("dspark", "dflash", "dflash2") @classmethod def get_instance(cls) -> "MtpManager": @@ -51,6 +51,10 @@ def get_decode_batch_multiplier(self, is_draft_model: bool) -> int: # Block draft models decode mtp_step rows per logical request. if spec_mode in self._BLOCK_DRAFT_MODES: + if spec_mode == "dflash2": + # DFlash2's physical block is [anchor, MASK...], while mtp_step + # counts only the proposal rows following the anchor. + return self.args.mtp_step + 1 return self.args.mtp_step return 1 @@ -83,7 +87,7 @@ def create_hidden_collector( if spec_mode is None: collector_type = NoopHiddenCollector elif model.is_mtp_draft_model: - if spec_mode == "dspark": + if spec_mode in ("dspark", "dflash2"): collector_type = MtpHeadOutputCollector elif spec_mode in self._BLOCK_DRAFT_MODES: collector_type = NoopHiddenCollector diff --git a/lightllm/common/req_manager/req_sampling_params.py b/lightllm/common/req_manager/req_sampling_params.py index 87271b6145..b514586b15 100644 --- a/lightllm/common/req_manager/req_sampling_params.py +++ b/lightllm/common/req_manager/req_sampling_params.py @@ -25,11 +25,14 @@ class ReqSamplingParamsManager: def __init__(self, max_request_num): # mode ["cpu_counter", "pin_mem_counter", "gpu_counter"] self.penalty_counter_mode = get_env_start_args().penalty_counter_mode + self.mtp_mode = get_env_start_args().mtp_mode self.vocab_size = get_vocab_size(get_env_start_args().model_dir) self.mtp_verify_width = get_env_start_args().mtp_step + 1 self.req_to_presence_penalty = torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda") self.req_to_frequency_penalty = torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda") self.req_to_repetition_penalty = torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda") + self.req_to_temperature = torch.ones(max_request_num + 1, dtype=torch.float32, device="cuda") + self.req_to_top_k = torch.ones(max_request_num + 1, dtype=torch.int32, device="cuda") self.req_to_next_token_ids = torch.zeros( (max_request_num + 1, self.mtp_verify_width), dtype=torch.int64, @@ -40,6 +43,8 @@ def __init__(self, max_request_num): if get_env_start_args().mtp_dynamic_verify else None ) + if self.mtp_mode == "dflash2": + self._init_dflash2_buffers(max_request_num) self.req_to_exponential_decay_length_penalty = torch.zeros( max_request_num + 1, dtype=torch.float32, device="cuda" @@ -54,6 +59,14 @@ def __init__(self, max_request_num): (max_request_num + 1, self.vocab_size), dtype=torch.int32, device="cpu", pin_memory=True ) + # DFlash2 的候选 token 和分布 q 需按请求跨轮保存，放在这里复用已有 token 状态的生命周期， + # 保证在 KV 容量估算前分配显存、请求槽复用时清零，并预留 HOLD_REQUEST_ID 槽位。 + def _init_dflash2_buffers(self, max_request_num: int) -> None: + # 当前 DFlash2 固定使用 16 个候选；Qwen3DFlash2Model._verify_params 校验配置值与此一致。 + shape = (max_request_num + 1, self.mtp_verify_width - 1, 16) + self.req_to_dflash2_candidate_ids = torch.zeros(shape, dtype=torch.int64, device="cuda") + self.req_to_dflash2_q_probs = torch.zeros(shape, dtype=torch.float32, device="cuda") + def init_req_sampling_params(self, req: "InferReq"): shm_param = req.sampling_param.shm_param self.req_to_next_token_ids[req.req_idx][0:1].fill_(req.get_last_gen_token()) @@ -63,6 +76,12 @@ def init_req_sampling_params(self, req: "InferReq"): self.req_to_presence_penalty[req.req_idx].fill_(shm_param.presence_penalty) self.req_to_frequency_penalty[req.req_idx].fill_(shm_param.frequency_penalty) self.req_to_repetition_penalty[req.req_idx].fill_(shm_param.repetition_penalty) + self.req_to_temperature[req.req_idx].fill_(shm_param.temperature) + self.req_to_top_k[req.req_idx].fill_(shm_param.top_k) + if self.mtp_mode == "dflash2": + self.req_to_dflash2_candidate_ids[req.req_idx].zero_() + self.req_to_dflash2_q_probs[req.req_idx].zero_() + exponential_decay_length_penalty = shm_param.exponential_decay_length_penalty.to_tuple() self.req_to_exponential_decay_length_penalty[req.req_idx].fill_(exponential_decay_length_penalty[1]) # 提前标记当前请求是否需要统计输出token的计数，因为这个统计可能会导致一些特定场景下后处理效率的下降 diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py index c7e9a59aad..6726c3119c 100644 --- a/lightllm/models/__init__.py +++ b/lightllm/models/__init__.py @@ -51,6 +51,7 @@ from lightllm.models.qwen3_5_moe_mtp.model import Qwen3_5MoeMTPModel from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel +from lightllm.models.qwen3_dflash2.model import Qwen3DFlash2Model from lightllm.models.qwen3_dspark.model import Qwen3DSparkModel from lightllm.models.qwen3_eagle.model import Qwen3EagleModel from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel diff --git a/lightllm/models/qwen3_dflash2/__init__.py b/lightllm/models/qwen3_dflash2/__init__.py new file mode 100644 index 0000000000..eab8cb8a4d --- /dev/null +++ b/lightllm/models/qwen3_dflash2/__init__.py @@ -0,0 +1,3 @@ +from lightllm.models.qwen3_dflash2.model import Qwen3DFlash2Model + +__all__ = ["Qwen3DFlash2Model"] diff --git a/lightllm/models/qwen3_dflash2/layer_infer/__init__.py b/lightllm/models/qwen3_dflash2/layer_infer/__init__.py new file mode 100644 index 0000000000..319b4d4701 --- /dev/null +++ b/lightllm/models/qwen3_dflash2/layer_infer/__init__.py @@ -0,0 +1,7 @@ +from lightllm.models.qwen3_dflash2.layer_infer.post_layer_infer import Qwen3DFlash2PostLayerInfer +from lightllm.models.qwen3_dflash2.layer_infer.transformer_layer_infer import Qwen3DFlash2TransformerLayerInfer + +__all__ = [ + "Qwen3DFlash2PostLayerInfer", + "Qwen3DFlash2TransformerLayerInfer", +] diff --git a/lightllm/models/qwen3_dflash2/layer_infer/post_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/post_layer_infer.py new file mode 100644 index 0000000000..72ced3147c --- /dev/null +++ b/lightllm/models/qwen3_dflash2/layer_infer/post_layer_infer.py @@ -0,0 +1,171 @@ +import torch +import torch.nn.functional as F + +from lightllm.distributed.communication_op import all_gather_into_tensor +from lightllm.models.qwen3_dflash.layer_infer.post_layer_infer import Qwen3DFlashPostLayerInfer +from lightllm.models.qwen3_dflash2.triton_kernel.selector_walk import selector_walk + + +class Qwen3DFlash2PostLayerInfer(Qwen3DFlashPostLayerInfer): + """Run DFlash2's top-k pair selector over a parallel draft block.""" + + def __init__(self, network_config): + super().__init__(network_config) + self.block_size_ = int(network_config["block_size"]) + self.selector_top_k_ = int(network_config["selector_top_k"]) + self.output_multiplier_ = float(network_config.get("output_multiplier", 1.0)) + softcap = network_config.get("final_logit_softcapping") + self.final_logit_softcapping_ = None if softcap is None else float(softcap) + if self.final_logit_softcapping_ is not None and self.final_logit_softcapping_ <= 0: + raise ValueError("final_logit_softcapping must be greater than 0") + + def _transform_unary_logits(self, logits): + logits = logits.float() + if self.output_multiplier_ != 1.0: + logits.mul_(self.output_multiplier_) + if self.final_logit_softcapping_ is not None: + logits.div_(self.final_logit_softcapping_).tanh_().mul_(self.final_logit_softcapping_) + return logits + + def _compute_candidates(self, hidden, infer_state, layer_weight): + """Select global top-k candidates without gathering full-vocabulary logits.""" + + token_num = hidden.shape[0] + lm_head = layer_weight.lm_head_weight_ + local_vocab_size = lm_head.weight.shape[0] + if local_vocab_size < self.selector_top_k_: + raise ValueError( + "DFlash2 selector_top_k exceeds the TP-local vocabulary size: " + f"top_k={self.selector_top_k_}, local_vocab_size={local_vocab_size}" + ) + + lm_head_input = hidden.permute(1, 0).contiguous() + local_logits = lm_head(input=lm_head_input, alloc_func=self.alloc_tensor) + local_logits = local_logits.permute(1, 0).contiguous() + local_values, local_ids = torch.topk(local_logits, k=self.selector_top_k_, dim=-1) + global_ids = local_ids.long().add_(lm_head.tp_vocab_start_id) + + if self.tp_world_size_ == 1: + return global_ids, self._transform_unary_logits(local_values) + + gathered_values = self.alloc_tensor( + (self.tp_world_size_ * token_num, self.selector_top_k_), + dtype=torch.float32, + ) + all_gather_into_tensor( + gathered_values, + local_values.float().contiguous(), + group=infer_state.dist_group, + async_op=False, + ) + gathered_ids = self.alloc_tensor( + (self.tp_world_size_ * token_num, self.selector_top_k_), + dtype=torch.int64, + ) + all_gather_into_tensor( + gathered_ids, + global_ids.contiguous(), + group=infer_state.dist_group, + async_op=False, + ) + + gathered_values = ( + gathered_values.view(self.tp_world_size_, token_num, self.selector_top_k_) + .permute(1, 0, 2) + .reshape(token_num, self.tp_world_size_ * self.selector_top_k_) + ) + gathered_ids = ( + gathered_ids.view(self.tp_world_size_, token_num, self.selector_top_k_) + .permute(1, 0, 2) + .reshape(token_num, self.tp_world_size_ * self.selector_top_k_) + ) + top_values, top_indexes = torch.topk(gathered_values, k=self.selector_top_k_, dim=-1) + candidate_ids = torch.gather(gathered_ids, dim=-1, index=top_indexes) + return candidate_ids, self._transform_unary_logits(top_values) + + def _select_path(self, hidden, candidate_ids, unary, anchor_token_ids, infer_state, layer_weight): + req_num, draft_width, _ = candidate_ids.shape + assert draft_width > 0 + assert hidden.shape[:2] == (req_num, draft_width) + assert unary.shape == candidate_ids.shape + + candidate_hidden = hidden.reshape(req_num * draft_width, -1) + gate = layer_weight.selector_hidden_projection_weight_.mm(candidate_hidden) + gate = gate.view(req_num, draft_width, -1) + + predecessor_codebook = layer_weight.selector_predecessor_codebook_weight_.weight + successor_codebook = layer_weight.selector_successor_codebook_weight_.weight + successor = F.embedding(candidate_ids, successor_codebook) + + anchor = F.embedding(anchor_token_ids, predecessor_codebook) + first_scores = unary[:, 0, :] + torch.sum( + anchor[:, None, :] * gate[:, 0, None, :] * successor[:, 0, :, :], + dim=-1, + ) + + pair_scores = None + if draft_width > 1: + predecessor = F.embedding(candidate_ids[:, :-1, :], predecessor_codebook) + conditioned_predecessor = predecessor * gate[:, 1:, None, :] + transitions = torch.matmul(conditioned_predecessor, successor[:, 1:, :, :].transpose(-1, -2)) + pair_scores = transitions + unary[:, 1:, None, :] + + score_lattice = first_scores[:, None, None, :].expand(-1, 1, self.selector_top_k_, -1) + if pair_scores is not None: + score_lattice = torch.cat((score_lattice, pair_scores), dim=1) + + request_ids = infer_state.b_req_idx.view(req_num, self.block_size_)[:, 0].long() + sampling_manager = infer_state.req_manager.req_sampling_params_manager + temperatures = sampling_manager.req_to_temperature.index_select(0, request_ids).clamp_min_(1e-5) + greedy_mask = sampling_manager.req_to_top_k.index_select(0, request_ids).eq(1) + # DFlash2 currently does not guarantee request-level determinism: draft selection uses the global RNG, + # while request seeds only control target verification. + uniforms = torch.rand( + (req_num, draft_width), + dtype=torch.float32, + device=candidate_ids.device, + ) + selected_ids, q_rows, _ = selector_walk( + scores=score_lattice, + candidate_ids=candidate_ids, + uniforms=uniforms, + temperatures=temperatures, + greedy_mask=greedy_mask, + ) + return selected_ids, q_rows + + def token_forward(self, input_embdings, infer_state, layer_weight): + if infer_state.is_prefill: + return super().token_forward(input_embdings, infer_state, layer_weight) + + last_input, token_num = self._slice_get_last_input(input_embdings, infer_state) + assert token_num % self.block_size_ == 0 + req_num = token_num // self.block_size_ + normed_hidden = self._norm(last_input, infer_state, layer_weight) + block_hidden = normed_hidden.view(req_num, self.block_size_, -1) + candidate_hidden = block_hidden[:, 1:, :].reshape(req_num * (self.block_size_ - 1), -1) + candidate_ids, unary = self._compute_candidates( + hidden=candidate_hidden, + infer_state=infer_state, + layer_weight=layer_weight, + ) + candidate_ids = candidate_ids.view(req_num, self.block_size_ - 1, self.selector_top_k_) + unary = unary.view(req_num, self.block_size_ - 1, self.selector_top_k_) + anchor_token_ids = infer_state.input_ids.view(req_num, self.block_size_)[:, 0] + draft_token_ids, draft_candidate_probs = self._select_path( + hidden=block_hidden[:, 1:, :], + candidate_ids=candidate_ids, + unary=unary, + anchor_token_ids=anchor_token_ids, + infer_state=infer_state, + layer_weight=layer_weight, + ) + infer_state.hidden_collector.add_mtp_outputs( + draft_token_ids=draft_token_ids, + draft_candidate_ids=candidate_ids, + draft_candidate_probs=draft_candidate_probs, + confidence_logits=None, + ) + # The proposer consumes selector outputs directly. Keep only a graph- + # compatible leading dimension instead of retaining full-vocabulary logits. + return unary.new_empty((token_num, 1)) diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..11989b9dda --- /dev/null +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -0,0 +1,103 @@ +from lightllm.common.basemodel.attention import AttControl +from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer +from lightllm.models.qwen3_dflash2.triton_kernel import grouped_dynamic_conv + + +class Qwen3DFlash2TransformerLayerInfer(Qwen3DFlashTransformerLayerInfer): + """DFlash2 layer with a grouped dynamic convolution around each sublayer.""" + + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + self.block_size_ = int(network_config["block_size"]) + self.conv_group_size_ = int(network_config["conv_group_size"]) + self.sliding_window_ = int(network_config.get("sliding_window", 0) or 0) + + def _token_attention_kernel(self, q, infer_state, layer_weight): + k, v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + k = self._reshape_storage_side_for_attention(k) + v = self._reshape_storage_side_for_attention(v) + q = q.view(-1, self.tp_q_head_num_, self.head_dim_) + use_sliding_window = self.sliding_window_ > 0 + # DFlash2 attention is non-causal, so its local window must retain both + # the visible prefix on the left and the complete draft block on the right. + window = self.sliding_window_ - 1 + output = infer_state.decode_att_state.decode_att( + q=q, + k=k, + v=v, + att_control=AttControl( + use_sliding_window=use_sliding_window, + sliding_window=(window, window) if use_sliding_window else (-1, -1), + ), + alloc_func=self.alloc_tensor, + ) + return output.view(-1, self.tp_q_head_num_ * self.head_dim_) + + def _reshape_storage_side_for_attention(self, cache): + draft_width = self.tp_k_head_num_ * self.head_dim_ + storage_width = cache.shape[1] * cache.shape[2] + assert draft_width == storage_width, ( + "DFlash2 draft and target KV must have equal flat widths: " + f"draft=({self.tp_k_head_num_}, {self.head_dim_}), " + f"storage=({cache.shape[1]}, {cache.shape[2]})" + ) + return cache.view(cache.shape[0], self.tp_k_head_num_, self.head_dim_) + + def _post_cache_kv(self, cache_kv, infer_state, layer_weight): + storage_head_num = infer_state.mem_manager.head_num + storage_head_dim = infer_state.mem_manager.head_dim + draft_width = self.tp_k_head_num_ * self.head_dim_ + storage_width = storage_head_num * storage_head_dim + assert draft_width == storage_width, ( + "DFlash2 draft and target KV must have equal flat widths: " + f"draft=({self.tp_k_head_num_}, {self.head_dim_}), " + f"storage=({storage_head_num}, {storage_head_dim})" + ) + storage_kv = cache_kv.view(cache_kv.shape[0], 2 * storage_head_num, storage_head_dim) + return super()._post_cache_kv(storage_kv, infer_state, layer_weight) + + def _run_dynamic_conv(self, hidden, dynamic, base_weight, side): + return grouped_dynamic_conv( + hidden=hidden.contiguous(), + dynamic=dynamic.contiguous(), + base_kernel=base_weight.weight.contiguous(), + block_size=self.block_size_, + group_size=self.conv_group_size_, + side=side, + ) + + def token_forward(self, input_embdings, infer_state, layer_weight): + attention_input = self._att_norm(input_embdings, infer_state, layer_weight) + attention_dynamic = layer_weight.attention_conv_projection_weight_.mm(attention_input) + attention_input = self._run_dynamic_conv( + attention_input, + attention_dynamic, + layer_weight.attention_conv_base_weight_, + side=0, + ) + attention_output = self.token_attention_forward(attention_input, infer_state, layer_weight) + attention_output = self._run_dynamic_conv( + attention_output, + attention_dynamic, + layer_weight.attention_conv_base_weight_, + side=1, + ) + input_embdings.add_(attention_output.view(-1, self.embed_dim_)) + + mlp_input = self._ffn_norm(input_embdings, infer_state, layer_weight) + mlp_dynamic = layer_weight.mlp_conv_projection_weight_.mm(mlp_input) + mlp_input = self._run_dynamic_conv( + mlp_input, + mlp_dynamic, + layer_weight.mlp_conv_base_weight_, + side=0, + ) + mlp_output = self._ffn(mlp_input, infer_state, layer_weight) + mlp_output = self._run_dynamic_conv( + mlp_output, + mlp_dynamic, + layer_weight.mlp_conv_base_weight_, + side=1, + ) + input_embdings.add_(mlp_output.view(-1, self.embed_dim_)) + return input_embdings diff --git a/lightllm/models/qwen3_dflash2/layer_weights/__init__.py b/lightllm/models/qwen3_dflash2/layer_weights/__init__.py new file mode 100644 index 0000000000..d97753ed53 --- /dev/null +++ b/lightllm/models/qwen3_dflash2/layer_weights/__init__.py @@ -0,0 +1,11 @@ +from lightllm.models.qwen3_dflash2.layer_weights.pre_and_post_layer_weight import ( + Qwen3DFlash2PreAndPostLayerWeight, +) +from lightllm.models.qwen3_dflash2.layer_weights.transformer_layer_weight import ( + Qwen3DFlash2TransformerLayerWeight, +) + +__all__ = [ + "Qwen3DFlash2PreAndPostLayerWeight", + "Qwen3DFlash2TransformerLayerWeight", +] diff --git a/lightllm/models/qwen3_dflash2/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen3_dflash2/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..e0407030a2 --- /dev/null +++ b/lightllm/models/qwen3_dflash2/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,65 @@ +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import ( + EmbeddingWeight, + LMHeadWeight, + ParameterWeight, + RMSNormWeight, + ROWMMWeight, +) +from lightllm.common.quantization import Quantcfg + + +class Qwen3DFlash2PreAndPostLayerWeight(PreAndPostLayerWeight): + """DFlash2 projection, normalization, and candidate-selector weights.""" + + def __init__(self, data_type, network_config, quant_cfg: Quantcfg): + super().__init__(data_type, network_config) + self.quant_cfg = quant_cfg + + hidden_size = network_config["hidden_size"] + vocab_size = network_config["vocab_size"] + target_layer_num = len(network_config["target_layer_ids"]) + selector_rank = int(network_config["selector_rank"]) + + # Published DFlash2 checkpoints share these two large weights with the + # target model. Qwen3DFlash2Model wires them up before loading weights. + self.wte_weight_: EmbeddingWeight = None + self.lm_head_weight_: LMHeadWeight = None + self.fc_weight_ = ROWMMWeight( + in_dim=hidden_size * target_layer_num, + out_dims=[hidden_size], + weight_names="fc.weight", + data_type=self.data_type_, + quant_method=self.quant_cfg.get_quant_method(0, "fc"), + tp_rank=0, + tp_world_size=1, + ) + self.hidden_norm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name="hidden_norm.weight", + data_type=self.data_type_, + ) + self.final_norm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name="norm.weight", + data_type=self.data_type_, + ) + self.selector_hidden_projection_weight_ = ROWMMWeight( + in_dim=hidden_size, + out_dims=[selector_rank], + weight_names="candidate_selector.hidden_projection.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.selector_predecessor_codebook_weight_ = ParameterWeight( + weight_name="candidate_selector.predecessor_codebook", + data_type=self.data_type_, + weight_shape=(vocab_size, selector_rank), + ) + self.selector_successor_codebook_weight_ = ParameterWeight( + weight_name="candidate_selector.successor_codebook", + data_type=self.data_type_, + weight_shape=(vocab_size, selector_rank), + ) diff --git a/lightllm/models/qwen3_dflash2/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen3_dflash2/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..de04220349 --- /dev/null +++ b/lightllm/models/qwen3_dflash2/layer_weights/transformer_layer_weight.py @@ -0,0 +1,46 @@ +from lightllm.common.basemodel.layer_weights.meta_weights import ParameterWeight, ROWMMWeight +from lightllm.models.qwen3_dflash.layer_weights.transformer_layer_weight import Qwen3DFlashTransformerLayerWeight + + +class Qwen3DFlash2TransformerLayerWeight(Qwen3DFlashTransformerLayerWeight): + """DFlash decoder weights plus the two dynamic convolutions.""" + + def __init__(self, layer_num, data_type, network_config, quant_cfg=None): + super().__init__(layer_num, data_type, network_config, quant_cfg) + + hidden_size = network_config["hidden_size"] + kernel_size = int(network_config["conv_kernel_size"]) + group_size = int(network_config["conv_group_size"]) + assert hidden_size % group_size == 0 + group_num = hidden_size // group_size + dynamic_size = 2 * kernel_size * group_num + weight_prefix = f"layers.{self.layer_num_}" + + self.attention_conv_projection_weight_ = ROWMMWeight( + in_dim=hidden_size, + out_dims=[dynamic_size], + weight_names=f"{weight_prefix}.attention_conv.kernel_projection.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.attention_conv_base_weight_ = ParameterWeight( + weight_name=f"{weight_prefix}.attention_conv.base_kernel", + data_type=self.data_type_, + weight_shape=(2, kernel_size, hidden_size), + ) + self.mlp_conv_projection_weight_ = ROWMMWeight( + in_dim=hidden_size, + out_dims=[dynamic_size], + weight_names=f"{weight_prefix}.mlp_conv.kernel_projection.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.mlp_conv_base_weight_ = ParameterWeight( + weight_name=f"{weight_prefix}.mlp_conv.base_kernel", + data_type=self.data_type_, + weight_shape=(2, kernel_size, hidden_size), + ) diff --git a/lightllm/models/qwen3_dflash2/model.py b/lightllm/models/qwen3_dflash2/model.py new file mode 100644 index 0000000000..2fb524e12a --- /dev/null +++ b/lightllm/models/qwen3_dflash2/model.py @@ -0,0 +1,81 @@ +from lightllm.models.draft_registry import DraftModelRegistry +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel +from lightllm.models.qwen3_dflash2.layer_infer.post_layer_infer import Qwen3DFlash2PostLayerInfer +from lightllm.models.qwen3_dflash2.layer_infer.transformer_layer_infer import Qwen3DFlash2TransformerLayerInfer +from lightllm.models.qwen3_dflash2.layer_weights.pre_and_post_layer_weight import ( + Qwen3DFlash2PreAndPostLayerWeight, +) +from lightllm.models.qwen3_dflash2.layer_weights.transformer_layer_weight import ( + Qwen3DFlash2TransformerLayerWeight, +) + + +@DraftModelRegistry(model_type="qwen3", spec_modes="dflash2") +class Qwen3DFlash2Model(Qwen3DFlashModel): + """Qwen3 DFlash2 draft model.""" + + pre_and_post_weight_class = Qwen3DFlash2PreAndPostLayerWeight + transformer_weight_class = Qwen3DFlash2TransformerLayerWeight + post_layer_infer_class = Qwen3DFlash2PostLayerInfer + transformer_layer_infer_class = Qwen3DFlash2TransformerLayerInfer + + def _init_config(self): + super()._init_config() + dflash_config = self.config.get("dflash_config", {}) + if not isinstance(dflash_config, dict): + raise ValueError("dflash_config must be an object in the DFlash2 checkpoint config") + self.config.update(dflash_config) + + rope_parameters = self.config.get("rope_parameters", {}) + if "rope_theta" in rope_parameters and "rope_theta" not in self.config: + self.config["rope_theta"] = rope_parameters["rope_theta"] + if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in self.config: + self.config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] + if rope_parameters and "rope_scaling" not in self.config: + self.config["rope_scaling"] = rope_parameters + + def _verify_params(self): + LlamaTpPartModel._verify_params(self) + assert not self.enable_tpsp_mix_mode, "Qwen3 DFlash2 draft model does not support TP-SP" + + if self.args.llm_kv_type == "fp8kv_sph": + raise NotImplementedError( + "DFlash2 sliding-window attention does not support fp8kv_sph; use --llm_kv_type None." + ) + + selector_top_k = self.config.get("selector_top_k") + if selector_top_k is None: + raise ValueError("selector_top_k is required in the DFlash2 checkpoint config") + selector_top_k = int(selector_top_k) + # 请求采样状态的候选缓冲区在 ReqSamplingParamsManager._init_dflash2_buffers 中固定分配为 16。 + if selector_top_k != 16: + raise ValueError(f"DFlash2 requires selector_top_k=16, got {selector_top_k}") + self.config["selector_top_k"] = selector_top_k + + physical_block_size = self.args.mtp_step + 1 + assert physical_block_size <= self.config["block_size"] + self.config["block_size"] = physical_block_size + + def _init_custom(self): + # The released Qwen3.8 drafter uses its own Qwen3 rotary layout. + LlamaTpPartModel._init_custom(self) + self.block_size = int(self.config["block_size"]) + self.mask_token_id = int(self.config["mask_token_id"]) + + def _init_mem_manager(self): + main_mem_manager = self.main_model.mem_manager + draft_head_num = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) + draft_width = draft_head_num * self.config["head_dim"] + target_width = main_mem_manager.head_num * main_mem_manager.head_dim + assert draft_width == target_width, ( + "DFlash2 draft and target KV must have equal flat widths: " + f"draft=({draft_head_num}, {self.config['head_dim']}), " + f"target=({main_mem_manager.head_num}, {main_mem_manager.head_dim})" + ) + self.mem_manager = main_mem_manager + + def _init_weights(self, start_layer_index=None): + super()._init_weights(start_layer_index=start_layer_index) + self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ + self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ diff --git a/lightllm/models/qwen3_dflash2/triton_kernel/__init__.py b/lightllm/models/qwen3_dflash2/triton_kernel/__init__.py new file mode 100644 index 0000000000..61291c94b9 --- /dev/null +++ b/lightllm/models/qwen3_dflash2/triton_kernel/__init__.py @@ -0,0 +1,3 @@ +from .grouped_dynamic_conv import grouped_dynamic_conv + +__all__ = ["grouped_dynamic_conv"] diff --git a/lightllm/models/qwen3_dflash2/triton_kernel/grouped_dynamic_conv.py b/lightllm/models/qwen3_dflash2/triton_kernel/grouped_dynamic_conv.py new file mode 100644 index 0000000000..7edae7b3a5 --- /dev/null +++ b/lightllm/models/qwen3_dflash2/triton_kernel/grouped_dynamic_conv.py @@ -0,0 +1,88 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _grouped_dynamic_conv_kernel( + hidden, + dynamic, + base_kernel, + output, + element_num, + hidden_size: tl.constexpr, + dynamic_stride: tl.constexpr, + block_size: tl.constexpr, + group_size: tl.constexpr, + group_num: tl.constexpr, + kernel_size: tl.constexpr, + side: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + row = offsets // hidden_size + channel = offsets % hidden_size + valid = offsets < element_num + group = channel // group_size + block_offset = row % block_size + + accumulator = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + for tap in tl.static_range(0, kernel_size): + has_predecessor = block_offset >= tap + value = tl.load( + hidden + (row - tap) * hidden_size + channel, + mask=valid & has_predecessor, + other=0.0, + ).to(tl.float32) + base_offset = (side * kernel_size + tap) * hidden_size + channel + dynamic_offset = row * dynamic_stride + (side * kernel_size + tap) * group_num + group + weight = tl.load(base_kernel + base_offset, mask=valid, other=0.0).to(tl.float32) + weight += tl.load(dynamic + dynamic_offset, mask=valid, other=0.0).to(tl.float32) + accumulator += value * weight + + tl.store(output + offsets, accumulator, mask=valid) + + +def grouped_dynamic_conv( + hidden: torch.Tensor, + dynamic: torch.Tensor, + base_kernel: torch.Tensor, + block_size: int, + group_size: int, + side: int, +) -> torch.Tensor: + """Apply one side of DFlash2's grouped dynamic causal convolution.""" + + assert hidden.ndim == 2 and hidden.is_contiguous() + assert dynamic.ndim == 2 and dynamic.is_contiguous() + assert base_kernel.ndim == 3 and base_kernel.is_contiguous() + assert hidden.shape[0] % block_size == 0 + assert hidden.shape[1] % group_size == 0 + assert side in (0, 1) + + token_num, hidden_size = hidden.shape + side_num, kernel_size, base_hidden_size = base_kernel.shape + group_num = hidden_size // group_size + assert side_num == 2 + assert base_hidden_size == hidden_size + assert dynamic.shape == (token_num, side_num * kernel_size * group_num) + + output = torch.empty_like(hidden) + element_num = token_num * hidden_size + block = 256 + _grouped_dynamic_conv_kernel[(triton.cdiv(element_num, block),)]( + hidden, + dynamic, + base_kernel, + output, + element_num, + hidden_size=hidden_size, + dynamic_stride=dynamic.shape[1], + block_size=block_size, + group_size=group_size, + group_num=group_num, + kernel_size=kernel_size, + side=side, + BLOCK_SIZE=block, + ) + return output diff --git a/lightllm/models/qwen3_dflash2/triton_kernel/selector_walk.py b/lightllm/models/qwen3_dflash2/triton_kernel/selector_walk.py new file mode 100644 index 0000000000..c42a294720 --- /dev/null +++ b/lightllm/models/qwen3_dflash2/triton_kernel/selector_walk.py @@ -0,0 +1,93 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _selector_walk_kernel( + scores_ptr, + candidate_ids_ptr, + uniforms_ptr, + temperatures_ptr, + greedy_mask_ptr, + tokens_ptr, + q_ptr, + path_indices_ptr, + SLOT_NUM: tl.constexpr, + TOP_K: tl.constexpr, +): + req_idx = tl.program_id(0) + offsets = tl.arange(0, TOP_K) + temperature = tl.load(temperatures_ptr + req_idx) + is_greedy = tl.load(greedy_mask_ptr + req_idx) != 0 + previous_index = 0 + + for slot_idx in range(SLOT_NUM): + score_offset = ((req_idx * SLOT_NUM + slot_idx) * TOP_K + previous_index) * TOP_K + scores = tl.load(scores_ptr + score_offset + offsets).to(tl.float32) + + if is_greedy: + best_score = tl.max(scores, axis=0) + selected_index = tl.min(tl.where(scores == best_score, offsets, TOP_K), axis=0) + probabilities = tl.where(offsets == selected_index, 1.0, 0.0) + else: + scaled_scores = scores / temperature + exponentials = tl.exp(scaled_scores - tl.max(scaled_scores, axis=0)) + probabilities = exponentials / tl.sum(exponentials, axis=0) + uniform = tl.load(uniforms_ptr + req_idx * SLOT_NUM + slot_idx) + selected_index = tl.sum( + tl.where(uniform >= tl.cumsum(probabilities, axis=0), 1, 0), + axis=0, + ) + selected_index = tl.minimum(selected_index, TOP_K - 1) + + output_offset = req_idx * SLOT_NUM + slot_idx + candidate_offset = output_offset * TOP_K + tl.store(q_ptr + candidate_offset + offsets, probabilities) + tl.store(tokens_ptr + output_offset, tl.load(candidate_ids_ptr + candidate_offset + selected_index)) + tl.store(path_indices_ptr + output_offset, selected_index) + previous_index = selected_index + + +@torch.no_grad() +def selector_walk( + scores: torch.Tensor, + candidate_ids: torch.Tensor, + uniforms: torch.Tensor, + temperatures: torch.Tensor, + greedy_mask: torch.Tensor, +): + """Sample one locally coherent candidate path and retain every conditional q row.""" + + req_num, slot_num, top_k, successor_top_k = scores.shape + assert top_k == successor_top_k + assert candidate_ids.shape == (req_num, slot_num, top_k) + assert uniforms.shape == (req_num, slot_num) + assert temperatures.shape == (req_num,) + assert greedy_mask.shape == (req_num,) + assert top_k == triton.next_power_of_2(top_k) + + scores = scores.contiguous() + candidate_ids = candidate_ids.contiguous() + uniforms = uniforms.contiguous() + temperatures = temperatures.contiguous() + greedy_mask = greedy_mask.contiguous() + tokens = torch.empty((req_num, slot_num), dtype=torch.int64, device=scores.device) + q_rows = torch.empty((req_num, slot_num, top_k), dtype=torch.float32, device=scores.device) + path_indices = torch.empty((req_num, slot_num), dtype=torch.int64, device=scores.device) + + _selector_walk_kernel[(req_num,)]( + scores, + candidate_ids, + uniforms, + temperatures, + greedy_mask, + tokens, + q_rows, + path_indices, + SLOT_NUM=slot_num, + TOP_K=top_k, + num_warps=1, + num_stages=1, + ) + return tokens, q_rows, path_indices diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index e6c077dbe4..10efc3c104 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -777,12 +777,13 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "eagle3", "dspark", "dflash", + "dflash2", None, ], default=None, help="""Speculative decoding mode. *_with_att and *_no_att select attention or non-attention draft models; - eagle3 uses autoregressive EAGLE-3 drafting; dflash uses block-diffusion drafting; + eagle3 uses autoregressive EAGLE-3 drafting; dflash and dflash2 use block-diffusion drafting; dspark uses semi-autoregressive parallel drafting.""", ) parser.add_argument( @@ -798,12 +799,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: type=int, default=0, help="""Number of additional draft tokens per request. - For DSpark and DFlash this value is derived from the draft checkpoint block_size.""", + For DSpark, DFlash, and DFlash2 this value is derived from the draft checkpoint block_size.""", ) parser.add_argument( "--mtp_dynamic_verify", action="store_true", - help="""Enable dynamic speculative scheduling.""", + help="""Enable dynamic speculative scheduling. Temporarily ignored for DFlash2, which uses fixed-width verification.""", ) parser.add_argument( "--kv_quant_calibration_config_path", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 6f8973425f..35160c7bb6 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -170,9 +170,15 @@ def _launch_subprocesses(args: StartArgs): "eagle3", "dspark", "dflash", + "dflash2", ), f"--mtp_draft_model_dir is required for {args.mtp_mode} mode" args.mtp_draft_model_dir = [args.model_dir] * args.mtp_step assert args.mtp_step > 0 + # TODO: DFlash2 dynamic verify 在 H200/Qwen3.8-27B 上未见稳定吞吐收益，暂保留固定宽度验证。 + # 后续降低动态调度/验证开销，并重新验证性能与正确性后再评估支持。 + if args.mtp_mode == "dflash2" and args.mtp_dynamic_verify: + logger.warning("DFlash2 currently uses fixed-width verification; disabling --mtp_dynamic_verify.") + args.mtp_dynamic_verify = False else: assert args.mtp_draft_model_dir is None assert args.mtp_step == 0 diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9c89975de7..d36053b9f9 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -194,6 +194,7 @@ class StartArgs: "eagle3", "dspark", "dflash", + "dflash2", None, ] }, diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 4d09476849..46e23b124f 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -9,7 +9,6 @@ prepare_prefill_inputs, prepare_decode_inputs, ) -from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample from lightllm.server.router.model_infer.infer_batch import g_infer_context from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine @@ -271,21 +270,20 @@ def decode_mtp( async_selected_row_mask_cpu.wait() selected_rows = async_selected_row_mask_cpu.tensor.tolist() run_reqs = [req for req, selected in zip(run_reqs, selected_rows) if selected] - next_token_ids, next_token_logprobs = sample( - model_output.logits, - run_reqs, - self.eos_id, - ) - next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) - b_req_mtp_start_loc = gen_b_req_mtp_start_loc(model_input.b_mtp_index, num_reqs=req_num) - mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( - backend=self, - next_token_ids=next_token_ids, + ( + next_token_ids, + next_token_logprobs, + mtp_accept_len, + accepted_index, + ) = spec_engine.sample_and_verify( + logits=model_output.logits, + run_reqs=run_reqs, b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, ) + next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", gpu_tensor=accepted_index, @@ -312,8 +310,7 @@ def decode_mtp( draft_step=spec_plan.draft_step, accept_len=mtp_accept_len, ) - mtp_utils.scatter_mtp_next_tokens( - backend=self, + spec_engine.prepare_next_verification_state( proposal=proposal, target_next_token_ids=next_token_ids, b_req_mtp_start_loc=b_req_mtp_start_loc, diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 9a81927bc1..cbd625cd8e 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -33,8 +33,8 @@ def __init__(self) -> None: # 在 mtp 模式下切换绑定的prefill 和 decode 函数 spec_mode = get_env_start_args().mtp_mode if spec_mode is not None: - if spec_mode in ("dspark", "dflash"): - raise NotImplementedError("DP backend does not support DFlash/DSpark parallel block drafting yet.") + if spec_mode in ("dspark", "dflash", "dflash2"): + raise NotImplementedError("DP backend does not support DFlash/DFlash2/DSpark parallel block drafting yet.") if self.enable_prefill_microbatch_overlap: self.prefill = self.prefill_overlap_mtp else: @@ -446,7 +446,6 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] sync_event.record() if req_num > 0: - # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 5b29ea0510..8538852b64 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -8,7 +8,7 @@ from lightllm.utils.envs_utils import get_env_start_args -def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): +def _prepare_sampling_probs(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int]): ( b_req_idx, b_temperatures, @@ -78,6 +78,64 @@ def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): logits.div_(b_temperatures.view((-1, 1))) probs = torch.softmax(logits, dim=-1) + return ( + probs, + b_top_ps, + b_top_ks, + is_all_greedy, + skip_top_k, + skip_top_p, + exist_req_use_random_seed, + ) + + +def build_sampling_probs( + logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] +) -> Tuple[torch.Tensor, torch.Tensor]: + """Return normalized sampling probabilities and pre-filter logprob probabilities. + + Reuses sample()'s in-place logits processing and the Triton sampling path's + top-k/top-p filtering, without drawing a token. This generic distribution + builder is currently used by DFlash2 rejection sampling; it contains no + draft-specific acceptance rules. Ordinary sampling keeps its fast paths. + """ + + ( + probs, + b_top_ps, + b_top_ks, + is_all_greedy, + skip_top_k, + skip_top_p, + _, + ) = _prepare_sampling_probs(logits, reqs, eos_id) + + if is_all_greedy: + token_ids = torch.argmax(logits, dim=-1, keepdim=True) + filtered_probs = torch.zeros_like(probs) + filtered_probs.scatter_(1, token_ids, 1.0) + elif skip_top_k and skip_top_p: + filtered_probs = probs + else: + sorted_probs, sorted_indices = _top_p_top_k(probs, b_top_ps, b_top_ks) + filtered_probs = torch.zeros_like(probs) + filtered_probs.scatter_(1, sorted_indices, sorted_probs) + filtered_probs.div_(filtered_probs.sum(dim=-1, keepdim=True).clamp_min_(1e-20)) + + return filtered_probs, probs + + +def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): + ( + probs, + b_top_ps, + b_top_ks, + is_all_greedy, + skip_top_k, + skip_top_p, + exist_req_use_random_seed, + ) = _prepare_sampling_probs(logits, reqs, eos_id) + if is_all_greedy: batch_next_token_ids = torch.argmax(logits, -1) batch_next_token_probs = torch.gather(probs, dim=1, index=batch_next_token_ids.view(-1, 1)) diff --git a/lightllm/server/router/model_infer/mtp_speculative/dflash2.py b/lightllm/server/router/model_infer/mtp_speculative/dflash2.py new file mode 100644 index 0000000000..ecdf89c918 --- /dev/null +++ b/lightllm/server/router/model_infer/mtp_speculative/dflash2.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List + +import torch + +if TYPE_CHECKING: + from lightllm.server.router.model_infer.infer_batch import InferReq + from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import SpecProposal + + +def save_dflash2_proposal_state( + backend: ModeBackend, + proposal: SpecProposal, + b_req_idx: torch.Tensor, + b_req_mtp_start_loc: torch.Tensor, +) -> None: + """Validate and save a DFlash2 proposal for the next verification step.""" + + from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import DFlash2SpecProposal + + if not isinstance(proposal, DFlash2SpecProposal): + raise TypeError(f"DFlash2 requires DFlash2SpecProposal, got {type(proposal).__name__}") + request_ids = b_req_idx.index_select(0, b_req_mtp_start_loc.long()).long() + sampling_manager = backend.model.req_manager.req_sampling_params_manager + assert sampling_manager.mtp_mode == "dflash2" + assert request_ids.ndim == 1 + expected_shape = (request_ids.shape[0], *sampling_manager.req_to_dflash2_candidate_ids.shape[1:]) + assert ( + proposal.candidate_ids.shape == proposal.q_probs.shape == expected_shape + ), "DFlash2 requires a complete fixed-width proposal" + # 按请求保存候选 token 和分布 q，供下一轮非贪心验证使用。 + sampling_manager.req_to_dflash2_candidate_ids[request_ids] = proposal.candidate_ids + sampling_manager.req_to_dflash2_q_probs[request_ids] = proposal.q_probs + + +def _request_uniforms( + request_reqs: List[InferReq], + width: int, + *, + device: torch.device, +) -> torch.Tensor: + """Draw row-wise uniforms while honoring optional per-request generators.""" + + uniforms = torch.rand( + (len(request_reqs), width), + dtype=torch.float32, + device=device, + ) + for row, req in enumerate(request_reqs): + if req.generator is not None: + uniforms[row].uniform_(generator=req.generator) + return uniforms + + +def sample_and_verify_dflash2_tokens( + backend: ModeBackend, + logits: torch.Tensor, + run_reqs: List[InferReq], + b_req_idx: torch.Tensor, + b_req_mtp_start_loc: torch.Tensor, +): + """Sample and verify a fixed-width DFlash2 block.""" + + from lightllm.server.router.model_infer.mode_backend.generic_post_process import build_sampling_probs + + req_num = int(b_req_mtp_start_loc.shape[0]) + verify_width = backend.model.req_manager.req_sampling_params_manager.mtp_verify_width + assert len(run_reqs) == logits.shape[0] + assert logits.shape[0] == req_num * verify_width, "DFlash2 requires fixed-width verification" + + sampling_probs, raw_probs = build_sampling_probs( + logits=logits, + reqs=run_reqs, + eos_id=backend.eos_id, + ) + sampling_manager = backend.model.req_manager.req_sampling_params_manager + request_ids = b_req_idx.index_select(0, b_req_mtp_start_loc.long()).long() + proposal_tokens = sampling_manager.req_to_next_token_ids.index_select(0, request_ids)[:, 1:] + candidate_ids = sampling_manager.req_to_dflash2_candidate_ids.index_select(0, request_ids) + q_rows = sampling_manager.req_to_dflash2_q_probs.index_select(0, request_ids) + + vocab_size = sampling_probs.shape[-1] + draft_width = verify_width - 1 + request_reqs = run_reqs[::verify_width] + return _rejection_sample_from_probs( + sampling_probs=sampling_probs.view(req_num, verify_width, vocab_size), + raw_probs=raw_probs.view(req_num, verify_width, vocab_size), + proposal_tokens=proposal_tokens[:, :draft_width], + candidate_ids=candidate_ids[:, :draft_width], + q_rows=q_rows[:, :draft_width], + request_reqs=request_reqs, + acceptance_uniforms=_request_uniforms( + request_reqs, + draft_width, + device=sampling_probs.device, + ), + ) + + +def _rejection_sample_from_probs( + sampling_probs: torch.Tensor, + raw_probs: torch.Tensor, + proposal_tokens: torch.Tensor, + candidate_ids: torch.Tensor, + q_rows: torch.Tensor, + request_reqs: List[InferReq], + acceptance_uniforms: torch.Tensor | None = None, +): + """Apply sequential speculative rejection sampling to one DFlash2 block.""" + + from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + _random_sample, + ) + + req_num, verify_width, _ = sampling_probs.shape + assert proposal_tokens.shape == (req_num, verify_width - 1) + assert candidate_ids.shape == q_rows.shape + assert candidate_ids.shape[:2] == proposal_tokens.shape + assert raw_probs.shape == sampling_probs.shape + if acceptance_uniforms is not None: + assert acceptance_uniforms.shape == proposal_tokens.shape + + draft_width = verify_width - 1 + row_ids = torch.arange(req_num, device=sampling_probs.device) + has_request_seed = any(req.generator is not None for req in request_reqs) + + target_proposal_probs = torch.gather( + sampling_probs[:, :draft_width], + dim=-1, + index=proposal_tokens.unsqueeze(-1), + ).squeeze(-1) + draft_proposal_probs = torch.where( + candidate_ids.eq(proposal_tokens.unsqueeze(-1)), + q_rows, + 0.0, + ).sum(dim=-1) + # 首轮 decode 尚无 proposal，q_rows 为零；必须拒绝占位 token， + # 让下方的 correction 从完整 target 分布采样。 + accept_probs = torch.where( + draft_proposal_probs > 0, + torch.minimum( + torch.ones_like(target_proposal_probs), + target_proposal_probs / draft_proposal_probs.clamp_min(1e-20), + ), + 0.0, + ) + if acceptance_uniforms is None: + acceptance_uniforms = torch.rand_like(accept_probs) + accepted_prefix = acceptance_uniforms.lt(accept_probs).to(torch.int32).cumprod(dim=-1) + accepted_draft_count = accepted_prefix.sum(dim=-1).to(torch.int32) + + # Only the first rejected position needs a residual sample. If all drafts + # are accepted, sample the target model's final bonus row instead. + selected_target_row = accepted_draft_count.long() + correction_probs = sampling_probs[row_ids, selected_target_row].clone() + is_rejection = accepted_draft_count.lt(draft_width) + rejected_slot = accepted_draft_count.clamp_max(draft_width - 1).long() + rejected_candidate_ids = candidate_ids[row_ids, rejected_slot] + rejected_q_rows = q_rows[row_ids, rejected_slot] * is_rejection[:, None] + correction_probs.scatter_add_(1, rejected_candidate_ids, -rejected_q_rows) + correction_probs.clamp_min_(0.0) + correction_mass = correction_probs.sum(dim=-1, keepdim=True) + target_fallback = sampling_probs[row_ids, selected_target_row] + correction_probs = torch.where( + correction_mass > 1e-20, + correction_probs / correction_mass.clamp_min(1e-20), + target_fallback, + ) + correction_token = _random_sample( + correction_probs, + request_reqs, + has_request_seed, + ) + + # Rows after the correction token are ignored by accepted_index, but keep + # them valid so raw-probability gathers and debug tooling are always safe. + output_ids = torch.zeros((req_num, verify_width), dtype=torch.int64, device=sampling_probs.device) + output_ids[:, :draft_width] = proposal_tokens + output_ids.scatter_(1, selected_target_row[:, None], correction_token[:, None]) + accept_lengths = accepted_draft_count + 1 + accepted_index = (torch.arange(verify_width, device=sampling_probs.device)[None, :] < accept_lengths[:, None]).to( + torch.int32 + ) + flat_output_ids = output_ids.reshape(-1) + flat_raw_probs = raw_probs.reshape(-1, raw_probs.shape[-1]) + output_logprobs = torch.log(torch.gather(flat_raw_probs, 1, flat_output_ids[:, None]).squeeze(1).clamp_min(1e-20)) + return flat_output_ids, output_logprobs, accept_lengths, accepted_index.reshape(-1) + + +__all__ = ["sample_and_verify_dflash2_tokens", "save_dflash2_proposal_state"] diff --git a/lightllm/server/router/model_infer/mtp_speculative/engine.py b/lightllm/server/router/model_infer/mtp_speculative/engine.py index 9d59afd8a7..5cd62e76e8 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/engine.py +++ b/lightllm/server/router/model_infer/mtp_speculative/engine.py @@ -21,11 +21,7 @@ class SpecEngine: - """Owns MTP planning and draft proposal generation. - - Target verification, request metrics, stream synchronization, and resource - cleanup are stateless operations exposed by ``mtp_speculative.utils``. - """ + """Owns MTP planning, proposal generation, and target verification.""" def __init__( self, @@ -34,6 +30,7 @@ def __init__( enable_dynmaic_mtp: bool, ) -> None: self.backend = backend + self.spec_mode = spec_mode self.proposer: BaseSpecProposer = build_spec_proposer( spec_mode=spec_mode, backend=backend, @@ -131,6 +128,85 @@ def propose_next( accept_len=accept_len, ) + # Target sampling and verification. + + def sample_and_verify( + self, + logits: torch.Tensor, + run_reqs: List, + b_req_idx: torch.Tensor, + b_req_mtp_start_loc: torch.Tensor, + b_mtp_index: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample + from lightllm.server.router.model_infer.mtp_speculative.utils import ( + update_mtp_state_after_verify, + verify_mtp_tokens, + ) + + if self.spec_mode == "dflash2": + assert logits.shape[0] == b_req_mtp_start_loc.shape[0] * ( + self.backend.max_draft_step + 1 + ), "DFlash2 requires fixed-width verification" + + # Greedy DFlash2 uses token equality and the regular sampler's logprobs, + # without reading proposal distributions or consuming rejection RNG. + if self.spec_mode == "dflash2" and any(req.sampling_param.shm_param.top_k != 1 for req in run_reqs): + from lightllm.server.router.model_infer.mtp_speculative.dflash2 import sample_and_verify_dflash2_tokens + + next_token_ids, next_token_logprobs, accept_lengths, accepted_index = sample_and_verify_dflash2_tokens( + backend=self.backend, + logits=logits, + run_reqs=run_reqs, + b_req_idx=b_req_idx, + b_req_mtp_start_loc=b_req_mtp_start_loc, + ) + update_mtp_state_after_verify( + backend=self.backend, + b_req_idx=b_req_idx, + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + ) + else: + next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.backend.eos_id) + accept_lengths, accepted_index = verify_mtp_tokens( + backend=self.backend, + next_token_ids=next_token_ids, + b_req_idx=b_req_idx, + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_mtp_index=b_mtp_index, + ) + return next_token_ids, next_token_logprobs, accept_lengths, accepted_index + + def prepare_next_verification_state( + self, + proposal: SpecProposal, + target_next_token_ids: torch.Tensor, + b_req_mtp_start_loc: torch.Tensor, + b_req_idx: torch.Tensor, + mtp_accept_len: torch.Tensor, + ) -> None: + from lightllm.server.router.model_infer.mtp_speculative.utils import scatter_mtp_next_tokens + + scatter_mtp_next_tokens( + backend=self.backend, + proposal=proposal, + target_next_token_ids=target_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_idx=b_req_idx, + mtp_accept_len=mtp_accept_len, + ) + if self.spec_mode == "dflash2": + from lightllm.server.router.model_infer.mtp_speculative.dflash2 import save_dflash2_proposal_state + + save_dflash2_proposal_state( + backend=self.backend, + proposal=proposal, + b_req_idx=b_req_idx, + b_req_mtp_start_loc=b_req_mtp_start_loc, + ) + # Planner runtime statistics. def update_planner_statics( diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/__init__.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/__init__.py index dd0d8d05ac..c6d98c1448 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/__init__.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/__init__.py @@ -14,6 +14,10 @@ def build_spec_proposer(*, spec_mode: str, backend: "ModeBackend", enable_dynmai from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import DFlashProposer return DFlashProposer(backend=backend, enable_dynmaic_mtp=enable_dynmaic_mtp) + if spec_mode == "dflash2": + from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash2 import DFlash2Proposer + + return DFlash2Proposer(backend=backend, enable_dynmaic_mtp=enable_dynmaic_mtp) if spec_mode == "eagle3": from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle3 import Eagle3Proposer diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py index 23f81f546a..c64895dc7b 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py @@ -145,6 +145,22 @@ def propose_next( draft_input.multimodal_params = [{"images": [], "audios": []} for _ in range(draft_input.batch_size)] draft_output = draft_model.forward(draft_input) + return self._build_proposal( + draft_output=draft_output, + req_num=req_num, + block_size=block_size, + draft_step=draft_step, + extra_mem_indexes_cpu=extra_mem_indexes_cpu, + ) + + def _build_proposal( + self, + draft_output: ModelOutput, + req_num: int, + block_size: int, + draft_step: int, + extra_mem_indexes_cpu: torch.Tensor, + ) -> DFlashSpecProposal: if self.enable_dynmaic_mtp: flat_draft_token_ids, flat_draft_token_probs = self.backend._gen_argmax_token_ids_and_prob(draft_output) else: diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash2.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash2.py new file mode 100644 index 0000000000..34f1102dfc --- /dev/null +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash2.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.server.router.model_infer.mtp_speculative.proposers.base import MtpMemIndexesToFree +from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import DFlashProposer +from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import DFlash2SpecProposal + + +class DFlash2Proposer(DFlashProposer): + """Fixed-width DFlash2 block proposer with selector-distribution outputs.""" + + def __init__(self, *, backend, enable_dynmaic_mtp: bool) -> None: + if enable_dynmaic_mtp: + raise ValueError("DFlash2 does not support dynamic MTP verification") + super().__init__(backend=backend, enable_dynmaic_mtp=False) + + def _build_proposal( + self, + draft_output: ModelOutput, + req_num: int, + block_size: int, + draft_step: int, + extra_mem_indexes_cpu: torch.Tensor, + ) -> DFlash2SpecProposal: + mtp_collector = draft_output.mtp_collector + selected_token_ids = mtp_collector.draft_token_ids + candidate_ids = mtp_collector.draft_candidate_ids + candidate_probs = mtp_collector.draft_candidate_probs + expected_token_shape = (req_num, block_size - 1) + assert selected_token_ids.shape == expected_token_shape + assert candidate_ids.shape == candidate_probs.shape, "candidate id/probability shapes must match" + assert candidate_ids.shape[:2] == expected_token_shape + + return DFlash2SpecProposal( + token_ids=selected_token_ids[:, :draft_step].contiguous(), + extra_mem_indexes_cpu=[MtpMemIndexesToFree(mem_indexes_cpu=extra_mem_indexes_cpu)], + candidate_ids=candidate_ids[:, :draft_step].contiguous(), + q_probs=candidate_probs[:, :draft_step].float().contiguous(), + ) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/proposal_type.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/proposal_type.py index ee70a26616..60e35634fa 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/proposal_type.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/proposal_type.py @@ -26,6 +26,14 @@ class DFlashSpecProposal(SpecProposal): schedule_scores: torch.Tensor | None = None +@dataclass +class DFlash2SpecProposal(SpecProposal): + """Fixed-width DFlash2 proposal with selector sampling state.""" + + candidate_ids: torch.Tensor | None = None + q_probs: torch.Tensor | None = None + + @dataclass class DSparkSpecProposal(SpecProposal): """DSpark proposal with GPU confidence scores and their CPU planner view.""" diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 935fe92d33..e6fadc9b01 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -34,6 +34,26 @@ def alloc_mem_indexes(token_count: int) -> torch.Tensor: return g_infer_context.req_manager.mem_manager.alloc(token_count) +def update_mtp_state_after_verify( + backend: ModeBackend, + b_req_idx: torch.Tensor, + b_req_mtp_start_loc: torch.Tensor, + b_mtp_index: torch.Tensor, + accepted_index: torch.Tensor, +) -> None: + """Select recurrent-state slots using the result of token or rejection verification.""" + + if backend.is_linear_att_mixed_model: + linear_att_mtp_state_index_update( + req_to_mtp_state_index=backend.model.req_manager.req_to_mtp_state_index, + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + verify_width=backend.max_draft_step + 1, + ) + + def verify_mtp_tokens( backend: ModeBackend, next_token_ids: torch.Tensor, @@ -49,15 +69,13 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - if backend.is_linear_att_mixed_model: - linear_att_mtp_state_index_update( - req_to_mtp_state_index=backend.model.req_manager.req_to_mtp_state_index, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=backend.max_draft_step + 1, - ) + update_mtp_state_after_verify( + backend=backend, + b_req_idx=b_req_idx, + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + ) return accept_lengths, accepted_index @@ -135,5 +153,6 @@ def free_mem_indexes( "free_mem_indexes", "record_request_mtp_metrics", "scatter_mtp_next_tokens", + "update_mtp_state_after_verify", "verify_mtp_tokens", ] diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index bded568bdc..d0178c3725 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -276,6 +276,8 @@ def get_added_mtp_kv_layer_num() -> int: return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) if mtp_mode == "dflash": return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) + if mtp_mode == "dflash2": + return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) raise ValueError(f"unsupported mtp_mode: {mtp_mode}") diff --git a/unit_tests/common/basemodel/test_mtp_manager.py b/unit_tests/common/basemodel/test_mtp_manager.py index c7ad04f56b..d49e434d52 100644 --- a/unit_tests/common/basemodel/test_mtp_manager.py +++ b/unit_tests/common/basemodel/test_mtp_manager.py @@ -42,6 +42,7 @@ def _decode_batch_multiplier(monkeypatch, spec_mode, *, is_draft_model, mtp_step ("eagle_no_att", True, 1), ("dspark", True, 7), ("dflash", True, 7), + ("dflash2", True, 8), ], ) def test_decode_batch_multiplier(monkeypatch, spec_mode, is_draft_model, expected): @@ -77,6 +78,7 @@ def test_decode_cuda_graph_grow_step_size(monkeypatch, dynamic_verify, is_draft_ ("vanilla_with_att", True, 0), ("dspark", True, 6), ("dflash", True, 6), + ("dflash2", True, 7), ], ) def test_decode_draft_step(monkeypatch, spec_mode, is_draft_model, expected): @@ -102,6 +104,7 @@ def test_get_instance_returns_singleton(monkeypatch): ("dspark", False, LayerHiddenCollector), ("eagle3", True, FinalHiddenCollector), ("dspark", True, MtpHeadOutputCollector), + ("dflash2", True, MtpHeadOutputCollector), ], ) def test_create_hidden_collector_selects_implementation(monkeypatch, spec_mode, is_draft_model, expected_type): diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py new file mode 100644 index 0000000000..13f45e86f8 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py @@ -0,0 +1,92 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.server.router.model_infer.mtp_speculative.dflash2 import ( + _rejection_sample_from_probs, + save_dflash2_proposal_state, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import ( + DFlash2SpecProposal, + SpecProposal, +) + + +@pytest.mark.parametrize( + "first,second,has_proposal,expected", + [ + ([1.0, 0, 0], [0, 1.0, 0], False, [0]), + ([1.0, 0, 0], [0, 1.0, 0], True, [0, 1, 2]), + ([0.5, 0.5, 0], [0, 1.0, 0], True, [1]), + ([1.0, 0, 0], [0, 0.5, 0.5], True, [0, 2]), + ], + ids=["first_round", "all_accepted", "first_rejected", "middle_rejected"], +) +def test_rejection_prefix_and_logprobs(first, second, has_proposal, expected): + probs = torch.tensor([[first, second, [0, 0, 1.0]]]) + raw = torch.tensor([[[0.2, 0.3, 0.5], [0.1, 0.2, 0.7], [0.3, 0.4, 0.3]]]) + q = torch.tensor([[[1.0, 0, 0], [0, 1.0, 0]]]) * has_proposal + tokens, logprobs, lengths, mask = _rejection_sample_from_probs( + sampling_probs=probs, + raw_probs=raw, + proposal_tokens=torch.tensor([[0, 1]]), + candidate_ids=torch.arange(3).expand(1, 2, 3), + q_rows=q, + request_reqs=[SimpleNamespace(generator=None)], + acceptance_uniforms=torch.tensor([[0.9, 0.9]]), + ) + assert lengths.tolist() == [len(expected)] + assert mask.tolist() == [int(i < len(expected)) for i in range(3)] + assert tokens[mask.bool()].tolist() == expected + expected_logprobs = raw[0, torch.arange(len(expected)), torch.tensor(expected)].log() + torch.testing.assert_close(logprobs[mask.bool()], expected_logprobs) + + +def test_rejection_recovers_target_distribution(): + # A non-delta q checks both p/q acceptance and residual sampling together. + count = 20000 + p = torch.tensor([0.6, 0.3, 0.1]) + q = torch.tensor([0.1, 0.3, 0.6]) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(123) + tokens, _, _, _ = _rejection_sample_from_probs( + sampling_probs=p.expand(count, 2, 3), + raw_probs=p.expand(count, 2, 3), + proposal_tokens=torch.multinomial(q, count, replacement=True).view(count, 1), + candidate_ids=torch.arange(3).expand(count, 1, 3), + q_rows=q.expand(count, 1, 3), + request_reqs=[SimpleNamespace(generator=None)] * count, + ) + frequencies = torch.bincount(tokens.view(count, 2)[:, 0], minlength=3).float() / count + torch.testing.assert_close(frequencies, p, atol=0.02, rtol=0) + + +@pytest.mark.parametrize("invalid", [None, "type", "width"]) +def test_proposal_state_preserves_request_slots(invalid): + manager = SimpleNamespace( + mtp_mode="dflash2", + req_to_dflash2_candidate_ids=torch.zeros(5, 2, 16, dtype=torch.int64), + req_to_dflash2_q_probs=torch.zeros(5, 2, 16), + ) + backend = SimpleNamespace(model=SimpleNamespace(req_manager=SimpleNamespace(req_sampling_params_manager=manager))) + candidates = torch.arange(64).reshape(2, 2, 16) + q = torch.zeros(2, 2, 16) + q[0, :, 0], q[1, :, 1] = 1, 1 + proposal = DFlash2SpecProposal(token_ids=candidates[:, :, 0], candidate_ids=candidates, q_probs=q) + req_ids, starts = torch.tensor([3, 3, 3, 1, 1, 1]), torch.tensor([0, 3]) + if invalid is not None: + if invalid == "type": + proposal = SpecProposal(token_ids=proposal.token_ids) + else: + proposal.candidate_ids, proposal.q_probs = candidates[:, :1], q[:, :1] + with pytest.raises(TypeError if invalid == "type" else AssertionError, match="DFlash2 requires"): + save_dflash2_proposal_state(backend, proposal, req_ids, starts) + assert not manager.req_to_dflash2_candidate_ids.any() + assert not manager.req_to_dflash2_q_probs.any() + return + save_dflash2_proposal_state(backend, proposal, req_ids, starts) + torch.testing.assert_close(manager.req_to_dflash2_candidate_ids[[3, 1]], candidates) + torch.testing.assert_close(manager.req_to_dflash2_q_probs[[3, 1]], q) + assert not manager.req_to_dflash2_candidate_ids[[0, 2, 4]].any() + assert not manager.req_to_dflash2_q_probs[[0, 2, 4]].any() diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py index 7baf34061b..0a967457e3 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py @@ -47,6 +47,9 @@ from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import ( DFlashProposer, ) +from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash2 import ( + DFlash2Proposer, +) from lightllm.server.router.model_infer.mtp_speculative.proposers.dspark import ( DSparkProposer, ) @@ -60,6 +63,7 @@ EagleWithAttProposer, ) from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import ( + DFlash2SpecProposal, DFlashSpecProposal, DSparkSpecProposal, EagleSpecProposal, @@ -167,7 +171,7 @@ def test_common_engine_delegates_empty_dp_batch_to_lightspec_planner(): ) -def test_spec_engine_only_exposes_planning_and_proposal_interfaces(): +def test_spec_engine_exposes_planning_proposal_and_verification_interfaces(): public_methods = { name for name, value in SpecEngine.__dict__.items() if callable(value) and not name.startswith("_") } @@ -177,6 +181,8 @@ def test_spec_engine_only_exposes_planning_and_proposal_interfaces(): "plan_decode", "prepare_decode_model_input", "propose_next", + "sample_and_verify", + "prepare_next_verification_state", "update_planner_statics", } @@ -196,6 +202,7 @@ def test_mode_proposals_own_their_schedule_metadata(): assert "schedule_scores_cpu" not in EagleSpecProposal.__dataclass_fields__ assert "schedule_scores_cpu" not in DFlashSpecProposal.__dataclass_fields__ assert "schedule_scores_cpu" in DSparkSpecProposal.__dataclass_fields__ + assert "schedule_scores" not in DFlash2SpecProposal.__dataclass_fields__ def test_scatter_mtp_next_tokens_consumes_mode_proposal(monkeypatch): @@ -307,6 +314,11 @@ def test_engine_routes_only_dspark_to_the_confidence_planner(): assert isinstance(dflash_planner, LightSpecPlanner) assert dflash_planner.draft_steps == (3,) + dflash2_planner = build_planner("dflash2", enable_dynmaic_mtp=False) + assert isinstance(dflash2_planner, FixedSpecPlanner) + with pytest.raises(ValueError, match="unsupported LightSpec mode: dflash2"): + build_planner("dflash2") + eagle_planner = build_planner("eagle3") assert isinstance(eagle_planner, LightSpecPlanner) assert eagle_planner.draft_steps == (1, 2, 3) @@ -376,6 +388,7 @@ def test_each_mode_proposer_inherits_its_expected_implementation_base(): for proposer_type in proposer_types: assert proposer_type.__bases__ == (BaseSpecProposer,) assert Eagle3Proposer.__bases__ == (EagleWithAttProposer,) + assert DFlash2Proposer.__bases__ == (DFlashProposer,) for proposer_type in dp_overlap_proposer_types: assert proposer_type.__bases__ == (BaseDpOverlapProposer,) assert DpOverlapEagle3Proposer.__bases__ == (DpOverlapEagleWithAttProposer,) @@ -390,6 +403,7 @@ def test_each_mtp_mode_builds_its_own_proposer(): "eagle_no_att": EagleNoAttProposer, "eagle3": Eagle3Proposer, "dflash": DFlashProposer, + "dflash2": DFlash2Proposer, "dspark": DSparkProposer, } diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py index aaa3e05268..6fc6a7b9c3 100644 --- a/unit_tests/utils/test_speculative_utils.py +++ b/unit_tests/utils/test_speculative_utils.py @@ -27,6 +27,7 @@ ("lightllm.models.qwen3_5_moe_mtp.model", "Qwen3_5MoeMTPModel"), ("lightllm.models.qwen3_eagle.model", "Qwen3EagleModel"), ("lightllm.models.qwen3_dflash.model", "Qwen3DFlashModel"), + ("lightllm.models.qwen3_dflash2.model", "Qwen3DFlash2Model"), ("lightllm.models.qwen3_5_dflash.model", "Qwen3_5DFlashModel"), ("lightllm.models.qwen3_dspark.model", "Qwen3DSparkModel"), ("lightllm.models.qwen3_5_dspark.model", "Qwen3_5DSparkModel"), @@ -54,6 +55,8 @@ def test_qwen3_eagle_uses_layers_checkpoint_prefix(): ("dspark", True, 7, True, False), ("dflash", False, 7, True, True), ("dflash", True, 7, True, False), + ("dflash2", False, 7, False, False), + ("dflash2", True, 7, False, False), ("vanilla_with_att", True, 7, True, False), ("vanilla_with_att", True, 0, True, False), ("eagle3", True, 0, True, False), @@ -90,6 +93,7 @@ def test_attention_backend_selects_dynamic_spec_layout( (None, False, True), ("dflash", False, True), ("dflash", True, False), + ("dflash2", True, False), ("dspark", True, False), ("eagle3", True, True), ("vanilla_with_att", True, True), @@ -210,6 +214,7 @@ def test_fa3_dynamic_decode_state_builds_group_markers(state_class): ("qwen3_5_moe", "vanilla_with_att", "Qwen3_5MoeMTPModel"), ("qwen3_5_moe_text", "eagle_with_att", "Qwen3_5MoeMTPModel"), ("qwen3", "dflash", "Qwen3DFlashModel"), + ("qwen3", "dflash2", "Qwen3DFlash2Model"), ("qwen3_5", "dflash", "Qwen3_5DFlashModel"), ("qwen3_5_text", "dflash", "Qwen3_5DFlashModel"), ("qwen3", "dspark", "Qwen3DSparkModel"), @@ -248,6 +253,7 @@ def test_draft_model_registry_rejects_unsupported_model_type(): ("qwen3_5_moe", "eagle3"), ("qwen3_5_moe", "dspark"), ("qwen3_5_moe", "dflash"), + ("qwen3_5", "dflash2"), ("qwen3", "eagle_no_att"), ], ) @@ -307,7 +313,8 @@ def test_fixed_added_mtp_kv_layer_num_by_mode(monkeypatch, mtp_mode, mtp_step, e assert envs_utils.get_added_mtp_kv_layer_num() == expected_layer_num -def test_dflash_added_kv_layers_come_from_draft_config(tmp_path): +@pytest.mark.parametrize("mtp_mode", ["dflash", "dflash2"]) +def test_dflash_added_kv_layers_come_from_draft_config(tmp_path, mtp_mode): config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"num_hidden_layers": 5})) @@ -315,7 +322,7 @@ def test_dflash_added_kv_layers_come_from_draft_config(tmp_path): envs_utils.get_added_mtp_kv_layer_num.cache_clear() envs_utils.set_env_start_args( { - "mtp_mode": "dflash", + "mtp_mode": mtp_mode, "mtp_step": 7, "mtp_dynamic_verify": False, "mtp_draft_model_dir": [str(tmp_path)], From 6ecfc6ad207281eee3980e6902125301058c4ef1 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Mon, 7 Sep 2026 03:54:48 +0000 Subject: [PATCH 2/2] fix format --- lightllm/server/api_cli.py | 2 +- .../model_infer/mode_backend/chunked_prefill/impl.py | 7 +------ .../router/model_infer/mode_backend/dp_backend/impl.py | 4 +++- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 10efc3c104..c6bdf76de5 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -804,7 +804,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--mtp_dynamic_verify", action="store_true", - help="""Enable dynamic speculative scheduling. Temporarily ignored for DFlash2, which uses fixed-width verification.""", + help="""Enable dynamic speculative scheduling.""", ) parser.add_argument( "--kv_quant_calibration_config_path", diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 46e23b124f..1bf428d901 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -271,12 +271,7 @@ def decode_mtp( selected_rows = async_selected_row_mask_cpu.tensor.tolist() run_reqs = [req for req, selected in zip(run_reqs, selected_rows) if selected] b_req_mtp_start_loc = gen_b_req_mtp_start_loc(model_input.b_mtp_index, num_reqs=req_num) - ( - next_token_ids, - next_token_logprobs, - mtp_accept_len, - accepted_index, - ) = spec_engine.sample_and_verify( + (next_token_ids, next_token_logprobs, mtp_accept_len, accepted_index,) = spec_engine.sample_and_verify( logits=model_output.logits, run_reqs=run_reqs, b_req_idx=model_input.b_req_idx, diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index cbd625cd8e..0729db4d2a 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -34,7 +34,9 @@ def __init__(self) -> None: spec_mode = get_env_start_args().mtp_mode if spec_mode is not None: if spec_mode in ("dspark", "dflash", "dflash2"): - raise NotImplementedError("DP backend does not support DFlash/DFlash2/DSpark parallel block drafting yet.") + raise NotImplementedError( + "DP backend does not support DFlash/DFlash2/DSpark parallel block drafting yet." + ) if self.enable_prefill_microbatch_overlap: self.prefill = self.prefill_overlap_mtp else: