ViperEkura
074642b6d2
- 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
44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
"""Inference module for continuous batching.
|
|
|
|
Subpackages:
|
|
- cache/: KV cache (buffers, strategies, pool)
|
|
- runtime/: Execution + sampling (executor, CUDA graph, sampling strategies)
|
|
- task/: Request lifecycle + performance metrics
|
|
- network/: HTTP protocol handling (server, protocol, OpenAI/Anthropic builders)
|
|
|
|
Modules:
|
|
- scheduler.py: Continuous batching loop
|
|
- workspace.py: Pre-allocated GPU buffers
|
|
- engine.py: Facade (InferenceEngine)
|
|
"""
|
|
|
|
from astrai.inference.engine import InferenceEngine, build_engine
|
|
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,
|
|
BatchedStreamCallback,
|
|
GenerationResult,
|
|
Task,
|
|
TaskManager,
|
|
TaskStatus,
|
|
)
|
|
|
|
__all__ = [
|
|
"InferenceEngine",
|
|
"build_engine",
|
|
"InferenceScheduler",
|
|
"BatchedStreamCallback",
|
|
"GenerationResult",
|
|
"Executor",
|
|
"STOP",
|
|
"Task",
|
|
"TaskManager",
|
|
"TaskStatus",
|
|
"sample",
|
|
"get_app",
|
|
"run_server",
|
|
]
|