perf: batch decode stream callbacks into one dispatch per step
- 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:
+72
-15
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user