fix(swa): graceful abort instead of crash when decode OOM after retract (#944)

When decode KV cache is exhausted in SWA hybrid mode, retract_decode
previously hit an assert crash when only one request remained but
memory was still insufficient. This changes the behavior to gracefully
abort the last request and return an error response to the client,
following upstream sglang's approach.

Author: cjx0709

python/sgl_jax/srt/managers/schedule_batch.py

@@ -976,7 +976,7 @@ def check_decode_mem(self, buf_multiplier=1, selected_indices: list[int] | None
         self._evict_tree_cache_if_needed(num_tokens)
         return self._is_available_size_sufficient(num_tokens)
 
-    def retract_decode(self, server_args: ServerArgs):
+    def retract_decode(self, server_args: ServerArgs) -> tuple[list[Req], float, list[Req]]:
         """Retract the decoding requests when there is not enough memory."""
         sorted_indices = list(range(len(self.reqs)))
 
@@ -990,19 +990,9 @@ def retract_decode(self, server_args: ServerArgs):
 
         retracted_reqs = []
         first_iter = True
-        while (not self.check_decode_mem(selected_indices=sorted_indices)) or first_iter:
+        while first_iter or (not self.check_decode_mem(selected_indices=sorted_indices)):
             if len(sorted_indices) == 1:
-                # Corner case: only one request left
-                if self.is_hybrid:
-                    full_available_size = self.token_to_kv_pool_allocator.full_available_size()
-                    swa_available_size = self.token_to_kv_pool_allocator.swa_available_size()
-                    assert (
-                        full_available_size > 0 and swa_available_size > 0
-                    ), f"No space left for only one request in SWA mode {full_available_size=}, {swa_available_size=}"
-                else:
-                    assert (
-                        self.token_to_kv_pool_allocator.available_size() > 0
-                    ), f"No space left for only one request, {self.token_to_kv_pool_allocator.available_size()=}"
+                # Keep at least one request in the loop; handle OOM below.
                 break
 
             first_iter = False
@@ -1011,11 +1001,24 @@ def retract_decode(self, server_args: ServerArgs):
             retracted_reqs.append(req)
             self.release_req(idx, len(sorted_indices), server_args)
 
-            if len(retracted_reqs) == 0:
-                # Corner case: only one request left
-                raise ValueError(
-                    "Failed to retract any request. No space left for only one request."
-                )
+        # If the last remaining request still can't fit, abort it gracefully
+        # instead of crashing the scheduler (follows upstream sglang).
+        reqs_to_abort: list[Req] = []
+        if len(sorted_indices) <= 1 and not self.check_decode_mem(selected_indices=sorted_indices):
+            last_idx = sorted_indices.pop()
+            last_req = self.reqs[last_idx]
+            last_req.to_finish = FINISH_ABORT(
+                "Out of memory even after retracting all other requests "
+                "in the decode batch. Aborting the last request.",
+                HTTPStatus.INTERNAL_SERVER_ERROR,
+                "InternalServerError",
+            )
+            reqs_to_abort.append(last_req)
+            self.release_req(last_idx, 0, server_args)
+            logger.warning(
+                "retract_decode: aborted last request %s due to OOM",
+                last_req.rid,
+            )
 
         self.filter_batch(keep_indices=sorted_indices)
 
@@ -1025,10 +1028,12 @@ def retract_decode(self, server_args: ServerArgs):
 
         new_estimate_ratio = (
             total_decoded_tokens + global_config.retract_decode_steps * len(self.reqs)
-        ) / total_max_new_tokens
+        ) / (
+            total_max_new_tokens + 1
+        )  # +1 to avoid zero division when all reqs aborted
         new_estimate_ratio = min(1.0, new_estimate_ratio)
 
-        return retracted_reqs, new_estimate_ratio
+        return retracted_reqs, new_estimate_ratio, reqs_to_abort
 
     def release_req(self, idx: int, remaing_req_count: int, server_args: ServerArgs):
         req = self.reqs[idx]

python/sgl_jax/srt/managers/scheduler.py

@@ -1298,13 +1298,22 @@ def update_running_batch(self, batch: ScheduleBatch) -> ScheduleBatch | None:
         ):
             old_ratio = self.new_token_ratio
 
-            retracted_reqs, new_token_ratio = batch.retract_decode(self.server_args)
+            retracted_reqs, new_token_ratio, reqs_to_abort = batch.retract_decode(self.server_args)
             num_retracted_reqs = len(retracted_reqs)
             self.new_token_ratio = new_token_ratio
 
+            # Send abort responses so clients get an error instead of a hung connection
+            for req in reqs_to_abort:
+                abort_out = AbortReq(rid=req.rid)
+                if self._comm_backend is not None:
+                    self._comm_backend.send_pyobj(abort_out)
+                else:
+                    self.send_to_tokenizer.send_pyobj(abort_out)
+
             logger.info(
-                "KV cache pool is full. Retract requests. #retracted_reqs: %d, #new_token_ratio: %.4f -> %.4f",
+                "KV cache pool is full. Retract requests. #retracted_reqs: %d, #aborted_reqs: %d, #new_token_ratio: %.4f -> %.4f",
                 num_retracted_reqs,
+                len(reqs_to_abort),
                 old_ratio,
                 self.new_token_ratio,
             )
@@ -1319,6 +1328,9 @@ def update_running_batch(self, batch: ScheduleBatch) -> ScheduleBatch | None:
         if batch.batch_size() < initial_bs:
             batch.batch_is_full = False
 
+        if batch.is_empty():
+            return batch
+
         # Update batch arrays
         batch.prepare_for_decode()
         return batch