fix(server): batch stream-chunk dispatch

atom/entrypoints/openai/api_server.py

@@ -287,13 +287,22 @@ def _prepare_multimodal_inputs(
     return inputs["input_ids"][0].tolist(), multimodal_data
 
 
+# ── Batched stream dispatch ──────────────────────────────────────────────
+# Per-seq `call_soon_threadsafe` floods the API event loop at high batch size
+# (one call per token). Instead the callback only decodes + buffers; the mgr
+# flushes a whole step with one scheduled call (see `flush_stream_batch`).
+import threading as _threading  # noqa: E402
+
+_stream_batch_tls = _threading.local()
+
+
 def _send_stream_chunk_direct(
     request_output: RequestOutput,
     request_id: str,
     stream_queue: asyncio.Queue,
     loop: AbstractEventLoop,
 ) -> None:
-    """Send stream chunk directly to the queue."""
+    """Decode the chunk and buffer it; dispatch happens in flush_stream_batch."""
     global tokenizer
 
     new_text = tokenizer.decode(request_output.output_tokens, skip_special_tokens=True)
@@ -309,7 +318,34 @@ def _send_stream_chunk_direct(
     }
     if getattr(request_output, "kv_transfer_params_output", None):
         chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output
-    loop.call_soon_threadsafe(stream_queue.put_nowait, chunk_data)
+
+    buf = getattr(_stream_batch_tls, "buf", None)
+    if buf is None:
+        buf = _stream_batch_tls.buf = []
+    buf.append((loop, stream_queue, chunk_data))
+
+
+def _drain_batch_into_queues(items: list) -> None:
+    """Runs ON the event loop: push each chunk into its per-request queue.
+    One scheduled call handles a whole step's worth of chunks."""
+    for _loop, q, chunk in items:
+        q.put_nowait(chunk)
+
+
+def flush_stream_batch() -> None:
+    """Flush a step's buffered chunks: one call_soon_threadsafe per loop
+    (normally one — all requests on a rank share the API loop)."""
+    buf = getattr(_stream_batch_tls, "buf", None)
+    if not buf:
+        return
+    _stream_batch_tls.buf = []
+    # Group by loop (normally a single loop). dict preserves insertion order
+    # so per-request chunk ordering within the step is maintained.
+    by_loop: Dict[AbstractEventLoop, list] = {}
+    for loop, q, chunk in buf:
+        by_loop.setdefault(loop, []).append((loop, q, chunk))
+    for loop, items in by_loop.items():
+        loop.call_soon_threadsafe(_drain_batch_into_queues, items)
 
 
 def _send_stream_chunk_tagged(

atom/model_engine/engine_core_mgr.py

@@ -45,6 +45,8 @@ def __init__(self, config: Config):
         self.stream_outputs_queue = queue.Queue()
         self.utility_response_queue = queue.Queue()
         self._seq_id_to_callback = {}
+        # Batched stream-flush hook, resolved lazily (avoids import cycle).
+        self._flush_stream_batch_fn = None
         self.engine_core_processes = []
         self.input_sockets = []
         self.output_sockets = []
@@ -230,27 +232,29 @@ def process_outputs_socket():
                             f"{self.label}: Received STREAM message with {len(stream_outputs)} outputs"
                         )
                         self.stream_outputs_queue.put_nowait(stream_outputs)
-                        # Also call callbacks if registered
+                        # Run per-seq callbacks (decode + buffer), then flush
+                        # the whole step in one scheduled call — avoids a
+                        # per-token call_soon_threadsafe storm on the API loop.
+                        any_callback = False
                         for seq_id, request_output in stream_outputs:
                             callback = self._seq_id_to_callback.get(seq_id)
-                            logger.debug(
-                                f"{self.label}: seq_id={seq_id}, callback={'found' if callback is not None else 'NOT FOUND'}, tokens={request_output.output_tokens}"
-                            )
                             if callback is not None:
+                                any_callback = True
                                 try:
                                     callback(request_output)
-                                    logger.debug(
-                                        f"{self.label}: Successfully called callback for seq_id={seq_id}"
-                                    )
                                 except Exception as e:
                                     logger.warning(
                                         f"Error calling stream_callback for sequence {seq_id}: {e}",
                                         exc_info=True,
                                     )
                             if request_output.finished:
                                 self._seq_id_to_callback.pop(seq_id, None)
-                                logger.debug(
-                                    f"{self.label}: Cleaned up callback for finished sequence {seq_id}"
+                        if any_callback:
+                            try:
+                                self._flush_stream_batch()
+                            except Exception:
+                                logger.exception(
+                                    f"{self.label}: flush_stream_batch failed"
                                 )
                     elif request_type == EngineCoreRequestType.UTILITY_RESPONSE:
                         self.utility_response_queue.put_nowait(data)
@@ -413,6 +417,20 @@ def add_request(self, seqs: List[Sequence]):
                         copy=False,
                     )
 
+    def _flush_stream_batch(self):
+        """Flush this step's buffered stream chunks (see flush_stream_batch).
+        Resolved lazily to avoid an import cycle; no-op without the entrypoint."""
+        fn = self._flush_stream_batch_fn
+        if fn is None:
+            try:
+                from atom.entrypoints.openai.api_server import flush_stream_batch
+
+                fn = self._flush_stream_batch_fn = flush_stream_batch
+            except Exception:
+                self._flush_stream_batch_fn = lambda: None  # resolve to no-op
+                return
+        fn()
+
     def get_stream_outputs(self):
         try:
             return self.stream_outputs_queue.get_nowait()