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

lightllm/common/basemodel/basemodel.py

@@ -23,6 +23,9 @@
 from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph
 from lightllm.common.quantization import Quantcfg
 from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed
+from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import (
+    is_vocab_parallel_topk_enabled,
+)
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.dist_utils import get_dp_world_size
 from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory
@@ -378,12 +381,22 @@ def forward(self, model_input: ModelInput):
         else:
             return self._decode(model_input)
 
+    def _is_cuda_graph_output_compatible(self, *model_inputs: ModelInput) -> bool:
+        """Whether inputs match the dense/sparse contract captured at startup."""
+
+        return (
+            self.is_mtp_draft_model
+            or not is_vocab_parallel_topk_enabled()
+            or all(model_input.use_vocab_parallel_topk for model_input in model_inputs)
+        )
+
     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.use_vocab_parallel_topk = self.is_mtp_draft_model or model_input.use_vocab_parallel_topk
         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 +547,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 +561,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
         )
@@ -643,9 +660,13 @@ def _decode(
         # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的
         # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。
         infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len)
-        use_cuda_graph = self.graph is not None and self.graph.can_run(
-            batch_size=infer_batch_size,
-            max_len_in_batch=infer_max_kv_seq_len,
+        use_cuda_graph = (
+            self._is_cuda_graph_output_compatible(model_input)
+            and self.graph is not None
+            and self.graph.can_run(
+                batch_size=infer_batch_size,
+                max_len_in_batch=infer_max_kv_seq_len,
+            )
         )
         need_capture = False
         if use_cuda_graph:
@@ -678,7 +699,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 +757,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,
         )
@@ -766,6 +787,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),
         )
 
@@ -897,7 +919,11 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1
         infer_batch_size = max(1, origin_batch_size0, origin_batch_size1)
         infer_batch_size = triton.cdiv(infer_batch_size, self.tp_world_size_) * self.tp_world_size_
 
-        if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch):
+        if (
+            self._is_cuda_graph_output_compatible(model_input0, model_input1)
+            and 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)
             padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size)
@@ -1020,11 +1046,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,
         )
@@ -1069,10 +1097,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

@@ -55,6 +55,10 @@ class ModelInput:
     # 的 draft 模型的输入
     mtp_draft_input_hiddens: Optional[torch.Tensor] = None
 
+    # The router sets this only when a target-model batch can sample directly
+    # from sparse candidates. Draft models always enable the same output form.
+    use_vocab_parallel_topk: bool = False
+
     def to_cuda(self):
         self.check_input()
 
@@ -200,10 +204,49 @@ 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
+
     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
 
     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()
+
+    @property
+    def has_vocab_parallel_logits(self) -> bool:
+        return self.logits_token_ids is not None
+
+    def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput":
+        """Select logit rows without dropping their vocabulary metadata."""
+
+        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
+            ),
+        )
+
+    @classmethod
+    def concat_logits_rows(cls, outputs: List["ModelOutput"]) -> "ModelOutput":
+        """Concatenate outputs that share the same dense or sparse layout."""
+
+        assert outputs
+        has_vocab_parallel_logits = outputs[0].has_vocab_parallel_logits
+        assert all(output.has_vocab_parallel_logits == has_vocab_parallel_logits for output in outputs)
+        return cls(
+            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_vocab_parallel_logits else None
+            ),
+        )

lightllm/common/basemodel/cuda_graph.py

@@ -9,6 +9,9 @@
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.distributed import dist_group_manager
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
+from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import (
+    is_vocab_parallel_topk_enabled,
+)
 from lightllm.utils.torch_memory_saver_utils import (
     TorchMemorySaverWrapper,
     MemoryTag,
@@ -279,6 +282,7 @@ def warmup(self, model):
                 b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"),
                 is_prefill=False,
                 multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)],
+                use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(),
                 **model._gen_special_model_input(batch_size),
             )
             model_output: ModelOutput = model.forward(model_input)
@@ -340,6 +344,7 @@ def warmup_overlap(self, model):
                     b_shared_radix_node_id=b_shared_radix_node_id,
                     b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"),
                     multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)],
+                    use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(),
                     **model._gen_special_model_input(batch_size),
                 )
                 decode_batches.append(micro_batch)

lightllm/common/basemodel/infer_struct.py

@@ -52,6 +52,8 @@ def __init__(self):
         self.mem_index: torch.Tensor = None
 
         self.return_all_prompt_logics: bool = False
+        self.use_vocab_parallel_topk: bool = False
+        self.logits_token_ids: Optional[torch.Tensor] = None
         # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个
         # token 位置的 logits，供后续回传 prompt logprobs 信息使用。
         # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。

lightllm/common/basemodel/prefill_cuda_graph.py

@@ -10,6 +10,9 @@
 from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
 from lightllm.distributed import dist_group_manager
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
+from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import (
+    is_vocab_parallel_topk_enabled,
+)
 from .infer_struct import InferStateInfo
 from .cuda_graph import CudaGraph
 
@@ -220,6 +223,7 @@ def warmup(self, model):
                 is_prefill=True,
                 b_prefill_has_output_cpu=[False],
                 multimodal_params=[{"images": [], "audios": []}],
+                use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(),
                 **model._gen_special_model_input(token_num=total_token_num),
             )
             model_output: ModelOutput = model.forward(model_input)
@@ -281,6 +285,7 @@ def warmup_overlap(self, model):
                     is_prefill=True,
                     b_prefill_has_output_cpu=[False],
                     multimodal_params=[{"images": [], "audios": []}],
+                    use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(),
                     **model._gen_special_model_input(token_num=total_token_num),
                 )
 

lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py

@@ -0,0 +1,94 @@
+"""Collect sparse candidates directly from tensor-parallel vocabulary shards."""
+
+import os
+
+import torch
+
+from lightllm.distributed.communication_op import all_gather_into_tensor
+from lightllm.utils.envs_utils import enable_env_vars
+
+
+VOCAB_PARALLEL_TOPK_ENV = "LIGHTLLM_VOCAB_PARALLEL_TOPK"
+VOCAB_PARALLEL_TOPK_SIZE_ENV = "LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE"
+DEFAULT_VOCAB_PARALLEL_TOPK = 128
+
+
+def is_vocab_parallel_topk_enabled() -> bool:
+    """Whether target-model greedy batches may use sparse vocabulary output."""
+
+    return enable_env_vars(VOCAB_PARALLEL_TOPK_ENV)
+
+
+def get_vocab_parallel_topk_size() -> int:
+    topk = int(os.getenv(VOCAB_PARALLEL_TOPK_SIZE_ENV, str(DEFAULT_VOCAB_PARALLEL_TOPK)))
+    assert topk > 0, f"{VOCAB_PARALLEL_TOPK_SIZE_ENV} must be positive, got {topk}"
+    return topk
+
+
+@torch.no_grad()
+def vocab_parallel_topk(
+    local_logits: torch.Tensor,
+    *,
+    vocab_size: int,
+    vocab_start_id: int,
+    topk: int,
+    tp_world_size: int,
+    group,
+    alloc_func,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather each TP rank's local top-k logits and their global token ids.
+
+    The returned width is ``tp_world_size * topk``. It intentionally keeps the
+    union of local candidates: greedy selection remains exact, while probability
+    calculations over the sparse result are an inexpensive approximation.
+    """
+
+    assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous()
+    local_vocab_size, token_num = local_logits.shape
+    # Collectives require every rank to contribute the same shape. Vocabulary
+    # shards can differ by one row, so cap against the smallest possible shard.
+    local_topk = min(topk, vocab_size // tp_world_size)
+    assert local_topk > 0
+    assert local_vocab_size >= local_topk
+
+    local_values, local_indexes = torch.topk(local_logits, k=local_topk, dim=0, sorted=False)
+    local_values = local_values.float()
+    local_token_ids = local_indexes.to(torch.int32).add_(int(vocab_start_id))
+
+    if tp_world_size == 1:
+        candidate_values = local_values.permute(1, 0).contiguous()
+        candidate_token_ids = local_token_ids.permute(1, 0).contiguous()
+    else:
+        # Values and ids are both four bytes. Bit-packing the ids into the FP32
+        # payload keeps the operation to one fixed-shape collective.
+        local_payload = alloc_func(
+            (local_topk * 2, token_num),
+            dtype=torch.float32,
+            device=local_logits.device,
+        )
+        local_payload[:local_topk].copy_(local_values)
+        local_payload[local_topk:].view(torch.int32).copy_(local_token_ids)
+
+        gathered_payload = alloc_func(
+            (tp_world_size, local_topk * 2, token_num),
+            dtype=torch.float32,
+            device=local_logits.device,
+        )
+        all_gather_into_tensor(
+            output_=gathered_payload,
+            input_=local_payload,
+            group=group,
+            async_op=False,
+        )
+        candidate_values = gathered_payload[:, :local_topk, :].permute(2, 0, 1).reshape(token_num, -1)
+        candidate_token_ids = (
+            gathered_payload[:, local_topk:, :].view(torch.int32).permute(2, 0, 1).reshape(token_num, -1)
+        )
+
+    output_logits = alloc_func(
+        candidate_values.shape,
+        dtype=torch.float32,
+        device=local_logits.device,
+    )
+    output_logits.copy_(candidate_values)
+    return output_logits, candidate_token_ids.to(torch.int64).contiguous()

lightllm/models/gemma4/layer_infer/post_layer_infer.py

@@ -5,18 +5,16 @@
 class Gemma4PostLayerInfer(LlamaPostLayerInfer):
     """
     Same final RMSNorm + tied lm_head path as Llama, with an extra tanh-based
-    logit softcap at the end: logits = softcap * tanh(logits / softcap).
+    transform before sampling: logits = softcap * tanh(logits / softcap).
     """
 
     def __init__(self, network_config):
         super().__init__(network_config)
         self.final_logit_softcapping = float(network_config.get("final_logit_softcapping"))
 
-    def token_forward(self, input_embdings, infer_state, layer_weight):
-        logits = super().token_forward(input_embdings, infer_state, layer_weight)
-        if self.final_logit_softcapping is not None and self.final_logit_softcapping > 0:
-            cap = self.final_logit_softcapping
-            logits = torch.tanh(logits / cap) * cap
-            if infer_state.prompt_logics is not None:
-                infer_state.prompt_logics = torch.tanh(infer_state.prompt_logics / cap) * cap
-        return logits
+    def _apply_logit_postprocessing(self, logits: torch.Tensor) -> torch.Tensor:
+        if self.final_logit_softcapping is None or self.final_logit_softcapping <= 0:
+            return logits
+        cap = self.final_logit_softcapping
+        # The historical path materializes FP32 logits before applying softcap.
+        return torch.tanh(logits.float() / cap) * cap