Optimize draft logits communication with sparse candidates

lightllm/common/basemodel/basemodel.py

@@ -389,6 +389,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0)
         infer_state.is_prefill = model_input.is_prefill
         infer_state.is_token_healing = self.is_token_healing
         infer_state.return_all_prompt_logics = self.return_all_prompt_logics
+        infer_state.is_mtp_draft_model = self.is_mtp_draft_model
         infer_state.batch_size = model_input.batch_size
         infer_state.total_token_num = model_input.total_token_num
         infer_state.max_q_seq_len = model_input.max_q_seq_len
@@ -539,6 +540,8 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba
             return model_output
         new_model_output = copy.copy(model_output)
         new_model_output.logits = new_model_output.logits[0:origin_batch_size]
+        if new_model_output.logits_token_ids is not None:
+            new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size]
         new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode(
             padded_batch_size=padded_batch_size,
             origin_batch_size=origin_batch_size,
@@ -551,6 +554,8 @@ def _create_unpad_prefill_model_output(
         new_model_output = copy.copy(padded_model_output)
         # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。
         new_model_output.logits = new_model_output.logits[0:origin_batch_size]
+        if new_model_output.logits_token_ids is not None:
+            new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size]
         new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill(
             origin_handle_token_num=origin_handle_token_num
         )
@@ -742,6 +747,7 @@ def prefill_func(input_tensors, _infer_state):
         hidden_collector.add_final_hidden(last_input_embs)
         model_output = ModelOutput(
             logits=predict_logits.contiguous(),
+            logits_token_ids=infer_state.logits_token_ids,
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
             prompt_logics=infer_state.prompt_logics,
         )
@@ -771,6 +777,7 @@ def _token_forward(self, infer_state: InferStateInfo):
         hidden_collector.add_final_hidden(last_input_embs)
         model_output = ModelOutput(
             logits=predict_logits.contiguous(),
+            logits_token_ids=infer_state.logits_token_ids,
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
         )
 
@@ -1025,11 +1032,13 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state
         hidden_collector1.add_final_hidden(last_input_embs1)
         model_output = ModelOutput(
             logits=predict_logits.contiguous(),
+            logits_token_ids=infer_state.logits_token_ids,
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
             prompt_logics=infer_state.prompt_logics,
         )
         model_output1 = ModelOutput(
             logits=predict_logits1.contiguous(),
+            logits_token_ids=infer_state1.logits_token_ids,
             mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1),
             prompt_logics=infer_state1.prompt_logics,
         )
@@ -1074,10 +1083,12 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1:
         hidden_collector1.add_final_hidden(last_input_embs1)
         model_output = ModelOutput(
             logits=predict_logits.contiguous(),
+            logits_token_ids=infer_state.logits_token_ids,
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
         )
         model_output1 = ModelOutput(
             logits=predict_logits1.contiguous(),
+            logits_token_ids=infer_state1.logits_token_ids,
             mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1),
         )
 

lightllm/common/basemodel/batch_objs.py

@@ -199,11 +199,32 @@ class ModelOutput:
     # 此时 logits 依然只保存每个请求最后一个位置的 logits，prompt_logics 为可选项，仅在
     # 需要返回 prompt logprobs 信息时才会非空。
     prompt_logics: Optional[torch.Tensor] = None
+    # Sparse draft logits map each candidate column back to a token id in the
+    # draft model's output vocabulary.  None means logits are dense and their
+    # column indexes are already token ids. Keep this optional field last so
+    # existing positional ModelOutput construction remains compatible.
+    logits_token_ids: Optional[torch.Tensor] = None
 
     def __post_init__(self) -> None:
         if self.mtp_collector is None:
             self.mtp_collector = ModelMtpOutputCollector()
+        if self.logits_token_ids is not None:
+            assert self.logits_token_ids.shape == self.logits.shape
+            assert self.logits_token_ids.dtype in (torch.int32, torch.int64)
+            assert self.logits_token_ids.device == self.logits.device
 
     def to_no_ref_tensor(self):
         self.logits = tensor_to_no_ref_tensor(self.logits)
+        if self.logits_token_ids is not None:
+            self.logits_token_ids = tensor_to_no_ref_tensor(self.logits_token_ids)
         self.mtp_collector.to_no_ref_tensor()
+
+    def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput":
+        """Select vocabulary-output rows while preserving sparse token ids."""
+
+        return ModelOutput(
+            logits=self.logits.index_select(0, index),
+            logits_token_ids=(
+                self.logits_token_ids.index_select(0, index) if self.logits_token_ids is not None else None
+            ),
+        )

lightllm/common/basemodel/infer_struct.py

@@ -53,6 +53,14 @@ def __init__(self):
 
         self.is_token_healing: bool = False
         self.return_all_prompt_logics: bool = False
+        # Draft models only need a small candidate set from the vocabulary
+        # projection.  The model marker enables the sparse logits path in the
+        # post layer without changing the target-model sampling interface.
+        self.is_mtp_draft_model: bool = False
+        # When logits are sparse, each column maps to the corresponding token
+        # id in the model's own output vocabulary.  Dense target logits leave
+        # this field as None and keep the historical column-index semantics.
+        self.logits_token_ids: Optional[torch.Tensor] = None
         # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个
         # token 位置的 logits，供后续回传 prompt logprobs 信息使用。
         # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。

lightllm/models/llama/layer_infer/post_layer_infer.py

@@ -7,15 +7,19 @@
 from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from lightllm.common.basemodel import PostLayerInferTpl
-from lightllm.distributed.communication_op import all_gather
+from lightllm.distributed.communication_op import all_gather, all_gather_into_tensor
 
 
 class LlamaPostLayerInfer(PostLayerInferTpl):
     """ """
 
+    DEFAULT_DRAFT_LOGITS_TOPK = 128
+
     def __init__(self, network_config):
         super().__init__(network_config)
         self.eps_ = network_config["rms_norm_eps"]
+        self.draft_logits_topk_ = int(os.getenv("LIGHTLLM_DRAFT_LOGITS_TOPK", str(self.DEFAULT_DRAFT_LOGITS_TOPK)))
+        assert self.draft_logits_topk_ > 0
         return
 
     def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor:
@@ -70,7 +74,13 @@ def _token_forward(
         input_embdings = None
 
         # 正常采样使用的 logits，始终只对应每个请求最后一个位置。
-        ans_logics = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state)
+        ans_logics = self._lm_head_and_gather(
+            last_input,
+            token_num,
+            layer_weight,
+            infer_state,
+            use_sparse_logits=getattr(infer_state, "is_mtp_draft_model", False),
+        )
         # 在 return_all_prompt_logics 模式下，prompt_logics 保存的是完整 prefill
         # 的 hidden state，需要在 norm/lm_head 之前取出来，避免被 input_embdings 置空。
         prompt_logics_hiddens = infer_state.prompt_logics
@@ -80,7 +90,11 @@ def _token_forward(
         if prompt_logics_hiddens is not None:
             prompt_token_num = prompt_logics_hiddens.shape[0]
             infer_state.prompt_logics = self._lm_head_and_gather(
-                prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state
+                prompt_logics_hiddens,
+                prompt_token_num,
+                layer_weight,
+                infer_state,
+                use_sparse_logits=False,
             )
 
         return ans_logics
@@ -91,12 +105,21 @@ def _lm_head_and_gather(
         token_num: int,
         layer_weight: LlamaPreAndPostLayerWeight,
         infer_state: LlamaInferStateInfo,
+        use_sparse_logits: bool = False,
     ) -> torch.Tensor:
         normed = self._norm(hidden, infer_state, layer_weight)
         normed = normed.permute(1, 0).view(-1, token_num)
         logic_batch = layer_weight.lm_head_weight_(input=normed, alloc_func=self.alloc_tensor)
         normed = None
 
+        if use_sparse_logits:
+            return self._gather_draft_topk_logits(
+                logic_batch=logic_batch,
+                token_num=token_num,
+                layer_weight=layer_weight,
+                infer_state=infer_state,
+            )
+
         vocab_size = layer_weight.lm_head_weight_.vocab_size
         if self.tp_world_size_ == 1:
             gather_data = logic_batch
@@ -116,6 +139,66 @@ def _lm_head_and_gather(
         gather_data = None
         return ans_logics
 
+    def _gather_draft_topk_logits(
+        self,
+        logic_batch: torch.Tensor,
+        token_num: int,
+        layer_weight: LlamaPreAndPostLayerWeight,
+        infer_state: LlamaInferStateInfo,
+    ) -> torch.Tensor:
+        """Gather each TP rank's local top-k instead of the full vocabulary."""
+
+        lm_head = layer_weight.lm_head_weight_
+        vocab_size = lm_head.vocab_size
+        local_topk = min(self.draft_logits_topk_, vocab_size)
+
+        # Every TP vocabulary shard is expected to contain at least top-k rows,
+        # so every rank contributes the same fixed-shape tensors directly.
+        assert logic_batch.shape[0] >= local_topk, (
+            f"draft logits top-k ({local_topk}) exceeds the local vocabulary shard " f"size ({logic_batch.shape[0]})"
+        )
+        local_values, local_indexes = torch.topk(logic_batch, k=local_topk, dim=0, sorted=False)
+        local_values = local_values.float()
+        local_indexes = local_indexes.to(dtype=torch.int32)
+        local_token_ids = local_indexes + int(lm_head.tp_vocab_start_id)
+
+        if self.tp_world_size_ == 1:
+            candidate_values = local_values.permute(1, 0)
+            candidate_token_ids = local_token_ids.permute(1, 0)
+        else:
+            # Pack FP32 logits and the bit representation of INT32 token ids
+            # into one FP32 payload. Both element types are four bytes, so the
+            # ids can be restored losslessly after a single all-gather.
+            local_payload = self.alloc_tensor(
+                (local_topk * 2, token_num),
+                dtype=torch.float32,
+                device=local_values.device,
+            )
+            local_payload[:local_topk].copy_(local_values)
+            local_payload[local_topk:].view(torch.int32).copy_(local_token_ids)
+            gathered_payload = self.alloc_tensor(
+                (self.tp_world_size_, local_topk * 2, token_num),
+                dtype=torch.float32,
+                device=local_values.device,
+            )
+            all_gather_into_tensor(
+                output_=gathered_payload,
+                input_=local_payload,
+                group=infer_state.dist_group,
+                async_op=False,
+            )
+            gathered_values = gathered_payload[:, :local_topk, :]
+            gathered_token_ids = gathered_payload[:, local_topk:, :].view(torch.int32)
+            candidate_values = gathered_values.permute(2, 0, 1).reshape(token_num, -1)
+            candidate_token_ids = gathered_token_ids.permute(2, 0, 1).reshape(token_num, -1)
+
+        candidate_token_ids = candidate_token_ids.to(dtype=torch.int64)
+        candidate_count = candidate_values.shape[1]
+        ans_logics = self.alloc_tensor((token_num, candidate_count), dtype=torch.float32, device=logic_batch.device)
+        ans_logics.copy_(candidate_values)
+        infer_state.logits_token_ids = candidate_token_ids.contiguous()
+        return ans_logics
+
     def token_forward(
         self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight
     ):

lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py

@@ -179,9 +179,20 @@ def token_forward(
             # Graph unpadding still uses the leading logits dimension when token ids are returned directly.
             return local_logits.new_empty((token_num, 1))
 
-        logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state)
+        logits = self._lm_head_and_gather(
+            last_input,
+            token_num,
+            layer_weight,
+            infer_state,
+            use_sparse_logits=getattr(infer_state, "is_mtp_draft_model", False),
+        )
         block_logits = logits.reshape(num_reqs, self.block_size_, -1)
-        sampled_tokens = torch.argmax(block_logits, dim=-1)
+        candidate_indexes = torch.argmax(block_logits, dim=-1)
+        if infer_state.logits_token_ids is None:
+            sampled_tokens = candidate_indexes
+        else:
+            block_token_ids = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_, -1)
+            sampled_tokens = block_token_ids.gather(-1, candidate_indexes.unsqueeze(-1)).squeeze(-1)
         confidence_logits = self.predict_confidence_logits(
             block_hidden,
             anchor_token_ids=anchor_token_ids,

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -840,14 +840,24 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]:
 
     def _gen_argmax_token_ids(self, model_output: ModelOutput):
         logits = model_output.logits
-        return torch.argmax(logits, dim=-1)
+        candidate_indexes = torch.argmax(logits, dim=-1)
+        return self._map_logits_indexes_to_token_ids(model_output, candidate_indexes)
 
     def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput):
         logits = model_output.logits
         probs = torch.softmax(logits, dim=-1)
-        max_probs, draft_next_token_ids_gpu = torch.max(probs, dim=-1)
+        max_probs, candidate_indexes = torch.max(probs, dim=-1)
+        draft_next_token_ids_gpu = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes)
         return draft_next_token_ids_gpu, max_probs
 
+    @staticmethod
+    def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexes: torch.Tensor):
+        """Map sparse-logits columns to vocabulary ids; dense logits are identity-mapped."""
+
+        if model_output.logits_token_ids is None:
+            return candidate_indexes
+        return model_output.logits_token_ids.gather(1, candidate_indexes.long().view(-1, 1)).view(-1).long()
+
     def _sample_and_scatter_token(
         self,
         logits: torch.Tensor,

lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py

@@ -141,7 +141,7 @@ def propose_next_overlap(
                 req_num_by_batch,
             )
         ):
-            accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows))
+            accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows)
             if self.enable_dynmaic_mtp:
                 draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output)
                 draft_token_probs = draft_token_probs.float()

lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py

@@ -81,7 +81,7 @@ def propose_next(
 
         # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行
         # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。
-        accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows))
+        accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows)
         if self.enable_dynmaic_mtp:
             draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output)
             schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1))

unit_tests/common/basemodel/test_model_output.py

@@ -11,23 +11,28 @@ def test_decode_unpad_slices_spec_output_with_logits():
     model = TpPartBaseModel.__new__(TpPartBaseModel)
     output = ModelOutput(
         logits=torch.arange(24).view(6, 4),
+        logits_token_ids=torch.arange(100, 124).view(6, 4),
         mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)),
     )
 
     unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=4)
 
     assert unpadded.logits.shape == (4, 4)
+    assert unpadded.logits_token_ids.shape == (4, 4)
+    torch.testing.assert_close(unpadded.logits_token_ids, output.logits_token_ids[:4])
     assert unpadded.mtp_collector.spec_hidden.shape == (4, 3)
     # Unpadding returns a shallow output copy and leaves the graph-owned
     # tensors on the original ModelOutput intact.
     assert output.logits.shape == (6, 4)
+    assert output.logits_token_ids.shape == (6, 4)
     assert output.mtp_collector.spec_hidden.shape == (6, 3)
 
 
 def test_prefill_unpad_uses_token_rows_for_spec_hidden():
     model = TpPartBaseModel.__new__(TpPartBaseModel)
     output = ModelOutput(
         logits=torch.arange(20).view(5, 4),
+        logits_token_ids=torch.arange(100, 120).view(5, 4),
         mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)),
         prompt_logics=torch.arange(32).view(8, 4),
     )
@@ -39,10 +44,27 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden():
     )
 
     assert unpadded.logits.shape == (3, 4)
+    assert unpadded.logits_token_ids.shape == (3, 4)
     assert unpadded.mtp_collector.spec_hidden.shape == (6, 3)
     assert unpadded.prompt_logics.shape == (6, 4)
 
 
+def test_index_select_logits_rows_preserves_sparse_token_mapping():
+    output = ModelOutput(
+        logits=torch.tensor([[1.0, 4.0], [7.0, 2.0], [3.0, 6.0]]),
+        logits_token_ids=torch.tensor([[10, 40], [70, 20], [30, 60]]),
+        mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))),
+    )
+
+    selected = output.index_select_logits_rows(torch.tensor([2, 0]))
+
+    torch.testing.assert_close(selected.logits, torch.tensor([[3.0, 6.0], [1.0, 4.0]]))
+    torch.testing.assert_close(selected.logits_token_ids, torch.tensor([[30, 60], [10, 40]]))
+    # The helper intentionally returns only the vocabulary outputs needed by
+    # the proposer; hidden collection remains owned by the original output.
+    assert selected.mtp_collector.spec_hidden is None
+
+
 def test_decode_unpad_restores_empty_output():
     model = TpPartBaseModel.__new__(TpPartBaseModel)
     output = ModelOutput(