faster

Author: Your Name

tests/beam/__init__.py

@@ -0,0 +1,28 @@
+import pytest
+
+from vllm.beam.beam import BeamScorer
+from vllm.entrypoints.openai.protocol import CompletionResponse, CompletionResponseChoice, EmbeddingResponse, UsageInfo
+
+classi_idx = {
+    "annotations_sexually_suggestive": 0,
+    "annotations_racist": 1,
+}
+
+@pytest.mark.asyncio
+async def test_beam_scorer():
+    responses = [CompletionResponse(
+        choices=[CompletionResponseChoice(text="Hello", index=0, logprobs=None, finish_reason="length", additional_heads=[[10000, 0, 0]],),],
+        model="test",
+        usage=UsageInfo(),
+    ),
+        CompletionResponse(
+            choices=[CompletionResponseChoice(text="Hello", index=0, logprobs=None, finish_reason="length",
+                                              additional_heads=[[-100, 0, 0]], ), ],
+            model="test",
+            usage=UsageInfo(),
+        )
+    ]
+
+    scorer = BeamScorer(classi_idx)
+    res = await scorer.pick_best_beam(responses)
+    assert res == responses[1]

tests/beam/test_beam.py

@@ -0,0 +1,30 @@
+import pytest
+
+from vllm.beam.beam import BeamScorer
+from vllm.beam.penalty import MEOW_CLASSI_IDX
+from vllm.entrypoints.openai.protocol import CompletionResponse, CompletionResponseChoice, EmbeddingResponse, UsageInfo
+
+@pytest.fixture()
+async def meow_random_beams():
+    return (
+        " Aizawa: You haven't given me your name, age, and quirk",
+
+    )
+@pytest.mark.asyncio
+async def test_beam_scorer():
+    responses = [CompletionResponse(
+        choices=[CompletionResponseChoice(text="Hello", index=0, logprobs=None, finish_reason="length", additional_heads=[[10000, 0, 0]],),],
+        model="test",
+        usage=UsageInfo(),
+    ),
+        CompletionResponse(
+            choices=[CompletionResponseChoice(text="Hello", index=0, logprobs=None, finish_reason="length",
+                                              additional_heads=[[-100, 0, 0]], ), ],
+            model="test",
+            usage=UsageInfo(),
+        )
+    ]
+
+    scorer = BeamScorer(MEOW_CLASSI_IDX)
+    res = await scorer.pick_best_beam(responses)
+    assert res == responses[1]

tests/beam/test_beam_meow.py

@@ -5,7 +5,7 @@
 from vllm.beam.penalty import PenaltyComputer
 import torch
 from vllm.beam.ranking import RankingComputer
-from vllm.entrypoints.openai.protocol import CompletionResponse, ErrorResponse
+from vllm.entrypoints.openai.protocol import CompletionResponse, ErrorResponse, CompletionResponseChoice
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
@@ -17,13 +17,13 @@ def __init__(self, classi_idx):
         self.ranking_computer = RankingComputer(classi_idx)
 
     async def pick_best_beam(self, responses: list[
-        Union[AsyncGenerator[str, None], CompletionResponse, ErrorResponse]]) -> Union[
-        AsyncGenerator[str, None], CompletionResponse, ErrorResponse]:
+        Union[AsyncGenerator[str, None], CompletionResponseChoice, ErrorResponse]]) -> Union[
+        AsyncGenerator[str, None], CompletionResponseChoice, ErrorResponse]:
         debug_info = [BeamDebugInfo() for _ in responses]
 
         scores = torch.zeros(len(responses), dtype=torch.float)
 
-        heads = [response.choices[0].additional_heads[0] for response in responses]
+        heads = [response.additional_heads[0] for response in responses]
         heads_tensor = torch.tensor(heads, dtype=torch.float)
         if len(heads_tensor) > 0:
             penalties = self.penalty_computer.compute(heads_tensor, debug_info)
@@ -36,7 +36,7 @@ async def pick_best_beam(self, responses: list[
 
         for i in range(len(responses)):
             debug_info[i].final_score = scores[i]
-            debug_info[i].content = responses[i].choices[0].text
+            debug_info[i].content = responses[i].text
 
         logger.debug('debug_info: %s', debug_info)
 

vllm/beam/beam.py

@@ -7,7 +7,7 @@
 from starlette.datastructures import MutableHeaders
 
 from vllm.entrypoints.openai.protocol import CompletionRequest, CompletionResponse, \
-    ErrorResponse
+    ErrorResponse, CompletionResponseChoice
 from vllm.logger import init_logger
 from vllm.utils import random_uuid
 
@@ -36,32 +36,32 @@ async def get_n_valid_beams(self, create_completion: Callable,
                                 request: CompletionRequest,
                                 chunk_num: int,
                                 raw_request: Optional[Request] = None) -> list[
-        Union[AsyncGenerator[str, None], CompletionResponse, ErrorResponse]]:
+        Union[AsyncGenerator[str, None], CompletionResponseChoice, ErrorResponse]]:
         request.stream = False
-        n = request.n if request.n > 1 else _DEFAULT_BEAM_SIZE
-        request.n = 1
-        # TODO(@tanuj): accept max tokens as a parameter
+        original_n = request.n
+        request.n = request.n if request.n > 1 else _DEFAULT_BEAM_SIZE
         request.max_tokens = _CHUNK_SIZE
         request.echo = True
         original_request_id = None
         if raw_request is not None:
             original_request_id = raw_request.headers.get("X-Request-Id", None)
-        
-        tasks = []
-        # TODO(@tanuj): deep copy request and raw_request?
-        for _ in range(n):
-            if original_request_id is not None:
-                mh = MutableHeaders(scope=raw_request.scope)
-                del mh["x-request-id"]
-                if hasattr(raw_request, "_headers"):
-                    delattr(raw_request, "_headers")
-
-            tasks.append(create_completion(
+
+        if original_request_id is not None:
+            mh = MutableHeaders(scope=raw_request.scope)
+            del mh["x-request-id"]
+            if hasattr(raw_request, "_headers"):
+                delattr(raw_request, "_headers")
+
+        raw_res = await create_completion(
                 request,
                 raw_request=raw_request,
-            ))
-        res = await asyncio.gather(*tasks)
-        request.n = n
+        )
+
+        if isinstance(raw_res, ErrorResponse):
+            return raw_res
+
+        res = raw_res.choices
+        request.n = original_n
         beam_validator_res = self.validate(res)
         if isinstance(beam_validator_res, ErrorResponse):
             return beam_validator_res
@@ -73,7 +73,7 @@ async def get_n_valid_beams(self, create_completion: Callable,
 
         return filtered_res
 
-    def validate(self, responses: list[AsyncGenerator],
+    def validate(self, responses: list[CompletionResponseChoice | ErrorResponse],
                  debug_infos_G: list[BeamDebugInfo] = None):
         error_responses = [r for r in responses if isinstance(r, ErrorResponse)]
         print(f"error_responses: {error_responses}")
@@ -86,7 +86,7 @@ def validate(self, responses: list[AsyncGenerator],
             )
 
         # TODO(@tanuj) - share this with the beam scorer
-        heads = [response.choices[0].additional_heads[0] for response in responses]
+        heads = [response.additional_heads[0] for response in responses]
         heads_tensor = torch.tensor(heads, dtype=torch.float)
         prob_GC = torch.sigmoid(heads_tensor)
         valid_G = torch.ones(prob_GC.shape[0], dtype=torch.bool)
@@ -99,8 +99,7 @@ def validate(self, responses: list[AsyncGenerator],
 
             if filtered:
                 valid_G[g] = False
-                for choice in responses[g].choices:
-                    choice.is_filtered = True
+                responses[g].is_filtered = True
 
         return valid_G
 

vllm/beam/filtering.py

@@ -104,7 +104,7 @@ async def _process_prefix(request: CompletionRequest):
         input_str_len = len(res.choices[0].text)
 
         async def _should_stop(final):
-            return final.choices[0].finish_reason == "stop" or final.choices[0].is_filtered
+            return final.finish_reason == "stop" or final.is_filtered
 
         max_chunks = math.ceil(request.max_tokens / _CHUNK_SIZE)
         async def _chunk_generator():
@@ -121,12 +121,12 @@ async def _chunk_generator():
                     break
 
                 final = await self.beam_scorer.pick_best_beam(beams)
-                request.prompt = final.choices[0].text
+                request.prompt = final.text
                 should_stop = await _should_stop(final)
-                final.choices[0].text = final.choices[0].text[input_str_len:]
-                output = final.choices[0].text
+                final.text = final.text[input_str_len:]
+                output = final.text
                 if self.request_logger:
-                    logger.info(f"yielding chunk {num_chunks} text: {final.choices[0].text}")
+                    logger.info(f"yielding chunk {num_chunks} text: {final.text}")
                 yield f"data: {final.model_dump_json()}\n\n"
 
                 if should_stop: