ViperEkura
8e39d9d8c9
- make decode steady state self-validating: gate the token fill and sampling reuse on the decode cache's own task signature, drop TaskCacheManager.last_task_signature_matches whose req-index signature recycles with slot reuse - extract PolicyVersionGuard (version protocol plus generation/weight mutex) and Stepper (shared one-token advancement) out of the scheduler, keeping its public API unchanged - pin the input staging buffer only on CUDA devices so CPU-only workspaces allocate
128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
"""One-token advancement primitive shared by every scheduling mode."""
|
|
|
|
from contextlib import nullcontext
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
from astrai.extension import AttentionBackend, attn_backend
|
|
from astrai.inference.cache import PagePool, TaskCacheManager
|
|
from astrai.inference.metrics import MetricsCollector
|
|
from astrai.inference.runtime.executor import Executor
|
|
from astrai.inference.task import Task, TaskStatus
|
|
|
|
|
|
class Stepper:
|
|
"""Advance every active task by one token (prefill + decode).
|
|
|
|
Single shared primitive for both the continuous-batching loop and the
|
|
synchronous ``run_batch`` path, so the two cannot drift.
|
|
|
|
Tasks must already be allocated in the KV cache. Tasks without output
|
|
are prefilled first and sample their first token from the final prompt
|
|
position. Tasks with output extend the cache by one position and decode
|
|
from their latest generated token.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pool: PagePool,
|
|
task_cache: TaskCacheManager,
|
|
executor: Executor,
|
|
metrics: MetricsCollector,
|
|
):
|
|
self._pool = pool
|
|
self._task_cache = task_cache
|
|
self._executor = executor
|
|
self._metrics = metrics
|
|
|
|
@staticmethod
|
|
def _task_backend_groups(tasks: List[Task]):
|
|
groups = {}
|
|
for task in tasks:
|
|
groups.setdefault(task.backend, (task.backend, []))[1].append(task)
|
|
return groups.values()
|
|
|
|
def step(
|
|
self, tasks: List[Task], return_logprobs: bool = False
|
|
) -> Tuple[List[Task], List[Task]]:
|
|
"""Advance ``tasks`` by one token.
|
|
|
|
Args:
|
|
tasks: Active tasks to advance by one token.
|
|
return_logprobs: Forwarded to the executor; per-token logprobs
|
|
are recorded on each task's ``output_logprobs``.
|
|
|
|
Returns:
|
|
``(decoded, aborted)``: tasks that produced a new token (its ID
|
|
already appended to ``output_ids``) and tasks that hit the
|
|
sequence cap and were marked ``ABORTED``.
|
|
"""
|
|
to_prefill = [t for t in tasks if not t.prefill_done and t.prompt_ids]
|
|
prefilled_ids = set()
|
|
produced: List[Task] = []
|
|
if to_prefill:
|
|
for t in to_prefill:
|
|
t.input_tokens = len(t.prompt_ids)
|
|
|
|
groups: Dict[Tuple[int, Optional[AttentionBackend]], List[Task]] = {}
|
|
for t in to_prefill:
|
|
start_pos = min(
|
|
self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1
|
|
)
|
|
groups.setdefault((start_pos, t.backend), []).append(t)
|
|
|
|
for (start_pos, _), group in groups.items():
|
|
backend = group[0].backend
|
|
backend_context = (
|
|
attn_backend(backend) if backend is not None else nullcontext()
|
|
)
|
|
with (
|
|
backend_context,
|
|
self._metrics.record([t.task_id for t in group], "prefill"),
|
|
):
|
|
prefilled, step_out = self._executor.execute_prefill(
|
|
group, start_pos=start_pos, return_logprobs=return_logprobs
|
|
)
|
|
|
|
for t, out in zip(prefilled, step_out):
|
|
t.output_ids.append(out[0] if return_logprobs else out)
|
|
t.output_tokens += 1
|
|
t.mark_prefill_done()
|
|
prefilled_ids.add(t.task_id)
|
|
produced.append(t)
|
|
|
|
start_logical_page = start_pos // self._pool.page_size
|
|
for t in group:
|
|
self._task_cache.task_record_hashes(
|
|
t.task_id, t.prompt_ids, start_logical_page
|
|
)
|
|
|
|
decoded: List[Task] = []
|
|
aborted: List[Task] = []
|
|
for t in tasks:
|
|
if t.task_id in prefilled_ids:
|
|
continue
|
|
if self._task_cache.task_extend(t.task_id, t.next_pos):
|
|
decoded.append(t)
|
|
else:
|
|
t.status = TaskStatus.ABORTED
|
|
aborted.append(t)
|
|
|
|
for backend, group in self._task_backend_groups(decoded):
|
|
backend_context = (
|
|
attn_backend(backend) if backend is not None else nullcontext()
|
|
)
|
|
with (
|
|
backend_context,
|
|
self._metrics.record([t.task_id for t in group], "decode"),
|
|
):
|
|
step_out = self._executor.execute_decode(
|
|
group, return_logprobs=return_logprobs
|
|
)
|
|
for t, out in zip(group, step_out):
|
|
t.output_ids.append(out[0] if return_logprobs else out)
|
|
t.output_tokens += 1
|
|
t.advance_kv()
|
|
produced.append(t)
|
|
|
|
return produced, aborted
|