feat: support DFlash2 for Qwen3.8

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(

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,15 +164,33 @@ 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)
 
     def unpad_decode(self, padded_batch_size: int, origin_batch_size: int) -> "ModelMtpOutputCollector":
         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

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
 

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

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的计数，因为这个统计可能会导致一些特定场景下后处理效率的下降

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

lightllm/models/qwen3_dflash2/__init__.py

@@ -0,0 +1,3 @@
+from lightllm.models.qwen3_dflash2.model import Qwen3DFlash2Model
+
+__all__ = ["Qwen3DFlash2Model"]

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",
+]

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))