[Bugfix] Fix delayed decoding bug for Bagel AR/DIT workflow (L3 test_bagel_img2img error) (#2422)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -95,54 +95,54 @@ class OmniARScheduler(VLLMScheduler):
             return False
 
         criteria_type = self.kv_transfer_criteria.get("type")
-        if (
-            self.kv_transfer_criteria.get("stop_after_transfer", True)
-            and request.request_id in self.transfer_triggered_requests
-        ):
-            # For split pipelines that only need the transferred KV
-            # snapshot, stop AR decode once KV extraction has completed.
-            # This frees stage-0 resources without requiring an
-            # orchestrator-side abort.
-            if request.request_id not in self.active_kv_transfers:
-                request.status = RequestStatus.FINISHED_STOPPED
-                return True
-            return False
+        stop_decode_on_trigger = self.kv_transfer_criteria.get("stop_after_transfer", True)
 
         if request.request_id in self.transfer_triggered_requests:
+            # Already triggered.  When stop_decode_on_trigger is True AND
+            # transfer was actually queued, the request was already stopped
+            # at trigger time (see below).  Any request that reaches this
+            # point either has stop_decode_on_trigger=False (continue
+            # decoding) or was not actually queued (should not be stopped).
             return False
 
         if criteria_type == "prefill_finished":
             if request.num_computed_tokens >= request.num_prompt_tokens:
-                logger.debug(f"[Omni] Request {request.request_id} triggered prefill_finished transfer (Non-Stop)")
                 self.transfer_triggered_requests.add(request.request_id)
                 self._mark_request_for_kv_transfer(request.request_id, request.num_computed_tokens)
+                actually_queued = request.request_id in self.requests_needing_kv_transfer
+
+                if stop_decode_on_trigger and actually_queued:
+                    # Stop immediately so the request is NOT scheduled in
+                    # the next step, freeing scheduling budget for companion
+                    # requests whose chunked-prefill boundaries must be
+                    # deterministic.  waiting_for_transfer_free keeps blocks
+                    # alive until the model runner finishes KV extraction.
+                    self.waiting_for_transfer_free.add(request.request_id)
+                    request.status = RequestStatus.FINISHED_STOPPED
+                    return True
 
-                # Return False means "Do NOT stop the request" -> Continue Decoding
                 return False
 
         elif criteria_type == "special_token":
             target_token_id = self.kv_transfer_criteria.get("token_id")
             if target_token_id is not None and target_token_id in new_token_ids:
-                logger.debug(f"[Omni] Request {request.request_id} triggered special_token criteria (Non-Stop)")
-
                 self.transfer_triggered_requests.add(request.request_id)
 
-                # Calculate precise snapshot length (trim to sentinel)
-                # Find the FIRST occurrence of the sentinel
                 try:
                     idx = new_token_ids.index(target_token_id)
-                    # seq_len = tokens_before_this_step + idx + 1 (include sentinel)
-                    # request.num_computed_tokens already includes ALL new_token_ids
-                    # so we subtract (len(new_token_ids) - (idx + 1))
                     tokens_to_exclude = len(new_token_ids) - (idx + 1)
                     snapshot_len = request.num_computed_tokens - tokens_to_exclude
                 except ValueError:
                     snapshot_len = request.num_computed_tokens
 
-                # Trigger Transfer
                 self._mark_request_for_kv_transfer(request.request_id, snapshot_len)
+                actually_queued = request.request_id in self.requests_needing_kv_transfer
+
+                if stop_decode_on_trigger and actually_queued:
+                    self.waiting_for_transfer_free.add(request.request_id)
+                    request.status = RequestStatus.FINISHED_STOPPED
+                    return True
 
-                # Do NOT stop request
                 return False
 
         return False
@@ -532,9 +532,12 @@ class OmniARScheduler(VLLMScheduler):
                     # It triggered but hasn't finished yet. We MUST wait.
                     logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
                     self.waiting_for_transfer_free.add(request_id)
-                    # We do NOT mark for transfer again, just wait.
-                    kv_xfer_params = None  # No new transfer params
+                    kv_xfer_params = None
                     return kv_xfer_params
+                elif request_id in self.waiting_for_transfer_free:
+                    # Stopped immediately by stop_decode_on_trigger; blocks are
+                    # held until KV extraction completes in a future step.
+                    return None
                 else:
                     logger.debug(
                         f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "