docs : clarify radix cache architecture

- document exact page-aligned radix prefix matching
- explain partial-page ownership and materialized KV boundaries
- remove bilingual wording from project overview

README.md

@@ -42,14 +42,14 @@
 
 ### Overview
 
-AstrAI is an end-to-end framework for building, training, evaluating, and serving bilingual Chinese-English Transformer models. It provides a compact PyTorch codebase for the complete model lifecycle, from declarative data preprocessing and distributed training to continuous-batching inference and OpenAI/Anthropic-compatible APIs.
+AstrAI is an end-to-end Transformer framework for building, training, evaluating, and serving models. It provides a compact PyTorch codebase for the complete model lifecycle, from declarative data preprocessing and distributed training to continuous-batching inference and OpenAI/Anthropic-compatible APIs.
 
 | Area | Capabilities |
 |---|---|
 | **Models** | Autoregressive language models and embedding models with GQA, MLA, MoE, RoPE, and extensible attention/FFN components |
 | **Training** | Pre-training (`seq`), supervised fine-tuning (`sft`), DPO, and GRPO with gradient accumulation, checkpointing, DDP, and FSDP |
 | **Data** | Declarative JSON preprocessing, configurable masking and packing, binary/JSONL storage, and streaming datasets |
-| **Inference** | Continuous batching, paged KV cache, prefix caching, streaming generation, and Torch/CUDA/FlashAttention backends |
+| **Inference** | Continuous batching, paged KV cache, radix prefix caching, streaming generation, and Torch/CUDA/FlashAttention backends |
 | **Serving** | FastAPI server with OpenAI and Anthropic chat completion protocols, including SSE streaming and tool calls |
 | **Evaluation** | Perplexity, MMLU, HumanEval, IFEval, IFD, and ROUGE evaluation tools |
 | **Extensibility** | Factory and registry architecture for models, datasets, training strategies, callbacks, kernels, and protocol components |

docs/README-zh-CN.md

@@ -48,14 +48,14 @@
 
 ### 项目概览
 
-AstrAI 是一个面向中英双语 Transformer 模型的端到端框架，覆盖模型构建、训练、评测与部署。项目以精简的 PyTorch 代码实现完整模型生命周期，包括声明式数据预处理、分布式训练、连续批处理推理，以及兼容 OpenAI 和 Anthropic 的服务接口。
+AstrAI 是一个覆盖模型构建、训练、评测与部署的端到端 Transformer 框架。项目以精简的 PyTorch 代码实现完整模型生命周期，包括声明式数据预处理、分布式训练、连续批处理推理，以及兼容 OpenAI 和 Anthropic 的服务接口。
 
 | 领域 | 能力 |
 |---|---|
 | **模型** | 自回归语言模型与嵌入模型，支持 GQA、MLA、MoE、RoPE，以及可扩展的 Attention/FFN 组件 |
 | **训练** | 预训练（`seq`）、监督微调（`sft`）、DPO 和 GRPO，支持梯度累积、检查点、DDP 与 FSDP |
 | **数据** | 声明式 JSON 预处理、可配置掩码与样本打包、二进制/JSONL 存储和流式数据集 |
-| **推理** | 连续批处理、分页 KV Cache、前缀缓存、流式生成，以及 Torch/CUDA/FlashAttention 后端 |
+| **推理** | 连续批处理、分页 KV Cache、Radix 前缀缓存、流式生成，以及 Torch/CUDA/FlashAttention 后端 |
 | **服务** | 基于 FastAPI 的 OpenAI 与 Anthropic 聊天补全协议，支持 SSE 流式输出和工具调用 |
 | **评测** | Perplexity、MMLU、HumanEval、IFEval、IFD 和 ROUGE 评测工具 |
 | **扩展** | 基于工厂与注册表扩展模型、数据集、训练策略、回调、内核和协议组件 |

docs/developer/architecture.md

@@ -870,12 +870,21 @@ classDiagram
             +ref_count(idx) int
         }
 
-        class PrefixCache {
+        class RadixNode {
+            +RadixNode parent
+            +Dict children
+            +Optional[int] page_idx
+            +Tuple tokens
+            +int lock_ref
+        }
+
+        class RadixCache {
             +int _page_size
             +evict(page_idx)
             +has_page(idx) bool
             +lookup(token_ids) List[int]
             +record(page_idx, token_ids, logical_page_idx)
+            +release(pages)
         }
 
         class KVStorage {
@@ -914,7 +923,7 @@ classDiagram
             -KVStorage _storage
             -ReqToTokenPool _req_pool
             -Allocator _alloc
-            -PrefixCache _prefix
+            -RadixCache _prefix
             +task_alloc(task_id, prompt_ids) bool
             +task_free(task_id)
             +task_extend(task_id, pos) bool
@@ -1321,7 +1330,8 @@ classDiagram
     PagePool *-- KVStorage
     PagePool *-- ReqToTokenPool
     PagePool *-- Allocator
-    PagePool *-- PrefixCache
+    PagePool *-- RadixCache
+    RadixCache *-- RadixNode
     InferenceEngine *-- InferenceScheduler
     InferenceScheduler *-- PagePool
     InferenceScheduler *-- Executor
@@ -1433,7 +1443,7 @@ classDiagram
 | **astrai.model** | ModelFactory, 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, InferenceWorkspace, 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.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, 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.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, attn_paged_prefill, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
 | **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 | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |

docs/developer/internals.md

@@ -168,9 +168,11 @@ 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.
+- **Allocator + RadixCache**: Paged-mode allocation with ref-counting, LRU eviction, and exact page-aligned prefix sharing when `page_size > 1`.
 
-`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. `bind_tasks()` returns a `KVCache` dataclass with `kv_indptr`, a prefix-sum index over sequence lengths computed once per step and shared across layers. Attention layers access buffers directly — no methods, no abstraction.
+`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand. `RadixCache` walks exact token-page edges from the root, preserving parent-prefix context instead of treating a page hash as a globally unique key. Only complete pages whose KV entries have been materialized are shared; partial pages remain request-private and are released at completion. The final sampled token is excluded because it has not yet been decoded into KV.
+
+`bind_tasks()` returns a `KVCache` dataclass with `kv_indptr`, a prefix-sum index over sequence lengths computed once per step and shared across layers. Attention layers access buffers directly — no methods, no abstraction.
 
 ### Attention Backend
 

docs/guides/inference.md

@@ -31,13 +31,17 @@ 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)
+  └── RadixCache             exact, page-aligned prefix matching (paged mode, page_size > 1)
 ```
 
 `PagePool` supports two modes:
 
 - **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.
+- **Paged** (`page_size=1` or `>1` with `n_tokens` set): shared token pool with on-demand allocation. `Allocator` provides ref-counted allocation and LRU eviction. When `page_size > 1`, `RadixCache` also enables prefix sharing.
+
+`RadixCache` indexes complete token pages as parent-linked radix edges. Lookup walks from the root and compares each page's exact token tuple, so an identical page can only be reused under the same parent prefix. Hash values are retained for introspection, but never determine a match.
+
+Only fully materialized KV pages enter the radix. A partial final page remains private to its request and is released when the request ends. On completion, the scheduler records the prompt plus generated tokens already decoded into KV; it excludes the final sampled token because that token has not yet passed through the model. A later request resumes prefill immediately after the longest complete-page hit.
 
 `bind_tasks()` returns a `KVCache` dataclass — pure data, no methods:
 
@@ -97,7 +101,7 @@ attention backends share the same rotary dispatch — it is backend-agnostic.
 `InferenceScheduler` runs a daemon thread with a 4-phase loop:
 
 ```
-1. Cleanup → Remove finished tasks, free KV cache slots/pages
+1. Cleanup → Record complete materialized pages, then release task-owned KV resources
 2. Refill  → Pop from waiting_queue, task_alloc resources, activate
 3. Prefill → Group by (prompt_len, start_pos), run full forward
 4. Decode  → Run single-token forward for each same-position group