A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[Bugfix]: fix issue with n>1 sampling on v1 requests overriding each other (#16863)
Signed-off-by: Jeffrey Li <jeffrey.dot.li@gmail.com>
Jeffrey Li committed
Apr 21, 2025 at 23:40 UTC
0e4254492f2367484ac04d8201a3ac35db742656
3 files changed
+104
-20
tests/v1/engine/test_output_processor.py
+81
@@ -921,3 +921,84 @@ async def test_request_output_collector():
921
# Cumulative logprobs should be the last one.
922
cumulative_logprob_expected = 1.0 * num_to_put
923
assert output.outputs[0].cumulative_logprob == cumulative_logprob_expected
924
+
925
+
926
+@pytest.mark.asyncio
927
+async def test_cumulative_output_collector_n():
928
+ """Test collector correctly handles multiple outputs by index."""
929
+ collector = RequestOutputCollector(RequestOutputKind.CUMULATIVE)
930
+ outputs = [
931
+ RequestOutput(
932
+ request_id="my-request-id",
933
+ prompt=None,
934
+ prompt_token_ids=[1, 2, 3],
935
+ prompt_logprobs=None,
936
+ outputs=[
937
+ CompletionOutput(
938
+ index=0,
939
+ text="a",
940
+ token_ids=[0],
941
+ cumulative_logprob=None,
942
+ logprobs=None,
943
+ finish_reason=None,
944
+ ),
945
+ CompletionOutput(
946
+ index=1,
947
+ text="b",
948
+ token_ids=[1],
949
+ cumulative_logprob=None,
950
+ logprobs=None,
951
+ finish_reason=None,
952
+ ),
953
+ ],
954
+ finished=False,
955
+ ),
956
+ RequestOutput(
957
+ request_id="my-request-id",
958
+ prompt=None,
959
+ prompt_token_ids=[1, 2, 3],
960
+ prompt_logprobs=None,
961
+ outputs=[
962
+ CompletionOutput(
963
+ index=0,
964
+ text="ab",
965
+ token_ids=[0, 1],
966
+ cumulative_logprob=None,
967
+ logprobs=None,
968
+ finish_reason=None,
969
+ ),
970
+ CompletionOutput(
971
+ index=2,
972
+ text="c",
973
+ token_ids=[2],
974
+ cumulative_logprob=None,
975
+ logprobs=None,
976
+ finish_reason=None,
977
+ ),
978
+ ],
979
+ finished=False,
980
+ ),
981
+ ]
982
+ for output in outputs:
983
+ collector.put(output)
984
+
985
+ # Get the output and check that the text and token_ids are correct.
986
+ result = await collector.get()
987
+ # We are expecting
988
+ # [{index: 0, text: "ab"}, {index: 1, text: "b"}, {index: 2, text: "c"}]
989
+ assert len(result.outputs) == 3
990
+ # First is the one where index is 0
991
+ first = [k for k in result.outputs if k.index == 0]
992
+ assert len(first) == 1
993
+ assert first[0].text == "ab"
994
+
995
+ # Second is the one where index is 1
996
+ second = [k for k in result.outputs if k.index == 1]
997
+ assert len(second) == 1
998
+ assert second[0].text == "b"
999
+ assert second[0].token_ids == [1]
1000
+
1001
+ # Third is the one where index is 2
1002
+ third = [k for k in result.outputs if k.index == 2]
1003
+ assert len(third) == 1
1004
+ assert third[0].text == "c"
vllm/outputs.py
+20
-14
@@ -134,26 +134,32 @@ class RequestOutput:
134
self.encoder_prompt_token_ids = encoder_prompt_token_ids
135
self.num_cached_tokens = num_cached_tokens
136
137
- def add(self, next_output: "RequestOutput") -> None:
137
+ def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
138
"""Merge subsequent RequestOutput into this one"""
139
140
self.finished |= next_output.finished
141
142
for next_completion in next_output.outputs:
143
- for completion in self.outputs:
143
+ for i, completion in enumerate(self.outputs):
144
if completion.index == next_completion.index:
145
- # Merge outputs with same index
146
- completion.text += next_completion.text
147
- if not isinstance(completion.token_ids, MutableSequence):
148
- completion.token_ids = list(completion.token_ids)
149
- completion.token_ids.extend(next_completion.token_ids)
150
- if next_completion.logprobs:
151
- assert completion.logprobs is not None
152
- completion.logprobs.extend(next_completion.logprobs)
153
- completion.cumulative_logprob = (
154
- next_completion.cumulative_logprob)
155
- completion.finish_reason = next_completion.finish_reason
156
- completion.stop_reason = next_completion.stop_reason
145
+ if aggregate:
146
+ # Merge outputs with same index
147
+ completion.text += next_completion.text
148
+ if not isinstance(completion.token_ids,
149
+ MutableSequence):
150
+ completion.token_ids = list(completion.token_ids)
151
+ completion.token_ids.extend(next_completion.token_ids)
152
+ if next_completion.logprobs:
153
+ assert completion.logprobs is not None
154
+ completion.logprobs.extend(
155
+ next_completion.logprobs)
156
+ completion.cumulative_logprob = (
157
+ next_completion.cumulative_logprob)
158
+ completion.finish_reason = next_completion.finish_reason
159
+ completion.stop_reason = next_completion.stop_reason
160
+ else:
161
+ # Replace the output with the new one
162
+ self.outputs[i] = next_completion
163
break
164
else:
165
self.outputs.append(next_completion)
vllm/v1/engine/output_processor.py
+3
-6
@@ -37,12 +37,9 @@ class RequestOutputCollector:
37
self.output = output
38
self.ready.set()
39
elif isinstance(self.output, RequestOutput):
40
- if self.aggregate:
41
- # Coalesce the outputs in delta case.
42
- self.output.add(output)
43
- else:
44
- # Just replace latest in non-delta case.
45
- self.output = output
40
+ # This ensures that request outputs with different request indexes
41
+ # (if n > 1) do not override each other.
42
+ self.output.add(output, aggregate=self.aggregate)
43
44
async def get(self) -> RequestOutput:
45
"""Get operation blocks on put event."""