perf: share vocab-parallel top-k logits for target and draft

docs/CN/source/tutorial/api_server_args.rst

@@ -641,6 +641,42 @@ MTP 多预测参数
     增加此值允许更多预测，但确保模型与指定的步数兼容。
     目前 deepseekv3/r1 模型仅支持 1 步
 
+词表并行采样
+------------
+
+.. option:: --disable_vocab_parallel_top1
+
+    关闭主模型和草稿模型的自动贪心归约。默认对满足 ``top_k=1``、
+    ``temperature=1`` 且无 logits 修改需求的主模型批次，汇总每卡最大值及
+    一个归一化统计量。token 与完整词表 argmax 一致，包括并列值规则；
+    主模型 logprob 使用完整词表分母，允许浮点归约的舍入误差。
+    固定贪心 draft 只取 top-1，不计算概率。
+
+.. option:: --enable_vocab_parallel_topk
+
+    显式接受主模型的近似采样，默认关闭。先截取每卡候选，再在汇总候选上
+    应用请求的 temperature、top-k 和 top-p。返回 logprob 使用请求 top-k/top-p
+    过滤前的候选 softmax，既不是完整词表 logprob，也不是最终过滤分布的
+    logprob。符合条件的 greedy 批次仍优先使用准确 top-1，除非已关闭它。
+
+.. option:: --vocab_parallel_topk_size
+
+    每个 TP rank 的候选数，正整数，默认 ``128``，不超过最小词表分片大小。
+    与请求的 ``top_k`` 不同；改变 TP 数量可能改变近似分布。
+    允许设置为 ``1``，但在 TP=1 时无法提供有区分度的概率置信度。
+
+动态 Vanilla/EAGLE/DFlash draft 默认使用上述候选数计算近似调度置信度；
+DSpark 使用独立 confidence head，可以保留 top-1。
+``--disable_vocab_parallel_top1`` 同时关闭 draft 的自动归约；配合不开启
+``--enable_vocab_parallel_topk``，即可让主模型和 draft 全部恢复完整输出。
+
+惩罚项、EOS/无效 token 屏蔽、约束采样、非 greedy 的显式 seed 及 RL rank
+需求回退完整 logits。不支持的模型 head、硬件平台，以及候选通信量不小于
+完整 logits 的情况也回退。Prompt logits 始终完整。Overlap 两侧使用兼容布局；
+decode 在启动时分别捕获启用布局的 graph，因此启动时间和 graph 显存可能增加。
+近似模式不保证与完整模式同 seed 的输出一致。
+
+
 DeepSeek 冗余专家参数
 ---------------------
 

docs/EN/source/tutorial/api_server_args.rst

@@ -644,6 +644,49 @@ MTP Multi-Prediction Parameters
     Increasing this value allows more predictions, but ensure the model is compatible with the specified number of steps.
     Currently deepseekv3/r1 models only support 1 step
 
+Vocabulary-Parallel Sampling
+----------------------------
+
+.. option:: --disable_vocab_parallel_top1
+
+    Disable automatic greedy reduction for target and draft models. By default,
+    eligible target batches with ``top_k=1`` and ``temperature=1`` gather one
+    maximum per rank and one normalization statistic. Selected tokens preserve
+    dense argmax tie-breaking; target logprobs use the full vocabulary denominator
+    (subject to floating-point reduction rounding). Fixed greedy drafts collect
+    top-1 without computing probabilities.
+
+.. option:: --enable_vocab_parallel_topk
+
+    Opt into approximate target sampling, disabled by default. Each rank retains
+    local candidates before applying request temperature, top-k and top-p to the
+    gathered candidate set. Returned logprobs use the candidate softmax before
+    request top-k/top-p filtering; they are not full-vocabulary logprobs or the
+    final filtered sampling distribution's logprobs. Eligible greedy batches
+    still prefer exact top-1 unless it is disabled.
+
+.. option:: --vocab_parallel_topk_size
+
+    Positive candidate count per TP rank, default ``128``, capped at the smallest
+    vocabulary shard. This is distinct from request ``top_k``. Changing TP size
+    can change the approximate distribution. Size ``1`` is allowed, but provides
+    no useful confidence signal at TP=1.
+
+Dynamic Vanilla/EAGLE/DFlash drafts use this candidate count for approximate
+scheduling confidence by default. DSpark uses its independent confidence head
+and can retain top-1. ``--disable_vocab_parallel_top1`` also disables automatic
+draft reduction; combined with leaving ``--enable_vocab_parallel_topk`` unset,
+it restores dense output for both target and draft models.
+
+Penalties, EOS/invalid-token masking, constrained sampling, non-greedy explicit
+seeds and non-greedy RL ranks retain dense output. Unsupported heads/platforms
+and candidate payloads no smaller than dense logits also retain the dense path.
+Prompt logits remain dense. Overlap microbatches share a compatible layout;
+decode graphs for enabled layouts are captured separately at startup, increasing
+capture time and potentially graph memory usage. Approximate mode does not
+promise seed-equivalent output to dense mode.
+
+
 DeepSeek Redundant Expert Parameters
 ------------------------------------
 

lightllm/common/basemodel/basemodel.py

@@ -378,12 +378,42 @@ def forward(self, model_input: ModelInput):
         else:
             return self._decode(model_input)
 
+    def _vocab_parallel_output_mode(self, model_input: ModelInput) -> tuple[int, bool]:
+        if self.args.hardware_platform != "cuda":
+            return 0, False
+        if self.is_mtp_draft_model:
+            if self.args.disable_vocab_parallel_top1:
+                return 0, False
+            needs_probs = self.args.mtp_dynamic_verify and self.args.mtp_mode != "dspark"
+            return (self.args.vocab_parallel_topk_size if needs_probs else 1), False
+        return (
+            model_input.vocab_parallel_topk,
+            model_input.vocab_parallel_topk == 1
+            and model_input.vocab_parallel_greedy
+            and not self.args.disable_vocab_parallel_top1,
+        )
+
+    def vocab_parallel_graph_modes(self):
+        if self.args.hardware_platform != "cuda":
+            return [(0, False)]
+        if self.is_mtp_draft_model:
+            return [self._vocab_parallel_output_mode(None)]
+        modes = [(0, False)]
+        if not self.args.disable_vocab_parallel_top1:
+            modes.append((1, True))
+        if self.args.enable_vocab_parallel_topk:
+            modes.append((self.args.vocab_parallel_topk_size, False))
+        return modes
+
     def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0):
         infer_state = self.infer_state_class()
         infer_state.hidden_collector = self.hidden_collector_prototype.new_instance()
         infer_state.input_ids = model_input.input_ids
         infer_state.is_prefill = model_input.is_prefill
         infer_state.return_all_prompt_logics = self.return_all_prompt_logics
+        infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy = self._vocab_parallel_output_mode(
+            model_input
+        )
         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
@@ -534,6 +564,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,
@@ -546,6 +578,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
         )
@@ -650,7 +684,7 @@ def _decode(
         need_capture = False
         if use_cuda_graph:
             infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size)
-            need_capture = self.graph.need_capture(infer_batch_size)
+            need_capture = self.graph.need_capture(infer_batch_size, self._vocab_parallel_output_mode(model_input))
 
         model_input = self._create_padded_decode_model_input(model_input=model_input, new_batch_size=infer_batch_size)
         infer_state = self._create_inferstate(model_input)
@@ -678,7 +712,6 @@ def _decode(
 
     @final
     def _context_forward(self, infer_state: InferStateInfo):
-
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"
@@ -737,6 +770,8 @@ 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,
+            logits_are_logprobs=infer_state.logits_are_logprobs,
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
             prompt_logics=infer_state.prompt_logics,
         )
@@ -766,6 +801,8 @@ 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,
+            logits_are_logprobs=infer_state.logits_are_logprobs,
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
         )
 
@@ -899,7 +936,8 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1
 
         if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch):
             infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size)
-            need_capture = self.graph.need_capture(infer_batch_size)
+            assert self._vocab_parallel_output_mode(model_input0) == self._vocab_parallel_output_mode(model_input1)
+            need_capture = self.graph.need_capture(infer_batch_size, self._vocab_parallel_output_mode(model_input0))
             padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size)
             padded_model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size)
             infer_state0 = self._create_inferstate(padded_model_input0, 0)
@@ -1020,11 +1058,15 @@ 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,
+            logits_are_logprobs=infer_state.logits_are_logprobs,
             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,
+            logits_are_logprobs=infer_state1.logits_are_logprobs,
             mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1),
             prompt_logics=infer_state1.prompt_logics,
         )
@@ -1069,10 +1111,14 @@ 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,
+            logits_are_logprobs=infer_state.logits_are_logprobs,
             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,
+            logits_are_logprobs=infer_state1.logits_are_logprobs,
             mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1),
         )
 

lightllm/common/basemodel/batch_objs.py

@@ -7,6 +7,12 @@
 from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
 
 
+def logits_indexes_to_token_ids(indexes: torch.Tensor, token_ids: Optional[torch.Tensor]) -> torch.Tensor:
+    if token_ids is None:
+        return indexes
+    return token_ids.gather(1, indexes.reshape(-1, 1)).reshape_as(indexes).long()
+
+
 @dataclass
 class ModelInput:
     # 通用变量
@@ -55,6 +61,10 @@ class ModelInput:
     # 的 draft 模型的输入
     mtp_draft_input_hiddens: Optional[torch.Tensor] = None
 
+    # 0: dense logits; positive: per-rank candidate count.
+    vocab_parallel_topk: int = 0
+    vocab_parallel_greedy: bool = False
+
     def to_cuda(self):
         self.check_input()
 
@@ -200,10 +210,52 @@ class ModelOutput:
     # 需要返回 prompt logprobs 信息时才会非空。
     prompt_logics: Optional[torch.Tensor] = None
 
+    # Sparse vocab-parallel outputs map every candidate column back to its
+    # global token id. None means logits are dense and column indexes are ids.
+    logits_token_ids: Optional[torch.Tensor] = None
+    # Exact target top-1 output already contains the selected full-vocab logprob.
+    logits_are_logprobs: bool = False
+
     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.ndim == 2
+            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
+        if self.logits_are_logprobs:
+            assert self.logits_token_ids is not None and self.logits.shape[1] == 1
 
     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 logit rows without dropping their vocabulary metadata."""
+
+        return ModelOutput(
+            logits_are_logprobs=self.logits_are_logprobs,
+            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
+            ),
+        )
+
+    @classmethod
+    def concat_logits_rows(cls, outputs: List["ModelOutput"]) -> "ModelOutput":
+        """Concatenate outputs that share the same dense or sparse layout."""
+
+        assert outputs
+        has_token_ids = outputs[0].logits_token_ids is not None
+        assert all((output.logits_token_ids is not None) == has_token_ids for output in outputs)
+        assert all(output.logits_are_logprobs == outputs[0].logits_are_logprobs for output in outputs)
+        return cls(
+            logits_are_logprobs=outputs[0].logits_are_logprobs,
+            logits=torch.cat([output.logits for output in outputs], dim=0),
+            logits_token_ids=(
+                torch.cat([output.logits_token_ids for output in outputs], dim=0) if has_token_ids else None
+            ),
+        )

lightllm/common/basemodel/cuda_graph.py

@@ -83,10 +83,10 @@ def __init__(
     def can_run(self, batch_size, max_len_in_batch):
         return batch_size <= self.max_batch_size and max_len_in_batch <= self.graph_max_len_in_batch
 
-    def need_capture(self, batch_size):
+    def need_capture(self, batch_size, output_mode=(0, False)):
         find_batch_size = self.find_closest_graph_batch_size(batch_size)
         if find_batch_size is not None:
-            return find_batch_size not in self.graph
+            return (find_batch_size, *output_mode) not in self.graph
         else:
             assert False, "dead code"
 
@@ -125,7 +125,11 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo):
 
         with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool):
             model_output = decode_func(infer_state)
-        self.graph[batch_size] = (graph_obj, infer_state, model_output)
+        self.graph[(batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy)] = (
+            graph_obj,
+            infer_state,
+            model_output,
+        )
         graph_obj.replay()
         self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size)
         return model_output
@@ -160,7 +164,7 @@ def _capture_decode_overlap(
 
         with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool):
             model_output, model_output1 = decode_func(infer_state, infer_state1)
-        self.graph[batch_size] = (
+        self.graph[(batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy)] = (
             graph_obj,
             infer_state,
             infer_state1,
@@ -212,7 +216,9 @@ def capture_decode(
 
     def _replay(self, infer_state: InferStateInfo):
         batch_size = infer_state.input_ids.shape[0]
-        graph_obj, graph_infer_state, graph_output = self.graph[batch_size]
+        graph_obj, graph_infer_state, graph_output = self.graph[
+            (batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy)
+        ]
         graph_infer_state.copy_for_cuda_graph(infer_state)
         graph_obj.replay()
         return graph_output
@@ -229,7 +235,11 @@ def _replay_overlap(
             graph_infer_state1,
             graph_model_output,
             graph_model_output1,
-        ) = self.graph[batch_size]
+        ) = self.graph[(batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy)]
+        assert (infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy) == (
+            infer_state1.vocab_parallel_topk,
+            infer_state1.vocab_parallel_greedy,
+        )
         graph_infer_state.copy_for_cuda_graph(infer_state)
         graph_infer_state1.copy_for_cuda_graph(infer_state1)
         graph_obj.replay()
@@ -281,8 +291,11 @@ def warmup(self, model):
                 multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)],
                 **model._gen_special_model_input(batch_size),
             )
-            model_output: ModelOutput = model.forward(model_input)
-            del model_output
+            for topk, greedy in model.vocab_parallel_graph_modes():
+                model_input.vocab_parallel_topk = topk
+                model_input.vocab_parallel_greedy = greedy
+                model_output: ModelOutput = model.forward(model_input)
+                del model_output
             del input_ids
             del mem_indexes
             del b_req_idx
@@ -350,7 +363,11 @@ def warmup_overlap(self, model):
                         del locals()[var_name]
                 torch.cuda.empty_cache()
 
-            _, _ = model.microbatch_overlap_decode(decode_batches[0], decode_batches[1])
+            for topk, greedy in model.vocab_parallel_graph_modes():
+                for batch in decode_batches:
+                    batch.vocab_parallel_topk = topk
+                    batch.vocab_parallel_greedy = greedy
+                _, _ = model.microbatch_overlap_decode(decode_batches[0], decode_batches[1])
 
             model.mem_manager.free_all()
             model.req_manager.free_all()