test: slim unit tests and repair chat logprob semantics

docs/CN/source/tutorial/function_calling.rst

@@ -53,8 +53,6 @@ Mistral
 
 **格式**: ``[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]``
 
-启用 chat logprobs 时暂不支持 ``mistral`` 工具调用解析器。
-
 DeepSeek-V3
 ~~~~~~~~~~~
 

docs/EN/source/tutorial/function_calling.rst

@@ -53,8 +53,6 @@ Mistral
 
 **Format**: ``[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]``
 
-The ``mistral`` tool-call parser is not supported when chat logprobs are enabled.
-
 DeepSeek-V3
 ~~~~~~~~~~~
 

lightllm/common/basemodel/triton_kernel/logprobs_capture.py

@@ -77,7 +77,10 @@ def _output_logprobs_partial(
     rows,
     candidates,
     sampled_ids,
-    higher_counts,
+    ranks,
+    partial_maxes,
+    partial_sums,
+    partial_posinf_counts,
     stride_row: tl.constexpr,
     vocab_size: tl.constexpr,
     tiles: tl.constexpr,
@@ -91,10 +94,17 @@ def _output_logprobs_partial(
     row_logits = logits + row * stride_row
     ids = tile * BLOCK + tl.arange(0, BLOCK)
     scores = tl.load(row_logits + ids, ids < vocab_size, -float("inf")).to(tl.float32)
-    sampled = tl.load(row_logits + tl.load(sampled_ids + row)).to(tl.float32)
-    # 只数严格更高的 logit。同分不占名次，采样 token 仍算在前 20 里。
-    higher = tl.sum(((ids < vocab_size) & (scores > sampled)).to(tl.int32), 0)
-    tl.store(higher_counts + index * tiles + tile, higher)
+    sampled = tl.load(row_logits + tl.load(sampled_ids + row))
+    rank = tl.sum(((ids < vocab_size) & (scores > sampled)).to(tl.int32), 0)
+    tl.store(ranks + index * tiles + tile, rank)
+    posinf_count = tl.sum(((ids < vocab_size) & (scores == float("inf"))).to(tl.int32), 0)
+    finite_scores = tl.where((scores != float("inf")) & (scores != -float("inf")), scores, -float("inf"))
+    tile_max = tl.max(finite_scores, 0)
+    tile_sum = tl.sum(tl.exp(finite_scores - tile_max), 0)
+    tile_sum = tl.where(tile_max > -float("inf"), tile_sum, 0.0)
+    tl.store(partial_maxes + index * tiles + tile, tile_max)
+    tl.store(partial_sums + index * tiles + tile, tile_sum)
+    tl.store(partial_posinf_counts + index * tiles + tile, posinf_count)
     # tl.topk 只能比较整数。高 32 位放可排序的分数，低 32 位放反转后的 token id。
     # 取最大的键时，分数高的在前；分数相同则 id 更小的在前。词表外的位置写成最小整数。
     bits = scores.to(tl.int32, bitcast=True)
@@ -115,7 +125,11 @@ def _output_logprobs_finish(
     candidates,
     token_ids,
     logprobs,
-    higher_counts,
+    partial_ranks,
+    sampled_ranks,
+    partial_maxes,
+    partial_sums,
+    partial_posinf_counts,
     stride_row: tl.constexpr,
     vocab_size: tl.constexpr,
     tiles: tl.constexpr,
@@ -127,7 +141,6 @@ def _output_logprobs_finish(
     row = tl.load(rows + index * 3).to(tl.int64)
     row_logits = logits + row * stride_row
     sampled = tl.load(row_logits + tl.load(sampled_ids + row))
-    sampled_logprob = tl.load(sampled_logprobs + row)
     width = tl.load(rows + index * 3 + 1)
     offset = tl.load(rows + index * 3 + 2)
     columns = tl.arange(0, TOPK)
@@ -137,16 +150,26 @@ def _output_logprobs_finish(
     # 键的低 32 位是反转后的 token id，这里还原。结果仍按分数从高到低排列。
     ids = 0xFFFFFFFF - (best & 0xFFFFFFFF)
     scores = tl.load(row_logits + ids, (columns < width) & (ids < vocab_size), -float("inf"))
-    stored_ids = tl.where(tl.abs(scores) < float("inf"), ids, -1).to(tl.int64)
-    # 候选和采样 token 的 logsumexp 相同，用采样 logprob 平移即可，不必再做整词表 softmax。
-    row_logprobs = scores - sampled + sampled_logprob
+    stored_ids = tl.where(scores == scores, ids, -1).to(tl.int64)
+    rank_tiles = tl.arange(0, TILE_BLOCK)
+    tile_maxes = tl.load(partial_maxes + index * tiles + rank_tiles, rank_tiles < tiles, -float("inf"))
+    tile_sums = tl.load(partial_sums + index * tiles + rank_tiles, rank_tiles < tiles, 0.0)
+    posinf_counts = tl.load(partial_posinf_counts + index * tiles + rank_tiles, rank_tiles < tiles, 0)
+    row_max = tl.max(tile_maxes, 0)
+    row_sum = tl.sum(tl.where(tile_maxes > -float("inf"), tile_sums * tl.exp(tile_maxes - row_max), 0.0), 0)
+    posinf_count = tl.sum(posinf_counts, 0)
+    finite_logprobs = scores - (row_max + tl.log(row_sum))
+    uniform_logprob = -tl.log(float(vocab_size))
+    row_logprobs = tl.where(
+        posinf_count > 0,
+        tl.where(scores == float("inf"), -tl.log(posinf_count.to(tl.float32)), -float("inf")),
+        tl.where(row_max > -float("inf"), finite_logprobs, uniform_logprob),
+    )
     tl.store(token_ids + offset + columns, stored_ids, columns < width)
     tl.store(logprobs + offset + columns, row_logprobs, columns < width)
-    # 候选平移用的是原始采样 logprob。至少 20 个 token 严格更高时，再把采样分数写成 -9999。
-    higher_tiles = tl.arange(0, TILE_BLOCK)
-    higher = tl.sum(tl.load(higher_counts + index * tiles + higher_tiles, higher_tiles < tiles, 0), 0)
-    outside = (higher >= 20) | ~(tl.abs(sampled) < float("inf"))
-    tl.store(sampled_logprobs + row, -9999.0, mask=outside)
+    rank = 1 + tl.sum(tl.load(partial_ranks + index * tiles + rank_tiles, rank_tiles < tiles, 0), 0)
+    rank = tl.where(sampled > -float("inf"), rank, vocab_size + 1)
+    tl.store(sampled_ranks + index, rank)
 
 
 def capture_output_logprobs(
@@ -156,26 +179,33 @@ def capture_output_logprobs(
     rows: torch.Tensor,
     max_topk: int,
     output_count: int,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """返回每行的 top-k token id 和 logprob。
 
     rows 每项是 (logits 行号, K, 写入起点)。写入起点按 token 个数累加，
-    第 r 行的候选是两个数组的 [起点, 起点 + K)。至少 20 个 token 的 logit
-    严格高于采样 token 时，把 sampled_logprobs 里对应位置写成 -9999。K=0 也判断。
+    第 r 行的候选是两个数组的 [起点, 起点 + K)。另返回每行采样 token 的
+    1-based rank（同分共享排名），不修改内部采样 logprob。K=0 仍计算排名。
     """
     block = 1024
     # 先按词表分段留下局部 top-k，再逐行合并成最终候选。
     tiles = triton.cdiv(logits.shape[1], block)
     # Triton 的 topk 块长必须是 2 的幂。请求的 K 最大 20，这里向上取整；真正写入仍按每行的 K 截断。
     topk = triton.next_power_of_2(max(2, max_topk))
     candidates = logits.new_empty((rows.shape[0], tiles, topk), dtype=torch.int64)
-    higher_counts = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32)
+    partial_ranks = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32)
+    sampled_ranks = logits.new_empty((rows.shape[0],), dtype=torch.int32)
+    partial_maxes = logits.new_empty((rows.shape[0], tiles), dtype=torch.float32)
+    partial_sums = logits.new_empty((rows.shape[0], tiles), dtype=torch.float32)
+    partial_posinf_counts = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32)
     _output_logprobs_partial[(rows.shape[0], tiles)](
         logits=logits,
         rows=rows,
         candidates=candidates,
         sampled_ids=sampled_ids,
-        higher_counts=higher_counts,
+        ranks=partial_ranks,
+        partial_maxes=partial_maxes,
+        partial_sums=partial_sums,
+        partial_posinf_counts=partial_posinf_counts,
         stride_row=logits.stride(0),
         vocab_size=logits.shape[1],
         tiles=tiles,
@@ -193,12 +223,16 @@ def capture_output_logprobs(
         candidates=candidates,
         token_ids=token_ids,
         logprobs=logprobs,
-        higher_counts=higher_counts,
+        partial_ranks=partial_ranks,
+        sampled_ranks=sampled_ranks,
+        partial_maxes=partial_maxes,
+        partial_sums=partial_sums,
+        partial_posinf_counts=partial_posinf_counts,
         stride_row=logits.stride(0),
         vocab_size=logits.shape[1],
         tiles=tiles,
         TOPK=topk,
         TILE_BLOCK=triton.next_power_of_2(tiles),
         num_warps=4,
     )
-    return token_ids, logprobs
+    return token_ids, logprobs, sampled_ranks