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

refactor: extract shared steady-state increment detection

Browse Source 
- add _BindState dataclass and _is_steady_increment() to cache.py
- replace _bind_sig/_bind_seq_lens dual fields with single _bind_state
- replace DecodeSteadyState bare tuple with named dataclass
- use _is_steady_increment() in both PagePool.bind_tasks and Executor.execute_decode
This commit is contained in:
ViperEkura
2026-08-07 23:00:25 +08:00
parent 02469887f5
commit 184fbbce5c
2 changed files with 57 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/core/cache.py
+33 -16
View File
@@ -13,9 +13,8 @@ Two modes:
"""
import threading
from collections import OrderedDict
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional
from typing import Callable, Dict, List, Optional, OrderedDict
import torch
from torch import Tensor
@@ -23,6 +22,30 @@ from torch import Tensor
from astrai.inference.core.workspace import InferenceWorkspace
@dataclass
class _BindState:
"""Cached bind metadata for steady-state decode increment detection."""
sig: tuple
seq_lens: List[int]
def _is_steady_increment(
prev_sig: Optional[tuple],
prev_vals: Optional[List[int]],
cur_sig: tuple,
cur_vals: List[int],
) -> bool:
"""True when the same ordered set has every value +1 from the previous step."""
return (
prev_sig is not None
and prev_vals is not None
and prev_sig == cur_sig
and len(prev_vals) == len(cur_vals)
and all(c == p + 1 for c, p in zip(cur_vals, prev_vals))
)
def page_hash(
token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
) -> int:
@@ -350,13 +373,10 @@ class PagePool:
self._task_pages: Dict[str, List[int]] = {}
self._lock = threading.Lock()
# Steady-state decode validation state: the ordered task set and its
# Python seq_lens mirror. When the same set advances every sequence
# by exactly one token per step, bind_tasks updates the stable
# buffers in-place (+=1 / +=inc) instead of re-cumsumming. Any
# task-set change is a miss and rebuilds.
self._bind_sig: Optional[tuple] = None
self._bind_seq_lens: Optional[List[int]] = None
# Steady-state decode validation: when the same ordered task set
# advances every sequence by exactly one token per step, bind_tasks
# updates the stable buffers in-place instead of re-cumsumming.
self._bind_state: Optional[_BindState] = None
# ---- task lifecycle ----
@@ -529,13 +549,11 @@ class PagePool:
inc_buf = workspace.inc
ocl_buf = workspace.out_cache_loc
prev = self._bind_state
incremental = (
start_pos is None
and self._bind_sig is not None
and self._bind_sig == sig
and self._bind_seq_lens is not None
and len(self._bind_seq_lens) == b
and all(s == p + 1 for s, p in zip(seq_lens, self._bind_seq_lens))
and prev is not None
and _is_steady_increment(prev.sig, prev.seq_lens, sig, seq_lens)
)
if incremental:
# Steady-state decode: advance the stable buffers in-place.
@@ -557,8 +575,7 @@ class PagePool:
req_pool_indices = rpi_buf[:b]
seq_lens_t = sl_buf[:b]
kv_indptr = kvp_buf[: b + 1]
self._bind_sig = sig
self._bind_seq_lens = list(seq_lens)
self._bind_state = _BindState(sig, list(seq_lens))
if start_pos is not None:
seq_len = seq_lens[0]
astrai/inference/core/executor.py
+24 -14
View File
@@ -14,7 +14,7 @@ from astrai.extension.attention_backend import (
attn_backend,
get_backend,
)
from astrai.inference.core.cache import PagePool
from astrai.inference.core.cache import PagePool, _is_steady_increment
from astrai.inference.core.graph import CudaGraphContext
from astrai.inference.core.task import Task
from astrai.inference.core.workspace import InferenceWorkspace
@@ -54,6 +54,19 @@ class SamplingBatchInfo:
has_freq: bool # any frequency_penalty != 0 (avoids per-step GPU .any())
@dataclass
class DecodeSteadyState:
"""Cached decode metadata for the steady-state case.
When the same ordered task set decodes one token per step, sampling
params and task signature are reused; only positions advance by 1.
"""
task_sig: tuple
positions: list[int]
sampling_info: SamplingBatchInfo
def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
pin = str(device).startswith("cuda")
freq_penalties = torch.tensor(
@@ -165,11 +178,10 @@ class Executor:
self.device = device or next(model.parameters()).device
self.dtype = dtype or next(model.parameters()).dtype
# Per-step decode cache for the steady-state case where the same
# ordered task set decodes one token per step. Sampling params are
# constant across steps; position_ids grows by exactly 1. Single-slot:
# any task-set change is a cache miss.
self._decode_cache: Optional[tuple] = None
# Per-step decode cache for the steady-state case (same ordered
# task set decodes one token per step). Sampling params stay
# constant; only positions advance.
self._decode_cache: Optional[DecodeSteadyState] = None
# Pre-allocated fixed-shape buffers for the decode hot path
# (input_ids, decode mask, KV bind metadata). Eagerly sized at init
@@ -340,20 +352,18 @@ class Executor:
sig = tuple(task_ids)
cached = self._decode_cache
if (
cached is not None
and cached[0] == sig
and cur_positions == [p + 1 for p in cached[1]]
):
info = cached[2]
prev_sig = cached.task_sig if cached is not None else None
prev_pos = cached.positions if cached is not None else None
if _is_steady_increment(prev_sig, prev_pos, sig, cur_positions):
info = cached.sampling_info
ws.position_ids[:b] += 1
self._decode_cache = (sig, cur_positions, info)
self._decode_cache = DecodeSteadyState(sig, cur_positions, info)
else:
info = _build_sampling_batch_info(tasks, self.device)
ws.position_ids[:b].copy_(
torch.tensor(cur_positions, dtype=torch.long, device=self.device)
)
self._decode_cache = (sig, cur_positions, info)
self._decode_cache = DecodeSteadyState(sig, cur_positions, info)
total_len = max(cur_positions) + 1
input_mask = ws.decode_mask(ws.position_ids[:b], total_len)