perf(sampling): fuse greedy token selection and logprob computation

lightllm/common/basemodel/triton_kernel/local_vocab_topk.py

@@ -108,3 +108,91 @@ def local_vocab_topk(
     )
     # 词表现在位于连续的最后一维；只需要候选集合，不额外要求候选按分数排序。
     return torch.topk(transposed_logits, k=top_k, dim=1, sorted=False)
+
+
+@triton.jit
+def _maximum_propagating_nan(a, b):
+    return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL)
+
+
+@triton.jit
+def _greedy_vocab_partial(
+    logits,
+    maxima,
+    sums,
+    indexes,
+    vocab: tl.constexpr,
+    stride_b: tl.constexpr,
+    chunks: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    b = tl.program_id(0)
+    v = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    scores = tl.load(logits + b * stride_b + v, v < vocab, -float("inf"))
+    maximum = tl.reduce(scores, axis=0, combine_fn=_maximum_propagating_nan)
+    # Fully masked chunks contribute zero instead of exp(-inf - -inf) = NaN.
+    total = tl.sum(tl.exp(scores - maximum), axis=0)
+    total = tl.where(maximum == -float("inf"), 0.0, total)
+    # Match torch.argmax: NaNs take precedence, then choose the first valid index.
+    valid = v < vocab
+    is_nan = scores != scores
+    is_maximum = scores == maximum
+    is_candidate = valid & (is_nan | is_maximum)
+    winner = tl.min(tl.where(is_candidate, v, 2147483647), axis=0)
+    offset = b * chunks + tl.program_id(1)
+    tl.store(maxima + offset, maximum)
+    tl.store(sums + offset, total)
+    tl.store(indexes + offset, winner)
+
+
+@triton.jit
+def _greedy_vocab_finish(maxima, sums, indexes, token_ids, logprobs, chunks: tl.constexpr, BLOCK: tl.constexpr):
+    b = tl.program_id(0)
+    c = tl.arange(0, BLOCK)
+    maximums = tl.load(maxima + b * chunks + c, c < chunks, -float("inf"))
+    partial_sums = tl.load(sums + b * chunks + c, c < chunks, 0)
+    ids = tl.load(indexes + b * chunks + c, c < chunks, 2147483647)
+    maximum = tl.reduce(maximums, axis=0, combine_fn=_maximum_propagating_nan)
+    total = tl.sum(partial_sums * tl.exp(maximums - maximum), axis=0)
+    valid = c < chunks
+    is_nan = maximums != maximums
+    is_maximum = maximums == maximum
+    is_candidate = valid & (is_nan | is_maximum)
+    winner = tl.min(tl.where(is_candidate, ids, 2147483647), axis=0)
+    tl.store(token_ids + b, winner.to(tl.int64))
+    tl.store(logprobs + b, -tl.log(total))
+
+
+def local_vocab_greedy_logprob(logits: torch.Tensor, alloc_func=torch.empty, run_config=None):
+    """Return greedy IDs and their normalized logprobs without a full probability matrix."""
+    assert logits.ndim == 2 and logits.dtype == torch.float32 and logits.stride(1) == 1
+    batch, vocab = logits.shape
+    if run_config is None:
+        run_config = {"BLOCK": 2048, "num_warps": 4}
+    chunks = triton.cdiv(vocab, run_config["BLOCK"])
+    maxima = alloc_func((batch, chunks), dtype=torch.float32, device=logits.device)
+    sums = alloc_func((batch, chunks), dtype=torch.float32, device=logits.device)
+    indexes = alloc_func((batch, chunks), dtype=torch.int32, device=logits.device)
+    token_ids = alloc_func((batch,), dtype=torch.int64, device=logits.device)
+    logprobs = alloc_func((batch,), dtype=torch.float32, device=logits.device)
+    _greedy_vocab_partial[(batch, chunks)](
+        logits,
+        maxima,
+        sums,
+        indexes,
+        vocab,
+        logits.stride(0),
+        chunks,
+        **run_config,
+    )
+    _greedy_vocab_finish[(batch,)](
+        maxima,
+        sums,
+        indexes,
+        token_ids,
+        logprobs,
+        chunks,
+        triton.next_power_of_2(chunks),
+        num_warps=4,
+    )
+    return token_ids, logprobs

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

@@ -60,7 +60,11 @@ def _fwd_kernel_apply_penalty(
         eos_id = tl.load(eos_ids + eos_index)
         cur_eos_logit_ptr = Logits + cur_batch * stride_logit_b + eos_id
         cur_eos_logit = tl.load(cur_eos_logit_ptr)
-        cur_eos_logit = cur_eos_logit + tl.abs(cur_eos_logit) * penalty_scale
+        # Preserve masked EOS and zero logits; neutral/overflowed scales must not create NaNs.
+        apply_length_penalty = (penalty_scale != 0.0) & (cur_eos_logit != 0.0) & (cur_eos_logit != -float("inf"))
+        cur_eos_logit = tl.where(
+            apply_length_penalty, cur_eos_logit + tl.abs(cur_eos_logit) * penalty_scale, cur_eos_logit
+        )
         cur_eos_logit = tl.where(mask_eos != 0, -10000000.0, cur_eos_logit)
         tl.store(cur_eos_logit_ptr, cur_eos_logit)
     return

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

@@ -75,7 +75,11 @@ def _eos_penalty(
         eos_id = tl.load(eos_ids + eos_index)
         cur_eos_logit_ptr = Logits + offs * stride_logit_b + eos_id
         cur_eos_logit = tl.load(cur_eos_logit_ptr, mask=mask, other=0.0)
-        cur_eos_logit = cur_eos_logit + tl.abs(cur_eos_logit) * penalty_scale
+        # Preserve masked EOS and zero logits; neutral/overflowed scales must not create NaNs.
+        apply_length_penalty = (penalty_scale != 0.0) & (cur_eos_logit != 0.0) & (cur_eos_logit != -float("inf"))
+        cur_eos_logit = tl.where(
+            apply_length_penalty, cur_eos_logit + tl.abs(cur_eos_logit) * penalty_scale, cur_eos_logit
+        )
         cur_eos_logit = tl.where(mask_eos != 0, -10000000.0, cur_eos_logit)
         tl.store(cur_eos_logit_ptr, cur_eos_logit, mask=mask)
     return

lightllm/server/core/objs/sampling_params.py

@@ -1,6 +1,7 @@
 import os
 import json
 import ctypes
+import math
 from typing import Optional, List, Tuple, Union
 from transformers import GenerationConfig
 from lightllm.server.httpserver.req_id_generator_server import MAX_BEST_OF
@@ -261,6 +262,8 @@ def initialize(self, inputs: Tuple[int, float]):
         self.item0 = inputs[0]
         assert self.item0 >= 0, "ExponentialDecayLengthPenalty item0 must be int >= 0"
         self.item1 = inputs[1]
+        if not math.isfinite(self.item1):
+            raise ValueError("exponential_decay_length_penalty[1] must be finite")
         assert self.item1 >= 1.0, "ExponentialDecayLengthPenalty item1 must be a float >= 1.0"
         return
 
@@ -454,6 +457,9 @@ def _cfg(key, default):
             pass
 
     def verify(self):
+        for name in ("presence_penalty", "frequency_penalty", "repetition_penalty", "temperature", "top_p"):
+            if not math.isfinite(getattr(self, name)):
+                raise ValueError(f"{name} must be finite, got {getattr(self, name)}")
         if self.best_of <= 0 or self.best_of > MAX_BEST_OF:
             raise ValueError(f"need 0 < best_of <= {MAX_BEST_OF}, but got {self.best_of}")
         if self.n != self.best_of:

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

@@ -4,6 +4,7 @@
 from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache
 from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids
 from lightllm.common.basemodel.triton_kernel.post_process.apply_constraint_mask import apply_constraint_mask
+from lightllm.common.basemodel.triton_kernel.local_vocab_topk import local_vocab_greedy_logprob
 from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context
 from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager
 from lightllm.utils.envs_utils import get_env_start_args
@@ -95,14 +96,12 @@ def sample(
             vocab_size=sampling_params_manager.vocab_size,
             b_mtp_index=b_mtp_index,
         )
-    probs = torch.softmax(logits, dim=-1)
-
     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))
-        return batch_next_token_ids.view(-1), torch.log(batch_next_token_probs).view(-1)
+        return local_vocab_greedy_logprob(logits)
+
+    probs = torch.softmax(logits, dim=-1)
 
-    elif skip_top_k and skip_top_p:
+    if skip_top_k and skip_top_p:
         # topk 等于整个词表，topp 等于1.0，等价于不进行topk topp过滤，直接进行随机采样，可以提升采样速度
         batch_next_token_ids = _random_sample(probs, reqs, exist_req_use_random_seed)
         batch_next_token_probs = torch.gather(probs, dim=1, index=batch_next_token_ids.view(-1, 1))