From 184fbbce5c713d21ebb8db3b97e17614ded15f32 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 7 Aug 2026 23:00:25 +0800 Subject: [PATCH] refactor: extract shared steady-state increment detection - 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 --- astrai/inference/core/cache.py | 49 +++++++++++++++++++++---------- astrai/inference/core/executor.py | 38 +++++++++++++++--------- 2 files changed, 57 insertions(+), 30 deletions(-) diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index 86caa70..dd7dcb8 100644 --- a/astrai/inference/core/cache.py +++ b/astrai/inference/core/cache.py @@ -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] diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index 99cad0a..6705172 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -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)