ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: batch decode stream callbacks into one dispatch per step

Browse Source 
- add BatchedStreamCallback sink type: TaskManager resolves a decode step's (task_id, token) events under one lock and delivers each sink a single list instead of one call per token
- keep the plain Callable[[str]] callback contract: per-token callbacks still receive one call per event, and invoke_callback/cancel_task wrap single events for batched sinks
- collect aborted, text, and finish STOP events in the scheduler decode loop and dispatch once per step instead of once per token
- register one _ResultSink per generate call (replacing per-task closures) so GenerateResult takes its lock and wakes waiters once per step, with late-bind replay for tasks that start decoding before add_task returns their id
- apply GenerateResult batches under a single condition hold via append_batch; append delegates to it
- update engine test fakes to the batched contract and add coverage for event grouping, single-event dispatch, cancel STOP, and late-bind replay

Benchmark: NVIDIA L20 (idle), CUDA 12.8, torch 2.11.0+cu128, 1.2B bf16 checkpoint, prompt 512, 256 greedy tokens, CUDA graph on, serving-level decode, 3 trials
- batch 32: 7.808 -> 7.506 ms/token (4098 -> 4263 batch tok/s, +4.0%)
- batch 1/8: unchanged within noise (3.768 -> 3.797 / 4.699 -> 4.607 ms/token)
- full suite: 896 passed
This commit is contained in:
ViperEkura
2026-09-05 00:03:36 +08:00
parent 1798474316
commit 074642b6d2
6 changed files with 265 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/__init__.py
+9 -1
View File
@@ -17,12 +17,20 @@ from astrai.inference.network import get_app, run_server
from astrai.inference.runtime.executor import Executor
from astrai.inference.runtime.sample import sample
from astrai.inference.scheduler import InferenceScheduler
from astrai.inference.task import STOP, GenerationResult, Task, TaskManager, TaskStatus
from astrai.inference.task import (
STOP,
BatchedStreamCallback,
GenerationResult,
Task,
TaskManager,
TaskStatus,
)
__all__ = [
"InferenceEngine",
"build_engine",
"InferenceScheduler",
"BatchedStreamCallback",
"GenerationResult",
"Executor",
"STOP",
astrai/inference/engine.py
+72 -15
View File
@@ -13,7 +13,7 @@ import torch.nn as nn
from astrai.extension import ATTN_BACKEND, AttentionBackend, get_backend
from astrai.inference.cache import PagePool
from astrai.inference.scheduler import InferenceScheduler
from astrai.inference.task import STOP
from astrai.inference.task import STOP, BatchedStreamCallback
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
@@ -33,15 +33,27 @@ class GenerateResult:
self._total = count
def append(self, token: str, idx: int = 0):
self.append_batch([(idx, token)])
def append_batch(self, items: List[Tuple[int, Any]]) -> None:
"""Append multiple ``(idx, token)`` events under one lock/notify.
Batched counterpart to :meth:`append` for per-step delivery: state
updates for every event happen under a single condition hold and
waiters are woken once per batch instead of once per token.
"""
if not items:
return
with self._cond:
self.tokens.append((idx, token))
if token is not STOP:
self.results[idx] += token
else:
if not self._done[idx]:
self._done[idx] = True
self._completed += 1
self._cond.notify_all()
for idx, token in items:
self.tokens.append((idx, token))
if token is STOP:
if not self._done[idx]:
self._done[idx] = True
self._completed += 1
self._cond.notify_all()
else:
self.results[idx] += token
self._event.set()
def pop_all(self) -> List[Tuple[int, str]]:
@@ -69,6 +81,47 @@ class GenerateResult:
return self.results.copy()
class _ResultSink(BatchedStreamCallback):
"""Batched stream channel from the scheduler into one GenerateResult.
Registered as the ``stream_callback`` for every task of a single
``generate`` call, so the scheduler's one dispatch per decode step
maps to one ``append_batch`` (one lock, one waiter wake). A task can
start decoding the moment ``add_task`` returns — before the engine
learns its id — so events for ids not yet bound are buffered and
replayed on ``bind``.
"""
def __init__(self, result: GenerateResult):
self._result = result
self._lock = threading.Lock()
self._index_of: Dict[str, int] = {}
self._pending: List[Tuple[str, Any]] = []
def bind(self, task_id: str, idx: int) -> None:
with self._lock:
self._index_of[task_id] = idx
replay = [(idx, token) for tid, token in self._pending if tid == task_id]
if replay:
self._pending = [
(tid, token) for tid, token in self._pending if tid != task_id
]
if replay:
self._result.append_batch(replay)
def __call__(self, events: List[Tuple[str, Any]]) -> None:
with self._lock:
items: List[Tuple[int, Any]] = []
for tid, token in events:
idx = self._index_of.get(tid)
if idx is None:
self._pending.append((tid, token))
else:
items.append((idx, token))
if items:
self._result.append_batch(items)
class InferenceEngine:
"""Unified inference engine backed by continuous-batching scheduler."""
@@ -147,6 +200,7 @@ class InferenceEngine:
) -> AsyncGenerator[str, None]:
request_backend = get_backend(use_default=False)
result = GenerateResult()
sink = _ResultSink(result)
task_id = self.scheduler.add_task(
prompt=prompt,
max_tokens=max_tokens,
@@ -156,8 +210,9 @@ class InferenceEngine:
frequency_penalty=frequency_penalty,
rep_window=rep_window,
backend=request_backend,
stream_callback=result.append,
stream_callback=sink,
)
sink.bind(task_id, 0)
async def _agen():
finished = False
@@ -191,8 +246,10 @@ class InferenceEngine:
n = len(prompts)
request_backend = get_backend(use_default=False)
result = GenerateResult(count=n)
task_ids = [
self.scheduler.add_task(
sink = _ResultSink(result)
task_ids = []
for i, p in enumerate(prompts):
task_id = self.scheduler.add_task(
prompt=p,
max_tokens=max_tokens,
temperature=temperature,
@@ -201,10 +258,10 @@ class InferenceEngine:
frequency_penalty=frequency_penalty,
rep_window=rep_window,
backend=request_backend,
stream_callback=lambda token, idx=i: result.append(token, idx),
stream_callback=sink,
)
for i, p in enumerate(prompts)
]
sink.bind(task_id, i)
task_ids.append(task_id)
if not stream:
try:
astrai/inference/scheduler.py
+7 -5
View File
@@ -264,17 +264,19 @@ class InferenceScheduler:
decoded, aborted = self._stepper.step(active)
for t in aborted:
self._task_mgr.invoke_callback(t.task_id, STOP)
# One dispatch per step: batch-aware sinks take their
# lock (and wake waiters) once instead of once per token.
events: List[Tuple[str, Any]] = [(t.task_id, STOP) for t in aborted]
for t in decoded:
if t.status == TaskStatus.ABORTED:
continue
new_text = t.decode_new_token(self._task_mgr.tokenizer)
if new_text:
self._task_mgr.invoke_callback(t.task_id, new_text)
events.append((t.task_id, new_text))
if t.is_finished(stop_ids):
self._task_mgr.invoke_callback(t.task_id, STOP)
events.append((t.task_id, STOP))
if events:
self._task_mgr.invoke_callbacks(events)
except Exception as e:
self._stop_event.set()
astrai/inference/task.py
+53 -3
View File
@@ -1,6 +1,7 @@
import threading
import time
import uuid
from abc import ABC, abstractmethod
from collections import deque
from dataclasses import dataclass
from enum import Enum
@@ -145,6 +146,22 @@ class Task:
return False
class BatchedStreamCallback(ABC):
"""Stream sink that receives a whole scheduler step's events in one call.
The scheduling loop dispatches once per decode step: every
``(task_id, token)`` event routed to the same sink object is delivered
as a single list, so batch-aware consumers take their lock and wake
waiters once per step instead of once per token. Plain per-token
callbacks keep the ``Callable[[str], None]`` contract.
"""
@abstractmethod
def __call__(self, events: List[Tuple[str, Any]]) -> None:
"""Consume ``[(task_id, token), ...]`` produced by one decode step."""
raise NotImplementedError
class TaskManager:
"""Thread-safe task queues and lifecycle transitions (no page ops)."""
@@ -256,7 +273,10 @@ class TaskManager:
immediate = [task]
if cancelled and callback is not None:
callback(STOP)
if isinstance(callback, BatchedStreamCallback):
callback([(task_id, STOP)])
else:
callback(STOP)
return immediate, cancelled
def remove_task(self, task_id: str) -> List[Task]:
@@ -264,10 +284,40 @@ class TaskManager:
immediate, _ = self.cancel_task(task_id)
return immediate
def invoke_callback(self, task_id: str, token: str):
def invoke_callback(self, task_id: str, token: Any):
with self._lock:
cb = self._callbacks.get(task_id)
if cb:
if isinstance(cb, BatchedStreamCallback):
cb([(task_id, token)])
elif cb:
cb(token)
def invoke_callbacks(self, events: List[Tuple[str, Any]]) -> None:
"""Dispatch one decode step's ``(task_id, token)`` events.
Callbacks resolve under a single lock acquisition; events aimed at
the same batched sink are delivered as one list (one consumer-side
lock/notify per step), while plain per-token callbacks receive one
call per event.
"""
grouped: Dict[int, Tuple[BatchedStreamCallback, List[Any]]] = {}
plain: List[Tuple[Callable[[str], None], Any]] = []
with self._lock:
for task_id, token in events:
cb = self._callbacks.get(task_id)
if cb is None:
continue
if isinstance(cb, BatchedStreamCallback):
entry = grouped.get(id(cb))
if entry is None:
grouped[id(cb)] = (cb, [(task_id, token)])
else:
entry[1].append((task_id, token))
else:
plain.append((cb, token))
for cb, batch in grouped.values():
cb(batch)
for cb, token in plain:
cb(token)
def get_stats(self) -> Dict[str, Any]: