feat: support batch response for mix instance.

Author: JimHsiung

xllm/core/scheduler/async_response_processor.cpp

@@ -41,7 +41,7 @@ AsyncResponseProcessor::AsyncResponseProcessor(
       role_(role.value_or(InstanceRole::DEFAULT)),
       enable_schedule_overlap_(enable_schedule_overlap),
       enable_decode_response_to_service_(enable_decode_response_to_service) {
-  if (role_ == InstanceRole::DECODE) {
+  if (role_ == InstanceRole::DECODE || role_ == InstanceRole::MIX) {
     enable_batch_response_ =
         util::get_bool_env("ENABLE_PD_DECODE_BATCH_RESPONSE", true);
   }
@@ -257,8 +257,9 @@ void AsyncResponseProcessor::process_stream_request(
 }
 
 void AsyncResponseProcessor::process_stream_requests(
-    const std::vector<std::shared_ptr<Request>>& requests) {
-  if (!enable_batch_response_) {
+    const std::vector<std::shared_ptr<Request>>& requests,
+    bool is_prefill) {
+  if (!enable_batch_response_ || is_prefill) {
     for (auto& req : requests) {
       process_stream_request(req);
     }

xllm/core/scheduler/async_response_processor.h

@@ -46,7 +46,8 @@ class AsyncResponseProcessor final {
 
   // in disagg pd mode, decode send requests' responses to prefill
   void process_stream_requests(
-      const std::vector<std::shared_ptr<Request>>& requests);
+      const std::vector<std::shared_ptr<Request>>& requests,
+      bool is_prefill);
 
   // wait for all responses in queue to be handled
   void wait_completion();

xllm/core/scheduler/continuous_scheduler.cpp

@@ -1037,7 +1037,8 @@ void ContinuousScheduler::process_batch_output(bool enable_schedule_overlap) {
     }
   }
   if (!stream_requests.empty()) {
-    response_processor_->process_stream_requests(stream_requests);
+    response_processor_->process_stream_requests(stream_requests,
+                                                 last_step_prefill_);
   }
 }
 

xllm/core/scheduler/disagg_pd_scheduler.cpp

@@ -250,7 +250,7 @@ void DisaggPDScheduler::step(const absl::Duration& timeout) {
   ContinuousScheduler::step(timeout);
   // send first generation token to decode instance
   // and remove the request from running_requests_ to remote_requests_map_
-  if (options_.instance_role() != InstanceRole::DECODE) {
+  if (options_.instance_role() != InstanceRole::DECODE && last_step_prefill_) {
     prefill_send_first_generation();
   }
 }