release: v1.3.4

- refactor: 分页 KV cache（PagedCache+CacheView）替换固定 slot，删除 PrefixCache
- refactor: 推理引擎控制逻辑重写，修复连续批处理核心缺陷、线程安全问题
- refactor: KV 缓存槽位下沉到注意力层，移除 _remap_kv / _writeback_kv
- refactor: 统一采样路径为 SamplingPipeline batch tensor，删除 apply_sampling_strategies
- refactor: 设计模式优化 inference 模块导入结构（cache/sampling 独立）
- feat: 推理引擎前缀缓存（KV cache 复用）
- feat: OpenAI 兼容 chat completion API（流式+非流式+usage）
- feat: Anthropic 兼容 /v1/messages API，移除旧版 /generate 端点
- feat: GRPO CLI 接入 + on-policy，OpenAI API top_k 参数化
- feat: Checkpoint 支持 extra 通用扩展数据
- feat: Docker Compose 一键部署（GPU/CPU 双模式）
- feat: GRPO 训练参数补充，批处理训练参数表
- fix: 调度器延迟优化 — 移除 5ms 睡眠，修复 refill 任务丢失
- fix: CLI 参数缺失/重复、device_ids 越界、generate 参数名不一致
- fix: 长对话截断方向错误，保留最新 token 而非最早
- fix: remove_task 未释放 KV cache slot 导致第二轮对话死锁
- fix: KV cache 槽位索引错位、版本校验缺失、注意力掩码
- fix: scheduler 越界 bug，SchedulerCallback 回调阶段修正
- perf: _Result 改用 Condition.wait_for 消除非流式 CPU 空转
- perf: decode 每步张量预分配；input_ids 改用一次构建代替逐元素赋值
- refactor: 移除 device_ids 参数，统一 CUDA_VISIBLE_DEVICES
- docs: 更新文档以匹配分页 KV cache 等代码重构
- docs: 修正多处文档错误、补充训练参数说明

.dockerignore

@@ -3,6 +3,7 @@
 
 # Allow necessary files
 !astrai/
-!scripts/tools/
+!scripts/
+!assets/
 !pyproject.toml
 !README.md

.github/workflows/docker.yml

@@ -0,0 +1,50 @@
+name: Build and Push Docker Image
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Extract metadata
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ghcr.io/${{ github.repository }}
+          tags: |
+            type=ref,event=tag
+            type=raw,value=latest
+
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          platforms: linux/amd64
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

.gitignore

@@ -15,6 +15,7 @@
 !/.gitattributes
 !/.dockerignore
 !/Dockerfile
+!/docker-compose.yml
 !/assets/**
 !/CONTRIBUTING.md
 !/LICENSE

Dockerfile

@@ -1,49 +1,54 @@
-# AstrAI Dockerfile
+# AstrAI Dockerfile - Multi-stage Build (Optimized)
-# Multi-stage build for optimized image size
 
-# Build stage
+# Build stage - use base image with minimal build tools
-FROM python:3.12-slim AS builder
+FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS builder
 
 WORKDIR /app
 
-# Install build dependencies
+# Install Python 3.12 and minimal build dependencies
-RUN apt-get update && apt-get install -y --no-install-recommends \
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
-    build-essential \
+    python3.12 \
+    python3.12-dev \
+    python3.12-venv \
+    gcc \
+    g++ \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy project files first for version extraction
+# Create isolated virtual environment
+RUN python3.12 -m venv --copies /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# Copy source code and install dependencies
 COPY astrai/ ./astrai/
 COPY pyproject.toml .
-
-# Install dependencies
 RUN pip install --no-cache-dir --upgrade pip \
-    && pip install --no-cache-dir .[dev]
+    && pip install --no-cache-dir . \
+    --extra-index-url https://download.pytorch.org/whl/cu126
 
 # Production stage
-FROM python:3.12-slim AS production
+FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production
 
 WORKDIR /app
 
-# Install runtime dependencies
+# Install Python 3.12 runtime
-RUN apt-get update && apt-get install -y --no-install-recommends \
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    python3.12 \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy installed packages from builder
+# Copy virtual environment from builder
-COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
+COPY --from=builder /opt/venv /opt/venv
-COPY --from=builder /usr/local/bin /usr/local/bin
+ENV PATH="/opt/venv/bin:$PATH"
 
 # Copy application code
 COPY astrai/ ./astrai/
-COPY scripts/tools/ ./scripts/tools/
+COPY scripts/ ./scripts/
+COPY assets/ ./assets/
 COPY pyproject.toml .
+COPY README.md .
 
 # Create non-root user
-RUN useradd -m -u 1000 astrai && chown -R astrai:astrai /app
+RUN useradd -m astrai && chown -R astrai:astrai /app
 USER astrai
 
-# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
-ENV PYTHONUNBUFFERED=1
+    PYTHONDONTWRITEBYTECODE=1
-ENV PYTHONDONTWRITEBYTECODE=1
-
-# Default command
-CMD ["python", "-m", "astrai.inference.server"]

README.md

@@ -27,9 +27,6 @@
 
 ## 📖 Table of Contents
 
-<details open>
-<summary><b>English</b></summary>
-
 - [Features](#features)
 - [Quick Start](#quick-start)
 - [Documentation](#documentation)
@@ -37,8 +34,6 @@
 - [Community](#community)
 - [License](#license)
 
-</details>
-
 ---
 
 <a id="english"></a>
@@ -51,7 +46,8 @@
 - 💡 **Easy to Use**: Simple API with comprehensive examples and demos.
 - 📦 **Lightweight**: Minimal dependencies, easy to deploy.
 - 🔬 **Research‑Friendly**: Modular design, easy to experiment with new ideas.
-- 🤗 **HuggingFace Integration**: Compatible with HuggingFace models and datasets.
+- 🤗 **HuggingFace-Style API**: AutoModel/AutoTokenizer APIs inspired by HuggingFace for easy model and tokenizer loading.
+- 🔌 **Dual API Compatibility**: Supports both OpenAI and Anthropic chat completion APIs out of the box.
 
 ### Quick Start
 
@@ -72,16 +68,109 @@ pip install -e ".[dev]"
 #### Train a Model
 
 ```bash
-python scripts/tools/train.py \
+CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
-  --train_type=seq \
+    --train_type seq \
-  --data_root_path=/path/to/dataset \
+    --data_root_path /path/to/dataset \
-  --param_path=/path/to/param_path
+    --param_path /path/to/model \
+    --batch_size 4 \
+    --accumulation_steps 8 \
+    --max_lr 3e-4 \
+    --warmup_steps 1000 \
+    --n_epoch 1
 ```
 
+Full reference at [Parameter Guide](assets/docs/params.md).
+
 #### Generate Text
 
 ```bash
-python scripts/tools/generate.py --param_path=/path/to/param_path
+python scripts/tools/generate.py \
+    --param_path /path/to/model \
+    --input_json_file /path/to/input.json \
+    --output_json_file /path/to/output.json
+```
+
+#### Docker
+
+Build and run with Docker (recommended for GPU environments):
+
+```bash
+# Build image
+docker build -t astrai:latest .
+
+# Run with GPU support
+docker run --gpus all -it astrai:latest
+
+# Run with specific GPUs
+docker run --gpus '"device=0,1"' -it astrai:latest
+
+# Run inference server
+docker run --gpus all -p 8000:8000 astrai:latest \
+  python -m scripts.tools.server --port 8000 --device cuda
+
+# Run with volume mount for data
+docker run --gpus all -v /path/to/data:/data -it astrai:latest
+
+# Docker Compose (GPU, default)
+docker compose up -d
+
+# Docker Compose (CPU only)
+docker compose --profile cpu up -d
+```
+
+> **Note**: `--gpus all` is required for CUDA support. Without it, `torch.cuda.is_available()` will return `False`.
+
+#### Start HTTP Server
+
+Start the inference server with OpenAI and Anthropic-compatible HTTP API:
+
+```bash
+python -m scripts.tools.server --port 8000 --device cuda
+```
+
+Make requests:
+
+```bash
+# OpenAI-compatible
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "Hello"}],
+    "max_tokens": 512
+  }'
+
+# OpenAI-compatible streaming
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "Tell a story"}],
+    "stream": true,
+    "max_tokens": 500
+  }'
+
+# Anthropic-compatible
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "astrai",
+    "system": "You are a helpful assistant.",
+    "messages": [{"role": "user", "content": "Hello"}],
+    "max_tokens": 512
+  }'
+
+# Anthropic-compatible streaming with stop sequences
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "astrai",
+    "messages": [{"role": "user", "content": "Write a story"}],
+    "max_tokens": 500,
+    "stream": true,
+    "stop_sequences": ["The end"]
+  }'
+
+# Health check
+curl http://localhost:8000/health
 ```
 
 #### Demo

assets/docs/README-zh-CN.md

@@ -52,7 +52,8 @@
 - 💡 **易用**: 简洁的 API 与丰富的示例、演示。
 - 📦 **轻量**: 依赖少，部署简单。
 - 🔬 **研究友好**: 模块化设计，便于实验新想法。
-- 🤗 **HuggingFace 集成**: 兼容 HuggingFace 模型与数据集。
+- 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
+- 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
 
 ### 快速开始
 
@@ -73,16 +74,109 @@ pip install -e ".[dev]"
 #### 训练模型
 
 ```bash
-python scripts/tools/train.py \
+CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
-  --train_type=seq \
+    --train_type seq \
-  --data_root_path=/path/to/dataset \
+    --data_root_path /path/to/dataset \
-  --param_path=/path/to/param_path
+    --param_path /path/to/model \
+    --batch_size 4 \
+    --accumulation_steps 8 \
+    --max_lr 3e-4 \
+    --warmup_steps 1000 \
+    --n_epoch 1
 ```
 
+完整参数列表见[参数说明](./params.md)。
+
 #### 文本生成
 
 ```bash
-python scripts/tools/generate.py --param_path=/path/to/param_path
+python scripts/tools/generate.py \
+    --param_path /path/to/model \
+    --input_json_file /path/to/input.json \
+    --output_json_file /path/to/output.json
+```
+
+#### Docker
+
+使用 Docker 构建和运行（推荐用于 GPU 环境）：
+
+```bash
+# 构建镜像
+docker build -t astrai:latest .
+
+# 启用 GPU 运行
+docker run --gpus all -it astrai:latest
+
+# 指定特定 GPU
+docker run --gpus '"device=0,1"' -it astrai:latest
+
+# 运行推理服务
+docker run --gpus all -p 8000:8000 astrai:latest \
+  python -m scripts.tools.server --port 8000 --device cuda
+
+# 挂载数据卷
+docker run --gpus all -v /path/to/data:/data -it astrai:latest
+
+# Docker Compose（GPU，默认）
+docker compose up -d
+
+# Docker Compose（仅 CPU）
+docker compose --profile cpu up -d
+```
+
+> **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。
+
+#### 启动 HTTP 服务
+
+启动推理服务器，支持 OpenAI 和 Anthropic 兼容的 HTTP API：
+
+```bash
+python -m scripts.tools.server --port 8000 --device cuda
+```
+
+发起请求：
+
+```bash
+# OpenAI 兼容
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "你好"}],
+    "max_tokens": 512
+  }'
+
+# OpenAI 兼容流式
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "讲个故事"}],
+    "stream": true,
+    "max_tokens": 500
+  }'
+
+# Anthropic 兼容
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "astrai",
+    "system": "你是一个乐于助人的助手。",
+    "messages": [{"role": "user", "content": "你好"}],
+    "max_tokens": 512
+  }'
+
+# Anthropic 兼容流式并设置停止序列
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "astrai",
+    "messages": [{"role": "user", "content": "写个故事"}],
+    "max_tokens": 500,
+    "stream": true,
+    "stop_sequences": ["结束"]
+  }'
+
+# 健康检查
+curl http://localhost:8000/health
 ```
 
 #### 演示

assets/docs/dataflow.md

@@ -7,13 +7,12 @@ This document describes the data flow of the AstrAI project (a training and infe
 AstrAI adopts a modular design with the following main components:
 - **Dataset Module** (`astrai/dataset/`): Dataset, sampler, serialization tools
 - **Model Module** (`astrai/model/`): AutoModel, Transformer model and its submodules
-- **Training Module** (`astrai/trainer/`): Trainer, training context, strategies, schedulers
+- **Training Module** (`astrai/trainer/`): Trainer, training context, strategies, schedulers, callbacks, metric utilities
 - **Inference Module** (`astrai/inference/`): Inference engine with continuous batching, streaming generation
-- **Config Module** (`astrai/config/`): Model, training, scheduler, and other configurations
+- **Config Module** (`astrai/config/`): ModelConfig, TrainConfig
 - **Factory Module** (`astrai/factory/`): Registry, BaseFactory for component registration
 - **Parallel Module** (`astrai/parallel/`): Distributed training support
-
+- **Serialization** (`astrai/serialization.py`): HDF5 data loading, checkpoint management
-The data flow can generally be divided into two main lines: **Training Data Flow** and **Inference Data Flow**.
 
 ## Data Flow Diagram
 
@@ -21,8 +20,8 @@ The data flow can generally be divided into two main lines: **Training Data Flow
 flowchart LR
     subgraph A[Data Preparation]
         direction TB
-        A1[Raw Text] --> A2[BpeTokenizer]
+        A1[Raw Text] --> A2[AutoTokenizer]
-        A2 --> A3[Serialize to .h5 files]
+        A2 --> A3[Tokenized .h5 files]
         A3 --> A4[BaseDataset]
         A4 --> A5[ResumableDistributedSampler]
         A5 --> A6[DataLoader]
@@ -30,30 +29,28 @@ flowchart LR
 
     subgraph B[Training]
         direction TB
-        B1[Batch Data] --> B2[TrainContextBuilder]
+        B1[DataLoader] --> B2[BaseStrategy]
-        B2 --> B3[TrainContext]
+        B2 --> B3[Transformer Forward]
-        B3 --> B4[BaseStrategy]
+        B3 --> B4[Loss + Backward]
-        B4 --> B5[Transformer]
+        B4 --> B5[Gradient Accumulation]
-        B5 --> B6[Compute Loss]
+        B5 -->|every accum_steps| B6[Optimizer Step]
-        B6 --> B7[Backward]
+        B6 --> B7[LR Scheduler]
-        B7 --> B8[Optimizer]
+        B7 -->|next batch| B2
-        B8 --> B9[LRScheduler]
+        B6 --> B8[CheckpointCallback]
-        B9 --> B10[CheckpointCallback]
     end
 
     subgraph C[Inference]
         direction TB
         C1[Checkpoint] --> C2[AutoModel]
-        C2 --> C3[Transformer + Tokenizer]
+        C1 --> C3[AutoTokenizer]
-        C3 --> C4[GenerationRequest + apply_chat_template]
+        C2 --> C4[InferenceEngine]
-        C4 --> C5[InferenceEngine]
+        C3 --> C4
-        C5 --> C6[InferenceScheduler]
+        C4 --> C5[InferenceScheduler]
-        C6 --> C7[apply_sampling_strategies]
+        C5 --> C6[Transformer Forward]
-        C7 --> C8[Transformer Forward]
+        C6 --> C7[sample]
-        C8 --> C9[KV Cache]
+        C7 --> C8{End?}
-        C9 --> C10{End Condition?}
+        C8 -->|No| C6
-        C10 -->|No| C8
+        C8 -->|Yes| C9[Generated Text]
-        C10 -->|Yes| C11[Output Text]
     end
 
     A --> B
@@ -62,167 +59,179 @@ flowchart LR
 
 ## Detailed Module Descriptions
 
-### 1. Dataset Module
+### 1. Serialization (`astrai/serialization.py`)
 
-#### 1.1 Tokenizer (`tokenizer.py`)
+- **`save_h5`**: Saves tensors by groups as HDF5 files (`.h5`), each key maps to a list of tensors
-- Implemented based on Byte-Level BPE (BPE)
+- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory
-- Supports special tokens: `<｜begin▁of▁sentence｜>`, `<｜end▁of▁sentence｜>`, `<｜▁pad▁｜>`, `<｜im▁start｜>`, `<｜im▁end｜>`
+- **`Checkpoint`**: Encapsulates model state dict + epoch + iteration; uses safetensors
-- Provides `encode`/`decode` methods for mutual conversion between text and token IDs
-- Learns vocabulary from corpus during training, saved as `.json` files
-- `BpeTrainer` class handles vocabulary training from corpus
 
-#### 1.2 Serialization (`serialization.py`)
+### 2. Dataset Module
-- **`save_h5`**: Saves multiple tensors by groups as HDF5 files (`.h5`), each key corresponds to a list of tensors
-- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory (`share_memory=True`)
-- **`Checkpoint` class**: Encapsulates model state dict, training epoch, iteration count; supports safetensors format for saving and loading
 
-#### 1.3 Dataset (`dataset.py`)
+#### 2.1 Dataset (`dataset.py`)
-- **`BaseDataset`**: Abstract base class, defines common logic for window sampling, stride, etc.
+- **`BaseDataset`**: Abstract base class for windowed sequence sampling
-- **`BaseSegmentFetcher`** and **`MultiSegmentFetcher`**: Efficiently fetch data from specified index ranges in multiple segments
+- **`BaseSegmentFetcher` / `MultiSegmentFetcher`**: Fetch tensor segments by index range
-- **`DatasetFactory`**: Factory pattern, supports dynamic registration of dataset types (`seq`, `sft`, `dpo`, `grpo`)
+- **`DatasetFactory`**: Creates dataset instances by `train_type` (`seq`, `sft`, `dpo`, `grpo`)
-- After dataset loading, multiple data keys (such as `"sequence"`, `"mask"`) are managed through `MultiSegmentFetcher`
+- Data keys: `"sequence"` (SEQ), `"loss_mask"` (SFT), `"chosen_mask"/"rejected_mask"` (DPO), `"masks"` (GRPO)
 
-#### 1.4 Sampler (`sampler.py`)
+#### 2.2 Sampler (`sampler.py`)
-- **`ResumableDistributedSampler`**: Resumable sampler supporting distributed training
+- **`ResumableDistributedSampler`**: Tracks `epoch` and `iter` for breakpoint resume; supports shuffle and drop_last
-- Records current epoch and iteration position, enabling training resume from breakpoints
-- Supports shuffle and drop_last options
 
-### 2. Model Module
+### 3. Model Module
 
-#### 2.1 Transformer / AutoModel (`transformer.py`, `automodel.py`)
+#### 3.1 Transformer / AutoModel
-- **`AutoModel`**: Base class for autoregressive language models with `from_pretrained()` and `save_pretrained()` methods
+- **`AutoModel`**: Base class with `from_pretrained()` / `save_pretrained()`
-- **`Transformer`**: Core autoregressive decoder architecture (registered via `@AutoModel.register('transformer')`)
+- **`Transformer`**: Decoder-only architecture, registered via `@AutoModel.register('transformer')`
-- Contains embedding layer, multi-layer `DecoderBlock`, RMSNorm, and linear output head
+- Embedding → N×DecoderBlock → RMSNorm → Linear lm_head
-- Supports weight tying (`tie_weight=True`) to reduce parameter count
+- RoPE position encoding, optional weight tying
-- Uses Rotary Position Embedding (RoPE) to inject position information
-- Supports loading from safetensors format with automatic model type detection from `config.json`
 
-#### 2.2 Submodules (`module.py`)
+#### 3.2 Submodules (`module.py`)
-- **`RotaryEmbedding`**: Generates RoPE cos/sin cache
+- **`DecoderBlock`**: GQA attention + residual + MLP + RMSNorm
-- **`DecoderBlock`**: Contains multi-head attention (supports GQA), feedforward network (FFN), residual connections
+- **`GQA`**: Grouped Query Attention (also `MLA` for multi-latent attention)
-- **`RMSNorm`**: Layer normalization variant
+- **`MLP`**: `SiLU(gate(x)) * up(x)` → down projection
-- **`Linear`**, **`Embedding`**: Custom linear layer and embedding layer, supporting parallelism wrappers
+- **`RotaryEmbedding`**: RoPE cos/sin cache
+- **`RMSNorm`**: Layer normalization
 
-### 3. Training Module
+### 4. Training Module
 
-#### 3.1 Training Context (`train_context.py`)
+#### 4.1 Training Context (`train_context.py`)
-- **`TrainContext`**: Data class encapsulating all components needed for training (model, optimizer, data loader, strategy, etc.)
+- **`TrainContext`**: Dataclass holding model, optimizer, dataloader, strategy, scheduler, checkpoint state
-- **`TrainContextBuilder`**: Builder pattern, progressively assembles training context, supports resume from checkpoint
+- **`TrainContextBuilder`**: Builder pattern — takes checkpoint for resume, builds all components
 
-#### 3.2 Trainer (`trainer.py`)
+#### 4.2 Trainer (`trainer.py`)
-- **`Trainer`**: Main training loop, manages callbacks (progress bar, checkpoint, metric logging, gradient clipping, scheduler)
-- Supports distributed training (launches multi-process via `spawn_parallel_fn`)
-- Training steps include:
-  1. `on_train_begin` → 2. `on_epoch_begin` → 3. `on_batch_begin` → 4. Forward/loss calculation → 5. `on_batch_end` → 6. Gradient accumulation → 7. `on_step_begin` → 8. Optimizer update → 9. `on_step_end` → 10. `on_epoch_end`
 
-#### 3.3 Strategy (`strategy.py`)
+The training loop is nested: **epoch** → **batch** (with step phase interspersed):
-- **`BaseStrategy`**: Defines training strategy interface (such as `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy`)
-- Strategy receives batch data, executes model forward pass, loss calculation, returns loss tensor
-- Created dynamically by `StrategyFactory` according to configuration
 
-#### 3.4 Scheduler (`schedule.py`)
+```
-- **`BaseScheduler`**: Abstract base class defining learning rate scheduling interface
+on_train_begin
-- **`SchedulerFactory`**: Factory pattern, supports registration of various schedulers (such as `cosine`, `sgdr`)
+  on_epoch_begin
-- Scheduler is automatically created according to configuration and bound to optimizer
+    for each batch:
+      if iteration % accumulation_steps == 0:        ← step phase
+        on_step_begin → optimizer.step() → zero_grad → on_step_end
+                                                      ← batch phase
+      on_batch_begin → strategy(batch) → loss → backward → on_batch_end
+      iteration += 1
 
-### 4. Factory Module
+    on_epoch_end
+on_train_end
+```
 
-#### 4.1 Registry and BaseFactory (`factory.py`)
+Key points:
-- **`Registry`**: Flexible registry for component classes with category and priority support
+- `on_step_*` wraps optimizer step (fires every `accumulation_steps` batches)
-- **`BaseFactory`**: Generic factory class for component registration and creation
+- `on_batch_*` wraps loss computation (fires every batch)
-- Supports decorator-based registration pattern for extensible components
+- `SchedulerCallback` fires on `on_batch_end` — LR scheduler steps every batch
-- Provides methods for registration, retrieval, and listing with filtering
+- `GradientClippingCallback` fires on `on_step_begin`
+
+#### 4.3 Strategy (`strategy.py`)
+- **`SEQStrategy`**: Next-token prediction, cross-entropy with label smoothing
+- **`SFTStrategy`**: Supervised fine-tuning with loss masking
+- **`DPOStrategy`**: Direct Preference Optimization with reference model
+- **`GRPOStrategy`**: Group Relative Policy Optimization with clipped ratio
+
+#### 4.4 Scheduler (`schedule.py`)
+- **`CosineScheduler`**: Cosine decay + linear warmup
+- **`SGDRScheduler`**: Cosine annealing with warm restarts
+- Created by `SchedulerFactory` and bound to optimizer
+
+#### 4.5 Callbacks
+- **`CheckpointCallback`**: Saves safetensors at `ckpt_interval` iterations
+- **`ProgressBarCallback`**: tqdm progress display
+- **`MetricLoggerCallback`**: Writes JSONL metrics to `{ckpt_dir}/logs/`
+- **`GradientClippingCallback`**: `clip_grad_norm_` on `on_step_begin`
+- **`SchedulerCallback`**: `scheduler.step()` on `on_batch_end`
 
 ### 5. Inference Module
 
 #### 5.1 Inference Engine (`engine.py`)
-- **`InferenceEngine`**: Unified inference interface, supports streaming and non-streaming generation
+- **`InferenceEngine`**: Facade over scheduler; provides `generate()`, `generate_with_request()`, `generate_async()`
-- **`InferenceScheduler`**: Continuous batching scheduler with dynamic batch composition
+- Accepts `prompt: str | List[str]`, returns generator (stream) or string (non-stream)
-- **`GenerationRequest`**: Encapsulates generation parameters (top_k, top_p, temperature, max_len, messages, etc.)
-- **`messages` format**: List of message dictionaries with `role` (system/user/assistant) and `content`
-- **`apply_chat_template`** (from `tokenizer.py`): Converts messages into prompt string using ChatML format
-- Provides streaming (`stream=True`) and non-streaming (`stream=False`) generation interfaces
-- Supports continuous batching with `max_batch_size` and `max_seq_len` parameters
-- Uses separate model and tokenizer initialization for flexibility
 
-#### 5.2 Scheduler (`scheduler.py`)
+#### 5.2 Scheduler 4-Phase Loop (`scheduler.py`)
-- **`Task`**: Individual generation task with state management (PENDING, RUNNING, FINISHED, ABORTED)
-- **`TaskStatus`**: Task state enumeration
-- **`apply_sampling_strategies`**: Applies temperature, top-k, top-p sampling to logits
-- Continuous batching: new requests can join at any time, completed requests are released immediately
 
-#### 5.3 Request (`engine.py`)
+Background thread runs continuously:
-- **`GenerationRequest`**: Encapsulates generation parameters (top_k, top_p, temperature, max_len, messages, etc.)
-- **`messages` format**: List of message dictionaries with `role` (system/user/assistant) and `content`
-- **`apply_chat_template`** (from `tokenizer.py`): Converts messages into prompt string using ChatML format
-- Provides streaming (`stream=True`) and non-streaming (`stream=False`) generation interfaces
 
-## Training Data Flow - Detailed Steps
+```
+1. Cleanup → Remove finished tasks, free KV cache pages
+2. Refill  → Pop from waiting_queue, alloc pages, add to active
+3. Prefill → Group active tasks by prompt_len, run full forward pass
+4. Decode  → Pick largest same-position group, run single-token forward
+```
+
+- **`Task`**: Tracks prompt_ids, output_ids, page_table, status (PENDING/RUNNING/FINISHED/ABORTED)
+- **`PagedCache`**: Bitmask-based page allocator with page-table-indirected read/write
+- **`CacheView`**: Batch view bundling cache + page table for attention layers
+- **`sample()`**: Temperature → top-k → top-p → multinomial
+
+#### 5.3 Server (`server.py`)
+- FastAPI with OpenAI `/v1/chat/completions` and Anthropic `/v1/messages` endpoints
+- Streaming via SSE, health check at `/health`, stats at `/stats`
+
+### 6. Tokenizer Module
+
+- **`AutoTokenizer`**: Wraps HuggingFace tokenizers (BBPE); `encode`/`decode`/`apply_chat_template`
+- **`ChatTemplate`**: Jinja2-based template rendering for multi-turn chat
+
+### 7. Factory & Parallel
+
+- **`Registry` / `BaseFactory`**: Decorator-based component registration
+- **`spawn_parallel_fn`**: Multi-process DDP launcher with NCCL backend
+- **`ParallelModel` / `ColumnParallelLinear` / `RowParallelLinear`**: Tensor model parallelism
+
+## Training Data Flow — Detailed Steps
 
 1. **Data Preparation**
-   - Raw text is converted to token ID sequences through BPE tokenizer
+   - Raw text → token IDs via `AutoTokenizer.encode()`
-   - Token ID sequences (possibly with masks, labels, etc.) are saved by groups as `.h5` files
+   - Save as `.h5` files (groups of tensor lists per data key)
-   - Files can contain multiple segments, each segment corresponds to a tensor
 
 2. **Dataset Loading**
-   - `BaseDataset`'s `load` method calls `load_h5`, obtaining `segments` dictionary
+   - `BaseDataset.load()` calls `load_h5()`, builds `MultiSegmentFetcher`
-   - Create `MultiSegmentFetcher` to manage data for multiple keys
+   - Sliding window of `window_size` with `stride` determines sample boundaries
-   - Calculate total sample count, and determine start/end indices for each sample based on window size and stride
 
-3. **Sampling and Batch Loading**
+3. **Sampling & Batching**
-   - `ResumableDistributedSampler` generates index sequence based on current epoch and iteration position
+   - `ResumableDistributedSampler` produces shuffled index sequences
-   - `DataLoader` uses sampler to get indices, calls dataset's `__getitem__` to get actual data
+   - `DataLoader` fetches `[batch_size, window_size]` tensors via `__getitem__`
-   - Batch data shape is `[batch_size, window_size]` (or varies according to specific dataset type)
 
-4. **Strategy Forward and Loss Calculation**
+4. **Strategy Forward**
-   - Batch data is passed to strategy (such as `SEQStrategy`)
+   - Strategy receives batch, calls `Transformer.forward()` for logits
-   - Strategy internally calls `Transformer` model, obtaining logits
+   - Computes task-specific loss (cross-entropy, DPO, GRPO)
-   - Calculate cross-entropy loss (or DPO loss, etc.) according to task type
-   - Return loss tensor
 
-5. **Backpropagation and Optimization**
+5. **Backward & Accumulation**
-   - Loss is normalized by dividing by accumulation steps, then `loss.backward()` is executed
+   - `loss = raw_loss / accumulation_steps`
-   - After accumulating `accumulation_steps` batches, optimizer `step()` and `zero_grad()` are executed
+   - `loss.backward()` accumulates gradients
-   - Learning rate scheduler updates learning rate after each step
+   - Every `accumulation_steps` batches: `optimizer.step()` → `zero_grad()`
+   - Every batch: `scheduler.step()` updates learning rate
 
-6. **Checkpoint Saving**
+6. **Checkpoint**
-   - `CheckpointCallback` saves checkpoints at set intervals
+   - `CheckpointCallback` saves `model.state_dict()` + metadata to safetensors at `ckpt_interval` iterations
-   - Checkpoints contain model state dict, current epoch, iteration, and other metadata
+   - Does NOT save optimizer/scheduler state (resume resets those)
-   - Saved in safetensors format, ensuring safety and efficiency
 
-## Inference Data Flow - Detailed Steps
+## Inference Data Flow — Detailed Steps
 
 1. **Model Loading**
-   - Load `Transformer` model from checkpoint via `AutoModel.from_pretrained()`
+   - `AutoModel.from_pretrained(path)` loads weights from safetensors
-   - Set model to evaluation mode (`model.eval()`), enable inference mode (`torch.inference_mode`)
+   - `torch.inference_mode()` wraps generation
 
-2. **Prompt Construction and Encoding**
+2. **Prompt Construction**
-   - User messages (list of dict with role and content) are converted to ChatML format string through `apply_chat_template` method in tokenizer
+   - Messages → `apply_chat_template(messages, tokenize=False)` → prompt string
-   - Tokenizer encodes prompt string to token ID sequence `input_ids`
+   - `tokenizer.encode(prompt)` → token IDs (truncated to `max_prompt_len`)
-   - For batch generation, use `pad_sequence` for padding
 
-3. **Autoregressive Generation Loop**
+3. **Continuous Batching Loop**
-   - Initialize KV cache (optional)
+   - **Cleanup**: Finished tasks → `stream_callback(STOP)`, free KV pages
-   - Loop until generating `max_len` tokens or encountering stop token:
+   - **Refill**: Pop from waiting queue, `PagedCache.alloc_n()` for prompt pages
-     - Input current `input_ids` (or cached new token) to model, obtain `logits`
+   - **Prefill**: Group by prompt length, run full forward with `start_pos=0`
-     - Apply `apply_sampling_strategies` (temperature, top-k, top-p) to `logits`
+   - **Decode**: Pick position group with most tasks, single-token forward:
-     - Sample next token ID from the processed distribution
+     - Model forward → `logits` → `sample()` → next token ID
-     - Append new token to `input_ids`, while updating KV cache
+     - Append to `output_ids`, update `output_tokens`
-     - For streaming generation, yield each token to caller immediately
+     - `_maybe_alloc_page()` grows page table as needed
+     - `stream_callback(token)` for streaming clients
 
-4. **Decoding and Output**
+4. **Output**
-   - Decode generated token ID sequence to text through tokenizer
+   - `tokenizer.decode(output_ids)` → text
-   - Remove special tokens, return plain text response
+   - Return to caller (streaming: token-by-token; non-streaming: complete string)
 
-## Checkpoint and Serialization
+## Checkpoint & Serialization
 
-- **Training Checkpoint**: Saves model parameters, optimizer state, scheduler state, current epoch and iteration
+- **Training Checkpoint**: safetensors weights + epoch/iteration metadata. Optimizer/scheduler state is NOT persisted.
-- **Model Parameters**: Supports safetensors format, automatically handles special logic like weight tying during loading
+- **Inference Loading**: `AutoModel.from_pretrained()` loads from the same safetensors format.
-- **Dataset Serialization**: HDF5 format supports efficient random access and shared memory, suitable for large-scale pre-training data
+- **Dataset Serialization**: HDF5 with shared memory support for large-scale pre-training data.
 
-## Summary
+> Document Update Time: 2026-05-09
-
-The data flow design of AstrAI reflects the characteristics of modularity, extensibility, and resumability. The training data flow supports large-scale distributed training through chunk loading, resumable sampling, gradient accumulation, and other mechanisms; the inference data flow achieves efficient text generation using KV cache and sampling strategies. Clear interfaces between modules facilitate customization and extension.
-
-> Document Update Time: 2026-04-05  
-> Corresponding Code Version: Refer to version number defined in `pyproject.toml`

assets/docs/design.md

@@ -8,7 +8,7 @@ Thus, the AstrAI project was born - 1B parameters, Chinese-English bilingual, su
 
 ```mermaid
 classDiagram
-    namespace astrai.config {
+    namespace config {
         class ModelConfig {
             +int vocab_size
             +int dim
@@ -50,23 +50,14 @@ classDiagram
             +str master_port
             +Callable parallel_wrapper
             +Callable state_dict_fn
-            +List[int] device_ids
             +str device_type
             +dict extra_kwargs
             +validate()
         }
 
-        class ModelParameter {
-            +nn.Module model
-            +BpeTokenizer tokenizer
-            +ModelConfig config
-            +save(instance, save_dir)
-            +load(load_dir, disable_init) ModelParameter
-            +to(*args, **kwargs)
-        }
     }
 
-    namespace astrai.dataset {
+    namespace dataset {
         class BaseDataset {
             +int window_size
             +int stride
@@ -93,22 +84,22 @@ classDiagram
         }
 
         class BaseSegmentFetcher {
-            +List~Tensor~ segments
+            +List[Tensor] segments
-            +List~int~ cum_lengths
+            +List[int] cum_lengths
             +int total_length
             +fetch_data(begin_idx, end_idx) Tensor
         }
 
         class MultiSegmentFetcher {
-            +Dict muti_fetchers
+            +Dict multi_fetchers
-            +List muti_keys
+            +List multi_keys
             +key_fetch(begin_idx, end_idx, keys) Dict
             +fetch_data(begin_idx, end_idx) Dict
         }
 
         class ResumableDistributedSampler {
-            +int start_epoch
+            +int epoch
-            +int start_iter
+            +int iter
         }
 
         class DatasetFactory {
@@ -117,7 +108,9 @@ classDiagram
             +create(train_type, window_size, stride) BaseDataset
             +load(train_type, load_path, window_size, stride) BaseDataset
         }
+    }
 
+    namespace serialization {
         class Checkpoint {
             +dict state_dict
             +int epoch
@@ -125,20 +118,12 @@ classDiagram
             +save(save_dir)
             +load(save_dir) Checkpoint
         }
-
-        class DataLoader {
-            +Dataset dataset
-            +int batch_size
-            +Sampler sampler
-            +__iter__()
-            +__len__()
-        }
     }
 
-    namespace astrai.model {
+    namespace model {
         class AutoModel {
             +ModelConfig config
-            +Dict _registry
+            +Registry _registry
             +register(model_type) decorator
             +get_model_class(model_type) Type
             +from_pretrained(path, disable_random_init) nn.Module
@@ -148,12 +133,12 @@ classDiagram
 
         class Transformer {
             +ModelConfig config
-            +RotaryEmbedding rotary_embeding
+            +RotaryEmbedding rotary_embedding
             +Embedding embed_tokens
             +ModuleList layers
             +RMSNorm norm
             +Linear lm_head
-            +forward(input_ids, input_mask, persistent_key_values, start_pos) Dict
+            +forward(input_ids, input_mask, paged_cache, start_pos) Dict
             +load_state_dict(state_dict)
             +state_dict()
         }
@@ -163,7 +148,7 @@ classDiagram
             +RMSNorm input_norm
             +MLP mlp
             +RMSNorm post_attention_norm
-            +forward(x, rotary_emb, attention_mask, kv_cache, start_pos) Tensor
+            +forward(x, rotary_emb, attention_mask, paged_cache, start_pos) Tensor
         }
 
         class GQA {
@@ -172,18 +157,20 @@ classDiagram
             +int head_dim
             +Linear q_proj, k_proj, v_proj, o_proj
             +RMSNorm q_norm, k_norm
-            +forward(x, rotary_emb, mask, kv_cache, start_pos) Tensor
+            +forward(x, rotary_emb, mask, paged_cache, start_pos) Tensor
         }
 
         class MLA {
             +int n_heads
             +int n_kv_heads
             +int head_dim
-            +Linear q_a_proj, q_b_proj, q_c_proj
+            +int kv_lora_rank
-            +Linear kv_a_proj, kv_b_proj, kv_c_proj
+            +int qk_nope_head_dim
+            +int qk_rope_head_dim
+            +Linear q_proj, kv_a_proj, kv_b_proj
             +Linear o_proj
-            +RMSNorm q_norm, k_norm
+            +RMSNorm kv_norm
-            +forward(x, rotary_emb, mask, kv_cache, start_pos) Tensor
+            +forward(x, rotary_emb, mask, paged_cache, start_pos) Tensor
         }
 
         class MLP {
@@ -207,7 +194,7 @@ classDiagram
             +int dim
             +int max_len
             +float base
-            +forward(x, start_pos) Tuple~Tensor, Tensor~
+            +forward(x, start_pos) Tuple[Tensor, Tensor]
         }
 
         class Embedding {
@@ -216,30 +203,52 @@ classDiagram
         }
     }
 
-    namespace astrai.tokenize {
+    namespace tokenize {
-        class Tokenizer {
+        class AutoTokenizer {
-            +encode(tokens, out_ids, add_special_tokens) List~int~
+            +List[int] stop_ids
-            +decode(tokens, skip_special_tokens) str
-            +__len__() int
-        }
-
-        class BpeTokenizer {
-            +List~str~ stop_ids
             +int bos_id
             +int eos_id
             +int pad_id
-            +encode(tokens, out_ids, add_special_tokens) List~int~
+            +vocab_size int
+            +encode(tokens, out_ids, add_special_tokens) List[int]
             +decode(tokens, skip_special_tokens) str
+            +apply_chat_template(messages, tokenize) Union[str, List[int]]
+            +set_chat_template(template)
+            +load(path)
+            +from_pretrained(path) AutoTokenizer
+            +save_pretrained(save_path)
+        }
+
+        class ChatTemplate {
+            +String template_str
+            +render(messages, system_prompt, **extra_variables) str
+            +from_string(template) ChatTemplate
         }
     }
 
-    namespace astrai.trainer {
+    namespace factory {
+        class Registry {
+            +Dict _entries
+            +register(name, component_cls, category, priority)
+            +get(name) Type
+            +list_names() List[str]
+        }
+
+        class BaseFactory {
+            +Registry _registry
+            +register(name, category, priority) decorator
+            +create(name, *args, **kwargs) T
+            +list_registered() list
+        }
+    }
+
+    namespace trainer {
         class Trainer {
             +TrainConfig train_config
-            +List~TrainCallback~ callbacks
+            +List[TrainCallback] callbacks
             +train(checkpoint)
             +_build_context(checkpoint) TrainContext
-            +_get_default_callbacks() List~TrainCallback~
+            +_get_default_callbacks() List[TrainCallback]
         }
 
         class TrainContext {
@@ -259,8 +268,6 @@ classDiagram
         class TrainContextBuilder {
             +TrainConfig config
             +with_checkpoint(checkpoint) TrainContextBuilder
-            +with_dataloader() TrainContextBuilder
-            +with_strategy() TrainContextBuilder
             +build() TrainContext
         }
 
@@ -302,7 +309,7 @@ classDiagram
         }
 
         class BaseScheduler {
-            +get_lr() List~float~
+            +get_lr() List[float]
             +step()
         }
 
@@ -337,6 +344,39 @@ classDiagram
             +on_error(context)
         }
 
+        class GradientClippingCallback {
+            +float max_grad_norm
+            +on_step_begin(context)
+        }
+
+        class SchedulerCallback {
+            +on_train_begin(context)
+            +on_batch_end(context)
+        }
+
+        class CheckpointCallback {
+            +str save_dir
+            +int interval
+            +_save_checkpoint(context)
+            +on_batch_end(context)
+            +on_train_end(context)
+            +on_error(context)
+        }
+
+        class ProgressBarCallback {
+            +int num_epoch
+            +on_epoch_begin(context)
+            +on_batch_end(context)
+            +on_epoch_end(context)
+        }
+
+        class MetricLoggerCallback {
+            +str log_dir
+            +int save_interval
+            +on_batch_end(context)
+            +on_train_end(context)
+        }
+
         class CallbackFactory {
             +Registry _registry
             +register(name) decorator
@@ -344,22 +384,28 @@ classDiagram
         }
     }
 
-    namespace astrai.inference {
+    namespace inference {
         class InferenceEngine {
-            +ModelParameter parameter
+            +nn.Module model
+            +AutoTokenizer tokenizer
             +InferenceScheduler scheduler
+            +int max_batch_size
+            +Optional int max_seq_len
             +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
             +generate_with_request(request) Union[Generator, str, List[str]]
+            +generate_async(prompt, max_tokens, temperature, top_p, top_k) AsyncGenerator
             +get_stats() Dict
             +shutdown()
         }
 
         class InferenceScheduler {
             +nn.Module model
-            +Tokenizer tokenizer
+            +AutoTokenizer tokenizer
-            +ModelConfig config
+            +PagedCache page_cache
-            +Tuple kv_cache
+            +int max_batch_size
-            +Tensor seq_mask
+            +int max_seq_len
+            +int max_prompt_len
+            +int page_size
             +List waiting_queue
             +List active_tasks
             +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
@@ -369,6 +415,28 @@ classDiagram
             +get_stats() Dict
         }
 
+        class PagedCache {
+            +int page_size
+            +int _free_mask
+            +List[int] _refs
+            +Tensor k_cache
+            +Tensor v_cache
+            +alloc() int
+            +alloc_n(n) List[int]
+            +free(idx)
+            +bind(page_table, total_len) CacheView
+            +write(layer_id, page_table, start_pos, k, v)
+            +gather(layer_id, page_table) Tuple[Tensor, Tensor]
+        }
+
+        class CacheView {
+            +PagedCache _cache
+            +Tensor _page_table
+            +int _total_len
+            +write(layer_id, start_pos, k, v)
+            +gather(layer_id) Tuple[Tensor, Tensor]
+        }
+
         class Task {
             +str task_id
             +List prompt_ids
@@ -380,47 +448,100 @@ classDiagram
             +List output_ids
             +int input_tokens
             +int output_tokens
-            +int slot
+            +List[int] page_table
+            +int n_pages
+            +float arrival_time
+            +float finish_time
             +Callable stream_callback
+            +int next_pos
             +is_finished(stop_ids) bool
         }
 
         class TaskStatus {
-            +str PENDING
+            <<enumeration>>
-            +str RUNNING
+            PENDING
-            +str FINISHED
+            RUNNING
-            +str ABORTED
+            FINISHED
-        }
+            ABORTED
-
-        class apply_sampling_strategies {
-            +Tensor logits
-            +float temperature
-            +int top_k
-            +float top_p
-            +forward() Tensor
-        }
-
-        class Server {
-            +start()
-            +predict(request)
         }
 
         class GenerationRequest {
+            +List[Dict] messages
+            +GenerationParams params
+            +bool stream
+        }
+
+        class GenerationParams {
+            <<value object>>
             +int top_k
             +float top_p
             +float temperature
-            +int max_len
+            +int max_tokens
-            +Union~str, List~str~~ query
+        }
-            +history Optional
+
-            +system_prompt Optional~str~
+        class BaseSamplingStrategy {
-            +stream bool
+            <<abstract>>
+            +apply(logits, filter_value) Tensor
+        }
+
+        class TemperatureStrategy {
+            +float temperature
+            +apply(logits, filter_value) Tensor
+        }
+
+        class TopKStrategy {
+            +int top_k
+            +apply(logits, filter_value) Tensor
+        }
+
+        class TopPStrategy {
+            +float top_p
+            +apply(logits, filter_value) Tensor
+        }
+
+        class SamplingPipeline {
+            +List strategies
+            +apply(logits, filter_value) Tensor
+            +sample(logits, filter_value) Tensor
+        }
+
+        class _Result {
+            +List[str] tokens
+            +List[str] results
+            +List[bool] _done
+            +append(token, idx)
+            +get_results() List[str]
+            +pop_all() List[str]
+            +wait(timeout) bool
+        }
+
+        class ChatMessage {
+            +str role
+            +str content
+        }
+
+        class ChatCompletionRequest {
+            +List[ChatMessage] messages
+            +float temperature
+            +float top_p
+            +int top_k
+            +int max_tokens
+            +bool stream
+            +Optional[str] stop
+            +Optional[int] n
         }
     }
 
-    namespace astrai.parallel {
+    namespace parallel {
-        class ParallelSetup {
+        class ParallelFunctions {
             +spawn_parallel_fn(fn, nprocs)
-            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type, device_ids)
+            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type)
+        }
+
+        class ParallelModel {
+            +dist.ProcessGroup process_group
+            +int rank
+            +int world_size
         }
 
         class ColumnParallelLinear {
@@ -433,73 +554,89 @@ classDiagram
     }
 
     %% Relationships
-    TrainConfig --> ModelConfig : contains
+    TrainConfig --> ModelConfig : uses
     TrainConfig --> BaseDataset : uses
-    TrainConfig --> Transformer : uses
+    TrainConfig --> StrategyFactory : selects
-    Trainer --> TrainConfig : configures
-    Trainer --> TrainContextBuilder : builds
-    Trainer --> TrainCallback : manages
-    TrainContextBuilder --> TrainContext : creates
-    TrainContext --> Checkpoint : manages
-    TrainContext --> BaseStrategy : uses
-    TrainContext --> BaseScheduler : uses
     StrategyFactory ..> BaseStrategy : creates
     BaseStrategy <|-- SEQStrategy
     BaseStrategy <|-- SFTStrategy
     BaseStrategy <|-- DPOStrategy
     BaseStrategy <|-- GRPOStrategy
-    DPOStrategy --> Transformer : creates ref_model
+    DPOStrategy --> Transformer : uses
-    GRPOStrategy --> Transformer : creates ref_model
+    GRPOStrategy --> Transformer : uses
+    Trainer --> TrainConfig : configures
+    Trainer --> TrainContextBuilder : builds
+    Trainer --> TrainCallback : manages
+    TrainContextBuilder --> TrainContext : creates
+    Checkpoint ..> Checkpoint : saves/loads
+    TrainContext --> Checkpoint : manages
+    TrainContext --> BaseStrategy : uses
+    TrainContext --> BaseScheduler : uses
     SchedulerFactory ..> BaseScheduler : creates
     BaseScheduler <|-- CosineScheduler
     BaseScheduler <|-- SGDRScheduler
     CallbackFactory ..> TrainCallback : creates
+    TrainCallback <|-- GradientClippingCallback
+    TrainCallback <|-- SchedulerCallback
+    TrainCallback <|-- CheckpointCallback
+    TrainCallback <|-- ProgressBarCallback
+    TrainCallback <|-- MetricLoggerCallback
     InferenceEngine --> InferenceScheduler : uses
+    InferenceEngine --> GenerationRequest : uses
+    GenerationRequest --> GenerationParams : contains
     InferenceScheduler --> Task : manages
+    Task --> TaskStatus : uses
     InferenceScheduler --> TaskStatus : uses
-    InferenceScheduler --> apply_sampling_strategies : uses
+    InferenceScheduler --> PagedCache : uses
     InferenceScheduler --> Transformer : uses
     InferenceEngine --> Transformer : uses
-    InferenceEngine --> GenerationRequest : uses
+    InferenceEngine --> _Result : uses
-    Server --> InferenceEngine : uses
+    BaseSamplingStrategy <|-- TemperatureStrategy
-    ParallelSetup --> Trainer : enables
+    BaseSamplingStrategy <|-- TopKStrategy
-    TrainConfig --> StrategyFactory : selects
+    BaseSamplingStrategy <|-- TopPStrategy
-    ModelParameter --> Transformer : contains
+    SamplingPipeline --> BaseSamplingStrategy : composes
-    ModelParameter --> BpeTokenizer : contains
-    ModelParameter --> ModelConfig : contains
     BaseDataset <|-- SEQDataset
     BaseDataset <|-- SFTDataset
     BaseDataset <|-- DPODataset
     BaseDataset <|-- GRPODataset
     DatasetFactory ..> BaseDataset : creates
-    BaseSegmentFetcher --> MultiSegmentFetcher : used by
+    MultiSegmentFetcher --> BaseSegmentFetcher : uses
-    MultiSegmentFetcher --> BaseDataset : used by
+    BaseDataset --> MultiSegmentFetcher : uses
     AutoModel <|-- Transformer
     AutoModel --> ModelConfig : contains
     Transformer --> DecoderBlock : uses
     Transformer --> RotaryEmbedding : uses
     Transformer --> Embedding : uses
     DecoderBlock --> GQA : uses
-    DecoderBlock --> MLA : uses
     DecoderBlock --> MLP : uses
     DecoderBlock --> RMSNorm : uses
-    BpeTokenizer --> Tokenizer : inherits
     TrainContextBuilder --> ResumableDistributedSampler : creates
-    DataLoader --> BaseDataset : uses
     ResumableDistributedSampler --> BaseDataset : samples
+    ParallelModel <|-- RowParallelLinear
+    ParallelModel <|-- ColumnParallelLinear
+    AutoTokenizer --> ChatTemplate : uses
+    TrainConfig --> DatasetFactory : selects
+    TrainConfig --> SchedulerFactory : selects
+    TrainConfig --> CallbackFactory : selects
+    AutoModel ..> AutoTokenizer : loads with
+    BaseFactory <|-- DatasetFactory
+    BaseFactory <|-- StrategyFactory
+    BaseFactory <|-- SchedulerFactory
+    BaseFactory <|-- CallbackFactory
 ```
 
 ### Module Overview
 
 | Module | Components | Description |
 |--------|------------|-------------|
-| **astrai.config** | ModelConfig, TrainConfig, ModelParameter | Configuration management |
+| **astrai.config** | ModelConfig, TrainConfig | Configuration management |
-| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory, Checkpoint, DataLoader | Dataset loading and management |
+| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
+| **astrai.serialization** | Checkpoint, save_h5, load_h5 | Model serialization and checkpoint management |
 | **astrai.model** | AutoModel, Transformer, DecoderBlock, GQA, MLA, MLP, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
-| **astrai.tokenize** | AutoTokenizer, BpeTokenizer, ChatTemplate, BpeTrainer | Tokenizer |
+| **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy, StrategyFactory, BaseScheduler, SchedulerFactory, TrainCallback, CallbackFactory | Training workflow management |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Task, TaskStatus, Server, GenerationRequest | Inference service with continuous batching |
+| **astrai.inference** | InferenceEngine, InferenceScheduler, PagedCache, CacheView, Task, TaskStatus, GenerationParams, GenerationRequest, BaseSamplingStrategy, TemperatureStrategy, TopKStrategy, TopPStrategy, SamplingPipeline, ChatMessage, ChatCompletionRequest | Inference service with continuous batching and paged KV cache |
-| **astrai.parallel** | ParallelSetup, ColumnParallelLinear, RowParallelLinear | Distributed parallel |
+| **astrai.parallel** | ParallelFunctions, ParallelModel, ColumnParallelLinear, RowParallelLinear | Distributed parallel |
 | **astrai.factory** | Registry, BaseFactory | Generic component registration |
 
 ### Design Patterns
@@ -510,20 +647,22 @@ classDiagram
 | **Builder** | `TrainContextBuilder` | Chain-building training context, step-by-step initialization of components |
 | **Factory** | `StrategyFactory`, `SchedulerFactory`, `DatasetFactory`, `CallbackFactory`, `BaseFactory` | Decorator registration mechanism, dynamically create training strategies, schedulers, datasets, and callbacks |
 | **Observer** | `TrainCallback`, `CallbackFactory` | Callback mechanism for training process monitoring (checkpoint, early stopping, metrics) |
-| **Singleton** | `TrainContext` | Training process global state management |
+| **Context** | `TrainContext` | Training process state container with model, optimizer, scheduler and checkpoint |
 | **Registry** | `BaseFactory`, `Registry` | Generic component registration with category and priority support |
+| **Object Pool** | `PagedCache` | Page-based KV cache with O(1) alloc/free via bitmask |
+| **Strategy (Sampling)** | `BaseSamplingStrategy`, `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations with temperature, top-k, top-p |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, `waiting_queue`, `active_tasks` | Continuous batching with dynamic task queue management |
 | **Event-Driven** | `threading.Event`, `_task_event` | Non-blocking wait mechanism for task scheduling using Python's `threading` module |
 | **AutoModel Registry** | `AutoModel`, `Transformer` | Model type registration and dynamic loading via decorator pattern |
-| **Generator Pattern** | `_StreamingResult`, `_NonStreamingResult` | Event-based result notification for streaming/non-streaming generation |
+| **Generator Pattern** | `_Result`, `GenerationRequest` | Event-based result notification for streaming/non-streaming generation |
 
 ### Core Relationships
 
 1. **Configuration → Training**: `TrainConfig` contains `ModelConfig`, holds model, dataset, optimizer and other references
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` to compute loss
 3. **Strategy Selection**: `StrategyFactory` creates corresponding strategy instance based on `train_type`
-4. **Inference Flow**: `Server` → `InferenceEngine` → `InferenceScheduler` → `Transformer`, supports continuous batching with streaming/non-streaming
+4. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `Transformer`, uses `PagedCache` for paged KV cache management and `SamplingPipeline` for efficient continuous batching with streaming/non-streaming
-5. **Distributed Support**: `ParallelSetup` provides multi-process training capability for `Trainer`
+5. **Distributed Support**: `spawn_parallel_fn` and `setup_parallel` provide multi-process training capability for `Trainer`
 6. **Dataset Loading**: `DatasetFactory` creates datasets (SEQDataset, SFTDataset, DPODataset, GRPODataset), supports HDF5 loading via `BaseSegmentFetcher` and `MultiSegmentFetcher`
 7. **Checkpoint Management**: `Checkpoint` handles model state serialization/deserialization with safetensors
 8. **Scheduler Support**: `SchedulerFactory` creates learning rate schedulers (CosineScheduler, SGDRScheduler)
@@ -567,12 +706,6 @@ $$
 L_{\text{GRPO}} = -\mathbb{E} \left[ \min\left( \frac{\pi_\theta(a|s)}{\pi_{\text{ref}}(a|s)} \cdot A, \text{clip}\left(\frac{\pi_\theta(a|s)}{\pi_{\text{ref}}(a|s)}, 1-\epsilon, 1+\epsilon\right) \cdot A \right) \right] + \lambda \cdot D_{KL}
 $$
 
-In this implementation, an off-policy approach is used ($\pi_\theta = \pi_{\text{ref}}$), and the policy loss simplifies to:
-
-$$
-L_{\text{policy}} = -\mathbb{E}[A]
-$$
-
 The KL divergence term uses mean squared error approximation:
 
 $$
@@ -582,3 +715,5 @@ $$
 The final loss is the sum of both: $L = L_{\text{policy}} + L_{KL}$
 
 Through the above three-stage progressive training, the model completes its evolution from a general language foundation to a specialized, highly-aligned dialogue intelligence.
+
+> Document Update Time: 2026-04-09

assets/docs/introduction.md

@@ -2,7 +2,7 @@
 
 ### 1. Model Architecture
 
-This model uses the Transformer architecture with GQA mechanism (q_head=24, kv_head=4), which saves KV cache memory compared to traditional MHA (although KV cache is not currently implemented). The model is built by stacking 32 layers of Transformer blocks, with 1.0 billion parameters. Transformer is an autoregressive model that calculates the relationship between all previous tokens to obtain the probability distribution of the next token.
+This model uses the Transformer architecture with GQA mechanism (q_head=24, kv_head=4), which saves KV cache memory compared to traditional MHA. The model is built by stacking 24 layers of Transformer blocks, with 1.0 billion parameters. Transformer is an autoregressive model that calculates the relationship between all previous tokens to obtain the probability distribution of the next token.
 
 The model now uses the **AutoModel** base class for flexible loading and saving:
 
@@ -48,14 +48,15 @@ flowchart TB
         S --> T[+]
         H --> T
         T --> U[RMSNorm]
-        U --> V[Linear]
+        U --> V["Linear (gate)"]
-        V --> W[SiLU]
+        U --> W["Linear (up)"]
-        V --> X[×]
+        V --> X[SiLU]
-        W --> X
+        X --> Y[×]
-        X --> Y[Linear]
+        W --> Y
-        Y --> Z[+]
+        Y --> Z["Linear (down)"]
-        T --> Z
+        Z --> AA[+]
-        Z --> AA[x']
+        T --> AA
+        AA --> BB[x']
     end
 
     classDef main fill:#e6f3ff,stroke:#0066cc;
@@ -168,8 +169,6 @@ from astrai.inference import InferenceEngine, GenerationRequest
 engine = InferenceEngine(
     model=model,
     tokenizer=tokenizer,
-    max_batch_size=8,
-    max_seq_len=4096,
 )
 
 # Use GenerationRequest with messages format
@@ -191,3 +190,145 @@ for token in engine.generate_with_request(request):
 ```
 
 The continuous batching feature allows dynamic batch composition where new requests can join at any time and completed requests are released immediately.
+
+## HTTP API Usage
+
+The inference server provides HTTP endpoints for remote inference. Start the server first:
+
+```bash
+python -m scripts.tools.server --port 8000
+```
+
+### OpenAI-Compatible Endpoint
+
+The server provides an OpenAI-compatible chat completion endpoint at `/v1/chat/completions`:
+
+```bash
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "system", "content": "You are a helpful assistant."},
+      {"role": "user", "content": "Hello, how are you?"}
+    ],
+    "temperature": 0.8,
+    "max_tokens": 2048,
+    "stream": false
+  }'
+```
+
+**Request Parameters:**
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `messages` | List[dict] | Required | Chat messages with role and content |
+| `temperature` | float | 1.0 | Sampling temperature (0.0-2.0) |
+| `top_p` | float | 1.0 | Nucleus sampling threshold |
+| `top_k` | int | 50 | Top-k sampling parameter |
+| `max_tokens` | int | 1024 | Maximum tokens to generate |
+| `stream` | bool | false | Enable streaming response |
+
+**Response (non-streaming):**
+```json
+{
+  "id": "chatcmpl-1234567890",
+  "object": "chat.completion",
+  "created": 1234567890,
+  "model": "astrai",
+  "choices": [
+    {
+      "index": 0,
+      "message": {"role": "assistant", "content": "Hello! I'm doing well..."},
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {
+    "prompt_tokens": 20,
+    "completion_tokens": 15,
+    "total_tokens": 35
+  }
+}
+```
+
+### Streaming Response
+
+Enable streaming for real-time token-by-token output:
+
+```bash
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "Write a story"}],
+    "stream": true,
+    "max_tokens": 500
+  }'
+```
+
+The server uses Server-Sent Events (SSE) with content type `text/event-stream`.
+
+### Anthropic-Compatible Endpoint
+
+The server also provides an Anthropic-compatible endpoint at `/v1/messages`:
+
+```bash
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "astrai",
+    "system": "You are a helpful assistant.",
+    "messages": [{"role": "user", "content": "Hello, how are you?"}],
+    "max_tokens": 2048
+  }'
+```
+
+Response:
+```json
+{
+  "id": "msg_abc123...",
+  "type": "message",
+  "role": "assistant",
+  "model": "astrai",
+  "content": [{"type": "text", "text": "Hello! I am doing well..."}],
+  "stop_reason": "end_turn",
+  "stop_sequence": null,
+  "usage": {"input_tokens": 20, "output_tokens": 15}
+}
+```
+
+Streaming:
+```bash
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "astrai",
+    "system": "You are a helpful assistant.",
+    "messages": [{"role": "user", "content": "Write a short poem"}],
+    "max_tokens": 500,
+    "stream": true
+  }'
+```
+
+Supports `stop_sequences` for early termination:
+```bash
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "astrai",
+    "messages": [{"role": "user", "content": "Write a story"}],
+    "max_tokens": 500,
+    "stop_sequences": ["The end", "THE END"]
+  }'
+```
+
+### Health Check
+
+Monitor server and model status:
+
+```bash
+curl http://localhost:8000/health
+# {"status": "ok", "model_loaded": true}
+
+curl http://localhost:8000/stats
+# {"total_tasks": 10, "total_tokens": 5000, "active_tasks": 1, "waiting_queue": 0}
+```
+
+> Document Update Time: 2026-04-09

assets/docs/params.md

@@ -4,70 +4,87 @@
 
 ### Basic Parameters
 
-| Parameter | Description | Default Value |
+| Parameter | Description | Default |
-|-----------|-------------|---------------|
+|-----------|-------------|---------|
-| `--train_type` | Training type (seq, sft, dpo, grpo) | required |
+| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`) | required |
-| `--model_type` | Model type for AutoModel loading (e.g., transformer) | transformer |
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
 | `--n_epoch` | Total training epochs | 1 |
-| `--batch_size` | Batch size | 4 |
+| `--batch_size` | Batch size | 1 |
-| `--accumulation_steps` | Gradient accumulation steps | 1 |
+| `--accumulation_steps` | Gradient accumulation steps between optimizer steps | 1 |
 
 ### Learning Rate Scheduling
 
-| Parameter | Description | Default Value |
+| Parameter | Description | Default |
-|-----------|-------------|---------------|
+|-----------|-------------|---------|
 | `--warmup_steps` | Warmup steps | 1000 |
-| `--max_lr` | Maximum learning rate (warmup + cosine decay) | 3e-4 |
+| `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
-| `--max_grad_norm` | Maximum gradient norm | 1.0 |
+| `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 
-### Checkpoint
+### Optimizer (AdamW)
 
-| Parameter | Description | Default Value |
+| Parameter | Description | Default |
-|-----------|-------------|---------------|
+|-----------|-------------|---------|
-| `--ckpt_interval` | Checkpoint save interval (iterations) | 5000 |
-| `--ckpt_dir` | Checkpoint save directory | checkpoint |
-| `--resume_dir` | Resume training from specified path | - |
-
-### Optimizer Parameters
-
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
 | `--adamw_beta1` | AdamW beta1 | 0.9 |
 | `--adamw_beta2` | AdamW beta2 | 0.95 |
 | `--adamw_weight_decay` | AdamW weight decay | 0.01 |
 
 ### Data Loading
 
-| Parameter | Description | Default Value |
+| Parameter | Description | Default |
-|-----------|-------------|---------------|
+|-----------|-------------|---------|
-| `--random_seed` | Random seed | 3407 |
+| `--window_size` | Max input sequence length | model config `max_len` |
-| `--num_workers` | DataLoader workers | 0 |
+| `--stride` | Stride for sliding window over sequences | None |
-| `--prefetch_factor` | Prefetch factor for dataloader | None |
+| `--random_seed` | Random seed for reproducibility | 3407 |
-| `--pin_memory` | Enable pin_memory | False |
+| `--num_workers` | DataLoader worker processes | 4 |
-| `--no_pin_memory` | Disable pin_memory | - |
+| `--no_pin_memory` | Disable pin_memory (enabled by default) | (flag) |
+
+### Checkpoint & Resume
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--ckpt_interval` | Iterations between checkpoints | 5000 |
+| `--ckpt_dir` | Checkpoint save directory | checkpoint |
+| `--start_epoch` | Resume from epoch (0 = from scratch) | 0 |
+| `--start_batch` | Resume from batch iteration | 0 |
 
 ### Distributed Training
 
-| Parameter | Description | Default Value |
+| Parameter | Description | Default |
-|-----------|-------------|---------------|
+|-----------|-------------|---------|
-| `--nprocs` | Number of GPUs | 1 |
+| `--nprocs` | Number of GPUs / processes | 1 |
-| `--device_type` | Device type (cuda/cpu) | cuda |
+| `--device_type` | Device type | cuda |
 
-### Other Parameters
+### Strategy-specific
 
-| Parameter | Description | Default Value |
+| Parameter | Description | Default | Used by |
-|-----------|-------------|---------------|
+|-----------|-------------|---------|---------|
-| `--window_size` | Maximum input sequence length | model config max_len |
+| `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
-| `--stride` | Input sequence stride | - |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.1 | `seq`, `sft` |
-| `--dpo_beta` | DPO beta value | 0.1 |
+| `--group_size` | GRPO group size | 4 | `grpo` |
-| `--grpo_clip_eps` | GRPO clip epsilon | 0.2 |
+| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
-| `--grpo_kl_coef` | GRPO KL coefficient | 0.01 |
+| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
-| `--grpo_group_size` | GRPO group size | 4 |
+| `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` |
-| `--label_smoothing` | Label smoothing parameter | 0.1 |
+
-| `--start_epoch` | Starting epoch | 0 |
+### Usage Example
-| `--start_batch` | Starting batch | 0 |
+
+```bash
+python scripts/tools/train.py \
+  --train_type seq \
+  --data_root_path /path/to/dataset \
+  --param_path /path/to/model \
+  --n_epoch 3 \
+  --batch_size 4 \
+  --accumulation_steps 8 \
+  --max_lr 3e-4 \
+  --warmup_steps 2000 \
+  --max_grad_norm 1.0 \
+  --ckpt_interval 5000 \
+  --ckpt_dir ./checkpoints \
+  --num_workers 4 \
+  --nprocs 1 \
+  --device_type cuda
+```
 
 ---
 
@@ -89,14 +106,14 @@
 ```python
 import torch
 from astrai.model import AutoModel
-from astrai.tokenize import Tokenizer
+from astrai.tokenize import AutoTokenizer
 from astrai.inference import InferenceEngine, GenerationRequest
 
 # Load model using AutoModel
 model = AutoModel.from_pretrained("your_model_dir")
 
 # Load tokenizer
-tokenizer = Tokenizer("your_model_dir")
+tokenizer = AutoTokenizer.from_pretrained("your_model_dir")
 
 # Create engine with separate model and tokenizer
 engine = InferenceEngine(
@@ -137,3 +154,5 @@ result = engine.generate(
 |------|-------------|
 | `stream=True` | Streaming output, yields token by token |
 | `stream=False` | Non-streaming output, returns complete result |
+
+> Document Update Time: 2026-04-09

astrai/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.3"
+__version__ = "1.3.4"
 __author__ = "ViperEkura"
 
 from astrai.config import (
@@ -12,18 +12,19 @@ from astrai.inference import (
     InferenceEngine,
 )
 from astrai.model import AutoModel, Transformer
-from astrai.tokenize import BpeTokenizer
+from astrai.tokenize import AutoTokenizer
-from astrai.trainer import SchedulerFactory, StrategyFactory, Trainer
+from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer
 
 __all__ = [
     "Transformer",
     "ModelConfig",
     "TrainConfig",
     "DatasetFactory",
-    "BpeTokenizer",
+    "AutoTokenizer",
     "GenerationRequest",
     "InferenceEngine",
     "Trainer",
+    "CallbackFactory",
     "StrategyFactory",
     "SchedulerFactory",
     "BaseFactory",

astrai/config/train_config.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Callable, List, Optional
+from typing import Callable, Optional
 
 import torch.nn as nn
 from torch.optim import Optimizer
@@ -74,9 +74,6 @@ class TrainConfig:
     )
 
     # others
-    device_ids: Optional[List[int]] = field(
-        default=None, metadata={"help": "Device ids for distributed training."}
-    )
     device_type: str = field(
         default="cuda", metadata={"help": "Device type for distributed training."}
     )

astrai/dataset/dataset.py

@@ -72,15 +72,16 @@ class MultiSegmentFetcher:
     Each key corresponds to a different type of data (e.g., "sequence", "mask").
     """
 
-    def __init__(self, muti_segments: Dict):
+    def __init__(self, multi_segments: Dict):
-        self.muti_keys = list(muti_segments.keys())
+        self.multi_keys = list(multi_segments.keys())
-        self.muti_fetchers = {
+        self.multi_fetchers = {
-            key: BaseSegmentFetcher(segments) for key, segments in muti_segments.items()
+            key: BaseSegmentFetcher(segments)
+            for key, segments in multi_segments.items()
         }
 
     def __len__(self) -> int:
         """Returns the minimum length across all fetchers."""
-        len_list = [len(seg) for seg in self.muti_fetchers.values()]
+        len_list = [len(seg) for seg in self.multi_fetchers.values()]
         return min(len_list)
 
     def key_fetch(
@@ -100,7 +101,7 @@ class MultiSegmentFetcher:
         keys = [keys] if isinstance(keys, str) else keys
 
         for key in keys:
-            fetcher = self.muti_fetchers[key]
+            fetcher = self.multi_fetchers[key]
             fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
             fetch_dict[key] = fetch_tensor
 
@@ -108,7 +109,7 @@ class MultiSegmentFetcher:
 
     def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
         """Fetch all keys."""
-        return self.key_fetch(begin_idx, end_idx, self.muti_keys)
+        return self.key_fetch(begin_idx, end_idx, self.multi_keys)
 
 
 class BaseDataset(Dataset, ABC):

astrai/inference/__init__.py

@@ -1,25 +1,46 @@
-"""Inference module for continuous batching."""
+"""Inference module for continuous batching.
+
+Layers:
+  - engine.py:    Facade (InferenceEngine), Value Object (GenerationParams, GenerationRequest)
+  - scheduler.py: Continuous-batching loop, Task state machine, TaskStatus enum
+  - cache.py:     PagedCache (page-table-indirected KV cache with alloc/free)
+  - sampling.py:  Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
+  - server.py:    FastAPI HTTP server (OpenAI-compatible endpoints)
+"""
 
 from astrai.inference.engine import (
+    GenerationParams,
     GenerationRequest,
     InferenceEngine,
 )
+from astrai.inference.sampling import (
+    BaseSamplingStrategy,
+    SamplingPipeline,
+    TemperatureStrategy,
+    TopKStrategy,
+    TopPStrategy,
+    sample,
+)
 from astrai.inference.scheduler import (
     InferenceScheduler,
     Task,
     TaskStatus,
-    apply_sampling_strategies,
 )
 
 __all__ = [
-    # Engine
+    # Engine / Requests
     "InferenceEngine",
+    "GenerationRequest",
+    "GenerationParams",
     # Scheduler
     "InferenceScheduler",
     "Task",
     "TaskStatus",
-    # Request
+    # Sampling (Strategy pattern)
-    "GenerationRequest",
+    "sample",
-    # Sampling
+    "BaseSamplingStrategy",
-    "apply_sampling_strategies",
+    "TemperatureStrategy",
+    "TopKStrategy",
+    "TopPStrategy",
+    "SamplingPipeline",
 ]

astrai/inference/cache.py

@@ -0,0 +1,174 @@
+"""Page-based KV cache with page-table-indirected read/write.
+
+Provides:
+  - PagedCache: paged KV cache combining page pool and tensor storage.
+"""
+
+from typing import Dict, List, Tuple
+
+import torch
+from torch import Tensor
+
+STOP = object()
+
+
+def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
+    start = page_idx * page_size
+    end = min(start + page_size, len(token_ids))
+    h = 0
+    for i in range(start, end):
+        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
+    return h
+
+
+class PagedCache:
+    """Paged KV cache with page-table-indirected read/write.
+
+    Combines:
+      - Page pool (ref-counted alloc/free via bitmask)
+      - KV tensor storage (k_cache, v_cache)
+      - Prefix-cache hash lookup (page_content_hash -> physical_page_idx)
+
+    Call :meth:`bind` to obtain a batch view for the attention layers.
+    """
+
+    def __init__(
+        self,
+        n_layers: int,
+        n_pages: int,
+        page_size: int,
+        n_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        self.page_size = page_size
+        self._free_mask = (1 << n_pages) - 1
+        self._refs: List[int] = [0] * n_pages
+        self.k_cache = torch.empty(
+            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
+            device=device,
+            dtype=dtype,
+        )
+        self.v_cache = torch.empty(
+            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
+            device=device,
+            dtype=dtype,
+        )
+        self._page_to_hash: Dict[int, int] = {}
+        self._hash_to_page: Dict[int, int] = {}
+
+    def record_page(
+        self, page_idx: int, token_ids: List[int], logical_page_idx: int
+    ) -> None:
+        h = page_hash(token_ids, logical_page_idx, self.page_size)
+        old_h = self._page_to_hash.pop(page_idx, None)
+        if old_h is not None:
+            self._hash_to_page.pop(old_h, None)
+        self._page_to_hash[page_idx] = h
+        self._hash_to_page[h] = page_idx
+
+    def lookup_prefix(self, token_ids: List[int]) -> List[int]:
+        full_pages = len(token_ids) // self.page_size
+        hits: List[int] = []
+        for i in range(full_pages):
+            h = page_hash(token_ids, i, self.page_size)
+            p = self._hash_to_page.get(h)
+            if p is None:
+                break
+            hits.append(p)
+        return hits
+
+    def inc_ref(self, idx: int) -> None:
+        self._refs[idx] += 1
+
+    def alloc(self) -> int:
+        lsb = self._free_mask & -self._free_mask
+        if lsb == 0:
+            return -1
+        idx = lsb.bit_length() - 1
+        self._free_mask ^= lsb
+        self._refs[idx] = 1
+        return idx
+
+    def alloc_n(self, n: int) -> List[int]:
+        pages = [self.alloc() for _ in range(n)]
+        if any(p < 0 for p in pages):
+            for p in pages:
+                if p >= 0:
+                    self.free(p)
+            return []
+        return pages
+
+    def free(self, idx: int) -> None:
+        self._refs[idx] -= 1
+        if self._refs[idx] == 0:
+            self._free_mask |= 1 << idx
+            h = self._page_to_hash.pop(idx, None)
+            if h is not None:
+                self._hash_to_page.pop(h, None)
+
+    def bind(self, page_table: Tensor, total_len: int = 0) -> "CacheView":
+        return CacheView(self, page_table, total_len)
+
+    def write(
+        self, layer_id: int, page_table: Tensor, start_pos: int, k: Tensor, v: Tensor
+    ) -> None:
+        seq_len = k.size(1)
+        if seq_len == 0:
+            return
+        page_size = self.page_size
+        written = 0
+        first_page = start_pos // page_size
+        last_page = (start_pos + seq_len - 1) // page_size
+        for pi in range(first_page, last_page + 1):
+            phys_pages = page_table[:, pi]
+            page_start = pi * page_size
+            write_start = max(page_start, start_pos)
+            write_end = min(page_start + page_size, start_pos + seq_len)
+            offset = write_start - page_start
+            chunk = write_end - write_start
+            self.k_cache[layer_id, phys_pages, offset : offset + chunk] = k[
+                :, written : written + chunk
+            ]
+            self.v_cache[layer_id, phys_pages, offset : offset + chunk] = v[
+                :, written : written + chunk
+            ]
+            written += chunk
+
+    def gather(self, layer_id: int, page_table: Tensor) -> Tuple[Tensor, Tensor]:
+        k_parts, v_parts = [], []
+        for pi in range(page_table.size(1)):
+            phys_pages = page_table[:, pi]
+            if not (phys_pages >= 0).any():
+                break
+            k_parts.append(self.k_cache[layer_id, phys_pages])
+            v_parts.append(self.v_cache[layer_id, phys_pages])
+        k = torch.cat(k_parts, dim=1)
+        v = torch.cat(v_parts, dim=1)
+        return k, v
+
+
+class CacheView:
+    """Per-batch view that bundles PagedCache + page_table + total_len.
+
+    Attention layers receive this as ``paged_cache`` and only see
+    ``write()`` / ``gather()``, never raw page tables or length params.
+    """
+
+    __slots__ = ("_cache", "_page_table", "_total_len")
+
+    def __init__(self, cache: PagedCache, page_table: Tensor, total_len: int = 0):
+        self._cache = cache
+        self._page_table = page_table
+        self._total_len = total_len
+
+    def write(self, layer_id: int, start_pos: int, k: Tensor, v: Tensor) -> None:
+        self._cache.write(layer_id, self._page_table, start_pos, k, v)
+
+    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
+        k, v = self._cache.gather(layer_id, self._page_table)
+        if self._total_len:
+            k = k[:, : self._total_len]
+            v = v[:, : self._total_len]
+        return k, v

astrai/inference/engine.py

@@ -1,21 +1,42 @@
-"""Unified inference engine."""
+"""Unified inference engine for continuous batching.
 
+Layers:
+  - GenerationParams:    Immutable value object for sampling parameters.
+  - GenerationRequest:   User-facing request DTO with validation.
+  - _Result:             Thread-safe token accumulator (Observer pattern).
+  - InferenceEngine:     Facade over InferenceScheduler + async wrapper.
+"""
+
+import asyncio
 import gc
-import logging
 import threading
-from typing import Any, Dict, Generator, List, Optional, Union
+from dataclasses import dataclass
+from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
 
 import torch
 import torch.nn as nn
 
+from astrai.inference.cache import STOP
 from astrai.inference.scheduler import InferenceScheduler
 from astrai.tokenize import AutoTokenizer
 
-logger = logging.getLogger(__name__)
+
+@dataclass(frozen=True)
+class GenerationParams:
+    """Immutable value object for sampling hyperparameters."""
+
+    top_k: int = 50
+    top_p: float = 1.0
+    temperature: float = 1.0
+    max_tokens: int = 1024
 
 
 class GenerationRequest:
-    """Request parameters for text generation."""
+    """Request parameters for text generation.
+
+    Encapsulates messages, sampling parameters (via GenerationParams),
+    and streaming preference for a single generation request.
+    """
 
     def __init__(
         self,
@@ -26,17 +47,44 @@ class GenerationRequest:
         max_len: int = 1024,
         stream: bool = False,
     ):
-        self.messages = messages
+        """Initializes a generation request.
-        self.top_k = top_k
-        self.top_p = top_p
-        self.temperature = temperature
-        self.max_len = max_len
-        self.stream = stream
 
+        Args:
+            messages: Conversation history as list of {"role": ..., "content": ...}.
+            top_k: Top-k sampling count (0 disables).
+            top_p: Nucleus sampling probability threshold.
+            temperature: Sampling temperature.
+            max_len: Maximum tokens to generate.
+            stream: Whether to return output as a token stream.
+        """
+        self.messages = messages
+        self.params = GenerationParams(
+            top_k=top_k,
+            top_p=top_p,
+            temperature=temperature,
+            max_tokens=max_len,
+        )
+        self.stream = stream
         self._validate()
 
+    @property
+    def top_k(self) -> int:
+        return self.params.top_k
+
+    @property
+    def top_p(self) -> float:
+        return self.params.top_p
+
+    @property
+    def temperature(self) -> float:
+        return self.params.temperature
+
+    @property
+    def max_len(self) -> int:
+        return self.params.max_tokens
+
     def _validate(self):
-        """Validate request parameters."""
+        """Validates sampling parameter ranges."""
         if not (isinstance(self.top_k, int) and self.top_k >= 0):
             raise ValueError("top_k must be a non-negative integer")
         if not (0.0 <= self.top_p <= 1.0):
@@ -45,66 +93,102 @@ class GenerationRequest:
             raise ValueError("temperature must be a non-negative number")
 
 
-class _StreamingResult:
+class _Result:
-    """Streaming result holder with event-based notification."""
+    """Thread-safe token accumulator for streaming and non-streaming modes.
 
-    def __init__(self):
+    Supports multiple concurrent generation tasks with per-index result tracking.
-        self.tokens: List[str] = []
+    Uses a threading.Condition for efficient completion notification
+    and a threading.Event for streaming wakeup.
+    """
+
+    def __init__(self, count: int = 1):
+        """Initializes the accumulator.
+
+        Args:
+            count: Number of concurrent generation tasks to track.
+        """
+        self._cond = threading.Condition()
         self._event = threading.Event()
-        self._lock = threading.Lock()
+        self.tokens: List[str] = []
+        self.results: List[str] = [""] * count
+        self._done: List[bool] = [False] * count
+        self._completed = 0
+        self._total = count
 
-    def append(self, token: str):
+    def append(self, token: str, idx: int = 0):
-        with self._lock:
+        """Appends a token to the result buffer.
+
+        In non-streaming mode, tokens are concatenated into results[idx].
+        The sentinel STOP marks a task as complete.
+
+        Args:
+            token: The decoded token string, or STOP sentinel.
+            idx: Index of the generation task this token belongs to.
+        """
+        with self._cond:
             self.tokens.append(token)
-        self._event.set()
+            if token is not STOP:
+                self.results[idx] += token
+            else:
+                if not self._done[idx]:
+                    self._done[idx] = True
+                    self._completed += 1
+                    self._cond.notify_all()
+            self._event.set()
 
     def pop_all(self) -> List[str]:
-        with self._lock:
+        """Returns and clears all accumulated tokens.
-            tokens = self.tokens.copy()
+
+        Returns:
+            List of token strings since the last call.
+        """
+        with self._cond:
+            out = self.tokens.copy()
             self.tokens.clear()
-            if not tokens:
+            if not out:
                 self._event.clear()
-            return tokens
+            return out
 
-    def wait(self, timeout: float = None) -> bool:
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        """Blocks until new tokens arrive or the timeout expires.
+
+        Args:
+            timeout: Maximum wait time in seconds (None = infinite).
+
+        Returns:
+            True if the event was set (new data available), False on timeout.
+        """
         return self._event.wait(timeout=timeout)
 
+    def wait_completion(self) -> None:
+        """Blocks until all tasks complete (non-streaming).
 
-class _NonStreamingResult:
+        Uses a Condition to sleep efficiently instead of busy-waiting.
-    """Non-streaming result holder with event-based completion notification."""
+        The calling thread is parked until a STOP signal arrives.
-
+        """
-    def __init__(self, count: int):
+        with self._cond:
-        self.results: List[str] = [""] * count
+            self._cond.wait_for(lambda: self._completed >= self._total)
-        self.done_flags: List[bool] = [False] * count
-        self._completed_count = 0
-        self._event = threading.Event()
-        self._lock = threading.Lock()
-
-    def append(self, idx: int, token: str):
-        with self._lock:
-            if token == "[DONE]":
-                if not self.done_flags[idx]:
-                    self.done_flags[idx] = True
-                    self._completed_count += 1
-                    if self._completed_count == len(self.results):
-                        self._event.set()
-            else:
-                self.results[idx] += token
-
-    def is_all_done(self) -> bool:
-        with self._lock:
-            return all(self.done_flags)
-
-    def wait(self, timeout: float = None) -> bool:
-        return self._event.wait(timeout=timeout)
 
     def get_results(self) -> List[str]:
-        with self._lock:
+        """Returns all accumulated results for non-streaming mode.
+
+        Returns:
+            List of complete generated strings, one per task index.
+        """
+        with self._cond:
             return self.results.copy()
 
 
 class InferenceEngine:
-    """Unified inference engine for continuous batching."""
+    """Unified inference engine backed by continuous-batching scheduler.
+
+    Usage:
+        with InferenceEngine(model, tokenizer) as engine:
+            for token in engine.generate("hello", stream=True):
+                print(token, end="")
+
+            text = engine.generate("hello")
+    """
 
     def __init__(
         self,
@@ -112,49 +196,37 @@ class InferenceEngine:
         tokenizer: AutoTokenizer,
         max_batch_size: int = 1,
         max_seq_len: Optional[int] = None,
+        max_prompt_len: int = 2048,
+        page_size: int = 128,
     ):
-        """
+        """Initializes the inference engine.
-        Initialize inference engine with separate model and tokenizer.
 
         Args:
-            model: The language model for inference (nn.Module, e.g., Transformer)
+            model: The model instance.
-            tokenizer: The tokenizer for encoding/decoding text
+            tokenizer: The tokenizer instance.
-            config: Model configuration
+            max_batch_size: Maximum number of concurrent tasks.
-            max_batch_size: Maximum batch size for continuous batching
+            max_seq_len: Maximum sequence length.
-            max_seq_len: Maximum sequence length (defaults to config.max_len)
+            max_prompt_len: Maximum prompt tokens.
+            compile: Whether to compile the model with torch.compile.
+            page_size: Number of tokens per KV cache page.
         """
         self.model = model
         self.tokenizer = tokenizer
-
-        # Get device and dtype from model parameters
-        try:
-            first_param = next(model.parameters())
-            device = first_param.device
-            dtype = first_param.dtype
-        except StopIteration:
-            # Model has no parameters, use default device/dtype
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            dtype = torch.float32
-
         self.scheduler = InferenceScheduler(
             model=self.model,
             tokenizer=self.tokenizer,
             max_batch_size=max_batch_size,
             max_seq_len=max_seq_len,
-            device=device,
+            max_prompt_len=max_prompt_len,
-            dtype=dtype,
+            page_size=page_size,
         )
 
-        self.kv_cache = self.scheduler.kv_cache
-        self.seq_mask = self.scheduler.seq_mask
-
         self.scheduler.start()
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        """Handle exceptions on exit."""
         self.shutdown()
         return False
 
@@ -166,46 +238,106 @@ class InferenceEngine:
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = 50,
-        abort_on_exception: bool = True,
     ) -> Union[Generator[str, None, None], str, List[str]]:
-        """Unified generation interface.
+        """Generates text from a prompt.
 
         Args:
-            abort_on_exception: If True, abort the generation when consumer
+            prompt: Single string or list of strings for batch generation.
-                stops iterating (GeneratorExit/StopIteration). Default: True.
+            stream: If True, returns a generator yielding tokens one by one.
+            max_tokens: Maximum number of tokens to generate.
+            temperature: Sampling temperature.
+            top_p: Nucleus sampling probability threshold.
+            top_k: Top-k sampling count (0 disables).
+
+        Returns:
+            Generator (stream=True), single string (non-stream, single prompt),
+            or list of strings (non-stream, batch prompts).
         """
         is_batch = isinstance(prompt, list)
         prompts = prompt if is_batch else [prompt]
 
         if stream:
             return self._generate_streaming(
-                prompts,
+                prompts, is_batch, max_tokens, temperature, top_p, top_k
-                is_batch,
-                max_tokens,
-                temperature,
-                top_p,
-                top_k,
-                abort_on_exception,
             )
         else:
             return self._generate_non_streaming(
                 prompts, is_batch, max_tokens, temperature, top_p, top_k
             )
 
+    def generate_async(
+        self,
+        prompt: str,
+        max_tokens: int = 1024,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+    ) -> AsyncGenerator[str, None]:
+        """Async streaming generator that does not block the event loop.
+
+        Runs the synchronous generator in a background thread pool executor,
+        yielding tokens to the async consumer as they arrive.
+
+        Args:
+            prompt: Input text to generate from.
+            max_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature.
+            top_p: Nucleus sampling threshold.
+            top_k: Top-k sampling count.
+
+        Yields:
+            Decoded token strings as they are generated.
+        """
+        sync_gen = self._generate_streaming(
+            [prompt], False, max_tokens, temperature, top_p, top_k
+        )
+
+        async def _agen():
+            loop = asyncio.get_event_loop()
+            while True:
+                token = await loop.run_in_executor(None, self._next_token, sync_gen)
+                if token is None:
+                    break
+                yield token
+
+        return _agen()
+
+    @staticmethod
+    def _next_token(gen: Generator) -> Optional[str]:
+        """Retrieves the next token from a synchronous generator.
+
+        Args:
+            gen: A synchronous generator yielding token strings.
+
+        Returns:
+            The next token, or None if the generator is exhausted.
+        """
+        try:
+            return next(gen)
+        except StopIteration:
+            return None
+
     def generate_with_request(
         self, request: GenerationRequest
     ) -> Union[Generator[str, None, None], str, List[str]]:
-        """Generate with GenerationRequest object."""
+        """Generates text from a structured GenerationRequest.
-        # Use tokenizer's chat template with messages
-        prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
 
+        Applies the chat template to the request's messages before generation.
+
+        Args:
+            request: A GenerationRequest with messages and parameters.
+
+        Returns:
+            Generator, string, or list of strings (see generate()).
+        """
+        prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
         return self.generate(
             prompt=prompt,
             stream=request.stream,
-            max_tokens=request.max_len,
+            max_tokens=request.params.max_tokens,
-            temperature=request.temperature,
+            temperature=request.params.temperature,
-            top_p=request.top_p,
+            top_p=request.params.top_p,
-            top_k=request.top_k,
+            top_k=request.params.top_k,
         )
 
     def _generate_streaming(
@@ -216,18 +348,27 @@ class InferenceEngine:
         temperature: float,
         top_p: float,
         top_k: int,
-        abort_on_exception: bool = True,
+    ) -> Generator[str, None, None]:
-    ) -> Union[Generator[str, None, None], List[Generator[str, None, None]]]:
+        """Internal streaming generator.
-        """Generate with streaming output.
+
+        Polls the _Result accumulator in a loop, yielding tokens as they arrive.
+        Cleans up the scheduler task on GeneratorExit.
 
         Args:
-            abort_on_exception: If True, abort the task when generator is
+            prompts: List of prompts (only first is used; batch not yet supported).
-                stopped early by consumer (GeneratorExit/StopIteration).
+            is_batch: If True, raises NotImplementedError.
+            max_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature.
+            top_p: Nucleus sampling threshold.
+            top_k: Top-k sampling count.
+
+        Yields:
+            Decoded token strings.
         """
         if is_batch:
-            raise NotImplementedError("Batch streaming is not implemented yet")
+            raise NotImplementedError("Batch streaming not yet supported")
 
-        result = _StreamingResult()
+        result = _Result()
 
         task_id = self.scheduler.add_task(
             prompt=prompts[0],
@@ -235,7 +376,7 @@ class InferenceEngine:
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            stream_callback=result.append,
+            stream_callback=lambda tok: result.append(tok, 0),
         )
 
         def gen():
@@ -243,17 +384,14 @@ class InferenceEngine:
                 while True:
                     tokens = result.pop_all()
                     for token in tokens:
-                        if token == "[DONE]":
+                        if token is STOP:
                             return
                         yield token
-                    result.wait(timeout=0.05)
+                    if not result.wait(timeout=0.05):
-            except Exception:
+                        pass
-                # Consumer stopped iterating - abort the task
+            finally:
-                if abort_on_exception:
+                self.scheduler.remove_task(task_id)
-                    self.scheduler.remove_task(task_id)
-                raise
 
-        gen.task_id = task_id
         return gen()
 
     def _generate_non_streaming(
@@ -265,36 +403,57 @@ class InferenceEngine:
         top_p: float,
         top_k: int,
     ) -> Union[str, List[str]]:
-        """Generate without streaming."""
+        """Internal non-streaming generator.
-        result = _NonStreamingResult(len(prompts))
+
+        Submits all prompts to the scheduler and waits for all to complete.
+
+        Args:
+            prompts: List of prompt strings.
+            is_batch: Whether multiple prompts were provided.
+            max_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature.
+            top_p: Nucleus sampling threshold.
+            top_k: Top-k sampling count.
+
+        Returns:
+            Single string for one prompt, list of strings for batch.
+        """
+        result = _Result(count=len(prompts))
+        task_ids = []
 
         for i, p in enumerate(prompts):
-            # Create closure to capture current index value using factory function
-            def make_callback(idx):
-                def callback(token):
-                    result.append(idx, token)
 
-                return callback
+            def make_cb(idx):
+                return lambda tok: result.append(tok, idx)
 
-            self.scheduler.add_task(
+            task_id = self.scheduler.add_task(
                 prompt=p,
                 max_tokens=max_tokens,
                 temperature=temperature,
                 top_p=top_p,
                 top_k=top_k,
-                stream_callback=make_callback(i),
+                stream_callback=make_cb(i),
             )
+            task_ids.append(task_id)
 
-        result.wait()
+        result.wait_completion()
-        results = result.get_results()
+
-        return results if is_batch else results[0]
+        for task_id in task_ids:
+            self.scheduler.remove_task(task_id)
+
+        res = result.get_results()
+        return res if is_batch else res[0]
 
     def get_stats(self) -> Dict[str, Any]:
-        """Get engine statistics."""
+        """Returns current engine statistics.
+
+        Returns:
+            Dict with total_tasks, total_tokens, active_tasks, waiting_queue.
+        """
         return self.scheduler.get_stats()
 
     def shutdown(self) -> None:
-        """Shutdown the engine and release all resources."""
+        """Shuts down the engine, stops the scheduler, and frees GPU memory."""
         self.scheduler.stop()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

astrai/inference/sampling.py

@@ -0,0 +1,178 @@
+"""Composable sampling strategies for logit transformation.
+
+Implements the Strategy pattern: each sampling technique
+(temperature, top-k, top-p) is a pluggable strategy that
+can be composed into a pipeline.
+
+All strategies accept both scalar and per-sample tensor
+parameters, so a single pipeline works for any batch size.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Union
+
+import torch
+from torch import Tensor
+
+
+class BaseSamplingStrategy(ABC):
+    """Abstract base for a logit transformation strategy."""
+
+    @abstractmethod
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
+        """Applies the strategy to logits.
+
+        Args:
+            logits: Raw logits tensor (batch, vocab_size).
+            filter_value: Value assigned to filtered-out positions.
+
+        Returns:
+            Transformed logits tensor.
+        """
+
+
+class TemperatureStrategy(BaseSamplingStrategy):
+    """Divides logits by temperature to control randomness.
+
+    Args:
+        temperature: Scalar or ``[batch]`` tensor.
+    """
+
+    def __init__(self, temperature: Union[float, Tensor] = 1.0):
+        self.temperature = temperature
+
+    def apply(self, logits, filter_value=-float("inf")):
+        t = self.temperature
+        if isinstance(t, Tensor):
+            if (t != 1.0).any():
+                logits = logits / t.to(logits.device, non_blocking=True).view(-1, 1)
+        elif t != 1.0:
+            logits = logits / t
+        return logits
+
+
+class TopKStrategy(BaseSamplingStrategy):
+    """Keeps only the top-k logits, setting the rest to filter_value.
+
+    Args:
+        top_k: Scalar or ``[batch]`` tensor (0 disables).
+    """
+
+    def __init__(self, top_k: Union[int, Tensor] = 0):
+        self.top_k = top_k
+
+    def apply(self, logits, filter_value=-float("inf")):
+        tk = self.top_k
+        if isinstance(tk, Tensor):
+            max_k = int(tk.max().item())
+            if max_k <= 0:
+                return logits
+            k = min(max_k, logits.size(-1))
+        elif tk > 0:
+            k = min(tk, logits.size(-1))
+        else:
+            return logits
+        thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
+        logits[logits < thresholds] = filter_value
+        return logits
+
+
+class TopPStrategy(BaseSamplingStrategy):
+    """Nucleus (top-p) filtering: keeps the smallest set of tokens whose
+    cumulative probability exceeds top_p.
+
+    Args:
+        top_p: Scalar or ``[batch]`` tensor (1.0 disables).
+    """
+
+    def __init__(self, top_p: Union[float, Tensor] = 1.0):
+        self.top_p = top_p
+
+    def _apply(self, logits, top_p, filter_value):
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+        remove = cum_probs > top_p
+        remove[..., 1:] = remove[..., :-1].clone()
+        remove[..., 0] = False
+        mask = torch.zeros_like(logits, dtype=torch.bool)
+        mask.scatter_(1, sorted_indices, remove)
+        logits[mask] = filter_value
+        return logits
+
+    def apply(self, logits, filter_value=-float("inf")):
+        tp = self.top_p
+        if isinstance(tp, Tensor):
+            tp = tp.to(logits.device, non_blocking=True)
+            if (tp < 1.0).any():
+                logits = self._apply(logits, tp.view(-1, 1), filter_value)
+        elif tp < 1.0:
+            logits = self._apply(logits, tp, filter_value)
+        return logits
+
+
+class SamplingPipeline(BaseSamplingStrategy):
+    """Composes multiple sampling strategies into a single transformation.
+
+    Strategies are applied sequentially in the order they are provided,
+    matching the original temperature -> top-k -> top-p ordering.
+
+    Usage::
+
+        pipeline = SamplingPipeline([
+            TemperatureStrategy(0.8),
+            TopKStrategy(50),
+            TopPStrategy(0.95),
+        ])
+        logits = pipeline.apply(logits)
+        token = pipeline.sample(logits)       # softmax + multinomial
+    """
+
+    def __init__(self, strategies: List[BaseSamplingStrategy]):
+        self.strategies = strategies
+
+    def apply(self, logits, filter_value=-float("inf")):
+        for strategy in self.strategies:
+            logits = strategy.apply(logits, filter_value)
+        return logits
+
+    @torch.no_grad()
+    def sample(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
+        """Apply strategies then sample (softmax + multinomial).
+
+        Args:
+            logits: Raw logits ``[batch, vocab_size]``.
+
+        Returns:
+            Sampled token IDs ``[batch]``.
+        """
+        return torch.multinomial(
+            torch.softmax(self.apply(logits, filter_value), dim=-1),
+            num_samples=1,
+        ).squeeze(-1)
+
+
+@torch.inference_mode()
+def sample(
+    logits: Tensor,
+    temperature: Union[float, Tensor] = 1.0,
+    top_k: Union[int, Tensor] = 0,
+    top_p: Union[float, Tensor] = 1.0,
+    filter_value: float = -float("inf"),
+) -> Tensor:
+    """Apply sampling strategies then sample (softmax + multinomial).
+
+    Shortcut for ``SamplingPipeline(...).sample(logits)``.
+
+    Args:
+        logits: Raw logits ``[batch, vocab_size]``.
+
+    Returns:
+        Sampled token IDs ``[batch]``.
+    """
+    return SamplingPipeline(
+        [
+            TemperatureStrategy(temperature),
+            TopKStrategy(top_k),
+            TopPStrategy(top_p),
+        ]
+    ).sample(logits, filter_value)

astrai/inference/scheduler.py

@@ -1,19 +1,25 @@
-"""Inference scheduler for continuous batching."""
+"""Inference scheduler for single-GPU continuous batching with paged KV cache."""
 
+import logging
 import threading
 import time
 import uuid
-from typing import Any, Callable, Dict, List, Optional
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
 from torch import Tensor
 
+from astrai.inference.cache import STOP, PagedCache
+from astrai.inference.sampling import sample
 from astrai.model.automodel import AutoModel
-from astrai.tokenize import AutoTokenizer
+from astrai.tokenize.tokenizer import AutoTokenizer
+
+logger = logging.getLogger(__name__)
 
 
-class TaskStatus:
+class TaskStatus(Enum):
-    """Task state for continuous batching."""
+    """Task states in the continuous batching lifecycle."""
 
     PENDING = "pending"
     RUNNING = "running"
@@ -22,7 +28,7 @@ class TaskStatus:
 
 
 class Task:
-    """Individual task for continuous batching."""
+    """Represents a single generation request with paged KV cache tracking."""
 
     def __init__(
         self,
@@ -45,58 +51,35 @@ class Task:
         self.output_ids: List[int] = []
         self.input_tokens: int = 0
         self.output_tokens: int = 0
-        self.slot: int = -1
+        self.page_table: List[int] = []
+        self.n_pages: int = 0
+        self._prefix_cached_tokens: int = 0
         self.arrival_time = time.time()
         self.finish_time: Optional[float] = None
-
         self.stream_callback = stream_callback
+        self._pages_freed: bool = False
+
+    @property
+    def next_pos(self) -> int:
+        return self.input_tokens + len(self.output_ids)
 
     def is_finished(self, stop_ids: List[int]) -> bool:
-        """Check if task is finished."""
+        if self.output_tokens >= self.max_tokens:
+            return True
         if self.output_ids and self.output_ids[-1] in stop_ids:
             return True
-        return self.output_tokens >= self.max_tokens
+        return False
-
-
-def apply_sampling_strategies(
-    logits: Tensor,
-    temperature: float,
-    top_k: int,
-    top_p: float,
-    filter_value: float = -float("inf"),
-) -> Tensor:
-    """Apply sampling strategies to the logits tensor."""
-    # Clone logits to avoid inplace updates on inference tensor
-    logits = logits.clone()
-
-    if temperature != 1.0:
-        logits = logits / temperature
-
-    if top_k > 0:
-        top_k = min(top_k, logits.size(-1))
-        indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
-        logits[indices_to_remove] = filter_value
-
-    if top_p < 1.0:
-        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
-        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-
-        sorted_indices_to_remove = cumulative_probs > top_p
-        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-        sorted_indices_to_remove[..., 0] = 0
-
-        indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
-        indices_to_remove.scatter_(
-            dim=1, index=sorted_indices, src=sorted_indices_to_remove
-        )
-
-        logits[indices_to_remove] = filter_value
-
-    return logits
 
 
 class InferenceScheduler:
-    """Inference scheduler with continuous batching support."""
+    """Continuous batching scheduler with paged KV cache.
+
+    Runs a background generation loop with four phases per iteration:
+      1. Cleanup finished tasks and release resources.
+      2. Refill active batch from the waiting queue.
+      3. Prefill newly activated tasks.
+      4. Decode the largest same-position group of active tasks.
+    """
 
     def __init__(
         self,
@@ -104,8 +87,10 @@ class InferenceScheduler:
         tokenizer: AutoTokenizer,
         max_batch_size: int = 16,
         max_seq_len: Optional[int] = None,
-        device: str = "cuda",
+        max_prompt_len: int = 512,
-        dtype: torch.dtype = torch.bfloat16,
+        page_size: int = 64,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
     ):
         config = model.config
 
@@ -113,38 +98,26 @@ class InferenceScheduler:
         self.tokenizer = tokenizer
         self.max_batch_size = max_batch_size
         self.max_seq_len = max_seq_len or config.max_len
+        self.max_prompt_len = max_prompt_len
+        self.page_size = page_size
         self.device = device or next(model.parameters()).device
         self.dtype = dtype or next(model.parameters()).dtype
 
-        num_kv_heads = config.n_kv_heads
+        n_kv_heads = config.n_kv_heads
         head_dim = config.dim // config.n_heads
         n_layers = config.n_layers
+        n_pages = (
+            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
+        ) // page_size
 
-        k_cache = torch.empty(
+        self.page_cache = PagedCache(
-            (
+            n_layers,
-                max_batch_size,
+            n_pages,
-                self.max_seq_len,
+            page_size,
-                n_layers,
+            n_kv_heads,
-                num_kv_heads,
+            head_dim,
-                head_dim,
+            self.device,
-            ),
+            self.dtype,
-            device=self.device,
-            dtype=self.dtype,
-        )
-        v_cache = torch.empty(
-            (
-                max_batch_size,
-                self.max_seq_len,
-                n_layers,
-                num_kv_heads,
-                head_dim,
-            ),
-            device=self.device,
-            dtype=self.dtype,
-        )
-        self.kv_cache = (k_cache, v_cache)
-        self.seq_mask = torch.ones(
-            (max_batch_size, self.max_seq_len), device=self.device, dtype=torch.bool
         )
 
         self.waiting_queue: List[Task] = []
@@ -157,6 +130,9 @@ class InferenceScheduler:
         self._total_tasks = 0
         self._total_tokens = 0
 
+    def _n_pages_for(self, n_tokens: int) -> int:
+        return (n_tokens + self.page_size - 1) // self.page_size
+
     def add_task(
         self,
         prompt: str,
@@ -166,9 +142,10 @@ class InferenceScheduler:
         top_k: int = 50,
         stream_callback: Optional[Callable[[str], None]] = None,
     ) -> str:
-        """Add a new task to the waiting queue."""
         task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
         prompt_ids = self.tokenizer.encode(prompt)
+        if len(prompt_ids) > self.max_prompt_len:
+            prompt_ids = prompt_ids[-self.max_prompt_len :]
 
         task = Task(
             task_id=task_id,
@@ -188,13 +165,28 @@ class InferenceScheduler:
         return task_id
 
     def remove_task(self, task_id: str) -> None:
-        """Remove a task from the scheduler."""
         with self._lock:
+            removed_active = [t for t in self.active_tasks if t.task_id == task_id]
             self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
             self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
 
+        for task in removed_active:
+            if not task._pages_freed:
+                self._free_pages(task.page_table)
+                task.page_table.clear()
+                task.n_pages = 0
+                task._pages_freed = True
+
+    def _free_pages(self, indices: List[int]) -> None:
+        for idx in indices:
+            self.page_cache.free(idx)
+
+    def _record_page_hashes(self, task: Task, start_logical_page: int = 0) -> None:
+        full_pages = len(task.prompt_ids) // self.page_size
+        for i in range(start_logical_page, full_pages):
+            self.page_cache.record_page(task.page_table[i], task.prompt_ids, i)
+
     def _remove_finished_tasks(self) -> None:
-        """Remove finished tasks from active batch."""
         finished = []
         for task in self.active_tasks:
             if task.is_finished(self.tokenizer.stop_ids):
@@ -204,198 +196,213 @@ class InferenceScheduler:
                 self._total_tokens += task.output_tokens
 
         for task in finished:
-            slot = task.slot
+            if not task._pages_freed:
-            if slot >= 0 and slot < len(self.active_tasks):
+                self._free_pages(task.page_table)
-                self.seq_mask[slot, :] = False
+                task.page_table.clear()
-            task.slot = -1
+                task.n_pages = 0
+                task._pages_freed = True
 
         self.active_tasks = [
             t for t in self.active_tasks if t.status != TaskStatus.FINISHED
         ]
 
     def _refill_active_batch(self) -> None:
-        """Refill active batch with waiting tasks."""
+        available = self.max_batch_size - len(self.active_tasks)
-        available_slots = self.max_batch_size - len(self.active_tasks)
+        if available <= 0:
-        if available_slots <= 0:
             return
 
+        to_add: List[Task] = []
         with self._lock:
-            to_add = []
+            n = min(available, len(self.waiting_queue))
-            for _ in range(min(available_slots, len(self.waiting_queue))):
+            for _ in range(n):
-                if self.waiting_queue:
+                to_add.append(self.waiting_queue.pop(0))
-                    task = self.waiting_queue.pop(0)
-                    task.status = TaskStatus.RUNNING
-                    to_add.append(task)
 
-            for task in to_add:
+        failed: List[Task] = []
-                for i in range(self.max_batch_size):
+        for task in to_add:
-                    if all(t.slot != i for t in self.active_tasks):
+            prompt_len = len(task.prompt_ids)
-                        task.slot = i
-                        break
-                self.active_tasks.append(task)
 
-    def _execute_prefill(self, tasks: List[Task]) -> None:
+            hit_pages = self.page_cache.lookup_prefix(task.prompt_ids)
-        """Execute Prefill phase."""
+            cached_tokens = len(hit_pages) * self.page_size
-        if not tasks:
+            for p in hit_pages:
-            return
+                self.page_cache.inc_ref(p)
 
-        tasks = sorted(tasks, key=lambda t: t.slot)
+            remaining = prompt_len - cached_tokens
+            n_new = self._n_pages_for(remaining) if remaining > 0 else 0
+            new_pages = self.page_cache.alloc_n(n_new) if n_new > 0 else []
 
-        prompt_lens = [len(task.prompt_ids) for task in tasks]
+            if remaining > 0 and not new_pages:
-        max_len = max(prompt_lens)
+                for p in hit_pages:
+                    self.page_cache.free(p)
+                failed.append(task)
+                continue
 
-        input_ids = torch.zeros(
+            task.page_table = hit_pages + new_pages
-            len(tasks), max_len, dtype=torch.long, device=self.device
+            task.n_pages = len(task.page_table)
-        )
+            task._prefix_cached_tokens = cached_tokens
-        for i, task in enumerate(tasks):
+            task.status = TaskStatus.RUNNING
-            if len(task.prompt_ids) > 0:
+            self.active_tasks.append(task)
-                input_ids[i, : len(task.prompt_ids)] = torch.tensor(
-                    task.prompt_ids, device=self.device
-                )
 
-        if self.tokenizer.pad_id is not None:
+        if failed:
-            input_mask = torch.ne(input_ids, self.tokenizer.pad_id)
+            with self._lock:
-        else:
+                self.waiting_queue[:0] = failed
-            input_mask = torch.ones(
+
-                input_ids.shape, dtype=torch.bool, device=self.device
+    def _execute_prefill(
+        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
+    ) -> None:
+        tasks = sorted(tasks, key=lambda t: t.task_id)
+        batch_sz = len(tasks)
+
+        seq_len = prompt_len - start_pos
+        input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
+        input_mask = torch.ones(batch_sz, seq_len, dtype=torch.bool, device=self.device)
+
+        for i, t in enumerate(tasks):
+            input_ids[i] = torch.tensor(
+                t.prompt_ids[start_pos:prompt_len], device=self.device
             )
 
+        page_tables = self._make_page_table_tensor(tasks)
+
         with torch.inference_mode():
             self.model(
                 input_ids,
                 input_mask=input_mask,
-                start_pos=0,
+                start_pos=start_pos,
-                persistent_key_values=self.kv_cache,
+                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
             )
 
-        for i, task in enumerate(tasks):
+        start_logical_page = start_pos // self.page_size
-            task.input_tokens = prompt_lens[i]
+        for t in tasks:
-            task.output_tokens = 0
+            self._record_page_hashes(t, start_logical_page=start_logical_page)
-
-        for task in tasks:
-            if task.slot >= 0:
-                self.seq_mask[task.slot, : task.input_tokens] = True
 
     def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
-        """Execute Decode phase."""
         if not tasks:
             return
 
-        tasks = sorted(tasks, key=lambda t: t.slot)
+        tasks = sorted(tasks, key=lambda t: t.task_id)
+        batch_sz = len(tasks)
 
-        input_ids = torch.zeros(len(tasks), dtype=torch.long, device=self.device)
+        for t in tasks:
-        for i, task in enumerate(tasks):
+            self._maybe_alloc_page(t, start_pos)
-            if task.output_ids:
-                input_ids[i] = task.output_ids[-1]
-            else:
-                input_ids[i] = task.prompt_ids[-1]
 
-        input_tensor = input_ids.unsqueeze(1)
+        input_ids = torch.tensor(
-        active_mask = torch.ones((len(tasks), 1), dtype=torch.bool, device=self.device)
+            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        active_mask = torch.ones((batch_sz, 1), dtype=torch.bool, device=self.device)
+
+        page_tables = self._make_page_table_tensor(tasks)
+        total_len = start_pos + 1
+
+        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
+        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
+        top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
 
         with torch.inference_mode():
             outputs = self.model(
-                input_tensor,
+                input_ids.unsqueeze(1),
                 input_mask=active_mask,
-                persistent_key_values=self.kv_cache,
+                paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
                 start_pos=start_pos,
             )
             logits = outputs["logits"][:, -1, :]
 
-        next_token_ids = []
+        next_tokens = sample(
-        for i, task in enumerate(tasks):
+            logits,
-            logit = logits[i : i + 1]
+            temperature=temperatures,
-            logit = apply_sampling_strategies(
+            top_k=top_ks,
-                logit,
+            top_p=top_ps,
-                task.temperature,
+        ).tolist()
-                task.top_k,
-                task.top_p,
-            )
-            probs = torch.softmax(logit, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            next_token_ids.append(next_token.item())
 
-        for task, next_token in zip(tasks, next_token_ids):
+        for t, ntok in zip(tasks, next_tokens):
-            task.output_ids.append(next_token)
+            t.output_ids.append(ntok)
-            task.output_tokens += 1
+            t.output_tokens += 1
+            pos = t.input_tokens + t.output_tokens
+            self._maybe_alloc_page(t, pos)
+            if t.stream_callback:
+                t.stream_callback(self.tokenizer.decode([ntok]))
 
-            pos = task.input_tokens + task.output_tokens
+        for t in tasks:
-            if task.slot >= 0 and pos < self.max_seq_len:
+            if t.is_finished(self.tokenizer.stop_ids):
-                self.seq_mask[task.slot, pos] = True
+                if t.stream_callback:
+                    t.stream_callback(STOP)
 
-            if task.stream_callback:
+    def _make_page_table_tensor(self, tasks: List[Task]) -> Tensor:
-                token_str = self.tokenizer.decode([next_token])
+        max_pages = max(t.n_pages for t in tasks)
-                task.stream_callback(token_str)
+        rows = [t.page_table + [-1] * (max_pages - t.n_pages) for t in tasks]
+        return torch.tensor(rows, dtype=torch.long, device=self.device)
 
-        for task in tasks:
+    def _maybe_alloc_page(self, task: Task, pos: int) -> None:
-            if task.output_tokens >= task.max_tokens or (
+        needed = self._n_pages_for(pos + 1)
-                task.output_ids and task.output_ids[-1] in self.tokenizer.stop_ids
+        while task.n_pages < needed:
-            ):
+            p = self.page_cache.alloc()
-                if task.stream_callback:
+            if p < 0:
-                    task.stream_callback("[DONE]")
+                break
+            task.page_table.append(p)
+            task.n_pages += 1
 
     def _run_generation_loop(self) -> None:
-        """Main generation loop."""
+        try:
-        while self._running:
+            while self._running:
-            self._remove_finished_tasks()
+                self._remove_finished_tasks()
-            self._refill_active_batch()
+                self._refill_active_batch()
 
-            if not self.active_tasks:
+                if not self.active_tasks and not self.waiting_queue:
-                self._task_event.wait(timeout=0.01)
+                    self._task_event.clear()
-                self._task_event.clear()
+                    self._task_event.wait(timeout=1.0)
-                continue
+                    continue
 
-            new_tasks = [t for t in self.active_tasks if t.output_tokens == 0]
+                to_prefill = [t for t in self.active_tasks if t.output_tokens == 0]
-            decode_tasks = [t for t in self.active_tasks if t.output_tokens > 0]
+                if to_prefill:
+                    for t in to_prefill:
+                        t.input_tokens = len(t.prompt_ids)
 
-            if decode_tasks:
+                    groups: Dict[Tuple[int, int], List[Task]] = {}
-                start_pos = max(t.input_tokens + t.output_tokens for t in decode_tasks)
+                    for t in to_prefill:
-            else:
+                        key = (len(t.prompt_ids), t._prefix_cached_tokens)
-                start_pos = 0
+                        groups.setdefault(key, []).append(t)
 
-            if new_tasks:
+                    for (prompt_len, start_pos), group in groups.items():
-                self._execute_prefill(new_tasks)
+                        if start_pos < prompt_len:
-                decode_tasks = new_tasks
+                            self._execute_prefill(group, prompt_len, start_pos)
-                start_pos = max(t.input_tokens for t in decode_tasks)
 
-            if decode_tasks:
+                pos_groups: Dict[int, List[Task]] = {}
-                self._execute_decode(decode_tasks, start_pos)
+                for t in self.active_tasks:
+                    pos_groups.setdefault(t.next_pos, []).append(t)
 
-            if not self.active_tasks and not self.waiting_queue:
+                if pos_groups:
-                self._task_event.wait(timeout=0.05)
+                    best_pos = max(pos_groups, key=lambda p: len(pos_groups[p]))
-                self._task_event.clear()
+                    self._execute_decode(pos_groups[best_pos], best_pos)
+        except Exception as e:
+            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
+            for task in self.active_tasks:
+                if task.stream_callback:
+                    task.stream_callback(STOP)
+            for task in self.waiting_queue:
+                if task.stream_callback:
+                    task.stream_callback(STOP)
+            raise
 
     def start(self) -> None:
-        """Start the generation loop."""
         if not self._running:
             self._running = True
-            self._loop_thread = threading.Thread(target=self._run_generation_loop)
+            t = threading.Thread(target=self._run_generation_loop, daemon=True)
-            self._loop_thread.daemon = True
+            t.start()
-            self._loop_thread.start()
+            self._loop_thread = t
 
     def stop(self) -> None:
-        """Stop the generation loop."""
         self._running = False
+        self._task_event.set()
         if hasattr(self, "_loop_thread"):
-            self._loop_thread.join(timeout=1.0)
+            self._loop_thread.join(timeout=2.0)
-
-        # Clear KV cache to free GPU memory
-        if self.kv_cache is not None:
-            k_cache, v_cache = self.kv_cache
-            if k_cache is not None:
-                k_cache.detach()
-            if v_cache is not None:
-                v_cache.detach()
-
-        # Clear seq mask
-        self.seq_mask.detach()
-
-        # Clear task lists
         self.waiting_queue.clear()
         self.active_tasks.clear()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
     def get_stats(self) -> Dict[str, Any]:
-        """Get scheduler statistics."""
         return {
             "total_tasks": self._total_tasks,
             "total_tokens": self._total_tokens,

astrai/inference/server.py

@@ -1,15 +1,14 @@
 """
-Inference Server with Continuous Batching Support
+OpenAI / Anthropic-compatible chat completion server backed by continuous-batching inference.
-
-FastAPI server for inference with continuous batching.
-Provides OpenAI-compatible chat completion endpoints.
 """
 
 import json
 import logging
+import time
+import uuid
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 import uvicorn
@@ -23,18 +22,63 @@ from astrai.tokenize import AutoTokenizer
 
 logger = logging.getLogger(__name__)
 
-# Global model parameter and engine (loaded once)
-_engine: Optional[InferenceEngine] = None
-_model_param: Optional[Any] = None
 _project_root = Path(__file__).parent.parent.parent
 
-# Server configuration (set before running server)
+
-_server_config: Dict[str, Any] = {
+class ServerState:
-    "device": "cuda",
+    def __init__(self):
-    "dtype": torch.bfloat16,
+        self.engine: Optional[InferenceEngine] = None
-    "param_path": None,
+        self.config: Dict[str, Any] = {
-    "max_batch_size": 16,
+            "device": "cuda",
-}
+            "dtype": torch.bfloat16,
+            "param_path": None,
+            "max_batch_size": 16,
+        }
+
+
+_state = ServerState()
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+
+class ChatCompletionRequest(BaseModel):
+    """OpenAI Chat Completion API request body."""
+
+    model: str = "astrai"
+    messages: List[ChatMessage]
+    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
+    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
+    top_k: Optional[int] = Field(default=50, ge=1)
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
+    max_tokens: Optional[int] = Field(default=2048, ge=1)
+    n: Optional[int] = Field(default=1, ge=1)
+    presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
+    frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
+    logit_bias: Optional[Dict[int, float]] = None
+    user: Optional[str] = None
+
+
+class AnthropicMessage(BaseModel):
+    role: str
+    content: Union[str, List[Dict[str, Any]]]
+
+
+class MessagesRequest(BaseModel):
+    """Anthropic Messages API request body."""
+
+    model: str = "astrai"
+    max_tokens: int = Field(default=1024, ge=1)
+    messages: List[AnthropicMessage]
+    system: Optional[str] = None
+    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
+    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
+    top_k: Optional[int] = Field(default=50, ge=1)
+    stream: Optional[bool] = False
+    stop_sequences: Optional[List[str]] = None
 
 
 def configure_server(
@@ -43,39 +87,29 @@ def configure_server(
     param_path: Optional[Path] = None,
     max_batch_size: int = 16,
 ):
-    """Configure server settings before starting.
+    _state.config.update(
-
+        device=device,
-    Args:
+        dtype=dtype,
-        device: Device to load model on (e.g., "cuda", "cpu", "cuda:0")
+        param_path=param_path,
-        dtype: Data type for model weights (e.g., torch.bfloat16, torch.float16)
+        max_batch_size=max_batch_size,
-        param_path: Path to model parameters directory
+    )
-        max_batch_size: Maximum batch size for continuous batching
-    """
-    _server_config["device"] = device
-    _server_config["dtype"] = dtype
-    _server_config["param_path"] = param_path
-    _server_config["max_batch_size"] = max_batch_size
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Lifespan context manager for startup and shutdown events."""
-    global _model_param, _engine
-    # Startup: Load model with configured settings
     try:
         load_model(
-            param_path=_server_config["param_path"],
+            param_path=_state.config["param_path"],
-            device=_server_config["device"],
+            device=_state.config["device"],
-            dtype=_server_config["dtype"],
+            dtype=_state.config["dtype"],
-            max_batch_size=_server_config["max_batch_size"],
+            max_batch_size=_state.config["max_batch_size"],
         )
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
         raise
     yield
-    # Shutdown: Cleanup engine
+    if _state.engine:
-    if _engine:
+        _state.engine.shutdown()
-        _engine.shutdown()
         logger.info("Inference engine shutdown complete")
 
 
@@ -88,270 +122,345 @@ def load_model(
     dtype: torch.dtype = torch.bfloat16,
     max_batch_size: int = 16,
 ):
-    """Load model parameters and initialize inference engine."""
-    global _model_param, _engine
     if param_path is None:
         param_path = _project_root / "params"
     if not param_path.exists():
         raise FileNotFoundError(f"Parameter directory not found: {param_path}")
 
-    # Load tokenizer separately
     tokenizer = AutoTokenizer.from_pretrained(param_path)
-    _model_param = AutoModel.from_pretrained(param_path)
+    model = AutoModel.from_pretrained(param_path)
-    _model_param.to(device=device, dtype=dtype)
+    model.to(device=device, dtype=dtype)
     logger.info(f"Model loaded on {device} with dtype {dtype}")
 
-    # Initialize inference engine with separate model and tokenizer
+    _state.engine = InferenceEngine(
-    _engine = InferenceEngine(
+        model=model,
-        model=_model_param,
         tokenizer=tokenizer,
         max_batch_size=max_batch_size,
     )
     logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
 
 
-# Pydantic models for API request/response
+def _get_engine() -> InferenceEngine:
-class ChatMessage(BaseModel):
+    if _state.engine is None:
-    role: str  # "user", "assistant", "system"
+        raise HTTPException(status_code=503, detail="Engine not initialized")
-    content: str
+    return _state.engine
 
 
-class ChatCompletionRequest(BaseModel):
+def _make_chunk(
-    messages: List[ChatMessage]
+    delta: Dict[str, str],
-    temperature: float = Field(0.8, ge=0.0, le=2.0)
+    finish_reason: Optional[str] = None,
-    top_p: float = Field(0.95, ge=0.0, le=1.0)
+    *,
-    top_k: int = Field(50, ge=0)
+    resp_id: str,
-    max_tokens: int = Field(2048, ge=1)
+    created: int,
-    stream: bool = False
+    model: str,
-    system_prompt: Optional[str] = None
+    index: int = 0,
-
-
-class CompletionResponse(BaseModel):
-    id: str = "chatcmpl-default"
-    object: str = "chat.completion"
-    created: int = 0
-    model: str = "astrai"
-    choices: List[Dict[str, Any]]
-
-
-class StreamCompletionResponse(BaseModel):
-    id: str = "chatcmpl-default"
-    object: str = "chat.completion.chunk"
-    created: int = 0
-    model: str = "astrai"
-    choices: List[Dict[str, Any]]
-
-
-def convert_messages_to_history(
-    messages: List[ChatMessage],
-) -> tuple[Optional[str], Optional[List[Tuple[str, str]]]]:
-    """Convert OpenAI-style messages to system_prompt and history."""
-    system_prompt = None
-    history: List[Tuple[str, str]] = []
-    user_buffer = []
-    assistant_buffer = []
-    for msg in messages:
-        if msg.role == "system":
-            system_prompt = msg.content
-        elif msg.role == "user":
-            if assistant_buffer:
-                # Flush previous pair
-                history.append(("".join(user_buffer), "".join(assistant_buffer)))
-                user_buffer = []
-                assistant_buffer = []
-            user_buffer.append(msg.content)
-        elif msg.role == "assistant":
-            assistant_buffer.append(msg.content)
-        else:
-            logger.warning(f"Unknown role {msg.role}")
-    return system_prompt, history if history else None
-
-
-def convert_messages_to_prompt(
-    messages: List[ChatMessage], engine: InferenceEngine = None
 ) -> str:
-    """Convert messages to prompt string.
+    """Build a single SSE ``data:`` chunk matching OpenAI streaming format."""
-
+    data = {
-    Args:
+        "id": resp_id,
-        messages: List of ChatMessage objects
+        "object": "chat.completion.chunk",
-        engine: InferenceEngine instance for accessing tokenizer
+        "created": created,
-
+        "model": model,
-    Returns:
+        "choices": [
-        str: Formatted prompt string
+            {
-    """
+                "index": index,
-    # Convert to dict format for chat template
+                "delta": delta,
-    msg_dicts = [{"role": m.role, "content": m.content} for m in messages]
+                "finish_reason": finish_reason,
-
+            }
-    # Extract system prompt if present
+        ],
-    system_prompt = None
+    }
-    filtered_messages = []
+    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
-    for msg in msg_dicts:
-        if msg["role"] == "system":
-            system_prompt = msg["content"]
-        else:
-            filtered_messages.append(msg)
-
-    # Use engine's tokenizer chat template if available
-    if engine is not None and engine.tokenizer is not None:
-        return engine.tokenizer.apply_chat_template(
-            filtered_messages, system_prompt=system_prompt, tokenize=False
-        )
-
-    # Fallback: simple concatenation (deprecated)
-    prompt_parts = []
-    for msg in filtered_messages:
-        prompt_parts.append(
-            f"<｜im▁start｜>{msg['role']}\n{msg['content']}<｜im▁end｜>"
-        )
-    return "\n".join(prompt_parts) + "\n<｜im▁start｜>assistant\n"
 
 
 @app.get("/health")
 async def health():
     return {
         "status": "ok",
-        "model_loaded": _model_param is not None,
+        "model_loaded": _state.engine is not None,
-        "engine_ready": _engine is not None,
     }
 
 
 @app.get("/stats")
 async def get_stats():
-    """Get inference engine statistics."""
+    return _get_engine().get_stats()
-    if _engine is None:
-        raise HTTPException(status_code=503, detail="Engine not initialized")
-    return _engine.get_stats()
 
 
-@app.post("/v1/chat/completions", response_model=CompletionResponse)
+@app.post("/v1/chat/completions")
 async def chat_completion(request: ChatCompletionRequest):
-    """OpenAI-compatible chat completion endpoint.
+    """OpenAI-compatible chat completion endpoint (streaming + non-streaming)."""
+    engine = _get_engine()
+    resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+    created = int(time.time())
+    model = request.model
 
-    Supports both streaming and non-streaming modes with continuous batching.
+    prompt = engine.tokenizer.apply_chat_template(
-    """
+        [{"role": m.role, "content": m.content} for m in request.messages],
-    if _engine is None:
+        tokenize=False,
-        raise HTTPException(status_code=503, detail="Engine not initialized")
+    )
-
+    prompt_tokens = len(engine.tokenizer.encode(prompt))
-    # Convert messages to prompt using engine's tokenizer
-    prompt = convert_messages_to_prompt(request.messages, engine=_engine)
 
     if request.stream:
-        # Streaming response (use synchronous generator)
+        agen = engine.generate_async(
-        generator = _engine.generate(
             prompt=prompt,
-            stream=True,
             max_tokens=request.max_tokens,
             temperature=request.temperature,
             top_p=request.top_p,
             top_k=request.top_k,
         )
 
-        def generate_stream():
+        async def event_stream():
-            for token in generator:
+            yield _make_chunk(
-                if token == "[DONE]":
+                {"role": "assistant"},
-                    break
+                finish_reason=None,
-                yield f"data: {json.dumps({'choices': [{'delta': {'content': token}}]})}\n\n"
+                resp_id=resp_id,
+                created=created,
+                model=model,
+            )
+
+            completion_tokens = 0
+            async for token in agen:
+                yield _make_chunk(
+                    {"content": token},
+                    finish_reason=None,
+                    resp_id=resp_id,
+                    created=created,
+                    model=model,
+                )
+                completion_tokens += 1
+
+            yield _make_chunk(
+                {},
+                finish_reason="stop",
+                resp_id=resp_id,
+                created=created,
+                model=model,
+            )
+
+            usage = {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            }
+            yield f"data: {json.dumps(usage, ensure_ascii=False)}\n\n"
             yield "data: [DONE]\n\n"
 
         return StreamingResponse(
-            generate_stream(),
+            event_stream(),
             media_type="text/event-stream",
             headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
         )
-    else:
+
-        # Non-streaming response
+    completion_tokens = 0
-        result = _engine.generate(
+    chunks: List[str] = []
+    agen = engine.generate_async(
+        prompt=prompt,
+        max_tokens=request.max_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        top_k=request.top_k,
+    )
+    async for token in agen:
+        chunks.append(token)
+        completion_tokens += 1
+    content = "".join(chunks)
+
+    return {
+        "id": resp_id,
+        "object": "chat.completion",
+        "created": created,
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": content},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+    }
+
+
+def _make_anthropic_sse(event: str, data: Dict[str, Any]) -> str:
+    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
+
+
+def _check_stop_sequence(text: str, stop_sequences: List[str]) -> Optional[str]:
+    for seq in stop_sequences:
+        if seq and seq in text:
+            return seq
+    return None
+
+
+def _extract_text_content(content: Union[str, List[Dict[str, Any]]]) -> str:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                return block.get("text", "")
+    return ""
+
+
+def _build_anthropic_messages(
+    messages: List[AnthropicMessage], system: Optional[str]
+) -> List[Dict[str, str]]:
+    result: List[Dict[str, str]] = []
+    if system:
+        result.append({"role": "system", "content": system})
+    for m in messages:
+        content = _extract_text_content(m.content)
+        if content:
+            result.append({"role": m.role, "content": content})
+    return result
+
+
+@app.post("/v1/messages")
+async def create_message(request: MessagesRequest):
+    """Anthropic-compatible Messages API endpoint (streaming + non-streaming)."""
+    engine = _get_engine()
+    resp_id = f"msg_{uuid.uuid4().hex[:24]}"
+    model = request.model
+
+    chat_messages = _build_anthropic_messages(request.messages, request.system)
+    prompt = engine.tokenizer.apply_chat_template(chat_messages, tokenize=False)
+    prompt_tokens = len(engine.tokenizer.encode(prompt))
+
+    stop_sequences = request.stop_sequences or []
+
+    if request.stream:
+        agen = engine.generate_async(
             prompt=prompt,
-            stream=False,
             max_tokens=request.max_tokens,
             temperature=request.temperature,
             top_p=request.top_p,
             top_k=request.top_k,
         )
 
-        # Build OpenAI-style response
+        async def event_stream():
-        import time
+            yield _make_anthropic_sse(
-
+                "message_start",
-        resp = CompletionResponse(
-            id=f"chatcmpl-{int(time.time())}",
-            created=int(time.time()),
-            choices=[
                 {
+                    "type": "message_start",
+                    "message": {
+                        "id": resp_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "model": model,
+                        "content": [],
+                        "usage": {"input_tokens": prompt_tokens},
+                    },
+                },
+            )
+
+            yield _make_anthropic_sse(
+                "content_block_start",
+                {
+                    "type": "content_block_start",
                     "index": 0,
-                    "message": {"role": "assistant", "content": result},
+                    "content_block": {"type": "text", "text": ""},
-                    "finish_reason": "stop",
+                },
-                }
+            )
-            ],
-        )
-        return resp
 
+            completion_tokens = 0
+            accumulated = ""
+            stopped_seq: Optional[str] = None
+            async for token in agen:
+                accumulated += token
+                completion_tokens += 1
 
-@app.post("/generate")
+                matched = _check_stop_sequence(accumulated, stop_sequences)
-async def generate(
+                if matched:
-    query: str,
+                    text = accumulated[: accumulated.rfind(matched)]
-    history: Optional[List[List[str]]] = None,
+                    stopped_seq = matched
-    temperature: float = 0.8,
+                    if text:
-    top_p: float = 0.95,
+                        yield _make_anthropic_sse(
-    top_k: int = 50,
+                            "content_block_delta",
-    max_len: int = 2048,
+                            {
-    stream: bool = False,
+                                "type": "content_block_delta",
-):
+                                "index": 0,
-    """Simple generation endpoint.
+                                "delta": {"type": "text_delta", "text": text},
+                            },
+                        )
+                    break
 
-    Args:
+                yield _make_anthropic_sse(
-        query: Input query string
+                    "content_block_delta",
-        history: Conversation history as list of [user, assistant] pairs
+                    {
-        temperature: Sampling temperature
+                        "type": "content_block_delta",
-        top_p: Top-p sampling parameter
+                        "index": 0,
-        top_k: Top-k sampling parameter
+                        "delta": {"type": "text_delta", "text": token},
-        max_len: Maximum tokens to generate
+                    },
-        stream: Enable streaming output
+                )
 
-    Returns:
+            yield _make_anthropic_sse(
-        dict: Generation result with response field
+                "content_block_stop",
-    """
+                {"type": "content_block_stop", "index": 0},
-    if _engine is None:
+            )
-        raise HTTPException(status_code=503, detail="Engine not initialized")
 
-    # Build messages for chat template
+            stop_reason = "stop_sequence" if stopped_seq else "end_turn"
-    messages = []
+            yield _make_anthropic_sse(
-    if history:
+                "message_delta",
-        # Convert history format: List[List[str]] -> List[Dict]
+                {
-        for h in history:
+                    "type": "message_delta",
-            if len(h) >= 2:
+                    "delta": {"stop_reason": stop_reason, "stop_sequence": stopped_seq},
-                messages.append({"role": "user", "content": h[0]})
+                    "usage": {"output_tokens": completion_tokens},
-                messages.append({"role": "assistant", "content": h[1]})
+                },
-    messages.append({"role": "user", "content": query})
+            )
 
-    # Use tokenizer's chat template
+            yield _make_anthropic_sse(
-    prompt = _engine.tokenizer.apply_chat_template(messages, tokenize=False)
+                "message_stop",
+                {"type": "message_stop"},
+            )
 
-    if stream:
+        return StreamingResponse(
-        # Synchronous streaming
+            event_stream(),
-        result = _engine.generate(
+            media_type="text/event-stream",
-            prompt=prompt,
+            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
-            stream=True,
-            max_tokens=max_len,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
         )
 
-        def stream_generator():
+    completion_tokens = 0
-            for token in result:
+    chunks: List[str] = []
-                yield token + "\n"
+    agen = engine.generate_async(
+        prompt=prompt,
+        max_tokens=request.max_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        top_k=request.top_k,
+    )
+    stopped_seq: Optional[str] = None
+    accumulated = ""
+    async for token in agen:
+        chunks.append(token)
+        completion_tokens += 1
+        accumulated += token
+        matched = _check_stop_sequence(accumulated, stop_sequences)
+        if matched:
+            stopped_seq = matched
+            break
 
-        return StreamingResponse(stream_generator(), media_type="text/plain")
+    content = "".join(chunks)
-    else:
+    if stopped_seq:
-        result = _engine.generate(
+        idx = content.rfind(stopped_seq)
-            prompt=prompt,
+        if idx != -1:
-            stream=False,
+            content = content[:idx]
-            max_tokens=max_len,
+
-            temperature=temperature,
+    return {
-            top_p=top_p,
+        "id": resp_id,
-            top_k=top_k,
+        "type": "message",
-        )
+        "role": "assistant",
-        return {"response": result}
+        "model": model,
+        "content": [{"type": "text", "text": content}],
+        "stop_reason": "stop_sequence" if stopped_seq else "end_turn",
+        "stop_sequence": stopped_seq,
+        "usage": {
+            "input_tokens": prompt_tokens,
+            "output_tokens": completion_tokens,
+        },
+    }
 
 
 def run_server(
@@ -363,17 +472,6 @@ def run_server(
     param_path: Optional[Path] = None,
     max_batch_size: int = 16,
 ):
-    """Run the FastAPI server with uvicorn.
-
-    Args:
-        host: Server host address
-        port: Server port number
-        reload: Enable auto-reload for development
-        device: Device to load model on (e.g., "cuda", "cpu", "cuda:0")
-        dtype: Data type for model weights (e.g., torch.bfloat16, torch.float16)
-        param_path: Path to model parameters directory
-        max_batch_size: Maximum batch size for continuous batching
-    """
     configure_server(
         device=device,
         dtype=dtype,

astrai/model/automodel.py

@@ -4,12 +4,13 @@ AutoModel base class for model loading and saving.
 
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Dict, Self, Type, Union
+from typing import Self, Type, Union
 
 import safetensors.torch as st
 import torch.nn as nn
 
 from astrai.config import ModelConfig
+from astrai.factory import Registry
 
 
 @contextmanager
@@ -44,8 +45,7 @@ class AutoModel(nn.Module):
     Provides model loading/saving and generation capabilities.
     """
 
-    # Model registry - stored as class attribute
+    _registry = Registry()
-    _registry: Dict[str, Type["AutoModel"]] = {}
 
     def __init__(self, config: ModelConfig):
         super().__init__()
@@ -63,7 +63,7 @@ class AutoModel(nn.Module):
         """
 
         def decorator(sub_cls: Type["AutoModel"]) -> Type["AutoModel"]:
-            cls._registry[model_type.lower()] = sub_cls
+            cls._registry.register(model_type.lower(), sub_cls)
             return sub_cls
 
         return decorator
@@ -72,18 +72,19 @@ class AutoModel(nn.Module):
     def get_model_class(cls, model_type: str) -> Type["AutoModel"]:
         """Get model class by model_type string."""
         model_type = model_type.lower()
-        if model_type not in cls._registry:
+        if not cls._registry.contains(model_type):
-            available = list(cls._registry.keys())
+            available = cls._registry.list_names()
             raise ValueError(
                 f"Unknown model_type: {model_type}. Available: {available}"
             )
-        return cls._registry[model_type]
+        return cls._registry.get(model_type)
 
     @classmethod
     def from_pretrained(
         cls,
         path: Union[str, Path],
         disable_random_init: bool = True,
+        strict: bool = True,
     ) -> nn.Module:
 
         model_path = Path(path)
@@ -96,14 +97,8 @@ class AutoModel(nn.Module):
         else:
             raise FileNotFoundError(f"Config file not found: {config_path}")
 
-        # If called from base class, use model_type to determine actual model class
+        model_type = config.model_type or "transformer"
-        if cls is AutoModel:
+        actual_cls = cls.get_model_class(model_type)
-            model_type = config.model_type or "transformer"
-            actual_cls = cls.get_model_class(model_type)
-        else:
-            raise ValueError(
-                f"Cannot call from_pretrained() on subclass {cls.__name__}"
-            )
 
         with _disable_random_init(enable=disable_random_init):
             model = actual_cls(config)
@@ -112,7 +107,7 @@ class AutoModel(nn.Module):
         weights_path = model_path / "model.safetensors"
         if weights_path.exists():
             state_dict = st.load_file(str(weights_path))
-            model.load_state_dict(state_dict, strict=False)
+            model.load_state_dict(state_dict, strict=strict)
 
         return model
 

astrai/model/module.py

@@ -5,17 +5,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 
+from astrai.inference.cache import CacheView
+
 
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
-    """
+    """Repeat KV heads n_rep times for GQA."""
-    Repeat k times along the dimension for attention heads.
-    Args:
-        x (Tensor): The input tensor.
-        n_rep (int): The number of repetitions.
-    Returns:
-        Tensor: The repeated tensor.
-    """
-
     bs, slen, n_heads, head_dim = x.shape
     if n_rep == 1:
         return x
@@ -30,49 +24,27 @@ def get_rotary_emb(
     dim: int,
     max_len: int,
     base: float = 10000,
+    device: Optional[torch.device] = None,
 ) -> Tuple[Tensor, Tensor]:
-    """
+    """Precompute cos/sin for RoPE."""
-    Get the rotary embedding for the given dimension and maximum length.
+    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
-    Args:
+    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
-        dim (int): The dimension of the input.
-        max_len (int): The maximum length of the input.
-        base (float, optional): The base for the frequency. Defaults to 10000.
-    Returns:
-        Tensor: The rotary embedding tensor.
-    """
-
-    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64) / dim)
-    t = torch.arange(0, max_len, dtype=torch.float64)
     freqs = torch.outer(t, theta)
-
     return torch.cos(freqs).float(), torch.sin(freqs).float()
 
 
 def apply_rotary_emb(x: torch.Tensor, rotary_emb: Tuple[Tensor, Tensor]) -> Tensor:
-    """
+    """Apply rotary embedding via cos/sin (shape-preserving)."""
-    Apply rotary embedding to the input tensor using cos/sin form.
-    Args:
-        x (Tensor): The input tensor (shape [..., seq_len, dim]).
-        rotary_emb (Tuple[Tensor, Tensor]): The rotary embedding (shape [seq_len, dim//2]).
-    Returns:
-        Tensor: The output tensor (rotated, same shape as input).
-    """
-
     dtype = x.dtype
     cos, sin = rotary_emb
-
+    cos = cos.unsqueeze(0).unsqueeze(2)
-    cos = cos.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, dim//2]
+    sin = sin.unsqueeze(0).unsqueeze(2)
-    sin = sin.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, dim//2]
+    x_real = x[..., 0::2]
-
+    x_imag = x[..., 1::2]
-    x_real = x[..., 0::2]  # [batch, seq_len, dim//2]
-    x_imag = x[..., 1::2]  # [batch, seq_len, dim//2]
-
     x_real_rot = x_real * cos - x_imag * sin
     x_imag_rot = x_real * sin + x_imag * cos
-
+    x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1)
-    x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1)  # [batch, seq_len, dim//2, 2]
+    x_out = x_out.view(*x_out.shape[:-2], -1)
-    x_out = x_out.view(*x_out.shape[:-2], -1)  # [batch, seq_len, dim]
-
     return x_out.to(dtype)
 
 
@@ -83,23 +55,20 @@ class RotaryEmbedding(nn.Module):
         self.max_len = max_len
         self.base = base
         self.max_len_cached = None
-        self._set_rotary_buffer(self.max_len)
+        self._set_rotary_buffer(self.max_len, None)
 
-    def _set_rotary_buffer(self, max_len: int):
+    def _set_rotary_buffer(self, max_len: int, device: Optional[torch.device] = None):
-        cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base)
+        cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base, device)
         self.register_buffer("cos_cached", cos_cached, persistent=False)
         self.register_buffer("sin_cached", sin_cached, persistent=False)
         self.max_len_cached = max_len
 
     def forward(self, x: Tensor, start_pos: int = 0) -> Tuple[Tensor, Tensor]:
         seq_len = x.size(1)
-
         if self.max_len_cached < seq_len + start_pos:
-            self._set_rotary_buffer(seq_len + start_pos)
+            self._set_rotary_buffer(self.max_len_cached * 2, x.device)
-
         cos = self.cos_cached[start_pos : start_pos + seq_len]
         sin = self.sin_cached[start_pos : start_pos + seq_len]
-
         return (cos, sin)
 
 
@@ -121,8 +90,7 @@ class RMSNorm(nn.Module):
         self.norm_eps = norm_eps
 
     def forward(self, x: Tensor) -> Tensor:
-        rms = F.rms_norm(x.float(), self.normalized_shape, self.weight, self.norm_eps)
+        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)
-        return rms.to(x.dtype)
 
 
 class MLP(nn.Module):
@@ -184,13 +152,13 @@ class GQA(nn.Module):
         x: Tensor,
         rotary_emb: Tuple[Tensor, Tensor],
         mask: Tensor = None,
-        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        paged_cache: Optional[CacheView] = None,
         start_pos: int = 0,
     ) -> Tensor:
         bsz, seq_len, _ = x.size()
         is_causal = mask is None
 
-        # x(bsz, seq_len, n_heads * head_dim) -> (bsz, seq_len, n_heads, head_dim)
+        # (bsz, seq_len, dim) -> (bsz, seq_len, n_heads, head_dim)
         q = self._split_heads(self.q_proj(x), self.n_heads)
         k = self._split_heads(self.k_proj(x), self.n_kv_heads)
         v = self._split_heads(self.v_proj(x), self.n_kv_heads)
@@ -199,22 +167,14 @@ class GQA(nn.Module):
         if self.use_qk_norm:
             q, k = self.q_norm(q), self.k_norm(k)
 
-        if kv_cache is not None:
+        if paged_cache is not None:
-            k_cache, v_cache = kv_cache
+            paged_cache.write(self.layer_id, start_pos, k, v)
-
+            k, v = paged_cache.gather(self.layer_id)
-            # copy to cache
-            k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
-            v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
-
-            # get cache
-            k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
-            v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
 
         k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
 
         # (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
         q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
-        # (bsz, n_heads, seq_len, head_dim) - > (bsz, seq_len, n_heads*head_dim)
         sdqa_out = (
             F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
             .permute(0, 2, 1, 3)
@@ -226,7 +186,6 @@ class GQA(nn.Module):
             sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
 
         out = self.o_proj(sdqa_out)
-
         return out
 
 
@@ -257,9 +216,9 @@ class MLA(nn.Module):
 
         self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
         self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
-        self.kv_norm = RMSNorm(kv_lora_rank, eps=norm_eps)
+        self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
 
-        # KV (k_nope, k_rope, v)
+        # fused KV: (k_nope, k_rope, v)
         self.kv_b_proj = Linear(
             kv_lora_rank,
             n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
@@ -275,7 +234,7 @@ class MLA(nn.Module):
         x: Tensor,
         rotary_emb: Tuple[Tensor, Tensor],
         mask: Tensor = None,
-        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        paged_cache: Optional[CacheView] = None,
         start_pos: int = 0,
     ) -> Tensor:
         bsz, seq_len, _ = x.size()
@@ -304,12 +263,9 @@ class MLA(nn.Module):
         q = torch.cat([q_nope, q_rope], dim=-1)
         k = torch.cat([k_nope, k_rope], dim=-1)
 
-        if kv_cache is not None:
+        if paged_cache is not None:
-            k_cache, v_cache = kv_cache
+            paged_cache.write(self.layer_id, start_pos, k, v)
-            k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
+            k, v = paged_cache.gather(self.layer_id)
-            v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
-            k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
-            v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
 
         q = q.permute(0, 2, 1, 3)
         k = k.permute(0, 2, 1, 3)
@@ -322,7 +278,6 @@ class MLA(nn.Module):
             attn_out = attn_out * F.sigmoid(self.gate(x))
 
         out = self.o_proj(attn_out)
-
         return out
 
 
@@ -357,18 +312,19 @@ class DecoderBlock(nn.Module):
         x: Tensor,
         rotary_emb: Tuple[Tensor, Tensor],
         attention_mask: Optional[Tensor] = None,
-        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        paged_cache: Optional[CacheView] = None,
         start_pos: int = 0,
     ) -> Tensor:
-        # attention
         attn_output = self.attention(
-            self.input_norm(x), rotary_emb, attention_mask, kv_cache, start_pos
+            self.input_norm(x),
+            rotary_emb,
+            attention_mask,
+            paged_cache,
+            start_pos,
         )
         x = attn_output + x
 
-        # feed forward
         x = self.mlp(self.post_attention_norm(x)) + x
-
         return x
 
 

astrai/model/transformer.py

@@ -1,10 +1,11 @@
-from typing import Any, Mapping, Optional, Tuple
+from typing import Any, Mapping, Optional
 
 import torch
 import torch.nn as nn
 from torch import Tensor
 
 from astrai.config.model_config import ModelConfig
+from astrai.inference.cache import CacheView
 from astrai.model.automodel import AutoModel
 from astrai.model.module import (
     DecoderBlock,
@@ -21,39 +22,25 @@ def process_attention_mask(
     start_pos: int = 0,
     is_causal: bool = False,
 ) -> Tensor:
-    """
+    """Build 4D attention mask from 2D seq_mask, with optional causal masking."""
-    Create attention mask for GQA
-    Args:
-        seq_mask (Tensor): A tensor indicating whether each position is valid or not.
-        input_tensor (Tensor): The input tensor.
-        start_pos (int): The starting position of the sequence.
-        is_causal (bool): Whether the attention is causal or not.
-    Returns:
-        Tensor: The attention mask tensor.
-    """
     device = input_tensor.device
     dtype = input_tensor.dtype
     seq_len = input_tensor.size(1)
 
     if seq_mask is None:
         if start_pos != 0:
-            # for single prompt chat
             seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
         else:
             return None
 
     if seq_mask.dim() > 2:
-        # shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos)
-        # if ndim > 2, it's 4D tensor
         return seq_mask
 
     batch_size = seq_mask.size(0)
     seq_mask = seq_mask[:, : start_pos + seq_len].to(device=device, dtype=torch.bool)
-    # (bsz, start_pos + seq_len)
     expanded_mask = seq_mask.unsqueeze(1).expand(
         batch_size, seq_len, start_pos + seq_len
     )
-    # (bsz, seq_len, start_pos + seq_len)
 
     if is_causal:
         expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)
@@ -62,16 +49,13 @@ def process_attention_mask(
     attention_mask = attention_mask.masked_fill_(
         ~expanded_mask, -torch.finfo(dtype).max / 2
     ).unsqueeze(1)
-    # (bsz, 1, seq_len, seq_len + start_pos)
 
     return attention_mask
 
 
 @AutoModel.register("transformer")
 class Transformer(AutoModel):
-    """
+    """Transformer language model with paged KV cache."""
-    Transformer language model.
-    """
 
     def __init__(self, config: ModelConfig):
         super().__init__(config)
@@ -114,18 +98,15 @@ class Transformer(AutoModel):
         lm_head_key = "lm_head.weight"
         embed_key = "embed_tokens.weight"
 
-        # Make a copy to avoid modifying the original state_dict
         state_dict = dict(state_dict)
 
         if self.config.tie_weight:
-            # same tensor
+            # same tensor for embed and lm_head
             if embed_key in state_dict:
                 state_dict[lm_head_key] = state_dict[embed_key]
         else:
-            # If lm_head.weight exists in checkpoint, use it directly
-            # If not, copy from embed_tokens.weight
             if lm_head_key not in state_dict and embed_key in state_dict:
-                # use clone to avoid sharing the same tensor
+                # clone to avoid sharing gradients
                 state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
 
         return super().load_state_dict(state_dict, strict, assign)
@@ -146,7 +127,7 @@ class Transformer(AutoModel):
         self,
         input_ids: Tensor,
         input_mask: Optional[Tensor] = None,
-        persistent_key_values: Optional[Tuple[Tensor, Tensor]] = None,
+        paged_cache: Optional[CacheView] = None,
         start_pos: int = 0,
     ) -> Tensor:
         assert input_ids.ndim == 2
@@ -157,7 +138,7 @@ class Transformer(AutoModel):
         attn_mask = process_attention_mask(input_mask, x, start_pos, is_causal=True)
 
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, persistent_key_values, start_pos)
+            x = layer(x, rotary_emb, attn_mask, paged_cache, start_pos)
 
         hidden_states = self.norm(x)
         logits = self.lm_head(hidden_states)

astrai/parallel/setup.py

@@ -1,7 +1,7 @@
 import os
 from contextlib import contextmanager
 from functools import wraps
-from typing import Callable, List, Optional
+from typing import Callable
 
 import torch
 import torch.distributed as dist
@@ -34,7 +34,6 @@ def setup_parallel(
     master_addr: str = "localhost",
     master_port: str = "29500",
     device_type: str = "cuda",
-    device_ids: Optional[List[int]] = None,
 ):
 
     if dist.is_available() and dist.is_initialized():
@@ -45,15 +44,10 @@ def setup_parallel(
         yield None
         return
 
-    if device_ids is None:
+    device_id = torch.device(device_type, rank)
-        device_ids = [i for i in range(world_size)]
-
-    rank = device_ids[rank % len(device_ids)]
-    device_id = torch.device(device_type, device_ids[rank])
 
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port
-
     os.environ["LOCAL_RANK"] = str(rank)
     os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["LOCAL_DEVICE"] = str(device_id)
@@ -103,7 +97,6 @@ def wrapper_spawn_func(
     master_addr: str,
     master_port: str,
     device_type: str,
-    device_ids: List[int],
     func: Callable,
     kwargs: dict,
 ):
@@ -115,7 +108,6 @@ def wrapper_spawn_func(
             master_addr=master_addr,
             master_port=master_port,
             device_type=device_type,
-            device_ids=device_ids,
         ):
             func(**kwargs)
 
@@ -131,7 +123,6 @@ def spawn_parallel_fn(
     master_addr: str = "localhost",
     master_port: str = "29500",
     device_type: str = "cuda",
-    device_ids: Optional[List[int]] = None,
     **kwargs,
 ):
     # clear environment variables
@@ -147,8 +138,9 @@ def spawn_parallel_fn(
             del os.environ[key]
 
     if world_size == 1:
-        device_ids = device_ids or [0]
+        device_id = torch.device(device_type, 0)
-        device_id = torch.device(device_type, device_ids[0])
+        os.environ["LOCAL_RANK"] = "0"
+        os.environ["WORLD_SIZE"] = "1"
         os.environ["LOCAL_DEVICE"] = str(device_id)
 
         func(**kwargs)
@@ -160,7 +152,6 @@ def spawn_parallel_fn(
         master_addr,
         master_port,
         device_type,
-        device_ids,
         func,
         kwargs,
     )

astrai/serialization.py

@@ -1,7 +1,7 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import h5py
 import safetensors.torch as st
@@ -54,10 +54,12 @@ class Checkpoint:
         state_dict: Dict[str, Any],
         epoch: int = 0,
         iteration: int = 0,
+        extra: Optional[Dict[str, Any]] = None,
     ):
         self.state_dict = state_dict
         self.epoch = epoch
         self.iteration = iteration
+        self.extra = extra or {}
 
     def save(
         self,
@@ -77,6 +79,8 @@ class Checkpoint:
                 json.dump(meta, f, indent=2)
 
             st.save_file(self.state_dict, save_path / "state_dict.safetensors")
+            if self.extra:
+                torch.save(self.extra, save_path / "extra.pt")
 
     @classmethod
     def load(
@@ -99,8 +103,14 @@ class Checkpoint:
 
         state_dict = st.load_file(save_path / "state_dict.safetensors")
 
+        extra = None
+        extra_path = save_path / "extra.pt"
+        if extra_path.exists():
+            extra = torch.load(extra_path, map_location="cpu", weights_only=False)
+
         return cls(
             state_dict=state_dict,
             epoch=meta["epoch"],
             iteration=meta["iteration"],
+            extra=extra,
         )

astrai/tokenize/__init__.py

@@ -1,15 +1,8 @@
 from astrai.tokenize.chat_template import ChatTemplate, MessageType
-from astrai.tokenize.tokenizer import (
+from astrai.tokenize.tokenizer import AutoTokenizer
-    AutoTokenizer,
-    BpeTokenizer,
-)
-from astrai.tokenize.trainer import BpeTrainer
 
 __all__ = [
     "AutoTokenizer",
-    "BpeTokenizer",
-    "BpeTrainer",
     "ChatTemplate",
     "MessageType",
-    "HistoryType",
 ]

astrai/tokenize/tokenizer.py

@@ -6,8 +6,7 @@ import json
 from pathlib import Path
 from typing import Dict, List, Optional, Union
 
-from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
+from tokenizers import Tokenizer
-from tokenizers.models import BPE
 
 from astrai.tokenize.chat_template import ChatTemplate
 
@@ -65,6 +64,11 @@ class AutoTokenizer:
             save_path: Path to save the tokenizer
         """
 
+        if self._tokenizer is None:
+            raise RuntimeError(
+                "Tokenizer not initialized. Load or create a tokenizer first."
+            )
+
         save_path = Path(save_path)
         save_path.mkdir(parents=True, exist_ok=True)
 
@@ -210,9 +214,9 @@ class AutoTokenizer:
 
         Args:
             messages: List of message dicts with 'role' and 'content'.
-            system_prompt: Optional system prompt string.
+            system_prompt: Optional system prompt string (auto-converted to first message).
             tokenize: Whether to return token IDs (True) or raw string (False).
-            add_generation_prompt: Whether to add the generation prompt (default: False).
+            add_generation_prompt: Whether to add the generation prompt (default: True).
             **kwargs: Additional variables to pass to the template.
 
         Returns:
@@ -226,10 +230,13 @@ class AutoTokenizer:
                 "Chat template not set. Use set_chat_template() to set a template first."
             )
 
+        # Auto-convert system_prompt to first message if provided
+        if system_prompt:
+            messages = [{"role": "system", "content": system_prompt}] + list(messages)
+
         # Render the template
         rendered = self._chat_template.render(
             messages=messages,
-            system_prompt=system_prompt,
             add_generation_prompt=add_generation_prompt,
             **kwargs,
         )
@@ -238,42 +245,3 @@ class AutoTokenizer:
             return self.encode(rendered)
 
         return rendered
-
-
-class BpeTokenizer(AutoTokenizer):
-    """BPE tokenizer implementation."""
-
-    def __init__(
-        self,
-        special_token_map: Dict[str, str] = None,
-        path: Optional[str] = None,
-        chat_template: Optional[str] = None,
-    ):
-        special_token_map = special_token_map or {
-            "bos": "<｜begin▁of▁sentence｜>",
-            "eos": "<｜end▁of▁sentence｜>",
-            "pad": "<｜▁pad▁｜>",
-            "im_start": "<｜im▁start｜>",
-            "im_end": "<｜im▁end｜>",
-        }
-        self._tokenizer = None
-        self._init_tokenizer()
-        super().__init__(
-            path, special_token_map=special_token_map, chat_template=chat_template
-        )
-
-    def _init_tokenizer(self):
-        """Initialize a new BPE tokenizer with default settings."""
-        model = BPE()
-        self._tokenizer = Tokenizer(model)
-        self._tokenizer.normalizer = normalizers.Sequence(
-            [normalizers.NFC(), normalizers.Strip()]
-        )
-        self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
-            [
-                pre_tokenizers.UnicodeScripts(),
-                pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True),
-            ]
-        )
-        self._tokenizer.decoder = decoders.ByteLevel()
-        self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)

astrai/tokenize/trainer.py

@@ -1,108 +0,0 @@
-"""
-BPE Tokenizer Trainer module.
-
-Provides training functionality for BPE tokenizers.
-"""
-
-from typing import List, Union
-
-from tokenizers import pre_tokenizers
-from tokenizers.trainers import BpeTrainer as BpeTrainerImpl
-
-
-class BpeTrainer:
-    """BPE tokenizer trainer."""
-
-    def __init__(self, tokenizer):
-        """Initialize trainer with a tokenizer instance.
-
-        Args:
-            tokenizer: A BpeTokenizer instance
-        """
-        self.tokenizer = tokenizer
-
-    def _prepare_trainer(
-        self,
-        vocab_size: int,
-        min_freq: int,
-        reserved_token_size: int,
-        max_token_length: int = 18,
-    ):
-        """Prepare the BPE trainer with proper configuration."""
-        assert reserved_token_size > len(self.tokenizer._special_tokens)
-        reserved_tokens = [
-            f"<|reserve{i:02d}|>"
-            for i in range(reserved_token_size - len(self.tokenizer._special_tokens))
-        ]
-        detail_vocab_size = vocab_size - (
-            len(reserved_tokens) + len(self.tokenizer._special_tokens)
-        )
-        alphabet = pre_tokenizers.ByteLevel.alphabet()
-        min_size = len(alphabet) + len(self.tokenizer._control_tokens)
-        assert detail_vocab_size > min_size
-
-        trainer = BpeTrainerImpl(
-            vocab_size=detail_vocab_size,
-            min_frequency=min_freq,
-            limit_alphabet=detail_vocab_size // 6,
-            max_token_length=max_token_length,
-            special_tokens=self.tokenizer._control_tokens,
-            initial_alphabet=alphabet,
-            show_progress=True,
-        )
-        return trainer, reserved_tokens
-
-    def train(
-        self,
-        files: Union[str, List[str]],
-        vocab_size: int,
-        min_freq: int,
-        reserved_token_size: int = 100,
-        **kwargs,
-    ):
-        """Train tokenizer from files.
-
-        Args:
-            files: Path or list of paths to training files
-            vocab_size: Target vocabulary size
-            min_freq: Minimum frequency for tokens
-            reserved_token_size: Number of reserved tokens
-            **kwargs: Additional arguments
-        """
-        trainer, reserved_tokens = self._prepare_trainer(
-            vocab_size, min_freq, reserved_token_size, **kwargs
-        )
-        self.tokenizer._tokenizer.train(files=files, trainer=trainer)
-        self.tokenizer._tokenizer.add_special_tokens(
-            self.tokenizer._special_tokens + reserved_tokens
-        )
-
-    def train_from_iterator(
-        self,
-        iterator,
-        vocab_size: int,
-        min_freq: int,
-        reserved_token_size: int = 100,
-        **kwargs,
-    ):
-        """Train tokenizer from iterator.
-
-        Args:
-            iterator: Iterator yielding training strings
-            vocab_size: Target vocabulary size
-            min_freq: Minimum frequency for tokens
-            reserved_token_size: Number of reserved tokens
-            **kwargs: Additional arguments
-        """
-        trainer, reserved_tokens = self._prepare_trainer(
-            vocab_size, min_freq, reserved_token_size, **kwargs
-        )
-        self.tokenizer._tokenizer.train_from_iterator(
-            iterator=iterator, trainer=trainer
-        )
-        self.tokenizer._tokenizer.add_special_tokens(
-            self.tokenizer._special_tokens + reserved_tokens
-        )
-
-
-__all__ = ["BpeTrainer"]

astrai/trainer/strategy.py

@@ -265,7 +265,9 @@ class DPOStrategy(BaseStrategy):
 class GRPOStrategy(BaseStrategy):
     """Group Relative Policy Optimization strategy.
 
-    Implements GRPO with clipping and KL penalty.
+    On-policy GRPO following DeepSeek-R1: the policy model is updated while
+    a frozen ref_model stores the old-policy log-probs.  ratio = exp(logπ_θ - logπ_ref),
+    clipped PPO objective.  Call ``sync_ref_model()`` after each data-generation round.
     """
 
     def __init__(
@@ -276,6 +278,7 @@ class GRPOStrategy(BaseStrategy):
         kl_coef: float = 0.01,
         group_size: int = 4,
         reduction: str = "mean",
+        sync_interval: int = 200,
         **kwargs,
     ):
         super().__init__(model, device, **kwargs)
@@ -284,8 +287,19 @@ class GRPOStrategy(BaseStrategy):
         self.kl_coef = kl_coef
         self.group_size = group_size
         self.reduction = reduction
+        self.sync_interval = sync_interval
+        self._step = 0
+
+    def sync_ref_model(self):
+        """Copy current model weights to ref model."""
+        ref_state = self.model.state_dict()
+        self.ref_model.load_state_dict(ref_state)
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        self._step += 1
+        if self._step % self.sync_interval == 0:
+            self.sync_ref_model()
+
         batch = move_to_device(batch, self.device)
         prompts = batch["prompts"]
         responses = batch["responses"]
@@ -297,7 +311,6 @@ class GRPOStrategy(BaseStrategy):
         masks_flat = masks.view(-1, response_len)
         prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
 
-        # Shape: (batch_size * group_size, seq_len + response_len)
         full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
         full_masks = torch.cat([torch.ones_like(prompt_expanded), masks_flat], dim=-1)
 
@@ -312,14 +325,13 @@ class GRPOStrategy(BaseStrategy):
             )
             log_probs_ref = log_probs_ref.view(batch_size, group_size)
 
-        # Compute advantages from rewards with normalization
         eps = torch.finfo(log_probs_policy.dtype).eps
         mean = rewards.mean(dim=-1, keepdim=True)
         std = rewards.std(dim=-1, keepdim=True)
         advantages = (rewards - mean) / (std + eps)
 
-        # PPO-style clipped surrogate objective
+        ratio = torch.exp(log_probs_policy - log_probs_ref)
-        ratio = torch.exp(0)  # Off-policy: policy_model = old_model
+
         surr1 = ratio * advantages
         surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
 

astrai/trainer/train_callback.py

@@ -121,11 +121,13 @@ class CheckpointCallback(TrainCallback):
         interval: int,
         weight_only: bool = False,
         state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
+        save_extra_fn: Optional[Callable[["TrainContext"], dict]] = None,
     ):
         self.save_dir = save_dir
         self.interval = interval
         self.weight_only = weight_only
         self.state_dict_fn = state_dict_fn
+        self.save_extra_fn = save_extra_fn
         self.last_ckpt_iter = 0
 
     @only_on_rank(0)
@@ -139,8 +141,12 @@ class CheckpointCallback(TrainCallback):
             else context.model.state_dict()
         )
 
+        extra = self.save_extra_fn(context) if self.save_extra_fn else None
         context.checkpoint = Checkpoint(
-            state_dict=state_dict, epoch=context.epoch, iteration=context.iteration
+            state_dict=state_dict,
+            epoch=context.epoch,
+            iteration=context.iteration,
+            extra=extra,
         )
 
         context.checkpoint.save(save_path)

astrai/trainer/train_context.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Optional, Self
+from typing import Callable, Optional, Self
 
 import torch.nn as nn
 from torch.optim import Optimizer
@@ -32,68 +32,70 @@ class TrainContext:
 
 
 class TrainContextBuilder:
-    def __init__(self, config: TrainConfig):
+    def __init__(
+        self,
+        config: TrainConfig,
+        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
+    ):
         self.config = config
-        self._context = TrainContext(
+        self._checkpoint: Optional[Checkpoint] = None
-            model=config.model,
+        self._load_extra_fn = load_extra_fn
+
+    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+        self._checkpoint = checkpoint
+        return self
+
+    def build(self) -> TrainContext:
+        context = TrainContext(
+            model=self.config.model,
             world_size=get_world_size(),
             rank=get_rank(),
         )
 
         device = get_current_device()
-        self._context.model = self._context.model.to(device=device)
+        context.model = context.model.to(device=device)
 
-        if self.config.nprocs > 1:
+        if self.config.nprocs > 1 and self.config.parallel_wrapper:
-            fn = self.config.parallel_wrapper
+            context.model = self.config.parallel_wrapper(context.model)
-            self._context.model = fn(self._context.model)
 
-        self._context.optimizer = self.config.optimizer_fn(self._context.model)
+        if self._checkpoint is not None:
-        self._context.scheduler = self.config.scheduler_fn(self._context.optimizer)
+            context.epoch = max(self._checkpoint.epoch, self.config.start_epoch)
-
+            context.iteration = max(self._checkpoint.iteration, self.config.start_batch)
-    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+            context.model.load_state_dict(self._checkpoint.state_dict)
-        if checkpoint is None:
+            context.checkpoint = self._checkpoint
-            checkpoint = Checkpoint(
-                state_dict=self._context.model.state_dict(),
-            )
         else:
-            # resume from the assigned checkpoint or assigned iteration
+            context.checkpoint = Checkpoint(
-            self._context.epoch = max(checkpoint.epoch, self.config.start_epoch)
+                state_dict=context.model.state_dict(),
-            self._context.iteration = max(checkpoint.iteration, self.config.start_batch)
+            )
-            self._context.model.load_state_dict(checkpoint.state_dict)
 
-        self._context.checkpoint = checkpoint
+        context.optimizer = self.config.optimizer_fn(context.model)
-        return self
+        context.scheduler = self.config.scheduler_fn(context.optimizer)
 
-    def with_dataloader(self) -> Self:
+        if self._checkpoint and self._checkpoint.extra and self._load_extra_fn:
-        # fix: change batch level iteration to sample level offset
+            self._load_extra_fn(self._checkpoint.extra, context)
-        config = self.config
+
-        sampler_offset = self._context.iteration * config.batch_size
+        cfg = self.config
-        resumeable_sampler = ResumableDistributedSampler(
+        sampler_offset = context.iteration * cfg.batch_size
-            data_source=config.dataset,
+        sampler = ResumableDistributedSampler(
-            start_epoch=self._context.epoch,
+            data_source=cfg.dataset,
+            start_epoch=context.epoch,
             start_iter=sampler_offset,
-            seed=config.random_seed,
+            seed=cfg.random_seed,
+        )
+        context.dataloader = DataLoader(
+            cfg.dataset,
+            batch_size=cfg.batch_size,
+            sampler=sampler,
+            num_workers=cfg.num_workers,
+            pin_memory=cfg.pin_memory,
+            prefetch_factor=cfg.prefetch_factor,
         )
 
-        dataloader = DataLoader(
+        context.strategy = StrategyFactory.create(
-            config.dataset,
+            model=context.model,
-            batch_size=config.batch_size,
-            sampler=resumeable_sampler,
-            num_workers=config.num_workers,
-            pin_memory=config.pin_memory,
-            prefetch_factor=config.prefetch_factor,
-        )
-        self._context.dataloader = dataloader
-        return self
-
-    def with_strategy(self) -> Self:
-        self._context.strategy = StrategyFactory.create(
-            model=self._context.model,
             train_type=self.config.strategy,
-            device=get_current_device(),
+            device=device,
             **self.config.extra_kwargs,
         )
-        return self
 
-    def build(self) -> TrainContext:
+        return context
-        return self._context

astrai/trainer/trainer.py

@@ -35,11 +35,7 @@ class Trainer:
 
     def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
         return (
-            TrainContextBuilder(self.train_config)
+            TrainContextBuilder(self.train_config).with_checkpoint(checkpoint).build()
-            .with_checkpoint(checkpoint)
-            .with_dataloader()
-            .with_strategy()
-            .build()
         )
 
     def _call_callbacks(self, method_name: str, context: TrainContext):
@@ -57,7 +53,6 @@ class Trainer:
             master_addr=config.master_addr,
             master_port=config.master_port,
             device_type=config.device_type,
-            device_ids=config.device_ids,
             checkpoint=checkpoint,
         )
 
@@ -72,8 +67,9 @@ class Trainer:
                 context.epoch = epoch
                 self._call_callbacks("on_epoch_begin", context)
 
+                accumulation_steps = max(self.train_config.accumulation_steps, 1)
                 for batch in context.dataloader:
-                    if context.iteration % self.train_config.accumulation_steps == 0:
+                    if context.iteration % accumulation_steps == 0:
                         # 2. step
                         self._call_callbacks("on_step_begin", context)
                         context.optimizer.step()
@@ -87,7 +83,7 @@ class Trainer:
                     context.iteration += 1
 
                     # to make the loss normalized by accumulation steps
-                    stand_loss = loss / self.train_config.accumulation_steps
+                    stand_loss = loss / accumulation_steps
                     stand_loss.backward()
 
                     self._call_callbacks("on_batch_end", context)

docker-compose.yml

@@ -0,0 +1,42 @@
+services:
+  server:
+    build: .
+    image: astrai:latest
+    ports:
+      - "8000:8000"
+    volumes:
+      - ./params:/app/params:ro
+      - ./checkpoints:/app/checkpoints
+    command: python -m scripts.tools.server --port 8000 --device cuda
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+      interval: 30s
+      timeout: 10s
+      retries: 3
+      start_period: 60s
+    restart: unless-stopped
+
+  server-cpu:
+    profiles: [cpu]
+    build: .
+    image: astrai:latest
+    ports:
+      - "8000:8000"
+    volumes:
+      - ./params:/app/params:ro
+      - ./checkpoints:/app/checkpoints
+    command: python -m scripts.tools.server --port 8000 --device cpu
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+      interval: 30s
+      timeout: 10s
+      retries: 3
+      start_period: 120s
+    restart: unless-stopped

scripts/demo/download.py

@@ -1,13 +1,41 @@
+import argparse
 from pathlib import Path
 
 from huggingface_hub import snapshot_download
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
-PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+DEFAULT_LOCAL_DIR = Path(PROJECT_ROOT, "params")
+DEFAULT_REPO_ID = "ViperEk/KHAOSZ"
 
 if __name__ == "__main__":
-    snapshot_download(
+    parser = argparse.ArgumentParser(
-        repo_id="ViperEk/KHAOSZ",
+        description="Download model parameters from HuggingFace"
-        local_dir=PARAMETER_ROOT,
-        force_download=True,
     )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default=DEFAULT_REPO_ID,
+        help=f"HuggingFace repo ID (default: {DEFAULT_REPO_ID})",
+    )
+    parser.add_argument(
+        "--local-dir",
+        type=Path,
+        default=DEFAULT_LOCAL_DIR,
+        help=f"Local directory to save model (default: {DEFAULT_LOCAL_DIR})",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Force download even if files exist",
+    )
+    args = parser.parse_args()
+
+    print(f"Downloading model from {args.repo_id} to {args.local_dir}")
+
+    snapshot_download(
+        repo_id=args.repo_id,
+        local_dir=args.local_dir,
+        force_download=args.force,
+    )
+
+    print("Download complete!")

scripts/demo/stream_chat.py

@@ -15,7 +15,7 @@ def chat():
     tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
     model.to(device="cuda", dtype=torch.bfloat16)
 
-    messages = []
+    messages = [{"role": "system", "content": "You are a helpful assistant."}]
     engine = InferenceEngine(model=model, tokenizer=tokenizer)
 
     while True:

scripts/tools/benchmark.py

@@ -1,9 +1,14 @@
+"""Benchmark Transformer with PagedCache (replaces old persistent_key_values)."""
+
 from dataclasses import dataclass
 from typing import Any, Dict
 
 import torch
+from torch import Tensor
 
-from astrai.model.transformer import ModelConfig, Transformer
+from astrai.config import ModelConfig
+from astrai.inference.cache import PagedCache
+from astrai.model.transformer import Transformer
 
 
 @dataclass
@@ -19,27 +24,25 @@ class GenerationBenchmark:
         self,
         config: ModelConfig,
         device: str = "cuda",
-        dtype: torch.dtype = torch.float16,
+        dtype: torch.dtype = torch.bfloat16,
+        page_size: int = 128,
     ):
         self.config = config
         self.device = device
         self.dtype = dtype
         self.model = Transformer(config).to(device=device, dtype=dtype)
         self.model.eval()
-
+        head_dim = config.dim // config.n_heads
-    def _initialize_kv_cache(self, batch_size: int) -> list:
+        n_pages = (config.max_len * 4 + page_size - 1) // page_size
-        """初始化KV缓存"""
+        self._page_cache = PagedCache(
-        config = self.config
-        shape = (
-            batch_size,
-            config.max_len,
             config.n_layers,
+            n_pages,
+            page_size,
             config.n_kv_heads,
-            config.dim // config.n_heads,
+            head_dim,
+            device,
+            dtype,
         )
-        k_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
-        v_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
-        return (k_cache, v_cache)
 
     def _prepare_inputs(self, batch_size: int, prompt_length: int, total_length: int):
         prompt_ids = torch.randint(
@@ -49,7 +52,6 @@ class GenerationBenchmark:
             device=self.device,
             dtype=torch.long,
         )
-
         gen_ids = torch.randint(
             low=0,
             high=self.config.vocab_size,
@@ -57,9 +59,11 @@ class GenerationBenchmark:
             device=self.device,
             dtype=torch.long,
         )
-
         return prompt_ids, gen_ids
 
+    def _make_mask(self, batch_size: int, seq_len: int) -> Tensor:
+        return torch.ones(batch_size, seq_len, dtype=torch.bool, device=self.device)
+
     @torch.inference_mode()
     def run_prefill_benchmark(
         self,
@@ -67,13 +71,11 @@ class GenerationBenchmark:
         prompt_length: int = 512,
         num_trials: int = 10,
     ) -> BenchmarkResult:
-
         for _ in range(3):
             prompt_ids, _ = self._prepare_inputs(
                 batch_size, prompt_length, prompt_length
             )
             _ = self.model(prompt_ids)
-
         torch.cuda.synchronize()
 
         total_time = 0.0
@@ -83,20 +85,20 @@ class GenerationBenchmark:
             prompt_ids, _ = self._prepare_inputs(
                 batch_size, prompt_length, prompt_length
             )
-            start_event = torch.cuda.Event(enable_timing=True)
+            start = torch.cuda.Event(enable_timing=True)
-            end_event = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
 
-            start_event.record()
+            start.record()
             _ = self.model(prompt_ids)
-            end_event.record()
+            end.record()
             torch.cuda.synchronize()
 
-            trial_time = start_event.elapsed_time(end_event) / 1000
+            trial_time = start.elapsed_time(end) / 1000
             total_time += trial_time
 
             print(
-                f"Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
+                f"  Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
-                f"({prompt_length / trial_time:.1f} tokens/s)"
+                f"({prompt_length / trial_time:.1f} tok/s)"
             )
 
         return BenchmarkResult(
@@ -107,7 +109,7 @@ class GenerationBenchmark:
                 "benchmark_type": "prefill",
                 "batch_size": batch_size,
                 "prompt_length": prompt_length,
-                "dtype": self.dtype,
+                "dtype": str(self.dtype),
                 "device": self.device,
             },
         )
@@ -120,41 +122,62 @@ class GenerationBenchmark:
         gen_length: int = 128,
         num_trials: int = 5,
     ) -> BenchmarkResult:
-
         total_time = 0.0
         total_tokens = batch_size * gen_length * num_trials
+        page_size = self._page_cache.page_size
 
         for trial in range(num_trials):
             prompt_ids, gen_ids = self._prepare_inputs(
-                batch_size, prompt_length, prompt_length + gen_length
+                batch_size,
+                prompt_length,
+                prompt_length + gen_length,
+            )
+
+            n_pages = (prompt_length + gen_length + page_size - 1) // page_size
+            pages = self._page_cache.alloc_n(n_pages * batch_size)
+            page_table = torch.tensor(
+                [pages[i * n_pages : (i + 1) * n_pages] for i in range(batch_size)],
+                dtype=torch.long,
+                device=self.device,
+            )
+
+            cv = self._page_cache.bind(page_table, total_len=prompt_length)
+            _ = self.model(
+                prompt_ids,
+                paged_cache=cv,
+                start_pos=0,
+                input_mask=self._make_mask(batch_size, prompt_length),
             )
-            kv_cache = self._initialize_kv_cache(batch_size)
-            _ = self.model(prompt_ids, persistent_key_values=kv_cache, start_pos=0)
 
             torch.cuda.synchronize()
 
-            start_event = torch.cuda.Event(enable_timing=True)
+            start = torch.cuda.Event(enable_timing=True)
-            end_event = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
-
-            start_event.record()
 
+            start.record()
             current_pos = prompt_length
             for i in range(gen_length):
                 input_token = gen_ids[:, i : i + 1]
+                cv = self._page_cache.bind(page_table, total_len=current_pos + 1)
                 _ = self.model(
-                    input_token, persistent_key_values=kv_cache, start_pos=current_pos
+                    input_token,
+                    paged_cache=cv,
+                    start_pos=current_pos,
+                    input_mask=self._make_mask(batch_size, 1),
                 )
                 current_pos += 1
-
+            end.record()
-            end_event.record()
             torch.cuda.synchronize()
 
-            trial_time = start_event.elapsed_time(end_event) / 1000
+            trial_time = start.elapsed_time(end) / 1000
             total_time += trial_time
 
+            for idx in pages:
+                self._page_cache.free(idx)
+
             print(
-                f"Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
+                f"  Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
-                f"({gen_length / trial_time:.1f} tokens/s)"
+                f"({gen_length / trial_time:.1f} tok/s)"
             )
 
         return BenchmarkResult(
@@ -166,31 +189,21 @@ class GenerationBenchmark:
                 "batch_size": batch_size,
                 "prompt_length": prompt_length,
                 "gen_length": gen_length,
-                "dtype": self.dtype,
+                "dtype": str(self.dtype),
                 "device": self.device,
             },
         )
 
 
 def print_benchmark_result(result: BenchmarkResult):
-    """打印基准测试结果"""
+    btype = result.metadata["benchmark_type"]
-    benchmark_type = result.metadata["benchmark_type"]
+    print(f"\n{' ' + btype.upper() + ' Benchmark ':-^80}")
-
-    print(f"\n{' ' + benchmark_type.upper().replace('_', ' ') + ' Benchmark ':-^80}")
     print(f"Total Tokens Processed: {result.total_tokens:,}")
     print(f"Time Consumed: {result.total_time:.3f}s")
-    print(f"Throughput: {result.tokens_per_second:,.1f} tokens/s")
+    print(f"Throughput: {result.tokens_per_second:,.1f} tok/s")
-
+    for k, v in result.metadata.items():
-    if benchmark_type == "prefill":
+        if k != "benchmark_type":
-        print(
+            print(f"{k.replace('_', ' ').title()}: {v}")
-            f"Batch Size: {result.metadata['batch_size']} | Prompt Length: {result.metadata['prompt_length']}"
-        )
-    elif benchmark_type == "decoding":
-        print(
-            f"Batch Size: {result.metadata['batch_size']} | Gen Length: {result.metadata['gen_length']}"
-        )
-
-    print(f"Device: {result.metadata['device']} | Dtype: {result.metadata['dtype']}")
     print("-" * 80)
 
 
@@ -209,15 +222,20 @@ if __name__ == "__main__":
     benchmark = GenerationBenchmark(config)
 
     print("=" * 80)
-    print("Running Transformer Generation Benchmark")
+    print("Running Transformer Generation Benchmark (PagedCache)")
     print("=" * 80)
 
     prefill_result = benchmark.run_prefill_benchmark(
-        batch_size=4, prompt_length=512, num_trials=5
+        batch_size=4,
+        prompt_length=512,
+        num_trials=5,
     )
     print_benchmark_result(prefill_result)
 
     gen_result = benchmark.run_decoding_benchmark(
-        batch_size=4, prompt_length=512, gen_length=128, num_trials=5
+        batch_size=4,
+        prompt_length=512,
+        gen_length=128,
+        num_trials=5,
     )
     print_benchmark_result(gen_result)

scripts/tools/generate.py

@@ -9,7 +9,7 @@ from astrai.tokenize import AutoTokenizer
 
 
 def processor(
-    model_dir: str,
+    param_path: str,
     input_json_file: str,
     output_json_file: str,
     temperature: float,
@@ -20,8 +20,8 @@ def processor(
     max_tokens: int,
 ):
     # Load model and tokenizer
-    model = AutoModel.from_pretrained(model_dir)
+    model = AutoModel.from_pretrained(param_path)
-    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
     model.to(device="cuda", dtype=torch.bfloat16)
 
     # Create inference engine
@@ -72,7 +72,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run generate with a Khaosz model.")
 
     parser.add_argument(
-        "--model_dir", type=str, required=True, help="Path to the model directory."
+        "--param_path", type=str, required=True, help="Path to the model directory."
     )
     parser.add_argument(
         "--input_json_file",

scripts/tools/train.py

@@ -23,7 +23,7 @@ def parse_args() -> argparse.Namespace:
         "--train_type",
         type=str,
         required=True,
-        choices=["seq", "sft", "dpo"],
+        choices=["seq", "sft", "dpo", "grpo"],
         help="Train type.",
     )
     parser.add_argument(
@@ -42,9 +42,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--n_epoch", type=int, default=1, help="Number of epochs to train."
     )
-    parser.add_argument(
+    parser.add_argument("--batch_size", type=int, default=1, help="Batch size per GPU.")
-        "--batch_size", type=int, default=1, help="Batch size for training."
-    )
     parser.add_argument(
         "--accumulation_steps",
         type=int,
@@ -55,7 +53,7 @@ def parse_args() -> argparse.Namespace:
         "--warmup_steps",
         type=int,
         default=1000,
-        help="Number of iters between warnings.",
+        help="Number of warmup steps for LR scheduler.",
     )
     parser.add_argument(
         "--max_lr", type=float, default=3e-4, help="Max learning rate for training."
@@ -100,12 +98,19 @@ def parse_args() -> argparse.Namespace:
         "--window_size",
         type=int,
         default=None,
-        help="the max length of the input sequence.",
+        help="Max length of the input sequence.",
     )
     parser.add_argument(
-        "--stride", type=int, default=None, help="the step size of the input sequence."
+        "--stride", type=int, default=None, help="Step size of the input sequence."
     )
     parser.add_argument("--dpo_beta", type=float, default=0.1, help="DPO beta value.")
+    parser.add_argument("--group_size", type=int, default=4, help="GRPO group size.")
+    parser.add_argument(
+        "--grpo_clip_eps", type=float, default=0.2, help="GRPO clipping epsilon."
+    )
+    parser.add_argument(
+        "--grpo_kl_coef", type=float, default=0.01, help="GRPO KL penalty coefficient."
+    )
     parser.add_argument(
         "--label_smoothing",
         type=float,
@@ -125,6 +130,12 @@ def parse_args() -> argparse.Namespace:
         default="checkpoint",
         help="Directory to save checkpoints.",
     )
+    parser.add_argument(
+        "--grpo_sync_interval",
+        type=int,
+        default=200,
+        help="GRPO ref model sync interval (steps).",
+    )
     parser.add_argument(
         "--start_epoch", type=int, default=0, help="Start epoch for training."
     )
@@ -144,7 +155,7 @@ def parse_args() -> argparse.Namespace:
 
 def ddp_wrap(model: nn.Module):
     local_rank = get_rank()
-    model = model.to(device=f"cuda:{local_rank}", dtype=torch.bfloat16)
+    model = model.to(dtype=torch.bfloat16)
     ddp_model = DDP(
         model,
         device_ids=[local_rank],
@@ -182,6 +193,10 @@ def train(
     ckpt_interval: int,
     ckpt_dir: str,
     dpo_beta: float,
+    grpo_clip_eps: float,
+    grpo_kl_coef: float,
+    group_size: int,
+    grpo_sync_interval: int,
     adamw_beta1: float,
     adamw_beta2: float,
     adamw_weight_decay: float,
@@ -195,7 +210,7 @@ def train(
     nprocs: int,
     device_type: str,
 ):
-    assert train_type in ["seq", "sft", "dpo"]
+    assert train_type in ["seq", "sft", "dpo", "grpo"]
     assert os.path.exists(param_path)
 
     # Load config
@@ -216,7 +231,14 @@ def train(
         state_dict = st.load_file(weights_path)
         model.load_state_dict(state_dict, strict=False)
 
-    strategy_kwargs = {"dpo_beta": dpo_beta, "label_smoothing": label_smoothing}
+    strategy_kwargs = {
+        "dpo_beta": dpo_beta,
+        "label_smoothing": label_smoothing,
+        "clip_eps": grpo_clip_eps,
+        "kl_coef": grpo_kl_coef,
+        "group_size": group_size,
+        "sync_interval": grpo_sync_interval,
+    }
 
     dataset = DatasetFactory.load(
         train_type=train_type,
@@ -234,13 +256,13 @@ def train(
         },
     )
 
-    toltal_steps = len(dataset) * n_epoch // (batch_size * nprocs)
+    total_steps = len(dataset) * n_epoch // (batch_size * nprocs)
     scheduler_fn = partial(
         create_scheduler,
         **{
             "schedule_type": "cosine",
             "warmup_steps": warmup_steps,
-            "lr_decay_steps": toltal_steps - warmup_steps,
+            "lr_decay_steps": total_steps - warmup_steps,
         },
     )
 

tests/conftest.py

@@ -7,12 +7,27 @@ import numpy as np
 import pytest
 import safetensors.torch as st
 import torch
-from tokenizers import pre_tokenizers
+from tokenizers import Tokenizer, models, pre_tokenizers, trainers
 from torch.utils.data import Dataset
 
 from astrai.config.model_config import ModelConfig
 from astrai.model.transformer import Transformer
-from astrai.tokenize import BpeTokenizer, BpeTrainer
+from astrai.tokenize import AutoTokenizer
+
+
+def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
+    """Create a simple tokenizer for testing purposes."""
+    tokenizer = Tokenizer(models.BPE())
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size, min_frequency=1, special_tokens=["<unk>", "<pad>"]
+    )
+    # Train on empty iterator with single character
+    tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer)
+    auto_tokenizer = AutoTokenizer()
+    auto_tokenizer._tokenizer = tokenizer
+    auto_tokenizer._special_token_map = {"unk_token": "<unk>", "pad_token": "<pad>"}
+    return auto_tokenizer
 
 
 class RandomDataset(Dataset):
@@ -109,7 +124,7 @@ def base_test_env(request: pytest.FixtureRequest):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     transformer_config = ModelConfig().load(config_path)
     model = Transformer(transformer_config).to(device=device)
-    tokenizer = BpeTokenizer()
+    tokenizer = create_test_tokenizer()
 
     yield {
         "device": device,
@@ -164,10 +179,7 @@ def test_env(request: pytest.FixtureRequest):
     with open(config_path, "w") as f:
         json.dump(config, f)
 
-    tokenizer = BpeTokenizer()
+    tokenizer = create_test_tokenizer(vocab_size=config["vocab_size"])
-    trainer = BpeTrainer(tokenizer)
-    sp_token_iter = iter(pre_tokenizers.ByteLevel.alphabet())
-    trainer.train_from_iterator(sp_token_iter, config["vocab_size"], 1)
     tokenizer.save(tokenizer_path)
 
     transformer_config = ModelConfig().load(config_path)

tests/inference/conftest.py

@@ -14,37 +14,32 @@ def client():
     return TestClient(app)
 
 
-@pytest.fixture
-def mock_model_param():
-    """Create a mock ModelParameter."""
-    mock_param = MagicMock()
-    mock_param.model = MagicMock()
-    mock_param.tokenizer = MagicMock()
-    mock_param.config = MagicMock()
-    mock_param.config.max_len = 100
-    mock_param.tokenizer.encode = MagicMock(return_value=[1, 2, 3])
-    mock_param.tokenizer.decode = MagicMock(return_value="mock response")
-    mock_param.tokenizer.stop_ids = []
-    mock_param.tokenizer.pad_id = 0
-    return mock_param
-
-
 @pytest.fixture
 def mock_engine():
     """Create a mock InferenceEngine."""
+
+    async def _async_gen():
+        yield "chunk1"
+        yield "chunk2"
+        yield "[DONE]"
+
     mock = MagicMock()
     mock.generate.return_value = "mock response"
+    mock.generate_async.return_value = _async_gen()
     mock.get_stats.return_value = {
         "total_tasks": 0,
         "total_tokens": 0,
         "active_tasks": 0,
         "waiting_queue": 0,
     }
+    mock.tokenizer.encode.return_value = [1, 2, 3]
+    mock.tokenizer.decode.return_value = "mock response"
+    mock.tokenizer.apply_chat_template.return_value = "mock prompt"
     return mock
 
 
 @pytest.fixture
-def loaded_model(mock_model_param, monkeypatch):
+def loaded_model(mock_engine, monkeypatch):
-    """Simulate that the model is loaded."""
+    """Simulate that the engine is loaded."""
-    monkeypatch.setattr("astrai.inference.server._model_param", mock_model_param)
+    monkeypatch.setattr("astrai.inference.server._state.engine", mock_engine)
-    return mock_model_param
+    return mock_engine

tests/inference/test_scheduler_concurrency.py

@@ -0,0 +1,177 @@
+"""Tests for scheduler concurrency."""
+
+import threading
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from astrai.inference.scheduler import InferenceScheduler
+
+
+@pytest.fixture
+def mock_model_and_tokenizer():
+    """Create mock model and tokenizer."""
+    mock_model = MagicMock()
+    mock_model.config = MagicMock()
+    mock_model.config.n_kv_heads = 8
+    mock_model.config.n_heads = 8
+    mock_model.config.dim = 128
+    mock_model.config.n_layers = 2
+    mock_model.config.max_len = 100
+    mock_model.parameters.return_value = iter(
+        [MagicMock(dtype=torch.float32, device=torch.device("cpu"))]
+    )
+
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.encode.return_value = [1, 2, 3, 4, 5]
+    mock_tokenizer.decode.return_value = "token"
+    mock_tokenizer.stop_ids = [0]
+    mock_tokenizer.pad_id = None
+
+    return mock_model, mock_tokenizer
+
+
+def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
+    """Test concurrent add_task operations."""
+    mock_model, mock_tokenizer = mock_model_and_tokenizer
+
+    with patch("astrai.inference.scheduler.AutoModel"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
+            scheduler = InferenceScheduler(
+                model=mock_model,
+                tokenizer=mock_tokenizer,
+                max_batch_size=4,
+                device="cpu",
+            )
+
+    results = {"task_ids": [], "errors": []}
+    lock = threading.Lock()
+
+    def add_task_worker(worker_id):
+        try:
+            for i in range(10):
+                task_id = scheduler.add_task(f"prompt from worker {worker_id}-{i}")
+                with lock:
+                    results["task_ids"].append(task_id)
+        except Exception as e:
+            results["errors"].append(str(e))
+
+    threads = [threading.Thread(target=add_task_worker, args=(i,)) for i in range(5)]
+
+    for t in threads:
+        t.start()
+
+    # Let some tasks be processed
+    time.sleep(0.1)
+
+    scheduler.stop()
+
+    for t in threads:
+        t.join()
+
+    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
+    assert len(results["task_ids"]) == 50
+
+
+def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
+    """Test concurrent add and remove task operations."""
+    mock_model, mock_tokenizer = mock_model_and_tokenizer
+
+    with patch("astrai.inference.scheduler.AutoModel"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
+            scheduler = InferenceScheduler(
+                model=mock_model,
+                tokenizer=mock_tokenizer,
+                max_batch_size=4,
+                device="cpu",
+            )
+
+    results = {"added": [], "removed": [], "errors": []}
+
+    def add_worker():
+        try:
+            for i in range(20):
+                task_id = scheduler.add_task(f"prompt {i}")
+                results["added"].append(task_id)
+                time.sleep(0.001)
+        except Exception as e:
+            results["errors"].append(f"Add: {str(e)}")
+
+    def remove_worker():
+        try:
+            time.sleep(0.05)  # Wait for some tasks to be added
+            for task_id in results["added"][:10]:
+                scheduler.remove_task(task_id)
+                results["removed"].append(task_id)
+        except Exception as e:
+            results["errors"].append(f"Remove: {str(e)}")
+
+    add_thread = threading.Thread(target=add_worker)
+    remove_thread = threading.Thread(target=remove_worker)
+
+    add_thread.start()
+    remove_thread.start()
+
+    time.sleep(0.2)
+    scheduler.stop()
+
+    add_thread.join()
+    remove_thread.join()
+
+    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
+    assert len(results["added"]) == 20
+
+
+def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
+    """Test concurrent get_stats operations."""
+    mock_model, mock_tokenizer = mock_model_and_tokenizer
+
+    with patch("astrai.inference.scheduler.AutoModel"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
+            scheduler = InferenceScheduler(
+                model=mock_model,
+                tokenizer=mock_tokenizer,
+                max_batch_size=4,
+                device="cpu",
+            )
+
+    results = {"stats": [], "errors": []}
+
+    def add_tasks():
+        try:
+            for i in range(20):
+                scheduler.add_task(f"prompt {i}")
+                time.sleep(0.001)
+        except Exception as e:
+            results["errors"].append(f"Add: {str(e)}")
+
+    def get_stats():
+        try:
+            for _ in range(50):
+                stats = scheduler.get_stats()
+                results["stats"].append(stats)
+                time.sleep(0.001)
+        except Exception as e:
+            results["errors"].append(f"Get stats: {str(e)}")
+
+    add_thread = threading.Thread(target=add_tasks)
+    stats_thread = threading.Thread(target=get_stats)
+
+    add_thread.start()
+    stats_thread.start()
+
+    time.sleep(0.3)
+    scheduler.stop()
+
+    add_thread.join()
+    stats_thread.join()
+
+    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
+    assert len(results["stats"]) == 50
+
+    # Verify stats are consistent
+    for stats in results["stats"]:
+        assert "total_tasks" in stats
+        assert stats["total_tasks"] >= 0

tests/inference/test_server.py

@@ -4,88 +4,38 @@ import pytest
 
 
 def test_health_no_model(client, monkeypatch):
-    """GET /health should return 200 even when model not loaded."""
+    """GET /health should return 200 even when engine not loaded."""
-    monkeypatch.setattr("astrai.inference.server._model_param", None)
+    monkeypatch.setattr("astrai.inference.server._state.engine", None)
-    monkeypatch.setattr("astrai.inference.server._engine", None)
     response = client.get("/health")
     assert response.status_code == 200
     data = response.json()
     assert data["status"] == "ok"
     assert not data["model_loaded"]
-    assert not data["engine_ready"]
 
 
-def test_health_with_model(client, loaded_model, mock_engine, monkeypatch):
+def test_health_with_model(client, loaded_model):
-    """GET /health should return 200 when model is loaded."""
+    """GET /health should return 200 when engine is loaded."""
-    monkeypatch.setattr("astrai.inference.server._engine", mock_engine)
     response = client.get("/health")
     assert response.status_code == 200
     data = response.json()
     assert data["status"] == "ok"
     assert data["model_loaded"] is True
-    assert data["engine_ready"] is True
 
 
-def test_generate_non_stream(client, loaded_model, mock_engine, monkeypatch):
+def test_chat_completions_non_stream(client, loaded_model, monkeypatch):
-    """POST /generate with stream=false should return JSON response."""
+    """POST /v1/chat/completions with stream=false returns OpenAI-style JSON."""
-    monkeypatch.setattr("astrai.inference.server._engine", mock_engine)
-    response = client.post(
-        "/generate",
-        params={
-            "query": "Hello",
-            "temperature": 0.8,
-            "top_p": 0.95,
-            "top_k": 50,
-            "max_len": 100,
-            "stream": False,
-        },
-    )
-    assert response.status_code == 200
-    data = response.json()
-    assert data["response"] == "mock response"
 
+    async def async_gen():
+        yield "Assistant reply"
 
-def test_generate_stream(client, loaded_model, mock_engine, monkeypatch):
+    mock_engine = loaded_model
-    """POST /generate with stream=true should return plain text stream."""
+    mock_engine.generate_async.return_value = async_gen()
-
+    monkeypatch.setattr("astrai.inference.server._state.engine", mock_engine)
-    # Create a streaming mock
-    def stream_gen():
-        yield "chunk1"
-        yield "chunk2"
-
-    mock_engine.generate.return_value = stream_gen()
-    monkeypatch.setattr("astrai.inference.server._engine", mock_engine)
-    response = client.post(
-        "/generate",
-        params={
-            "query": "Hello",
-            "temperature": 0.8,
-            "top_p": 0.95,
-            "top_k": 50,
-            "max_len": 100,
-            "stream": True,
-        },
-        headers={"Accept": "text/plain"},
-    )
-    assert response.status_code == 200
-    assert response.headers["content-type"] == "text/plain; charset=utf-8"
-    # The stream yields lines ending with newline
-    content = response.content.decode("utf-8")
-    assert "chunk1" in content
-    assert "chunk2" in content
-
-
-def test_chat_completions_non_stream(client, loaded_model, mock_engine, monkeypatch):
-    """POST /v1/chat/completions with stream=false returns OpenAI‑style JSON."""
-    mock_engine.generate.return_value = "Assistant reply"
-    monkeypatch.setattr("astrai.inference.server._engine", mock_engine)
     response = client.post(
         "/v1/chat/completions",
         json={
             "messages": [{"role": "user", "content": "Hello"}],
             "temperature": 0.8,
-            "top_p": 0.95,
-            "top_k": 50,
             "max_tokens": 100,
             "stream": False,
         },
@@ -94,57 +44,120 @@ def test_chat_completions_non_stream(client, loaded_model, mock_engine, monkeypa
     data = response.json()
     assert data["object"] == "chat.completion"
     assert len(data["choices"]) == 1
-    assert data["choices"][0]["message"]["content"] == "Assistant reply"
+    assert "usage" in data
+    assert "prompt_tokens" in data["usage"]
 
 
-def test_chat_completions_stream(client, loaded_model, mock_engine, monkeypatch):
+def test_chat_completions_stream(client, loaded_model, monkeypatch):
     """POST /v1/chat/completions with stream=true returns SSE stream."""
 
-    # Simulate a streaming generator that yields cumulative responses
+    async def async_gen():
-    def stream_gen():
         yield "cumulative1"
         yield "cumulative2"
-        yield "[DONE]"
 
-    mock_engine.generate.return_value = stream_gen()
+    mock_engine = loaded_model
-    monkeypatch.setattr("astrai.inference.server._engine", mock_engine)
+    mock_engine.generate_async.return_value = async_gen()
+    monkeypatch.setattr("astrai.inference.server._state.engine", mock_engine)
     response = client.post(
         "/v1/chat/completions",
         json={
             "messages": [{"role": "user", "content": "Hello"}],
             "temperature": 0.8,
-            "top_p": 0.95,
-            "top_k": 50,
             "max_tokens": 100,
             "stream": True,
         },
         headers={"Accept": "text/event-stream"},
     )
     assert response.status_code == 200
-    assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
-    # Parse SSE lines
     lines = [
         line.strip() for line in response.content.decode("utf-8").split("\n") if line
     ]
-    # Should contain data lines and a final [DONE]
     assert any("cumulative1" in line for line in lines)
     assert any("cumulative2" in line for line in lines)
+    assert any("[DONE]" in line for line in lines)
 
 
-def test_generate_with_history(client, loaded_model, mock_engine, monkeypatch):
+def test_messages_non_stream(client, loaded_model, monkeypatch):
-    """POST /generate with history parameter."""
+    """POST /v1/messages with stream=false returns Anthropic-style JSON."""
-    monkeypatch.setattr("astrai.inference.server._engine", mock_engine)
+
+    async def async_gen():
+        yield "Assistant reply"
+
+    mock_engine = loaded_model
+    mock_engine.generate_async.return_value = async_gen()
+    monkeypatch.setattr("astrai.inference.server._state.engine", mock_engine)
     response = client.post(
-        "/generate",
+        "/v1/messages",
-        params={
+        json={
-            "query": "Hi",
+            "messages": [{"role": "user", "content": "Hello"}],
-            "history": [["user1", "assistant1"], ["user2", "assistant2"]],
+            "temperature": 0.8,
+            "max_tokens": 100,
             "stream": False,
         },
     )
     assert response.status_code == 200
-    # Verify the engine.generate was called
+    data = response.json()
-    mock_engine.generate.assert_called_once()
+    assert data["type"] == "message"
+    assert data["role"] == "assistant"
+    assert len(data["content"]) == 1
+    assert data["content"][0]["type"] == "text"
+    assert "usage" in data
+    assert "input_tokens" in data["usage"]
+
+
+def test_messages_stream(client, loaded_model, monkeypatch):
+    """POST /v1/messages with stream=true returns Anthropic SSE stream."""
+
+    async def async_gen():
+        yield "cumulative1"
+        yield "cumulative2"
+
+    mock_engine = loaded_model
+    mock_engine.generate_async.return_value = async_gen()
+    monkeypatch.setattr("astrai.inference.server._state.engine", mock_engine)
+    response = client.post(
+        "/v1/messages",
+        json={
+            "messages": [{"role": "user", "content": "Hello"}],
+            "temperature": 0.8,
+            "max_tokens": 100,
+            "stream": True,
+        },
+        headers={"Accept": "text/event-stream"},
+    )
+    assert response.status_code == 200
+    content = response.content.decode("utf-8")
+    assert "message_start" in content
+    assert "content_block_start" in content
+    assert "content_block_delta" in content
+    assert "cumulative1" in content
+    assert "cumulative2" in content
+    assert "content_block_stop" in content
+    assert "message_delta" in content
+    assert "message_stop" in content
+
+
+def test_messages_with_system(client, loaded_model, monkeypatch):
+    """POST /v1/messages with system prompt."""
+
+    async def async_gen():
+        yield "Reply"
+
+    mock_engine = loaded_model
+    mock_engine.generate_async.return_value = async_gen()
+    monkeypatch.setattr("astrai.inference.server._state.engine", mock_engine)
+    response = client.post(
+        "/v1/messages",
+        json={
+            "messages": [{"role": "user", "content": "Hello"}],
+            "system": "You are a helpful assistant.",
+            "max_tokens": 100,
+            "stream": False,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["type"] == "message"
 
 
 if __name__ == "__main__":

tests/trainer/test_train_strategy.py

@@ -72,6 +72,7 @@ def test_schedule_factory_random_configs():
 
             # Test scheduler step functionality
             initial_lr = scheduler.get_last_lr()
+            optimizer.step()
             scheduler.step()
             new_lr = scheduler.get_last_lr()
 
@@ -112,6 +113,7 @@ def test_schedule_factory_edge_cases():
 
         # Test multiple steps
         for _ in range(10):
+            optimizer.step()
             scheduler.step()
 
 
@@ -136,6 +138,7 @@ def test_schedule_factory_state_persistence():
 
     # Take a few steps
     for _ in range(5):
+        optimizer.step()
         scheduler.step()
 
     # Save state