perf(api-server): stop keeping finished sequences, keep prompt rendering and tokenizing off the event loop

atom/config.py

@@ -2108,6 +2108,9 @@ class Config:
     compilation_config: CompilationConfig = field(default_factory=CompilationConfig)
     quant_config: QuantizationConfig = field(init=False)
     asyncio_mode: bool = False
+    # Queue finished sequences on outputs_queue; only its readers turn it on
+    # (LLMEngine.add_request for step / generate, and asyncio_mode).
+    keep_finished_outputs: bool = False
     mark_trace: bool = False
     load_dummy: str | None = None
     enable_expert_parallel: bool = False

atom/entrypoints/openai/api_server.py

@@ -51,6 +51,7 @@
 from atom.model_engine.arg_utils import EngineArgs
 from atom.model_engine.llm_engine import load_tokenizer
 from atom.model_engine.request import RequestOutput
+from atom.model_engine.segment_tokenizer import PrefixEncoding
 from atom.model_engine.sequence import new_token_ids
 from atom.multimodal.processing import prepare_multimodal_inputs
 from atom.multimodal.registry import get_native_multimodal_processor
@@ -853,7 +854,7 @@ async def _submit_engine_requests(seqs):
 
 
 async def generate_async(
-    prompt_or_tokens: str | list[int],
+    prompt_or_tokens: str | list[int] | PrefixEncoding,
     sampling_params: SamplingParams,
     request_id: str,
     kv_transfer_params: dict[str, Any] | None = None,
@@ -1306,6 +1307,50 @@ def _prompt_headroom(
     return headroom
 
 
+def _prepare_prompt(
+    messages,
+    template_kwargs: dict[str, Any],
+    *,
+    tools,
+    preserve_tool_call_arguments: bool,
+    max_tokens: int,
+    kv_transfer_params: dict[str, Any] | None,
+) -> tuple[PrefixEncoding | list[int], str]:
+    """Render, tokenize the engine's way and cut to the context, in one
+    executor step for ``/v1/messages`` and ``/v1/responses``.
+
+    Returns ``(engine_input, text)``: what preprocess takes in place of the
+    prompt (no second tokenize), and the text the engine runs, which seeds the
+    reasoning channel. A prompt longer than the context minus the output
+    headroom is cut from the end; a context with no room for any prompt is
+    refused (ValueError, a 400)."""
+    prompt = apply_chat_template(
+        tokenizer,
+        custom_message_encoder,
+        messages,
+        tools=tools,
+        preserve_tool_call_arguments=preserve_tool_call_arguments,
+        **template_kwargs,
+    )
+    encoded = engine.io_processor.encode_prompt(prompt)
+    max_ctx = int(_get_engine_max_model_len() or 30720)
+    cap = max_ctx - _prompt_headroom(max_tokens, max_ctx, kv_transfer_params)
+    if cap < 1:
+        raise ValueError(
+            f"max_tokens={max_tokens} leaves no room for a prompt in this "
+            f"model's {max_ctx}-token context"
+        )
+    token_ids = encoded.token_ids if isinstance(encoded, PrefixEncoding) else encoded
+    if len(token_ids) <= cap:
+        return encoded, prompt
+    logger.warning("Prompt too long (%d > %d), truncating", len(token_ids), cap)
+    token_ids = list(token_ids[:cap])
+    text = tokenizer.decode(token_ids, skip_special_tokens=False)
+    if not text:
+        raise ValueError("The prompt is empty after truncation to the context")
+    return token_ids, text
+
+
 def _with_kv_transfer_params(response: dict, final_output: dict) -> dict:
     """Return the prefill's KV location to the PD router, as chat does."""
     if "kv_transfer_output_meta_info" in final_output:
@@ -1314,7 +1359,7 @@ def _with_kv_transfer_params(response: dict, final_output: dict) -> dict:
 
 
 async def setup_streaming_request(
-    prompt_or_tokens: str | list[int],
+    prompt_or_tokens: str | list[int] | PrefixEncoding,
     sampling_params: SamplingParams,
     request_id: str,
     kv_transfer_params: dict[str, Any] | None = None,
@@ -1924,7 +1969,8 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
             prompt_or_tokens = pretokenized
             kv_transfer_params = _engine_kv_transfer_params(kv_transfer_params)
         else:
-            prompt_or_tokens = apply_chat_template(
+            prompt_or_tokens = await asyncio.to_thread(
+                apply_chat_template,
                 tokenizer,
                 custom_message_encoder,
                 messages,
@@ -2351,15 +2397,6 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req
         # no such spelling, and are not answered here or on the chat path.
         if forbids_tool_calls(request.tool_choice):
             merged_kwargs["tool_choice"] = "none"
-        prompt = apply_chat_template(
-            tokenizer,
-            custom_message_encoder,
-            messages,
-            tools=anthropic_to_openai_tools(request.tools),
-            preserve_tool_call_arguments=preserve_tool_call_arguments,
-            **merged_kwargs,
-        )
-
         generation_config = engine.config.generation_config
         model_temperature = getattr(generation_config, "temperature", None)
         model_top_p = getattr(generation_config, "top_p", None)
@@ -2388,47 +2425,23 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req
         request_id = uuid.uuid4().hex[:24]
         set_request_id(request_id)
         kv_transfer_params = _engine_kv_transfer_params(request.kv_transfer_params)
-        input_tokens = len(tokenizer.encode(prompt))
-
-        max_ctx = None
-        for _path in (
-            lambda: engine.config.max_model_len,
-            lambda: engine.model_config.max_model_len,
-            lambda: engine.scheduler.max_model_len,
-            lambda: engine.max_model_len,
-        ):
-            try:
-                _v = _path()
-                if _v:
-                    max_ctx = int(_v)
-                    break
-            except Exception:
-                logger.debug("max_ctx candidate rejected", exc_info=True)
-                continue
-        if not max_ctx:
-            max_ctx = 30720
-        logger.warning(f"[anthropic] resolved max_ctx={max_ctx}")
-        headroom = _prompt_headroom(
-            sampling_params.max_tokens, max_ctx, kv_transfer_params
+        engine_input, prompt = await asyncio.to_thread(
+            _prepare_prompt,
+            messages,
+            merged_kwargs,
+            tools=anthropic_to_openai_tools(request.tools),
+            preserve_tool_call_arguments=preserve_tool_call_arguments,
+            max_tokens=sampling_params.max_tokens,
+            kv_transfer_params=kv_transfer_params,
         )
-        max_input = max_ctx - headroom
-        if input_tokens > max_input:
-            logger.warning(
-                f"Prompt too long ({input_tokens} > {max_input}), truncating"
-            )
-            token_ids = tokenizer.encode(prompt)[:max_input]
-            prompt = tokenizer.decode(token_ids, skip_special_tokens=False)
-            input_tokens = max_input
 
         if request.stream:
             # Streaming response
-            seq_id, stream_collector, _num_prompt_tokens = (
-                await setup_streaming_request(
-                    prompt,
-                    sampling_params,
-                    request_id,
-                    kv_transfer_params=kv_transfer_params,
-                )
+            seq_id, stream_collector, input_tokens = await setup_streaming_request(
+                engine_input,
+                sampling_params,
+                request_id,
+                kv_transfer_params=kv_transfer_params,
             )
 
             async def generate_anthropic_stream():
@@ -2600,7 +2613,7 @@ async def generate_anthropic_stream():
         # Non-streaming response
         final_output = await _run_nonstream_with_disconnect(
             generate_async(
-                prompt,
+                engine_input,
                 sampling_params,
                 request_id,
                 kv_transfer_params=kv_transfer_params,
@@ -2610,6 +2623,7 @@ async def generate_anthropic_stream():
         )
         if final_output is None:
             raise RuntimeError("No output generated")
+        input_tokens = final_output["num_tokens_input"]
 
         raw_text = final_output["text"]
         # Separating is unconditional -- the same call the chat path makes,
@@ -2633,7 +2647,7 @@ async def generate_anthropic_stream():
         if drop_reasoning:
             events = [e for e in events if e[0] != "reasoning"]
         _, tool_calls = flatten_tool_events(events)
-        output_tokens = len(tokenizer.encode(raw_text))
+        output_tokens = final_output["num_tokens_output"]
         cache_read_input_tokens = final_output.get("num_cached_tokens", 0)
         response = build_anthropic_response(
             request_id=request_id,
@@ -2783,15 +2797,6 @@ async def responses_create(raw_request: Request):
             rules=rules,
         )
 
-        prompt = apply_chat_template(
-            tokenizer,
-            custom_message_encoder,
-            messages,
-            tools=openai_tools,
-            preserve_tool_call_arguments=preserve_tool_call_arguments,
-            **merged_kwargs,
-        )
-
         generation_config = engine.config.generation_config
         model_temperature = getattr(generation_config, "temperature", None)
         model_top_p = getattr(generation_config, "top_p", None)
@@ -2828,30 +2833,15 @@ async def responses_create(raw_request: Request):
         schema = schema_from_format((text_config or {}).get("format"), responses=True)
         request_id = f"resp_{uuid.uuid4().hex[:24]}"
         set_request_id(request_id)
-        input_tokens = len(tokenizer.encode(prompt))
-
-        max_ctx = 30720
-        for obj in (
-            getattr(engine, "config", None),
-            getattr(engine, "model_config", None),
-            getattr(engine, "scheduler", None),
-            engine,
-        ):
-            value = getattr(obj, "max_model_len", None) if obj is not None else None
-            if value:
-                max_ctx = int(value)
-                break
-        headroom = _prompt_headroom(
-            sampling_params.max_tokens, max_ctx, kv_transfer_params
+        engine_input, prompt = await asyncio.to_thread(
+            _prepare_prompt,
+            messages,
+            merged_kwargs,
+            tools=openai_tools,
+            preserve_tool_call_arguments=preserve_tool_call_arguments,
+            max_tokens=sampling_params.max_tokens,
+            kv_transfer_params=kv_transfer_params,
         )
-        max_input = max_ctx - headroom
-        if input_tokens > max_input:
-            logger.warning(
-                f"Prompt too long ({input_tokens} > {max_input}), truncating"
-            )
-            token_ids = tokenizer.encode(prompt)[:max_input]
-            prompt = tokenizer.decode(token_ids, skip_special_tokens=False)
-            input_tokens = max_input
 
         created_at = int(time.time())
         _reasoning = reasoning_channel(
@@ -2869,13 +2859,11 @@ async def responses_create(raw_request: Request):
         _reasoning = structured_response_channel(_reasoning, schema)
 
         if request.stream:
-            seq_id, stream_collector, _num_prompt_tokens = (
-                await setup_streaming_request(
-                    prompt,
-                    sampling_params,
-                    request_id,
-                    kv_transfer_params=kv_transfer_params,
-                )
+            seq_id, stream_collector, input_tokens = await setup_streaming_request(
+                engine_input,
+                sampling_params,
+                request_id,
+                kv_transfer_params=kv_transfer_params,
             )
 
             async def generate_responses_stream():
@@ -2984,7 +2972,7 @@ async def generate_responses_stream():
 
         final_output = await _run_nonstream_with_disconnect(
             generate_async(
-                prompt,
+                engine_input,
                 sampling_params,
                 request_id,
                 kv_transfer_params=kv_transfer_params,
@@ -2994,6 +2982,7 @@ async def generate_responses_stream():
         )
         if final_output is None:
             raise RuntimeError("No output generated")
+        input_tokens = final_output["num_tokens_input"]
 
         raw_text = final_output["text"]
         events = read_whole_blocks(
@@ -3003,7 +2992,7 @@ async def generate_responses_stream():
             openai_tools,
             suppress_calls=forbids_tool_calls(request.tool_choice),
         )
-        output_tokens = len(tokenizer.encode(raw_text))
+        output_tokens = final_output["num_tokens_output"]
         cache_read_input_tokens = final_output.get("num_cached_tokens", 0)
         response = build_responses_response(
             request_id=request_id,

atom/model_engine/engine_core_mgr.py

@@ -231,6 +231,7 @@ def _init_shared_state(
         )
         self.ctx = zmq.Context(io_threads=2)
         self.outputs_queue = queue.Queue[list[Sequence]]()
+        self.keep_finished_outputs = config.keep_finished_outputs or config.asyncio_mode
         # Utility replies are matched to their caller by request id, not by
         # position on a shared queue, where a late or unasked-for reply became
         # the next caller's answer.
@@ -834,7 +835,8 @@ def process_outputs_socket():
                                     f"{self.label}: flush_stream_batch failed: {e}",
                                     exc_info=True,
                                 )
-                        self.outputs_queue.put_nowait(seqs)
+                        if self.keep_finished_outputs:
+                            self.outputs_queue.put_nowait(seqs)
                         # These completions freed slots on their ranks — re-feed
                         # from the pending pool so a fast rank never drains to
                         # empty while peers still have work. No-op when
@@ -1306,7 +1308,23 @@ def _bp_fail_stalled_run(self) -> None:
         ``LLMEngine.generate`` waits while the engine is alive and any request
         is unfinished. A dropped seq never finishes, so that wait never ends.
         The seqs are still in ``_bp_pending``; this only surfaces the error.
+        A pending seq with a stream callback is not reading that queue, so it
+        is taken out and ended through its callback, as abort does.
         """
+        with self._lb_lock:
+            waiting = [s for s in self._bp_pending if s.id in self._seq_id_to_callback]
+            for seq in waiting:
+                self._bp_pending.remove(seq)
+        delivered = False
+        for seq in waiting:
+            if not getattr(seq, "leave_reason", None):
+                seq.leave_reason = "aborted"
+            delivered |= self._deliver_terminal_output(seq)
+        if delivered and self._flush_stream_batch_fn is not None:
+            try:
+                self._flush_stream_batch_fn()
+            except Exception as e:  # noqa: BLE001 - the error below must still go out
+                logger.warning("%s: flush_stream_batch failed: %s", self.label, e)
         pending = len(getattr(self, "_bp_pending", ()))
         logger.error(
             "%s: backpressure send failed with nothing in flight; "