diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 2104a40da6..7862ffe8b8 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -885,7 +885,12 @@ async def stream_results() -> AsyncGenerator[bytes, None]: model=request.model, choices=[final_choice], ) - yield f"data: {_serialize_sse_chunk(final_chunk, _final_choice_nulls)}\n\n" + final_logprobs = None + if logprob_buffer is not None: + content_tokens, reasoning_tokens = logprob_buffer.message_logprobs() + if content_tokens or reasoning_tokens: + final_logprobs = {"content": content_tokens, "reasoning_content": reasoning_tokens} + yield f"data: {_serialize_sse_chunk(final_chunk, _final_choice_nulls, logprobs=final_logprobs)}\n\n" reasoning_parser = get_env_start_args().reasoning_parser completion_tokens_details = None diff --git a/lightllm/server/chat_logprob_buffer.py b/lightllm/server/chat_logprob_buffer.py index 2f41449b9c..ad19ce03c4 100644 --- a/lightllm/server/chat_logprob_buffer.py +++ b/lightllm/server/chat_logprob_buffer.py @@ -12,7 +12,8 @@ class ChatLogprobBuffer: """按 token 序列识别分隔符，把可见文本的 logprob 放进正文或思维链。 分隔符可以是一个或多个 token。凑齐才切换通道，分隔符本身不输出。 - 空增量的分数先累加，加到后面第一个非空文本上。生成结束时还没拼上的空增量丢掉。 + 每个采样位置单独一条 logprob。未完成的 UTF-8 字节文本为空，字符文本落在补齐的那个 + token 上，分数不相加。特殊 token 的空文本不输出。 """ def __init__( @@ -28,7 +29,6 @@ def __init__( self._separate_reasoning = separate_reasoning self._channel = "content" # 当前输出通道：content、reasoning，或 commentary 的 drop。 self._in_tool = False # 是否处在工具调用开始和结束分隔符之间。 - self._pending_logprob = 0.0 # 尚未拼成字符的空增量分数，等下一个非空文本一起输出。 self._reasoning: list[dict[str, Any]] = [] # 分开思考时的思维链 logprob。 self._content: list[dict[str, Any]] = [] # 正文 logprob；不分开思考时思维链也放这里。 self._sequences: list[tuple[tuple[int, ...], str]] = [] # 分隔符 token id 序列和切换动作。 @@ -47,7 +47,6 @@ def record(self, text: str, metadata: dict[str, Any], finish_status: FinishStatu self._resolve_held(closing=False) if finish_status.is_finished(): self._resolve_held(closing=True) - self._pending_logprob = 0.0 def reasoning_logprobs(self) -> dict[str, Any] | None: tokens = self._drain(items=self._reasoning) @@ -193,16 +192,11 @@ def _release(self, item: _Held) -> None: self._save(text=text, score=score, metadata=metadata, token_id=token_id) def _save(self, text: str, score: float, metadata: dict[str, Any], token_id: int) -> None: - if not text: - # 未完成的 UTF-8 字节并进后面的字符。特殊 token 的空文本直接丢掉。 - if not self._is_special(metadata=metadata, token_id=token_id): - self._pending_logprob += score + if not text and self._is_special(metadata=metadata, token_id=token_id): return - logprob = self._pending_logprob + score - self._pending_logprob = 0.0 token = { "token": text, - "logprob": logprob, + "logprob": score, "bytes": None if "\ufffd" in text else list(text.encode("utf-8")), "top_logprobs": [ _candidate_logprob(tokenizer=self.tokenizer, token_id=candidate, logprob=value) diff --git a/test.md b/test.md deleted file mode 100644 index f1a6298685..0000000000 --- a/test.md +++ /dev/null @@ -1,58 +0,0 @@ -# ReplaySSM 折回：精度确认与两种时机的对比 - -测量对象是当前工作区的实现：普通 decode 放不下时，折回发生在每一层正在算线性注意力的 `replay` kernel 开头。容量够时只重放，不写 canonical SSM。Qwen3.5-0.8B，思考关闭，GSM8K 1000 题，`max_tokens=512`。 - -## 1. 1000 题上，现在这条 kernel 没有原先的精度崩溃 - -崩溃的样子是答案尾部刷出长串 `!` 后被 512 token 截断。同一套协议下，现在三条 replay 配置的 `!` 循环都是 0，strict 正确数贴着 native fp32。 - -| 配置 | strict | flexible | 截断 | `!` 循环 | 与 native 最终数字一致 | -|---|---:|---:|---:|---:|---:| -| native fp32 | 519 | 542 | 74 | 0 | 1000 | -| inline bf16，cache 4 | 511 | 534 | 81 | 0 | 807 | -| inline bf16，cache 16 | 515 | 536 | 63 | 0 | 810 | -| inline fp32，cache 16 | 518 | 533 | 67 | 0 | 803 | - -同一批题上，改层内折回之前： - -| 配置 | strict | 截断 | `!` 循环 | -|---|---:|---:|---:| -| inline bf16，cache 4 | 336 | 403 | 383 | -| inline bf16，cache 16 | 444 | 134 | 60 | -| inline fp32，cache 16 | 481 | 126 | 58 | - -和 native 对不上的大约两成，是回放和 fused kernel 的正常差别。cache 4 的正确数从 336 回到 511，已经贴着 native 的 519。这次没有重跑 precompute 和 KDA。 - -## 2. 整步结束后再折回，和层内折回的差别 - -「整步后再折回」指旧路径：每层 replay 先把 token 算完，最后一层 `maybe_compress_prefix` 再用独立 kernel `_fold_prefix` 折回，并在该 kernel 内把 cursor 写成 0。「层内折回」指现在的路径：放不下时，折回在每一层的 replay kernel 开头完成，共享 cursor 等全部层结束才在外面更新。 - -### 精度 - -独立折回如果等全部线程块写完再改 cursor，和层内折回是同一个结果。bf16 和 fp32、cache 4 和 cache 16，一步边界 token 的输出和 SSM 最大误差都是 0。 - -旧路径不是这样。`WHEN_FULL` 由其中一个线程块在 kernel 内部把 cursor 写成 0，其他线程块可能还没读到原来的长度，对应的层或头就跳过折回。和正确折回相比，状态最大误差： - -| 批量 | cache | 20 次里对不上 | 状态最大误差 | -|---|---:|---:|---:| -| 1 | 4 | 1 | 0.125 | -| 1 | 16 | 11 | 0.132 | -| 16 | 4 | 20 | 0.218 | -| 16 | 16 | 20 | 0.268 | - -这个误差会留在 SSM 里。cache 4 的 1000 题里有 530 题最终数字因此改变，其中 379 题改到与 native 一致，44 题离开 native。 - -### 性能 - -形状是 Qwen3.5-0.8B 的 18 个线性层、16 头、维度 128、bf16，CUDA graph。时间是 18 层加在一起。普通步是前缀已满但还没触发折回；折回步是刚好放不下的那一步。 - -| 批量 | cache | 普通回放 | 层内折回再回放 | 整步后的折回 kernel | 每步都有的空折回 | -|---|---:|---:|---:|---:|---:| -| 1 | 4 | 106 µs | 178 µs | 12 µs | 3.3 µs | -| 1 | 16 | 140 µs | 361 µs | 31 µs | 3.4 µs | -| 16 | 4 | 180 µs | 309 µs | 161 µs | 13 µs | -| 16 | 16 | 283 µs | 732 µs | 400 µs | 13 µs | - -单条请求时，折回串在每一层的 replay 里。cache 4 的那一步大约多 70 µs；独立 kernel 把 18 层一起折完只要 12 µs。摊到每个 token 大约多 10 µs。批量 16 时，独立折回本身要 160–400 µs，放进已经要启动的层 kernel 里，每个 token 大约少 10 µs。 - -旧路径在容量够的时候仍会每步启动一次空折回，批量 1 约 3 µs，批量 16 约 13 µs。现在这条路径不再发这个 kernel。相对完整 decode（还有 full attention 和 MLP），这几十微秒不是端到端的大头。1000 题墙钟时间现在是 55–61 秒，native 是 54 秒。 diff --git a/unit_tests/server/test_chat_output_logprobs.py b/unit_tests/server/test_chat_output_logprobs.py index 18d4e64b11..f3bf4513e9 100644 --- a/unit_tests/server/test_chat_output_logprobs.py +++ b/unit_tests/server/test_chat_output_logprobs.py @@ -107,8 +107,11 @@ async def generate(prompt, params, *args, **kwargs): assert [entry["token"] for entry in reasoning] == (["think"] if separate_reasoning else []) expected_content = "你 yes" if separate_reasoning else "think你 yes" assert bytes(byte for entry in content for byte in entry["bytes"]).decode() == expected_content - assert [entry["token"] for entry in content] == (["你", " yes"] if separate_reasoning else ["think", "你", " yes"]) - assert [entry["logprob"] for entry in content] == ([-0.75, -0.25] if separate_reasoning else [-0.25, -0.75, -0.25]) + assert [entry["token"] for entry in content] == ( + ["", "", "你", " yes"] if separate_reasoning else ["think", "", "", "你", " yes"] + ) + expected_logprobs = [-0.25, -0.25, -0.25, -0.25] if separate_reasoning else [-0.25] * 5 + assert [entry["logprob"] for entry in content] == expected_logprobs assert [entry["logprob"] for entry in reasoning] == ([-0.25] if separate_reasoning else []) expected = [{"token": " yes", "logprob": -0.125, "bytes": [32, 121, 101, 115]}] if topk > 0 else [] assert all(entry["top_logprobs"] == expected for entry in content + reasoning) @@ -287,13 +290,13 @@ async def test_chat_keeps_the_sampled_logprob(chat_response, stream, topk, rank) @pytest.mark.anyio @pytest.mark.parametrize("stream", [False, True]) -async def test_empty_increments_fuse_into_the_following_text(chat_response, stream): +async def test_empty_increments_keep_their_own_logprob(chat_response, stream): tokenizer = make_tokenizer(["<0xE4>", "<0xBD>", "<0xA0>"]) text, _, content, _ = await chat_response(tokenizer, [(0, ""), (1, ""), (2, "你")], stream=stream) assert text == "你" - assert [entry["token"] for entry in content] == ["你"] - assert content[0]["logprob"] == pytest.approx(-2.401099 * 3) - assert content[0]["bytes"] == list("你".encode()) + assert [entry["token"] for entry in content] == ["", "", "你"] + assert [entry["logprob"] for entry in content] == pytest.approx([-2.401099, -2.401099, -2.401099]) + assert [entry["bytes"] for entry in content] == [[], [], list("你".encode())] @pytest.mark.anyio @@ -424,7 +427,7 @@ async def test_tool_delimiter_fused_into_another_token_is_rejected(chat_response await chat_response(tokenizer, list(enumerate(pieces)), stream=False, tools=tools) -def test_empty_increments_fuse_and_trailing_ones_are_dropped(): +def test_empty_increments_keep_each_position_including_a_trailing_one(): from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: "") @@ -432,12 +435,14 @@ def test_empty_increments_fuse_and_trailing_ones_are_dropped(): for token_id, text in enumerate(["", "", "你"]): buffer.record(text, {"id": token_id, "logprob": -0.25}, FinishStatus()) content, _ = buffer.message_logprobs() - assert [token["token"] for token in content] == ["你"] - assert content[0]["logprob"] == -0.75 - assert content[0]["bytes"] == list("你".encode()) + assert [token["token"] for token in content] == ["", "", "你"] + assert [token["logprob"] for token in content] == [-0.25, -0.25, -0.25] + assert content[2]["bytes"] == list("你".encode()) buffer.record("", {"id": 3, "logprob": -0.5}, FinishStatus(FinishStatus.FINISHED_STOP)) - assert buffer.message_logprobs() == (None, None) + content, reasoning = buffer.message_logprobs() + assert [(token["token"], token["logprob"]) for token in content] == [("", -0.5)] + assert reasoning is None def test_reasoning_delimiter_must_be_one_token():