feat(grpc): support batched completion prompts in the pipeline

e2e_test/completions/test_basic.py

@@ -261,3 +261,126 @@ def test_streaming_echo_max_tokens_zero(self, model, api_client):
         assert full_text == prompt, f"Expected echoed prompt, got: {full_text!r}"
         assert len(finish_reasons) == 1
         assert finish_reasons[0] in ("stop", "length")
+
+
+@pytest.mark.engine("sglang", "vllm", "tokenspeed")
+@pytest.mark.gpu(1)
+@pytest.mark.model("meta-llama/Llama-3.1-8B-Instruct")
+@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
+class TestCompletionBatch:
+    """Tests for batched prompt arrays on /v1/completions."""
+
+    PROMPTS = ["The capital of France is", "The capital of Germany is"]
+
+    @staticmethod
+    def _collect_stream_by_index(stream):
+        """Consume a stream, returning per-index text, finish reasons, usage, and usage-chunk count."""
+        texts = {}
+        finish_reasons = {}
+        usage = None
+        usage_chunks = 0
+        for chunk in stream:
+            assert chunk.object == "text_completion"
+            if chunk.usage is not None:
+                usage = chunk.usage
+                usage_chunks += 1
+            for choice in chunk.choices:
+                if choice.text:
+                    texts.setdefault(choice.index, []).append(choice.text)
+                if choice.finish_reason:
+                    finish_reasons.setdefault(choice.index, []).append(choice.finish_reason)
+        return (
+            {index: "".join(parts) for index, parts in texts.items()},
+            finish_reasons,
+            usage,
+            usage_chunks,
+        )
+
+    def test_batch_non_streaming(self, model, api_client):
+        """Test that a prompt array returns one choice per prompt with global indices."""
+
+        response = api_client.completions.create(
+            model=model,
+            prompt=self.PROMPTS,
+            max_tokens=20,
+            temperature=0,
+        )
+
+        assert len(response.choices) == len(self.PROMPTS)
+        assert sorted(choice.index for choice in response.choices) == [0, 1]
+        for choice in response.choices:
+            assert isinstance(choice.text, str)
+            assert len(choice.text) > 0
+            assert choice.finish_reason in ("stop", "length")
+
+        assert response.usage is not None
+        assert response.usage.prompt_tokens > 0
+        assert response.usage.completion_tokens > 0
+        assert response.usage.total_tokens == (
+            response.usage.prompt_tokens + response.usage.completion_tokens
+        )
+
+    def test_batch_non_streaming_with_n(self, model, api_client):
+        """Test prompt-major global indices with n > 1: index = prompt_index * n + i."""
+
+        n = 2
+        response = api_client.completions.create(
+            model=model,
+            prompt=self.PROMPTS,
+            max_tokens=20,
+            temperature=0.7,
+            n=n,
+            echo=True,
+        )
+
+        assert len(response.choices) == len(self.PROMPTS) * n
+        assert [choice.index for choice in response.choices] == [0, 1, 2, 3]
+        for choice in response.choices:
+            assert choice.text.startswith(self.PROMPTS[choice.index // n])
+
+    def test_batch_echo_maps_prompts_to_choices(self, model, api_client):
+        """Test that echo=True prepends each prompt to its own choice."""
+
+        response = api_client.completions.create(
+            model=model,
+            prompt=self.PROMPTS,
+            max_tokens=10,
+            temperature=0,
+            echo=True,
+        )
+
+        choices = {choice.index: choice for choice in response.choices}
+        assert len(choices) == len(self.PROMPTS)
+        for prompt_index, prompt in enumerate(self.PROMPTS):
+            assert choices[prompt_index].text.startswith(prompt)
+
+    def test_batch_streaming(self, model, api_client):
+        """Test streaming with a prompt array and n > 1: global per-choice deltas
+        and exactly one aggregated usage chunk."""
+
+        n = 2
+        stream = api_client.completions.create(
+            model=model,
+            prompt=self.PROMPTS,
+            max_tokens=20,
+            temperature=0.7,
+            n=n,
+            stream=True,
+            stream_options={"include_usage": True},
+        )
+
+        texts, finish_reasons, usage, usage_chunks = self._collect_stream_by_index(stream)
+
+        expected_indices = list(range(len(self.PROMPTS) * n))
+        assert sorted(texts) == expected_indices, (
+            f"Expected deltas for all choices, got {sorted(texts)}"
+        )
+        assert sorted(finish_reasons) == expected_indices
+        for index in expected_indices:
+            assert len(texts[index]) > 0
+            assert finish_reasons[index] in (["stop"], ["length"])
+
+        assert usage_chunks == 1, f"Expected exactly one usage chunk, got {usage_chunks}"
+        assert usage is not None
+        assert usage.prompt_tokens > 0
+        assert usage.completion_tokens > 0

model_gateway/src/routers/grpc/common/response_collection.rs

@@ -67,6 +67,14 @@ pub(crate) async fn collect_responses(
                 "Embedding result encountered in response collection",
             ));
         }
+        // Batches are split into per-prompt results by the completion processor
+        // before collection.
+        ExecutionResult::Batch { .. } => {
+            return Err(error::internal_error(
+                "invalid_execution_mode",
+                "Batch result encountered in response collection",
+            ));
+        }
     };
 
     if all_responses.is_empty() {

model_gateway/src/routers/grpc/common/stages/request_execution.rs

@@ -4,7 +4,7 @@ use std::time::Instant;
 
 use async_trait::async_trait;
 use axum::response::Response;
-use futures::future::join_all;
+use futures::future::{join_all, try_join_all};
 use tracing::{debug, error, info_span, Instrument};
 
 use super::PipelineStage;
@@ -15,8 +15,8 @@ use crate::{
         grpc::{
             common::stages::encode::EncodeDispatchPlan,
             context::{
-                ClientSelection, ExecutionPlan, ExecutionResult, LoadGuards, PdTiming,
-                RequestContext, WorkerSelection,
+                ClientSelection, ExecutionPlan, ExecutionPlanKind, ExecutionResult, LoadGuards,
+                PdTiming, RequestContext, WorkerSelection,
             },
             proto_wrapper::{
                 ProtoEmbedRequest, ProtoGenerateRequest, ProtoRequest, ProtoResponseVariant,
@@ -167,7 +167,15 @@ impl PipelineStage for RequestExecutionStage {
             )
         })?;
 
-        ctx.state.load_guards = Some(LoadGuards::new(workers, ctx.input.headers.as_ref()));
+        let sub_requests = match &execution_plan {
+            ExecutionPlan::Batch { requests, .. } => requests.len(),
+            _ => 1,
+        };
+        ctx.state.load_guards = Some(LoadGuards::scaled(
+            workers,
+            ctx.input.headers.as_ref(),
+            sub_requests,
+        ));
 
         // Extract dispatch metadata for tracing span
         let dispatch = ctx.state.dispatch.as_ref();
@@ -204,6 +212,10 @@ impl PipelineStage for RequestExecutionStage {
                     self.execute_epd_dispatch(request, clients, workers, model, encode_dispatch)
                         .await
                 }
+                ExecutionPlan::Batch { kind, requests, .. } => {
+                    self.execute_batch_dispatch(kind, requests, clients, workers, model)
+                        .await
+                }
             }
         }
         .instrument(span)
@@ -329,6 +341,37 @@ impl RequestExecutionStage {
         });
     }
 
+    /// Dispatch one backend request per batched prompt concurrently, preserving
+    /// prompt order. Fail-fast: the first failed dispatch fails the batch and
+    /// drops the remaining streams (abort-on-drop reclaims them backend-side).
+    async fn execute_batch_dispatch(
+        &self,
+        kind: ExecutionPlanKind,
+        requests: Vec<ProtoGenerateRequest>,
+        clients: &ClientSelection,
+        workers: &WorkerSelection,
+        model: &str,
+    ) -> Result<ExecutionResult, Response> {
+        let dispatches = requests.into_iter().map(|request| {
+            let mut clients = clients.clone();
+            async move {
+                match kind {
+                    ExecutionPlanKind::Single => {
+                        self.execute_single(request, &mut clients, workers).await
+                    }
+                    // Completion EPD carries no encode jobs; sub-requests dispatch as PD.
+                    ExecutionPlanKind::PrefillDecode | ExecutionPlanKind::EncodePrefillDecode => {
+                        self.execute_pd_dispatch(request, &mut clients, workers, model)
+                            .await
+                    }
+                }
+            }
+        });
+
+        let results = try_join_all(dispatches).await?;
+        Ok(ExecutionResult::Batch { results })
+    }
+
     async fn execute_single(
         &self,
         mut proto_request: ProtoGenerateRequest,