[Perf][Frontend] Defer reasoning usage recounts for non-continuous chat streams (#56067)
Signed-off-by: Cheng Rui <286040359@qq.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>
Cheng Rui committed
Sep 26, 2026 at 05:08 UTC
8b47e8b22caf0022b00df3b421cf2d41419a04a7
2 files changed
+341
-29
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+327
-25
@@ -798,52 +798,354 @@ async def test_chat_streaming_metrics_ride_on_usage_chunk():
798
799
800
@pytest.mark.asyncio
801
-async def test_streaming_reasoning_usage_counts_across_deltas():
802
- serving = _build_minimal_metrics_serving_chat(enable_per_request_metrics=False)
803
- serving._include_reasoning_tokens_details = True
804
- serving.model_config = None
805
-
806
- parser = MagicMock()
807
- parser.parse_delta.side_effect = [
808
- DeltaMessage(reasoning="reasoning"),
809
- DeltaMessage(content="answer"),
810
- ]
811
- parser.count_reasoning_tokens.side_effect = lambda token_ids: sum(
812
- token_id == 20 for token_id in token_ids
801
+@pytest.mark.parametrize("n", [1, 2])
802
+@pytest.mark.parametrize("include_reasoning", [True, False])
803
+@pytest.mark.parametrize("include_details", [True, False])
804
+@pytest.mark.parametrize("separate_finish", [True, False])
805
+@pytest.mark.parametrize(
806
+ "stream_options, force_usage, include_usage, continuous_usage",
807
+ [
808
+ (None, False, False, False),
809
+ ({"include_usage": True}, False, True, False),
810
+ ({"include_usage": True, "continuous_usage_stats": True}, False, True, True),
811
+ ({"include_usage": False, "continuous_usage_stats": True}, False, False, False),
812
+ (None, True, True, True),
813
+ ],
814
+)
815
+async def test_streaming_reasoning_usage_counts_across_deltas(
816
+ n,
817
+ include_reasoning,
818
+ include_details,
819
+ separate_finish,
820
+ stream_options,
821
+ force_usage,
822
+ include_usage,
823
+ continuous_usage,
824
+):
825
+ """Final usage needs one recount per choice, not per streamed delta."""
826
+ serving = _build_minimal_metrics_serving_chat(
827
+ enable_per_request_metrics=False, enable_force_include_usage=force_usage
828
)
814
- serving.parser_cls = MagicMock(return_value=parser)
829
+ serving._include_reasoning_tokens_details = include_details
830
+ serving.model_config = None
831
816
- first = _make_metrics_request_output(metrics=None, token_ids=(10, 20))
817
- first.outputs[0].text = "<think>reasoning"
818
- first.outputs[0].finish_reason = None
819
- second = _make_metrics_request_output(metrics=None, token_ids=(11, 30))
820
- second.outputs[0].text = "</think>answer"
832
+ parsers = [MagicMock() for _ in range(n)]
833
+ request_outputs = []
834
+ num_steps = 3 if separate_finish else 2
835
+ for parser in parsers:
836
+ parser.parse_delta.side_effect = [
837
+ DeltaMessage(reasoning="reasoning") if include_reasoning else None,
838
+ DeltaMessage(content="answer"),
839
+ DeltaMessage(),
840
+ ]
841
+ parser.count_reasoning_tokens.side_effect = lambda token_ids: sum(
842
+ token_id == 20 for token_id in token_ids
843
+ )
844
+ serving.parser_cls = MagicMock(side_effect=parsers)
845
+ for step in range(num_steps):
846
+ for i in reversed(range(n)):
847
+ tokens = ((10,) + (20,) * (i + 1), (11, 20, 30), ())[step]
848
+ result = _make_metrics_request_output(metrics=None, token_ids=tokens)
849
+ result.outputs[0].index = i
850
+ result.outputs[0].text = ("<think>reasoning", "</think>answer", "")[step]
851
+ finished = step == num_steps - 1
852
+ result.outputs[0].finish_reason = "stop" if finished else None
853
+ result.finished = finished and i == 0
854
+ request_outputs.append(result)
855
856
request = ChatCompletionRequest(
857
model="test-model",
858
messages=[{"role": "user", "content": "Test prompt"}],
859
max_tokens=10,
860
stream=True,
827
- stream_options={"include_usage": True},
861
+ n=n,
862
+ include_reasoning=include_reasoning,
863
+ return_token_ids=True,
864
+ stream_options=stream_options,
865
)
866
+ metadata = RequestResponseMetadata(request_id="chatcmpl-test-id")
867
chunks: list[dict[str, Any]] = []
868
async for line in serving.chat_completion_stream_generator(
869
request,
832
- _stream_request_outputs(first, second),
870
+ _stream_request_outputs(*request_outputs),
871
"chatcmpl-test-id",
872
"test-model",
873
conversation=[{"role": "user", "content": "Test"}],
874
tokenizer=MagicMock(),
837
- request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"),
875
+ request_metadata=metadata,
876
):
877
payload = line.removeprefix("data: ").strip()
878
if payload != "[DONE]":
879
chunks.append(json.loads(payload))
880
843
- usage_chunks = [chunk for chunk in chunks if chunk.get("usage")]
844
- assert usage_chunks[-1]["usage"]["completion_tokens_details"] == {
845
- "reasoning_tokens": 1
846
- }
881
+ assert all("error" not in chunk for chunk in chunks)
882
+ expected_count = sum(i + 2 for i in range(n))
883
+ expected_details = {"reasoning_tokens": expected_count} if include_details else None
884
+ assert metadata.final_usage_info is not None
885
+ assert metadata.final_usage_info.model_dump()["completion_tokens_details"] == (
886
+ expected_details
887
+ )
888
+ final_chunks = [chunk for chunk in chunks if not chunk["choices"]]
889
+ assert len(final_chunks) == int(include_usage)
890
+ if include_usage:
891
+ assert final_chunks[0]["usage"].get("completion_tokens_details") == (
892
+ expected_details
893
+ )
894
+ for i, parser in enumerate(parsers):
895
+ expected_calls = (
896
+ (num_steps if continuous_usage else 1) if include_details else 0
897
+ )
898
+ assert parser.count_reasoning_tokens.call_count == expected_calls
899
+ if include_details:
900
+ assert parser.count_reasoning_tokens.call_args.args == (
901
+ (10,) + (20,) * (i + 1) + (11, 20, 30),
902
+ )
903
+ choices = [
904
+ chunk
905
+ for chunk in chunks
906
+ if chunk["choices"] and chunk["choices"][0]["index"] == i
907
+ ]
908
+ assert (
909
+ "".join(c["choices"][0]["delta"].get("content") or "" for c in choices)
910
+ == "answer"
911
+ )
912
+ if not include_reasoning:
913
+ assert all(c["choices"][0].get("token_ids") is None for c in choices)
914
+ if continuous_usage and include_details:
915
+ counts = [
916
+ c["usage"]["completion_tokens_details"]["reasoning_tokens"]
917
+ for c in choices
918
+ ]
919
+ expected_counts = [0, i + 1] if include_reasoning else [0]
920
+ assert counts == expected_counts + [i + 2] * (num_steps - 1)
921
+
922
+
923
+def _reasoning_lifecycle_output(
924
+ index: int = 0,
925
+ tokens: tuple[int, ...] = (20, 30),
926
+ text: str = "answer",
927
+ finish: str | None = None,
928
+) -> RequestOutput:
929
+ result = _make_metrics_request_output(None, tokens)
930
+ result.outputs[0].index = index
931
+ result.outputs[0].text = text
932
+ result.outputs[0].finish_reason = finish
933
+ result.finished = finish is not None
934
+ return result
935
+
936
+
937
+_EMPTY_OUTPUTS = _reasoning_lifecycle_output(tokens=(), text="")
938
+_EMPTY_OUTPUTS.outputs = []
939
+
940
+_REASONING_LIFECYCLE_CASES = [
941
+ ("empty-generator", 1, [], 0, 0),
942
+ ("empty-outputs", 1, [_EMPTY_OUTPUTS], 0, 0),
943
+ (
944
+ "empty-finish",
945
+ 1,
946
+ [_reasoning_lifecycle_output(tokens=(), text="", finish="stop")],
947
+ 0,
948
+ 0,
949
+ ),
950
+ (
951
+ "empty-prefill",
952
+ 1,
953
+ [
954
+ _reasoning_lifecycle_output(tokens=(), text=""),
955
+ _reasoning_lifecycle_output(finish="stop"),
956
+ ],
957
+ 2,
958
+ 1,
959
+ ),
960
+ (
961
+ "one-empty-choice",
962
+ 2,
963
+ [_reasoning_lifecycle_output(index=1, finish="stop")],
964
+ 2,
965
+ 1,
966
+ ),
967
+ (
968
+ "uneven-finish",
969
+ 2,
970
+ [
971
+ _reasoning_lifecycle_output(index=1, finish="stop"),
972
+ _reasoning_lifecycle_output(),
973
+ _reasoning_lifecycle_output(tokens=(20,), finish="stop"),
974
+ ],
975
+ 5,
976
+ 3,
977
+ ),
978
+ (
979
+ "empty-final-delta",
980
+ 1,
981
+ [
982
+ _reasoning_lifecycle_output(),
983
+ _reasoning_lifecycle_output(tokens=(), text="", finish="stop"),
984
+ ],
985
+ 2,
986
+ 1,
987
+ ),
988
+ (
989
+ "ignore-finished-choice",
990
+ 2,
991
+ [
992
+ _reasoning_lifecycle_output(index=1, finish="stop"),
993
+ _reasoning_lifecycle_output(index=1, tokens=(20, 20, 20)),
994
+ _reasoning_lifecycle_output(finish="stop"),
995
+ ],
996
+ 4,
997
+ 2,
998
+ ),
999
+ ("exception-before", 1, [RuntimeError("fixture engine error")], None, None),
1000
+ (
1001
+ "exception-after",
1002
+ 1,
1003
+ [_reasoning_lifecycle_output(), RuntimeError("fixture engine error")],
1004
+ None,
1005
+ None,
1006
+ ),
1007
+ (
1008
+ "generation-error",
1009
+ 1,
1010
+ [
1011
+ _reasoning_lifecycle_output(),
1012
+ _reasoning_lifecycle_output(finish="error"),
1013
+ ],
1014
+ None,
1015
+ None,
1016
+ ),
1017
+ (
1018
+ "cancel-after",
1019
+ 1,
1020
+ [_reasoning_lifecycle_output(), asyncio.CancelledError()],
1021
+ None,
1022
+ None,
1023
+ ),
1024
+ (
1025
+ "close-after",
1026
+ 1,
1027
+ [
1028
+ _reasoning_lifecycle_output(),
1029
+ _reasoning_lifecycle_output(finish="stop"),
1030
+ ],
1031
+ None,
1032
+ None,
1033
+ ),
1034
+]
1035
+
1036
+
1037
+@pytest.mark.asyncio
1038
+@pytest.mark.parametrize("usage_mode", ["none", "final", "continuous"])
1039
+@pytest.mark.parametrize(
1040
+ "name, n, outputs, expected_tokens, expected_reasoning",
1041
+ _REASONING_LIFECYCLE_CASES,
1042
+ ids=[case[0] for case in _REASONING_LIFECYCLE_CASES],
1043
+)
1044
+async def test_streaming_reasoning_usage_lifecycle_edges(
1045
+ usage_mode,
1046
+ name,
1047
+ n,
1048
+ outputs,
1049
+ expected_tokens,
1050
+ expected_reasoning,
1051
+):
1052
+ """Cancel, errors, and empty choices still match baseline stream usage."""
1053
+ serving = _build_minimal_metrics_serving_chat(enable_per_request_metrics=False)
1054
+ serving.model_config = None
1055
+ serving._include_reasoning_tokens_details = True
1056
+ parsers = [MagicMock() for _ in range(n)]
1057
+ for parser in parsers:
1058
+ parser.parse_delta.side_effect = lambda **kwargs: DeltaMessage(
1059
+ content=kwargs["delta_text"]
1060
+ )
1061
+ parser.count_reasoning_tokens.side_effect = lambda ids: ids.count(20)
1062
+ serving.parser_cls = MagicMock(side_effect=parsers)
1063
+ stream_options = (
1064
+ None
1065
+ if usage_mode == "none"
1066
+ else {
1067
+ "include_usage": True,
1068
+ "continuous_usage_stats": usage_mode == "continuous",
1069
+ }
1070
+ )
1071
+ request = ChatCompletionRequest(
1072
+ model="test-model",
1073
+ messages=[{"role": "user", "content": "Test prompt"}],
1074
+ stream=True,
1075
+ n=n,
1076
+ return_token_ids=True,
1077
+ stream_options=stream_options,
1078
+ )
1079
+ metadata = RequestResponseMetadata(request_id="chatcmpl-lifecycle")
1080
+
1081
+ async def source():
1082
+ for result in outputs:
1083
+ if isinstance(result, BaseException):
1084
+ raise result
1085
+ yield result
1086
+
1087
+ iterator = source()
1088
+ stream = serving.chat_completion_stream_generator(
1089
+ request,
1090
+ iterator,
1091
+ "chatcmpl-lifecycle",
1092
+ "test-model",
1093
+ conversation=[{"role": "user", "content": "Test"}],
1094
+ tokenizer=MagicMock(),
1095
+ request_metadata=metadata,
1096
+ )
1097
+ chunks: list[object] = []
1098
+ outcome = "exhausted"
1099
+ try:
1100
+ async for line in stream:
1101
+ raw = line.removeprefix("data: ").strip()
1102
+ if raw == "[DONE]":
1103
+ chunks.append(raw)
1104
+ continue
1105
+ item = json.loads(raw)
1106
+ chunks.append(item)
1107
+ if name == "close-after" and any(
1108
+ choice["delta"].get("content") for choice in item.get("choices", [])
1109
+ ):
1110
+ await stream.aclose()
1111
+ outcome = "closed"
1112
+ break
1113
+ except asyncio.CancelledError:
1114
+ outcome = "cancelled"
1115
+ finally:
1116
+ await stream.aclose()
1117
+ await iterator.aclose()
1118
+
1119
+ error_chunks = [c for c in chunks if isinstance(c, dict) and "error" in c]
1120
+ call_count = sum(parser.count_reasoning_tokens.call_count for parser in parsers)
1121
+ if expected_tokens is not None:
1122
+ assert outcome == "exhausted" and chunks[-1] == "[DONE]" and not error_chunks
1123
+ assert metadata.final_usage_info is not None
1124
+ assert metadata.final_usage_info.completion_tokens == expected_tokens
1125
+ assert (
1126
+ metadata.final_usage_info.completion_tokens_details.reasoning_tokens
1127
+ == expected_reasoning
1128
+ )
1129
+ nonempty_choices = sum(
1130
+ bool(parser.count_reasoning_tokens.call_args) for parser in parsers
1131
+ )
1132
+ if usage_mode == "continuous":
1133
+ assert call_count >= nonempty_choices
1134
+ else:
1135
+ assert call_count == nonempty_choices
1136
+ return
1137
+
1138
+ assert metadata.final_usage_info is None
1139
+ if name in {"cancel-after", "close-after"}:
1140
+ assert outcome == ("cancelled" if name == "cancel-after" else "closed")
1141
+ assert "[DONE]" not in chunks and not error_chunks
1142
+ else:
1143
+ assert outcome == "exhausted" and chunks[-1] == "[DONE]"
1144
+ assert len(error_chunks) == 1
1145
+ if usage_mode != "continuous" or name == "exception-before":
1146
+ assert call_count == 0
1147
+ else:
1148
+ assert call_count >= 1
1149
1150
1151
@pytest.mark.asyncio
vllm/entrypoints/openai/chat_completion/serving.py
+14
-4
@@ -672,11 +672,14 @@ class OpenAIServingChat(GenerateBaseServing):
672
673
# set the previous values for the next iteration
674
previous_num_tokens[i] += len(output.token_ids)
675
- if parser is not None:
675
+ if parser is not None and self._include_reasoning_tokens_details:
676
generated_token_ids[i].extend(output.token_ids)
677
- previous_reasoning_tokens[i] = parser.count_reasoning_tokens(
678
- tuple(generated_token_ids[i])
679
- )
677
+ if include_continuous_usage:
678
+ previous_reasoning_tokens[i] = (
679
+ parser.count_reasoning_tokens(
680
+ tuple(generated_token_ids[i])
681
+ )
682
+ )
683
684
# if the message delta is None (e.g. because it was a
685
# "control token" for tool calls or the parser otherwise
@@ -820,6 +823,13 @@ class OpenAIServingChat(GenerateBaseServing):
823
data = chunk.model_dump_json(exclude_unset=True)
824
yield f"data: {data}\n\n"
825
826
+ if self._include_reasoning_tokens_details and not include_continuous_usage:
827
+ for i, parser in enumerate(parsers):
828
+ if parser is not None and generated_token_ids[i]:
829
+ previous_reasoning_tokens[i] = parser.count_reasoning_tokens(
830
+ tuple(generated_token_ids[i])
831
+ )
832
+
833
# once the final token is handled, if stream_options.include_usage
834
# is sent, send the usage
835
if include_usage: