From 372477d2d9890381dbc794cf66ef86ff989f04e0 Mon Sep 17 00:00:00 2001
From: junuxyz <216036880+junuxyz@users.noreply.github.com>
Date: Sat, 12 Sep 2026 15:22:34 +0900
Subject: [PATCH] [Bugfix][Examples] Launch prefill and decode concurrently in
 NixlPushConnector demo (#55088)

Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com>
---
 .../disagg_proxy_pushconnector_demo.py        | 28 ++++++++++++-------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
index 672162f8ba..a35a44fbba 100644
--- a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
+++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
@@ -31,6 +31,7 @@ disagg_proxy_pushconnector_demo.py \
 """
 
 import argparse
+import asyncio
 import contextlib
 import ipaddress
 import itertools
@@ -259,17 +260,24 @@ class PushProxy:
         decode_instance = self.schedule(self.decode_cycler)
         headers = self._common_headers(request_id)
 
-        # Fire prefill; we don't read its body but must drain the
-        # connection so the upstream server can free its slot.
-        async for _ in self.forward_request(
-            f"http://{prefill_instance}{path}", prefill_request, headers
-        ):
-            continue
+        async def drain_prefill():
+            async for _ in self.forward_request(
+                f"http://{prefill_instance}{path}", prefill_request, headers
+            ):
+                continue
 
-        generator = self.forward_request(
-            f"http://{decode_instance}{path}", decode_request, headers
-        )
-        return StreamingResponse(generator, media_type="application/json")
+        prefill_task = asyncio.create_task(drain_prefill())
+
+        async def stream_decode():
+            try:
+                async for chunk in self.forward_request(
+                    f"http://{decode_instance}{path}", decode_request, headers
+                ):
+                    yield chunk
+            finally:
+                await prefill_task
+
+        return StreamingResponse(stream_decode(), media_type="application/json")
 
     async def create_completion(self, raw_request: Request):
         try: