diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index a710cfc..94f8217 100644 --- a/astrai/inference/core/cache.py +++ b/astrai/inference/core/cache.py @@ -155,10 +155,10 @@ class ReqToTokenPool: class KVStorage: - """Token-level flat KV cache storage with NHD layout. + """Token-level KV cache storage. Buffers: [n_layers, size, n_kv_heads, head_dim]. Each token occupies - one contiguous row. Logical ordering is determined by ReqToTokenPool. + one slot indexed by ReqToTokenPool. """ def __init__( @@ -185,8 +185,8 @@ class KVStorage: return self.v_buffer[layer_id] def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None: - self.k_buffer[layer_id][loc] = k - self.v_buffer[layer_id][loc] = v + self.k_buffer[layer_id, loc] = k + self.v_buffer[layer_id, loc] = v @dataclass diff --git a/astrai/model/components/attention.py b/astrai/model/components/attention.py index 1fb0acd..4f4a0b3 100644 --- a/astrai/model/components/attention.py +++ b/astrai/model/components/attention.py @@ -87,8 +87,8 @@ class GQA(nn.Module): q, k = self.q_norm(q), self.k_norm(k) if kv_cache is not None: - kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k - kv_cache.v_buffer[self.layer_id][kv_cache.out_cache_loc] = v + kv_cache.k_buffer[self.layer_id, kv_cache.out_cache_loc] = k + kv_cache.v_buffer[self.layer_id, kv_cache.out_cache_loc] = v max_len = kv_cache.seq_lens.max() indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len] @@ -97,8 +97,8 @@ class GQA(nn.Module): < kv_cache.seq_lens[:, None] ) indices = torch.where(pos_mask, indices, torch.zeros_like(indices)) - k = kv_cache.k_buffer[self.layer_id][indices] - v = kv_cache.v_buffer[self.layer_id][indices] + k = kv_cache.k_buffer[self.layer_id, indices] + v = kv_cache.v_buffer[self.layer_id, indices] k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep) @@ -204,8 +204,8 @@ class MLA(nn.Module): k = self.k_norm(k) if kv_cache is not None: - kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k - kv_cache.v_buffer[self.layer_id][kv_cache.out_cache_loc] = v + kv_cache.k_buffer[self.layer_id, kv_cache.out_cache_loc] = k + kv_cache.v_buffer[self.layer_id, kv_cache.out_cache_loc] = v max_len = kv_cache.seq_lens.max() indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len] @@ -214,8 +214,8 @@ class MLA(nn.Module): < kv_cache.seq_lens[:, None] ) indices = torch.where(pos_mask, indices, torch.zeros_like(indices)) - k = kv_cache.k_buffer[self.layer_id][indices] - v = kv_cache.v_buffer[self.layer_id][indices] + k = kv_cache.k_buffer[self.layer_id, indices] + v = kv_cache.v_buffer[self.layer_id, indices] q = q.permute(0, 2, 1, 3) k = k.permute(0, 2, 1, 3) diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index be8d2bc..3158307 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -277,7 +277,7 @@ classDiagram +ModuleList layers +RMSNorm norm +Linear lm_head - +forward(input_ids, input_mask, paged_cache, position_ids) Dict[str, Tensor] + +forward(input_ids, input_mask, kv_cache, position_ids) Dict[str, Tensor] +load_state_dict(state_dict, strict, assign) +state_dict() } @@ -299,7 +299,7 @@ classDiagram +RMSNorm input_norm +nn.Module mlp # MLP or DeepSeekMoE via FFNFactory +RMSNorm post_attention_norm - +forward(x, rotary_emb, attention_mask, paged_cache) Tensor + +forward(x, rotary_emb, attention_mask, kv_cache) Tensor } class GQA { @@ -314,7 +314,7 @@ classDiagram +Linear q_proj, k_proj, v_proj, o_proj +Linear gate # only if use_gated_attention +RMSNorm q_norm, k_norm # only if use_qk_norm - +forward(x, rotary_emb, attn_mask, paged_cache) Tensor + +forward(x, rotary_emb, attn_mask, kv_cache) Tensor } class MLA { @@ -334,7 +334,7 @@ classDiagram +Linear gate # only if use_gated_attention +RMSNorm kv_norm +RMSNorm q_norm, k_norm # only if use_qk_norm - +forward(x, rotary_emb, attn_mask, paged_cache) Tensor + +forward(x, rotary_emb, attn_mask, kv_cache) Tensor } class MLP { @@ -824,75 +824,46 @@ classDiagram +record(page_idx, token_ids, logical_page_idx) } - class Storage { - +int page_size - +Tensor k_cache - +Tensor v_cache - +write(layer_id, page_table, start_pos, k, v) - +gather(layer_id, page_table, total_len) Tuple[Tensor, Tensor] + class KVStorage { + +int size + +Tensor k_buffer + +Tensor v_buffer + +get_key_buffer(layer_id) Tensor + +get_value_buffer(layer_id) Tensor + +set_kv_buffer(layer_id, loc, k, v) + } + + class ReqToTokenPool { + +int size + +int max_context_len + +Tensor req_to_token + +alloc(num_reqs) List[int] + +free(req_indices) + +write(indices, values) } class KVCache { - <> - +task_alloc(task_id, prompt_ids) bool - +task_free(task_id) - +task_extend(task_id, pos) bool - +task_cached(task_id) int - +task_record_hashes(task_id, prompt_ids, start_logical_page) - +bind_tasks(task_ids, total_len, device) CacheView + +Tensor k_buffer + +Tensor v_buffer + +Tensor req_to_token + +Tensor req_pool_indices + +Tensor seq_lens + +Tensor out_cache_loc } - class PageCache { + class PagePool { +int page_size - -PagePool _pool - -Storage _storage - -TaskTable _table + +bool contiguous + -KVStorage _storage + -ReqToTokenPool _req_pool + -Allocator _alloc + -PrefixCache _prefix +task_alloc(task_id, prompt_ids) bool +task_free(task_id) +task_extend(task_id, pos) bool +task_cached(task_id) int +task_record_hashes(task_id, prompt_ids, start_logical_page) - +bind_tasks(task_ids, total_len, device) PageCacheView - } - - class ContiguousCache { - +int max_seq_len - +Tensor k, v - +task_alloc(task_id, prompt_ids) bool - +task_free(task_id) - +task_extend(task_id, pos) bool - +bind_tasks(task_ids, total_len, device) ContiguousCacheView - } - - class CacheView { - <> - +write(layer_id, k, v) - +gather(layer_id) Tuple[Tensor, Tensor] - } - - class PageCacheView { - -Storage _storage - +Tensor _page_table - +int _total_len - +write(layer_id, k, v) - +gather(layer_id) Tuple[Tensor, Tensor] - } - - class ContiguousCacheView { - -ContiguousCache _cache - +Tensor _batch_indices - +int _total_len - +write(layer_id, k, v) - +gather(layer_id) Tuple[Tensor, Tensor] - } - - class TaskTable { - +set(task_id, page_table, cached) - +get(task_id) List[int] - +get_cached(task_id) int - +get_ref(task_id) List[int] - +pop(task_id) Tuple[List[int], int] - +table_tensor(task_ids, device) Tensor + +bind_tasks(task_ids, seq_lens, device, start_pos) KVCache } class Task { @@ -1314,17 +1285,13 @@ classDiagram RawRollout <|-- RolloutResult LaunchStrategy <|-- TorchrunStrategy LaunchStrategy <|-- LocalStrategy - KVCache <|-- PageCache - KVCache <|-- ContiguousCache - CacheView <|-- PageCacheView - CacheView <|-- ContiguousCacheView - %% --- Composition (strong ownership, part destroyed with whole) --- - PageCache *-- PagePool - PageCache *-- Storage - PageCache *-- TaskTable + PagePool *-- KVStorage + PagePool *-- ReqToTokenPool + PagePool *-- Allocator + PagePool *-- PrefixCache InferenceEngine *-- InferenceScheduler - InferenceScheduler *-- KVCache + InferenceScheduler *-- PagePool InferenceScheduler *-- Executor InferenceScheduler *-- TaskManager AutoRegressiveLM *-- DecoderBlock @@ -1352,8 +1319,6 @@ classDiagram TrainContext o-- BaseScheduler TrainContext o-- Checkpoint TrainContext o-- BaseExecutor - PageCacheView o-- Storage - ContiguousCacheView o-- ContiguousCache SamplingPipeline o-- BaseSamplingStrategy BaseDataset o-- Store Pipeline o-- PipelineConfig @@ -1398,8 +1363,7 @@ classDiagram TrainContextBuilder ..> RDSampler : creates Checkpoint ..> Checkpoint : serializes CheckpointCallback ..> Checkpoint : creates - PageCache ..> PageCacheView : binds - ContiguousCache ..> ContiguousCacheView : binds + PagePool ..> KVCache : binds InferenceEngine ..> GenerationRequest : uses InferenceEngine ..> GenerateResult : creates OpenAIResponseBuilder ..> ChatCompletionRequest : receives @@ -1436,7 +1400,7 @@ classDiagram | **astrai.model** | AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model | | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template | | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow | -| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–ContiguousCache/PageCache, CacheView–ContiguousCacheView/PageCacheView, Allocator–Storage, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service | +| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, PrefixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service | | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation | | **astrai.factory** | BaseFactory | Component registration | | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers | diff --git a/docs/developer/internals.md b/docs/developer/internals.md index a2ec20d..2d98bc5 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -143,10 +143,15 @@ The cache stores $k_j$ and $v_j$ for all previous positions. At each decode step If RoPE were applied after caching, the rotation factors would be inconsistent between cached and new tokens. -### Cache Implementations +### Cache Architecture -- **ContiguousCache**: Each task gets a fixed slot of `[max_seq_len, num_key_value_heads, head_dim]`. Simple, efficient for small-to-medium batch sizes. -- **PageCache**: Paged KV cache with prefix sharing. Uses `PagePool` (allocator + LRU + prefix matching) and `Storage` (page tensors). Enables sharing of common prompt prefixes across requests. +Three-layer separation (SGLang-inspired): + +- **KVStorage**: Flat token-level buffers `[n_layers, size, n_kv_heads, head_dim]`. +- **ReqToTokenPool**: Index table `[req_idx, pos] → physical token slot`, shared across all layers. +- **Allocator + PrefixCache**: Paged-mode slot allocation with ref-counting, LRU eviction, and hash-based prefix sharing. + +`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. Attention layers access buffers directly via `KVCache` dataclass — no methods, no abstraction. ## Mask Algorithm Internals diff --git a/docs/guides/inference.md b/docs/guides/inference.md index 8eb2a0c..98b852d 100644 --- a/docs/guides/inference.md +++ b/docs/guides/inference.md @@ -23,30 +23,33 @@ RoPE is applied **before** KV cache write, not after — otherwise position enco ## KVCache System -Seven classes working together, with two concrete cache implementations: - -### ContiguousCache (default) +Three-layer separation (SGLang-inspired): storage, index table, allocator. ``` -ContiguousCache (simple contiguous per-slot cache) - ├── ContiguousCacheView bundles k/v tensors + slot indices for attention layers +PagePool (top-level manager, orchestrates all layers) + ├── KVStorage k_buffer / v_buffer [n_layers, size, n_kv_heads, head_dim] + ├── ReqToTokenPool req_to_token [num_reqs, max_ctx_len] → physical token slot + ├── Allocator bitmask-based page allocator + ref-count + LRU (paged mode only) + └── PrefixCache hash-based prefix matching (paged mode only) ``` -Created by default when no cache is passed to `InferenceScheduler`. Each task occupies a fixed slot of `[max_seq_len, num_key_value_heads, head_dim]`. Simple and efficient for small-to-medium batch sizes. +`PagePool` supports two modes: -### PageCache (paged with prefix sharing) +- **Contiguous (default)**: pre-allocates `max_batch_size * max_seq_len` token slots. `req_to_token` is a trivial linear mapping (`slot = req_idx * max_seq_len + pos`). No dynamic allocation. +- **Paged** (`page_size=1` or `>1` with `n_tokens` set): shared token pool with on-demand allocation. Allocator + PrefixCache enable prefix sharing and LRU eviction. + +`bind_tasks()` returns a `KVCache` dataclass — pure data, no methods: ``` -PageCache (paged KV cache with prefix sharing, alternative) - ├── PagePool orchestrates page allocation + prefix matching - │ ├── Allocator bitmask-based page allocator + ref-count + LRU - │ └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) - ├── TaskTable maps task_id → page_table + cached token count - ├── Storage k_cache / v_cache tensors (num_hidden_layers × n_pages × page_size × num_key_value_heads × head_dim) - └── PageCacheView bundles Storage + page_table + total_len for attention layers +KVCache + ├── k_buffer, v_buffer [n_layers, size, n_kv_heads, head_dim] + ├── req_to_token [num_reqs, max_ctx_len] + ├── req_pool_indices [batch_size] + ├── seq_lens [batch_size] + └── out_cache_loc [batch, seq_len] — write indices for this forward ``` -`isinstance(cache, KVCache)` checks dispatch to the correct view. Both implement the abstract `KVCache` interface used by `Executor` and `InferenceScheduler`. +Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather. ## Continuous Batching diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index 4cd5820..6cd08db 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -239,12 +239,12 @@ def test_page_pool_contiguous_bind_roundtrip(): kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0) k = torch.randn(1, 4, 2, 4) v = torch.randn(1, 4, 2, 4) - kv.k_buffer[0][kv.out_cache_loc] = k - kv.v_buffer[0][kv.out_cache_loc] = v + kv.k_buffer[0, kv.out_cache_loc] = k + kv.v_buffer[0, kv.out_cache_loc] = v indices = kv.req_to_token[kv.req_pool_indices, :4] - gathered_k = kv.k_buffer[0][indices] - gathered_v = kv.v_buffer[0][indices] + gathered_k = kv.k_buffer[0, indices] + gathered_v = kv.v_buffer[0, indices] assert torch.allclose(gathered_k, k) assert torch.allclose(gathered_v, v) @@ -301,11 +301,11 @@ def test_page_pool_paged_bind_roundtrip(): kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0) k = torch.randn(1, 4, 2, 4) v = torch.randn(1, 4, 2, 4) - kv.k_buffer[0][kv.out_cache_loc] = k - kv.v_buffer[0][kv.out_cache_loc] = v + kv.k_buffer[0, kv.out_cache_loc] = k + kv.v_buffer[0, kv.out_cache_loc] = v indices = kv.req_to_token[kv.req_pool_indices, :4] - gathered_k = kv.k_buffer[0][indices] + gathered_k = kv.k_buffer[0, indices] assert torch.allclose(gathered_k, k) @@ -352,9 +352,9 @@ def test_page_pool_paged_ps64_bind_roundtrip(): kv = pool.bind_tasks(["t1"], [128], torch.device("cpu"), start_pos=0) k = torch.randn(1, 128, 2, 4) v = torch.randn(1, 128, 2, 4) - kv.k_buffer[0][kv.out_cache_loc] = k - kv.v_buffer[0][kv.out_cache_loc] = v + kv.k_buffer[0, kv.out_cache_loc] = k + kv.v_buffer[0, kv.out_cache_loc] = v indices = kv.req_to_token[kv.req_pool_indices, :128] - gathered_k = kv.k_buffer[0][indices] + gathered_k = kv.k_buffer[0, indices] assert torch.allclose(gathered_k, k)