refactor: harden inference cache state and attention dispatch
- split KVCache into phase-specific PrefillKVCache/DecodeKVCache types selected by start_pos - unify steady-state detection in TaskCacheManager - guard decode steady-state reuse with the cached task signature so recycled req slots cannot replay a prior generation's tokens and positions - collapse attention backend fwd_decode/fwd_prefill into a single subclass-owned forward with a shared _check_fwd guard - fix thread-safety gap in weight update and validate prefill inputs before KV allocation - centralize magic constants in InferenceConfig and align docs with behavior
This commit is contained in:
Vendored
+40
-13
@@ -2,7 +2,10 @@
|
||||
|
||||
Layer 1 — ``KVStorage``: flat token-level K/V GPU buffers [n_layers, size, n_kv_heads, head_dim]
|
||||
Layer 2 — ``ReqToTokenPool``: index table [req_idx, pos] → physical token slot
|
||||
Layer 3 — ``KVCache``: pure dataclass passed to the model for direct buffer access
|
||||
Layer 3 — ``BaseKVCache``: shared fields for all cache modes
|
||||
``PrefillKVCache``: prefill-specific layout
|
||||
``DecodeKVCache``: decode-specific layout
|
||||
``KVCache``: union type for backward compatibility
|
||||
|
||||
These classes have no knowledge of tasks, allocation policies, or scheduling.
|
||||
They are the "dumb" physical storage layer.
|
||||
@@ -10,7 +13,7 @@ They are the "dumb" physical storage layer.
|
||||
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
@@ -74,8 +77,8 @@ class KVStorage:
|
||||
|
||||
|
||||
@dataclass
|
||||
class KVCache:
|
||||
"""Pure data struct passed to model for KV cache I/O.
|
||||
class BaseKVCache:
|
||||
"""Shared fields for all KV cache modes.
|
||||
|
||||
The attention layer does raw buffer indexing — no methods, no abstraction.
|
||||
"""
|
||||
@@ -85,12 +88,36 @@ class KVCache:
|
||||
req_to_token: Tensor
|
||||
req_pool_indices: Tensor
|
||||
seq_lens: Tensor
|
||||
out_cache_loc: Tensor
|
||||
max_len: int = 0
|
||||
kv_indptr: Optional[Tensor] = None
|
||||
qo_indptr: Optional[Tensor] = None
|
||||
q_tile_to_batch: Optional[Tensor] = None
|
||||
q_tile_to_index: Optional[Tensor] = None
|
||||
decode_o_part: Optional[Tensor] = None
|
||||
decode_ml_part: Optional[Tensor] = None
|
||||
decode_out: Optional[Tensor] = None
|
||||
max_len: int
|
||||
kv_indptr: Tensor # Always present in both modes
|
||||
|
||||
|
||||
@dataclass
|
||||
class PrefillKVCache(BaseKVCache):
|
||||
"""Prefill-specific KV cache layout.
|
||||
|
||||
Handles packed ragged batching where prompts have variable lengths.
|
||||
"""
|
||||
|
||||
out_cache_loc: Tensor # [total_q_tokens] - flattened write locations
|
||||
qo_indptr: Tensor # [B+1] - prefix sum of q_lens for unpacking
|
||||
q_tile_to_batch: Tensor # [num_q_tiles] - maps Q tiles to batch indices
|
||||
q_tile_to_index: Tensor # [num_q_tiles] - maps Q tiles to local indices
|
||||
|
||||
|
||||
@dataclass
|
||||
class DecodeKVCache(BaseKVCache):
|
||||
"""Decode-specific KV cache layout.
|
||||
|
||||
Single-token incremental generation with split-KV partial results.
|
||||
"""
|
||||
|
||||
out_cache_loc: Tensor # [B] - one write position per request
|
||||
qo_indptr: Tensor # [B+1] - sequential [0, 1, 2, ..., B]
|
||||
decode_o_part: Tensor # [B, max_q_heads, MAX_SPLITS, head_dim]
|
||||
decode_ml_part: Tensor # [B, max_q_heads, MAX_SPLITS, 2]
|
||||
decode_out: Tensor # [B, max_q_heads, head_dim]
|
||||
|
||||
|
||||
# Backward compatibility: union type for existing code
|
||||
KVCache = Union[PrefillKVCache, DecodeKVCache]
|
||||
|
||||
Vendored
+54
-19
@@ -16,7 +16,13 @@ from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
|
||||
from astrai.inference.cache.buffer import (
|
||||
DecodeKVCache,
|
||||
KVCache,
|
||||
KVStorage,
|
||||
PrefillKVCache,
|
||||
ReqToTokenPool,
|
||||
)
|
||||
from astrai.inference.cache.strategy import (
|
||||
AllocationStrategy,
|
||||
Allocator,
|
||||
@@ -31,6 +37,8 @@ from astrai.inference.workspace import Q_TILE_ROWS, InferenceWorkspace
|
||||
# continues to work unchanged after the file split.
|
||||
__all__ = [
|
||||
"KVCache",
|
||||
"PrefillKVCache",
|
||||
"DecodeKVCache",
|
||||
"KVStorage",
|
||||
"ReqToTokenPool",
|
||||
"Allocator",
|
||||
@@ -232,7 +240,20 @@ class PagePool:
|
||||
)
|
||||
q_tile_to_batch = workspace.q_tile_to_batch[:n_tiles]
|
||||
q_tile_to_index = workspace.q_tile_to_index[:n_tiles]
|
||||
decode_o_part = decode_ml_part = decode_out = None
|
||||
|
||||
return PrefillKVCache(
|
||||
k_buffer=self._storage.k_buffer,
|
||||
v_buffer=self._storage.v_buffer,
|
||||
req_to_token=self._req_pool.req_to_token,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=seq_lens_t,
|
||||
max_len=max(seq_lens),
|
||||
kv_indptr=kv_indptr,
|
||||
out_cache_loc=out_cache_loc,
|
||||
qo_indptr=qo_indptr,
|
||||
q_tile_to_batch=q_tile_to_batch,
|
||||
q_tile_to_index=q_tile_to_index,
|
||||
)
|
||||
else:
|
||||
# ---- decode: out_cache_loc is a single column (last position) ----
|
||||
write_pos = seq_lens_t - 1
|
||||
@@ -241,27 +262,24 @@ class PagePool:
|
||||
out_cache_loc = ocl_buf[:b].reshape(-1)
|
||||
workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1])
|
||||
qo_indptr = workspace.qo_indptr[: b + 1]
|
||||
q_tile_to_batch = q_tile_to_index = None
|
||||
decode_o_part = getattr(workspace, "decode_o_part", None)
|
||||
decode_ml_part = getattr(workspace, "decode_ml_part", None)
|
||||
decode_out = getattr(workspace, "decode_out", None)
|
||||
|
||||
return KVCache(
|
||||
k_buffer=self._storage.k_buffer,
|
||||
v_buffer=self._storage.v_buffer,
|
||||
req_to_token=self._req_pool.req_to_token,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=seq_lens_t,
|
||||
out_cache_loc=out_cache_loc,
|
||||
max_len=max(seq_lens),
|
||||
kv_indptr=kv_indptr,
|
||||
qo_indptr=qo_indptr,
|
||||
q_tile_to_batch=q_tile_to_batch,
|
||||
q_tile_to_index=q_tile_to_index,
|
||||
decode_o_part=decode_o_part,
|
||||
decode_ml_part=decode_ml_part,
|
||||
decode_out=decode_out,
|
||||
)
|
||||
return DecodeKVCache(
|
||||
k_buffer=self._storage.k_buffer,
|
||||
v_buffer=self._storage.v_buffer,
|
||||
req_to_token=self._req_pool.req_to_token,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=seq_lens_t,
|
||||
max_len=max(seq_lens),
|
||||
kv_indptr=kv_indptr,
|
||||
out_cache_loc=out_cache_loc,
|
||||
qo_indptr=qo_indptr,
|
||||
decode_o_part=decode_o_part,
|
||||
decode_ml_part=decode_ml_part,
|
||||
decode_out=decode_out,
|
||||
)
|
||||
|
||||
|
||||
class TaskCacheManager:
|
||||
@@ -385,8 +403,25 @@ class TaskCacheManager:
|
||||
|
||||
@property
|
||||
def bind_was_steady(self) -> bool:
|
||||
"""True if the last bind was a steady-state increment (same tasks, +1 seq_lens)."""
|
||||
return self._bind_was_steady
|
||||
|
||||
def last_task_signature_matches(self, task_ids: List[str]) -> bool:
|
||||
"""Check if task_ids match the previous bind's signature.
|
||||
|
||||
Used by Executor to detect steady-state decode for device-to-device
|
||||
token copy optimization.
|
||||
"""
|
||||
if self._bind_state is None:
|
||||
return False
|
||||
prev_sig = self._bind_state.sig
|
||||
# sig is tuple of req_indices, need to map task_ids to req_indices
|
||||
try:
|
||||
current_sig = tuple(self._states[tid].req_idx for tid in task_ids)
|
||||
return prev_sig == current_sig
|
||||
except KeyError:
|
||||
return False
|
||||
|
||||
# -- internals --
|
||||
|
||||
def _rollback(self, state: TaskCacheState, task_id: str):
|
||||
|
||||
Reference in New Issue
Block a user