refactor: use single-index access and update docs for cache architecture

- Replace all buffer[layer_id][loc] double indexing with buffer[layer_id, loc] single advanced indexing in cache.py and attention.py
- Revert KVStorage buffers back to 4D [n_layers, size, n_kv_heads, head_dim], remove leftover 3D reshape/view in MLA path
- Update docs/guides/inference.md, docs/developer/internals.md, docs/developer/architecture.md to reflect new PagePool/KVStorage/ReqToTokenPool/KVCache classes

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.k_buffer[layer_id, loc] = k
-        self.v_buffer[layer_id][loc] = v
+        self.v_buffer[layer_id, loc] = v
 
 
 @dataclass

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.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.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]
+            k = kv_cache.k_buffer[self.layer_id, indices]
-            v = kv_cache.v_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.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.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]
+            k = kv_cache.k_buffer[self.layer_id, indices]
-            v = kv_cache.v_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)

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 {
+        class KVStorage {
-            +int page_size
+            +int size
-            +Tensor k_cache
+            +Tensor k_buffer
-            +Tensor v_cache
+            +Tensor v_buffer
-            +write(layer_id, page_table, start_pos, k, v)
+            +get_key_buffer(layer_id) Tensor
-            +gather(layer_id, page_table, total_len) Tuple[Tensor, 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 {
-            <<abstract>>
+            +Tensor k_buffer
-            +task_alloc(task_id, prompt_ids) bool
+            +Tensor v_buffer
-            +task_free(task_id)
+            +Tensor req_to_token
-            +task_extend(task_id, pos) bool
+            +Tensor req_pool_indices
-            +task_cached(task_id) int
+            +Tensor seq_lens
-            +task_record_hashes(task_id, prompt_ids, start_logical_page)
+            +Tensor out_cache_loc
-            +bind_tasks(task_ids, total_len, device) CacheView
         }
 
-        class PageCache {
+        class PagePool {
             +int page_size
-            -PagePool _pool
+            +bool contiguous
-            -Storage _storage
+            -KVStorage _storage
-            -TaskTable _table
+            -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
+            +bind_tasks(task_ids, seq_lens, device, start_pos) KVCache
-        }
-
-        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 {
-            <<abstract>>
-            +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
         }
 
     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
+    PagePool *-- KVStorage
-    PageCache *-- Storage
+    PagePool *-- ReqToTokenPool
-    PageCache *-- TaskTable
+    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
+    PagePool ..> KVCache : binds
-    ContiguousCache ..> ContiguousCacheView : 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 |

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.
+Three-layer separation (SGLang-inspired):
-- **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.
+
+- **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
 

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:
+Three-layer separation (SGLang-inspired): storage, index table, allocator.
-
-### ContiguousCache (default)
 
 ```
-ContiguousCache (simple contiguous per-slot cache)
+PagePool (top-level manager, orchestrates all layers)
-  ├── ContiguousCacheView  bundles k/v tensors + slot indices for attention 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)
+KVCache
-  ├── PagePool               orchestrates page allocation + prefix matching
+  ├── k_buffer, v_buffer     [n_layers, size, n_kv_heads, head_dim]
-  │     ├── Allocator         bitmask-based page allocator + ref-count + LRU
+  ├── req_to_token           [num_reqs, max_ctx_len]
-  │     └── PrefixCache       hash-based prefix matching (page_hash via polynomial hash)
+  ├── req_pool_indices       [batch_size]
-  ├── TaskTable              maps task_id → page_table + cached token count
+  ├── seq_lens               [batch_size]
-  ├── Storage                k_cache / v_cache tensors (num_hidden_layers × n_pages × page_size × num_key_value_heads × head_dim)
+  └── out_cache_loc          [batch, seq_len] — write indices for this forward
-  └── PageCacheView          bundles Storage + page_table + total_len for attention layers
 ```
 
-`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
 

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.k_buffer[0, kv.out_cache_loc] = k
-    kv.v_buffer[0][kv.out_cache_loc] = v
+    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]
-    gathered_v = kv.v_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.k_buffer[0, kv.out_cache_loc] = k
-    kv.v_buffer[0][kv.out_cache_loc] = v
+    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.k_buffer[0, kv.out_cache_loc] = k
-    kv.v_buffer[0][kv.out_cache_loc] = v
+    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)