From ca0f0865d03ae4196dcf6eb50c9820e7289ef06c Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Mon, 2 Nov 2026 17:49:23 +0800 Subject: [PATCH 1/6] perf(sampling): fuse greedy vocabulary selection and logprob --- .../triton_kernel/local_vocab_topk.py | 72 +++++++++++++++++++ .../mode_backend/generic_post_process.py | 11 ++- 2 files changed, 77 insertions(+), 6 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py index 5fcbb66ee9..2dbb5c24b2 100644 --- a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py +++ b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py @@ -108,3 +108,75 @@ def local_vocab_topk( ) # 词表现在位于连续的最后一维；只需要候选集合，不额外要求候选按分数排序。 return torch.topk(transposed_logits, k=top_k, dim=1, sorted=False) + + +@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.max(scores, axis=0) + total = tl.sum(tl.exp(scores - maximum), axis=0) + winner = tl.min(tl.where(scores == maximum, 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.max(maximums, axis=0) + total = tl.sum(partial_sums * tl.exp(maximums - maximum), axis=0) + winner = tl.min(tl.where(maximums == maximum, 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 diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 90c96ecf2f..585f1b274e 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/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)) From 944386e0b6e913f646f8fdf6a2bea328e2f25601 Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Tue, 3 Nov 2026 14:34:18 +0800 Subject: [PATCH 2/6] =?UTF-8?q?fix(sampling):=20=E9=81=BF=E5=85=8D?= =?UTF-8?q?=E6=95=B4=E5=9D=97=E5=B1=8F=E8=94=BD=E5=AF=BC=E8=87=B4=20greedy?= =?UTF-8?q?=20logprob=20=E4=B8=BA=20NaN?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- lightllm/common/basemodel/triton_kernel/local_vocab_topk.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py index 2dbb5c24b2..415ad87d3e 100644 --- a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py +++ b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py @@ -125,7 +125,8 @@ def _greedy_vocab_partial( v = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) scores = tl.load(logits + b * stride_b + v, v < vocab, -float("inf")) maximum = tl.max(scores, axis=0) - total = tl.sum(tl.exp(scores - maximum), axis=0) + # Fully masked chunks contribute zero instead of exp(-inf - -inf) = NaN. + total = tl.sum(tl.exp(scores - tl.where(maximum == -float("inf"), 0.0, maximum)), axis=0) winner = tl.min(tl.where(scores == maximum, v, 2147483647), axis=0) offset = b * chunks + tl.program_id(1) tl.store(maxima + offset, maximum) From c4570b3ec4652ab108f12d27662591a3752d895f Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Tue, 3 Nov 2026 14:40:28 +0800 Subject: [PATCH 3/6] =?UTF-8?q?refactor(sampling):=20=E5=88=86=E5=BC=80?= =?UTF-8?q?=E8=AE=A1=E7=AE=97=E9=83=A8=E5=88=86=E5=92=8C=E4=B8=8E=E5=B1=8F?= =?UTF-8?q?=E8=94=BD=E5=9D=97=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- lightllm/common/basemodel/triton_kernel/local_vocab_topk.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py index 415ad87d3e..c125379cb4 100644 --- a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py +++ b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py @@ -126,7 +126,8 @@ def _greedy_vocab_partial( scores = tl.load(logits + b * stride_b + v, v < vocab, -float("inf")) maximum = tl.max(scores, axis=0) # Fully masked chunks contribute zero instead of exp(-inf - -inf) = NaN. - total = tl.sum(tl.exp(scores - tl.where(maximum == -float("inf"), 0.0, maximum)), axis=0) + total = tl.sum(tl.exp(scores - maximum), axis=0) + total = tl.where(maximum == -float("inf"), 0.0, total) winner = tl.min(tl.where(scores == maximum, v, 2147483647), axis=0) offset = b * chunks + tl.program_id(1) tl.store(maxima + offset, maximum) From f7745932a5e1462ef0f81ce3f15c214c383f8518 Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Tue, 3 Nov 2026 15:05:41 +0800 Subject: [PATCH 4/6] =?UTF-8?q?fix(sampling):=20=E6=8B=92=E7=BB=9D?= =?UTF-8?q?=E9=9D=9E=E6=9C=89=E9=99=90=E9=87=87=E6=A0=B7=E5=8F=82=E6=95=B0?= =?UTF-8?q?=E5=B9=B6=E5=AF=B9=E9=BD=90=20NaN=20=E9=80=89=E8=AF=8D=E8=A1=8C?= =?UTF-8?q?=E4=B8=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../basemodel/triton_kernel/local_vocab_topk.py | 13 +++++++++---- lightllm/server/core/objs/sampling_params.py | 4 ++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py index c125379cb4..8cfc9e9224 100644 --- a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py +++ b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py @@ -110,6 +110,11 @@ def local_vocab_topk( return torch.topk(transposed_logits, k=top_k, dim=1, sorted=False) +@triton.jit +def _greedy_maximum(a, b): + return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL) + + @triton.jit def _greedy_vocab_partial( logits, @@ -124,11 +129,11 @@ def _greedy_vocab_partial( 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.max(scores, axis=0) + maximum = tl.reduce(scores, axis=0, combine_fn=_greedy_maximum) # 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) - winner = tl.min(tl.where(scores == maximum, v, 2147483647), axis=0) + winner = tl.min(tl.where((v < vocab) & ((scores == maximum) | (scores != scores)), v, 2147483647), axis=0) offset = b * chunks + tl.program_id(1) tl.store(maxima + offset, maximum) tl.store(sums + offset, total) @@ -142,9 +147,9 @@ def _greedy_vocab_finish(maxima, sums, indexes, token_ids, logprobs, chunks: tl. 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.max(maximums, axis=0) + maximum = tl.reduce(maximums, axis=0, combine_fn=_greedy_maximum) total = tl.sum(partial_sums * tl.exp(maximums - maximum), axis=0) - winner = tl.min(tl.where(maximums == maximum, ids, 2147483647), axis=0) + winner = tl.min(tl.where((c < chunks) & ((maximums == maximum) | (maximums != maximums)), ids, 2147483647), axis=0) tl.store(token_ids + b, winner.to(tl.int64)) tl.store(logprobs + b, -tl.log(total)) diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py index 778007ecee..663a5e1b50 100644 --- a/lightllm/server/core/objs/sampling_params.py +++ b/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 @@ -454,6 +455,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: From 58fcee4658f2c21657f6db9c1724b2d46bef2f9b Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Tue, 3 Nov 2026 15:14:00 +0800 Subject: [PATCH 5/6] =?UTF-8?q?refactor(sampling):=20=E5=B1=95=E5=BC=80=20?= =?UTF-8?q?greedy=20=E5=80=99=E9=80=89=E5=88=A4=E6=96=AD=E4=BB=A5=E6=8F=90?= =?UTF-8?q?=E9=AB=98=E5=8F=AF=E8=AF=BB=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../triton_kernel/local_vocab_topk.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py index 8cfc9e9224..7516c32f26 100644 --- a/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py +++ b/lightllm/common/basemodel/triton_kernel/local_vocab_topk.py @@ -111,7 +111,7 @@ def local_vocab_topk( @triton.jit -def _greedy_maximum(a, b): +def _maximum_propagating_nan(a, b): return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL) @@ -129,11 +129,16 @@ def _greedy_vocab_partial( 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=_greedy_maximum) + 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) - winner = tl.min(tl.where((v < vocab) & ((scores == maximum) | (scores != scores)), v, 2147483647), axis=0) + # 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) @@ -147,9 +152,13 @@ def _greedy_vocab_finish(maxima, sums, indexes, token_ids, logprobs, chunks: tl. 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=_greedy_maximum) + maximum = tl.reduce(maximums, axis=0, combine_fn=_maximum_propagating_nan) total = tl.sum(partial_sums * tl.exp(maximums - maximum), axis=0) - winner = tl.min(tl.where((c < chunks) & ((maximums == maximum) | (maximums != maximums)), ids, 2147483647), 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)) From 3eca9aa80b01c4ecc923d9adff45e4d01802136c Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Tue, 3 Nov 2026 15:25:02 +0800 Subject: [PATCH 6/6] =?UTF-8?q?fix(sampling):=20=E4=BB=8E=E5=8F=82?= =?UTF-8?q?=E6=95=B0=E5=85=A5=E5=8F=A3=E5=92=8C=20EOS=20penalty=20?= =?UTF-8?q?=E8=BF=90=E7=AE=97=E6=B6=88=E9=99=A4=20NaN=20=E6=9D=A5=E6=BA=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../basemodel/triton_kernel/post_process/apply_penalty.py | 6 +++++- .../triton_kernel/post_process/apply_penalty_gpu_cache.py | 6 +++++- lightllm/server/core/objs/sampling_params.py | 2 ++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty.py b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty.py index f31b3994dd..197763ab40 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty.py +++ b/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 diff --git a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py index a46fd7436c..bcd7b1041f 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py +++ b/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 diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py index 663a5e1b50..c157e23997 100644 --- a/lightllm/server/core/objs/sampling_params.py +++ b/lightllm/server/core/objs/sampling_params.py @@ -262,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