[Bugfix][PD] set max_completion_tokens=1 if req has this value (vllm-project#21841)

Signed-off-by: Abirdcfly <[email protected]>
Signed-off-by: Diego-Castan <[email protected]>

Author: Abirdcfly

examples/online_serving/disaggregated_serving/disagg_proxy_demo.py

@@ -293,6 +293,8 @@ async def create_chat_completion(self, raw_request: Request):
             # add params to request
             kv_prepare_request = request.copy()
             kv_prepare_request["max_tokens"] = 1
+            if "max_completion_tokens" in kv_prepare_request:
+                kv_prepare_request["max_completion_tokens"] = 1
 
             # prefill stage
             prefill_instance = self.schedule(self.prefill_cycler)

examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_proxy_p2p_nccl_xpyd.py

@@ -128,6 +128,8 @@ async def handle_request():
         prefill_request = original_request_data.copy()
         # change max_tokens = 1 to let it only do prefill
         prefill_request["max_tokens"] = 1
+        if "max_completion_tokens" in prefill_request:
+            prefill_request["max_completion_tokens"] = 1
 
         global count
         global prefill_instances