[V1][WIP] V1 sampler implements parallel sampling (PR 1/N for parallel sampling support)

vllm/v1/engine/async_llm.py

@@ -167,7 +167,7 @@ async def add_request(
     # requests we don't need to send multiple messages to core proc,
     # and so we don't need multiple streams which then get
     # re-multiplexed in the API server anyhow.
-    async def generate(
+    async def _generate(
         self,
         prompt: PromptType,
         sampling_params: SamplingParams,
@@ -238,6 +238,43 @@ async def generate(
             await self.abort(request_id)
             raise
 
+    async def _parallel_sampling_batch(
+        self,
+        prompt: PromptType,
+        sampling_params: SamplingParams,
+        request_id: str,
+        lora_request: Optional[LoRARequest] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
+        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        priority: int = 0,
+    ) -> AsyncGenerator[RequestOutput, None]:
+        pass
+
+    async def generate(
+        self,
+        prompt: PromptType,
+        sampling_params: SamplingParams,
+        request_id: str,
+        lora_request: Optional[LoRARequest] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
+        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        priority: int = 0,
+    ) -> AsyncGenerator[RequestOutput, None]:
+        n = sampling_params.n
+        if n is None or sampling_params.n == 1:
+            generator = self._generate(prompt, sampling_params, request_id,
+                                       lora_request, trace_headers,
+                                       prompt_adapter_request, priority)
+        else:
+            generator = self._parallel_sampling_batch(prompt, sampling_params,
+                                                      request_id, lora_request,
+                                                      trace_headers,
+                                                      prompt_adapter_request,
+                                                      priority)
+
+        async for output in generator:
+            yield output
+
     async def _run_output_handler(self):
         """Background loop: pulls from EngineCore and pushes to AsyncStreams."""
 

vllm/v1/engine/parallel_sampling.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from copy import copy
+from typing import Any, Dict, Optional
+
+from vllm.outputs import RequestOutput
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+
+
+class ParentRequestState:
+    sampling_params: SamplingParams
+    request_output: Optional[RequestOutput] = None
+
+    def get_child_sampling_params(
+        self,
+        kwargs: Dict[str, Any] = {},
+    ) -> SamplingParams:
+        sampling_params = copy(self.sampling_params)
+        for kw in kwargs:
+            setattr(sampling_params, kw, kwargs[kw])
+        return sampling_params
+
+    def add_output(
+        self,
+        child_req_output: RequestOutput,
+    ) -> None:
+        if self.output_kind != RequestOutputKind.DELTA:
+            pass
+
+    @property
+    def n(self) -> int:
+        return self.sampling_params.n
+
+    @property
+    def logprobs(self) -> Optional[int]:
+        return self.sampling_params.logprobs
+
+    @property
+    def prompt_logprobs(self) -> Optional[int]:
+        return self.sampling_params.prompt_logprobs
+
+    @property
+    def output_kind(self) -> RequestOutputKind:
+        return self.sampling_params.output_kind