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
This commit is contained in:
2026-08-07 23:00:25 +08:00
parent 02469887f5
commit 184fbbce5c
2 changed files with 57 additions and 30 deletions
+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]