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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user