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: 修正多处文档错误、补充训练参数说明

.gitignore

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

README.md

@@ -46,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
 
@@ -67,44 +68,28 @@ 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/model \
+    --param_path /path/to/model \
-  --n_epoch=3 \
+    --batch_size 4 \
-  --batch_size=4 \
+    --accumulation_steps 8 \
-  --accumulation_steps=8 \
+    --max_lr 3e-4 \
-  --max_lr=3e-4 \
+    --warmup_steps 1000 \
-  --warmup_steps=2000 \
+    --n_epoch 1
-  --ckpt_interval=5000 \
-  --ckpt_dir=./checkpoints
 ```
 
+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
 ```
 
-#### Training Parameters
-
-| Parameter | Description | Default |
-|-----------|-------------|---------|
-| `--train_type` | Training type (`seq`, `sft`, `dpo`) | required |
-| `--data_root_path` | Dataset root directory | required |
-| `--param_path` | Model / checkpoint path | required |
-| `--n_epoch` | Training epochs | 1 |
-| `--batch_size` | Batch size | 1 |
-| `--accumulation_steps` | Gradient accumulation steps | 1 |
-| `--max_lr` | Peak learning rate (cosine decay) | 3e-4 |
-| `--warmup_steps` | LR warmup steps | 1000 |
-| `--ckpt_interval` | Checkpoint interval (iters) | 5000 |
-| `--ckpt_dir` | Checkpoint directory | checkpoint |
-| `--num_workers` | DataLoader workers | 4 |
-| `--nprocs` | Number of GPUs | 1 |
-
-Full reference at [Parameter Guide](./assets/docs/params.md#training-parameters).
-
 #### Docker
 
 Build and run with Docker (recommended for GPU environments):
@@ -125,13 +110,19 @@ docker run --gpus all -p 8000:8000 astrai:latest \
 
 # 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-compatible HTTP API:
+Start the inference server with OpenAI and Anthropic-compatible HTTP API:
 
 ```bash
 python -m scripts.tools.server --port 8000 --device cuda
@@ -140,7 +131,7 @@ python -m scripts.tools.server --port 8000 --device cuda
 Make requests:
 
 ```bash
-# Chat API (OpenAI compatible)
+# OpenAI-compatible
 curl -X POST http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
@@ -148,7 +139,7 @@ curl -X POST http://localhost:8000/v1/chat/completions \
     "max_tokens": 512
   }'
 
-# Streaming response
+# OpenAI-compatible streaming
 curl -X POST http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
@@ -157,6 +148,27 @@ curl -X POST http://localhost:8000/v1/chat/completions \
     "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
 ```

assets/docs/README-zh-CN.md

@@ -52,7 +52,8 @@
 - 💡 **易用**: 简洁的 API 与丰富的示例、演示。
 - 📦 **轻量**: 依赖少，部署简单。
 - 🔬 **研究友好**: 模块化设计，便于实验新想法。
-- 🤗 **HuggingFace 集成**: 兼容 HuggingFace 模型与数据集。
+- 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
+- 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
 
 ### 快速开始
 
@@ -73,44 +74,28 @@ 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/model \
+    --param_path /path/to/model \
-  --n_epoch=3 \
+    --batch_size 4 \
-  --batch_size=4 \
+    --accumulation_steps 8 \
-  --accumulation_steps=8 \
+    --max_lr 3e-4 \
-  --max_lr=3e-4 \
+    --warmup_steps 1000 \
-  --warmup_steps=2000 \
+    --n_epoch 1
-  --ckpt_interval=5000 \
-  --ckpt_dir=./checkpoints
 ```
 
+完整参数列表见[参数说明](./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
 ```
 
-#### 训练参数
-
-| 参数 | 说明 | 默认值 |
-|------|------|--------|
-| `--train_type` | 训练类型（`seq`, `sft`, `dpo`） | 必填 |
-| `--data_root_path` | 数据集根目录 | 必填 |
-| `--param_path` | 模型参数或断点路径 | 必填 |
-| `--n_epoch` | 训练轮数 | 1 |
-| `--batch_size` | 批次大小 | 1 |
-| `--accumulation_steps` | 梯度累积步数 | 1 |
-| `--max_lr` | 峰值学习率（余弦衰减） | 3e-4 |
-| `--warmup_steps` | 预热步数 | 1000 |
-| `--ckpt_interval` | 检查点间隔（迭代步） | 5000 |
-| `--ckpt_dir` | 检查点保存目录 | checkpoint |
-| `--num_workers` | 数据加载线程数 | 4 |
-| `--nprocs` | GPU 数量 | 1 |
-
-完整参数列表见[参数说明](./params.md#training-parameters)。
-
 #### Docker
 
 使用 Docker 构建和运行（推荐用于 GPU 环境）：
@@ -131,13 +116,19 @@ docker run --gpus all -p 8000:8000 astrai:latest \
 
 # 挂载数据卷
 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 兼容的 HTTP API：
+启动推理服务器，支持 OpenAI 和 Anthropic 兼容的 HTTP API：
 
 ```bash
 python -m scripts.tools.server --port 8000 --device cuda
@@ -146,7 +137,7 @@ python -m scripts.tools.server --port 8000 --device cuda
 发起请求：
 
 ```bash
-# Chat API（OpenAI 兼容）
+# OpenAI 兼容
 curl -X POST http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
@@ -154,7 +145,7 @@ curl -X POST http://localhost:8000/v1/chat/completions \
     "max_tokens": 512
   }'
 
-# 流式响应
+# OpenAI 兼容流式
 curl -X POST http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
@@ -163,6 +154,27 @@ curl -X POST http://localhost:8000/v1/chat/completions \
     "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

@@ -9,13 +9,11 @@ AstrAI adopts a modular design with the following main components:
 - **Model Module** (`astrai/model/`): AutoModel, Transformer model and its submodules
 - **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
 
 ```mermaid
@@ -23,38 +21,36 @@ flowchart LR
     subgraph A[Data Preparation]
         direction TB
         A1[Raw Text] --> A2[AutoTokenizer]
-        A2 --> A3[Serialize to .h5 files]
+        A2 --> A3[Tokenized .h5 files]
         A3 --> A4[BaseDataset]
         A4 --> A5[ResumableDistributedSampler]
-        A5 --> A6[PyTorch DataLoader]
+        A5 --> A6[DataLoader]
     end
 
     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]
+        C5 --> C6[Transformer Forward]
         C6 --> C7[sample]
-        C7 --> C8[Transformer Forward]
+        C7 --> C8{End?}
-        C8 --> C9[Paged KV Cache]
+        C8 -->|No| C6
-        C9 --> C10{End Condition?}
+        C8 -->|Yes| C9[Generated Text]
-        C10 -->|No| C8
-        C10 -->|Yes| C11[Output Text]
     end
 
     A --> B
@@ -65,215 +61,177 @@ flowchart LR
 
 ### 1. Serialization (`astrai/serialization.py`)
 
-- **`save_h5`**: Saves multiple tensors by groups as HDF5 files (`.h5`), each key corresponds to a list of tensors
+- **`save_h5`**: Saves tensors by groups as HDF5 files (`.h5`), each key maps to a list of tensors
-- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory (`share_memory=True`)
+- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory
-- **`Checkpoint` class**: Encapsulates model state dict, training epoch, iteration count; supports safetensors format for saving and loading
+- **`Checkpoint`**: Encapsulates model state dict + epoch + iteration; uses safetensors
 
 ### 2. Dataset Module
 
 #### 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)
 
 #### 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
 
 ### 3. Model Module
 
-#### 3.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`
 
 #### 3.2 Submodules (`module.py`)
-- **`RotaryEmbedding`**: Generates RoPE cos/sin cache
+- **`DecoderBlock`**: GQA attention + residual + MLP + RMSNorm
-- **`DecoderBlock`**: Contains multi-head attention (supports GQA and MLA), feedforward network (FFN), residual connections
+- **`GQA`**: Grouped Query Attention (also `MLA` for multi-latent attention)
-- **`GQA`**: Grouped Query Attention implementation
+- **`MLP`**: `SiLU(gate(x)) * up(x)` → down projection
-- **`MLA`**: Multi-Latent Attention implementation (like Qwen2-VL)
+- **`RotaryEmbedding`**: RoPE cos/sin cache
-- **`MLP`**: Feed-forward network with SiLU activation and gated mechanism
+- **`RMSNorm`**: Layer normalization
-- **`RMSNorm`**: Layer normalization variant
-- **`Linear`**, **`Embedding`**: Custom linear layer and embedding layer, supporting parallelism wrappers
 
 ### 4. Training Module
 
 #### 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
 
 #### 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`)
+The training loop is nested: **epoch** → **batch** (with step phase interspersed):
-- 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`
+```
+on_train_begin
+  on_epoch_begin
+    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
+
+    on_epoch_end
+on_train_end
+```
+
+Key points:
+- `on_step_*` wraps optimizer step (fires every `accumulation_steps` batches)
+- `on_batch_*` wraps loss computation (fires every batch)
+- `SchedulerCallback` fires on `on_batch_end` — LR scheduler steps every batch
+- `GradientClippingCallback` fires on `on_step_begin`
 
 #### 4.3 Strategy (`strategy.py`)
-- **`BaseStrategy`**: Defines training strategy interface
+- **`SEQStrategy`**: Next-token prediction, cross-entropy with label smoothing
-- **`SEQStrategy`**: Standard next-token prediction training
+- **`SFTStrategy`**: Supervised fine-tuning with loss masking
-- **`SFTStrategy`**: Supervised Fine-tuning with loss masking
+- **`DPOStrategy`**: Direct Preference Optimization with reference model
-- **`DPOStrategy`**: Direct Preference Optimization
+- **`GRPOStrategy`**: Group Relative Policy Optimization with clipped ratio
-- **`GRPOStrategy`**: Group Relative Policy Optimization
-- Strategy receives batch data, executes model forward pass, loss calculation, returns loss tensor
-- Created dynamically by `StrategyFactory` according to configuration
 
 #### 4.4 Scheduler (`schedule.py`)
-- **`BaseScheduler`**: Abstract base class defining learning rate scheduling interface
+- **`CosineScheduler`**: Cosine decay + linear warmup
-- **`CosineScheduler`**: Cosine decay scheduler with warmup
+- **`SGDRScheduler`**: Cosine annealing with warm restarts
-- **`SGDRScheduler`**: Stochastic Gradient Descent with Warm Restarts
+- Created by `SchedulerFactory` and bound to optimizer
-- **`SchedulerFactory`**: Factory pattern, supports registration of various schedulers
-- Scheduler is automatically created according to configuration and bound to optimizer
 
-#### 4.5 Callbacks (`train_callback.py`)
+#### 4.5 Callbacks
-- **`TrainCallback`**: Protocol interface for trainer callbacks
+- **`CheckpointCallback`**: Saves safetensors at `ckpt_interval` iterations
-- **`CheckpointCallback`**: Saves model checkpoints at configurable intervals
+- **`ProgressBarCallback`**: tqdm progress display
-- **`ProgressBarCallback`**: Displays training progress
+- **`MetricLoggerCallback`**: Writes JSONL metrics to `{ckpt_dir}/logs/`
-- **`MetricLoggerCallback`**: Logs training metrics to JSON files
+- **`GradientClippingCallback`**: `clip_grad_norm_` on `on_step_begin`
-- **`GradientClippingCallback`**: Clips gradient norms
+- **`SchedulerCallback`**: `scheduler.step()` on `on_batch_end`
-- **`SchedulerCallback`**: Steps learning rate scheduler
 
-#### 4.6 Metric Utility (`metric_util.py`)
+### 5. Inference Module
-- **`MetricTracker`**: Tracks and aggregates training metrics across epochs
-- **`get_learning_rate`**: Utility to extract current learning rates from optimizer param groups
 
-### 5. Factory Module
+#### 5.1 Inference Engine (`engine.py`)
+- **`InferenceEngine`**: Facade over scheduler; provides `generate()`, `generate_with_request()`, `generate_async()`
+- Accepts `prompt: str | List[str]`, returns generator (stream) or string (non-stream)
 
-#### 5.1 Registry and BaseFactory (`factory.py`)
+#### 5.2 Scheduler 4-Phase Loop (`scheduler.py`)
-- **`Registry`**: Flexible registry for component classes with category and priority support
-- **`BaseFactory`**: Generic factory class for component registration and creation
-- Supports decorator-based registration pattern for extensible components
-- Provides methods for registration, retrieval, and listing with filtering
 
-### 6. Parallel Module
+Background thread runs continuously:
 
-#### 6.1 Setup (`setup.py`)
+```
-- **`spawn_parallel_fn`**: Spawns multiple processes for distributed training using PyTorch multiprocessing
+1. Cleanup → Remove finished tasks, free KV cache pages
-- **`setup_parallel`**: Context manager for initializing distributed process group (NCCL/CCL backend)
+2. Refill  → Pop from waiting_queue, alloc pages, add to active
-- **`only_on_rank`**: Decorator to execute functions only on specific ranks
+3. Prefill → Group active tasks by prompt_len, run full forward pass
-- **`get_rank`**: Returns current process rank in distributed group
+4. Decode  → Pick largest same-position group, run single-token forward
-- **`get_world_size`**: Returns total number of processes in distributed group
+```
-- **`get_current_device`**: Returns current device from environment
 
-#### 6.2 Parallel Layers (`module.py`)
+- **`Task`**: Tracks prompt_ids, output_ids, page_table, status (PENDING/RUNNING/FINISHED/ABORTED)
-- **`ParallelModel`**: Base class for parallel models with process group
+- **`PagedCache`**: Bitmask-based page allocator with page-table-indirected read/write
-- **`ColumnParallelLinear`**: Column-parallel linear layer with input splitting and output gathering
+- **`CacheView`**: Batch view bundling cache + page table for attention layers
-- **`RowParallelLinear`**: Row-parallel linear layer with output reduction
+- **`sample()`**: Temperature → top-k → top-p → multinomial
 
-### 7. Inference Module
+#### 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`
 
-#### 7.1 Inference Engine (`engine.py`)
+### 6. Tokenizer Module
-- **`InferenceEngine`**: Unified inference interface, supports streaming, async streaming, and non-streaming generation
-- **`InferenceScheduler`**: Continuous batching scheduler with paged KV cache
-- **`GenerationRequest`**: Encapsulates generation parameters (top_k, top_p, temperature, max_len, messages, etc.)
-- **`GenerationParams`**: Immutable value object for sampling hyperparameters
-- **`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`), async streaming (`generate_async`), 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
 
-#### 7.2 Cache (`cache.py`)
+- **`AutoTokenizer`**: Wraps HuggingFace tokenizers (BBPE); `encode`/`decode`/`apply_chat_template`
-- **`PagedCache`**: Page-based KV cache with page-table-indirected read/write; uses bitmask for O(1) page allocation/deallocation
+- **`ChatTemplate`**: Jinja2-based template rendering for multi-turn chat
-- **`CacheView`**: Per-batch view bundling a `PagedCache` with its page table for attention layer access
 
-#### 7.3 Scheduler (`scheduler.py`)
+### 7. Factory & Parallel
-- **`Task`**: Individual generation task with state management (PENDING, RUNNING, FINISHED, ABORTED)
-- **`TaskStatus`**: Task state enumeration
-- **`sample`** (from `sampling.py`): Applies temperature, top-k, top-p sampling to logits via composable `SamplingPipeline`
-- Uses `PagedCache` for paged KV cache management with page table indirection
-- Continuous batching: new requests can join at any time, completed requests release pages immediately
 
-#### 7.4 Server (`server.py`)
+- **`Registry` / `BaseFactory`**: Decorator-based component registration
-- FastAPI-based HTTP inference server
+- **`spawn_parallel_fn`**: Multi-process DDP launcher with NCCL backend
-- OpenAI-compatible `/v1/chat/completions` endpoint
+- **`ParallelModel` / `ColumnParallelLinear` / `RowParallelLinear`**: Tensor model parallelism
-- Health check and statistics endpoints
-- Supports both streaming and non-streaming responses
 
-### 8. Tokenizer Module
+## Training Data Flow — Detailed Steps
-
-#### 8.1 Tokenizer (`tokenizer.py`)
-- Implemented based on HuggingFace tokenizers library (Byte-Level BPE)
-- **`AutoTokenizer`**: Auto-loading tokenizer class
-- Supports special tokens: `<｜begin▁of▁sentence｜>`, `<｜end▁of▁sentence｜>`, `<｜▁pad▁｜>`, `<｜im▁start｜>`, `<｜im▁end｜>`
-- Provides `encode`/`decode` methods for mutual conversion between text and token IDs
-- Uses `AutoTokenizer` for loading pre-trained tokenizers
-
-#### 8.2 Chat Template (`chat_template.py`)
-- **`ChatTemplate`**: Jinja2-based chat template with rendering support
-- Handles multi-role message formatting (system, user, assistant)
-- Supports dynamic prompts and generation prompts
-
-## Training Data Flow - Detailed Steps
 
 1. **Data Preparation**
-   - Raw text is converted to token ID sequences through AutoTokenizer
+   - 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
-   - PyTorch `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**
-   - Scheduler allocates pages via `PagedCache.alloc_n()` for each task's prompt
+   - **Cleanup**: Finished tasks → `stream_callback(STOP)`, free KV pages
-   - Prefill phase: runs full prompt through model with `PagedCache.bind()` to fill initial KV cache pages
+   - **Refill**: Pop from waiting queue, `PagedCache.alloc_n()` for prompt pages
-   - Decode phase: loops until generating `max_len` tokens or encountering stop token:
+   - **Prefill**: Group by prompt length, run full forward with `start_pos=0`
-     - Input last token ID to model, obtain `logits`
+   - **Decode**: Pick position group with most tasks, single-token forward:
-     - Apply `sample()` (temperature, top-k, top-p) to `logits`
+     - Model forward → `logits` → `sample()` → next token ID
-     - Sample next token ID from the processed distribution
+     - Append to `output_ids`, update `output_tokens`
-     - Write new KV entries into paged cache; allocate additional pages as needed
+     - `_maybe_alloc_page()` grows page table as needed
-     - For streaming generation, yield each token to caller immediately via `stream_callback`
+     - `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 paged KV cache, continuous batching, and composable sampling strategies. Clear interfaces between modules facilitate customization and extension.
-
-> Document Update Time: 2026-04-09

assets/docs/design.md

@@ -50,7 +50,6 @@ classDiagram
             +str master_port
             +Callable parallel_wrapper
             +Callable state_dict_fn
-            +List[int] device_ids
             +str device_type
             +dict extra_kwargs
             +validate()
@@ -99,8 +98,8 @@ classDiagram
         }
 
         class ResumableDistributedSampler {
-            +int start_epoch
+            +int epoch
-            +int start_iter
+            +int iter
         }
 
         class DatasetFactory {
@@ -124,7 +123,7 @@ classDiagram
     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
@@ -139,7 +138,7 @@ classDiagram
             +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()
         }
@@ -149,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 {
@@ -158,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 {
@@ -204,7 +205,7 @@ classDiagram
 
     namespace tokenize {
         class AutoTokenizer {
-            +List[str] stop_ids
+            +List[int] stop_ids
             +int bos_id
             +int eos_id
             +int pad_id
@@ -220,7 +221,7 @@ classDiagram
 
         class ChatTemplate {
             +String template_str
-            +render(messages, add_generation_prompt) str
+            +render(messages, system_prompt, **extra_variables) str
             +from_string(template) ChatTemplate
         }
     }
@@ -267,8 +268,6 @@ classDiagram
         class TrainContextBuilder {
             +TrainConfig config
             +with_checkpoint(checkpoint) TrainContextBuilder
-            +with_dataloader() TrainContextBuilder
-            +with_strategy() TrainContextBuilder
             +build() TrainContext
         }
 
@@ -454,7 +453,7 @@ classDiagram
             +float arrival_time
             +float finish_time
             +Callable stream_callback
-            +next_pos() int
+            +int next_pos
             +is_finished(stop_ids) bool
         }
 
@@ -506,15 +505,10 @@ classDiagram
             +sample(logits, filter_value) Tensor
         }
 
-        class Server {
-            +start()
-            +predict(request)
-        }
-
         class _Result {
             +List[str] tokens
             +List[str] results
-            +List[bool] done_flags
+            +List[bool] _done
             +append(token, idx)
             +get_results() List[str]
             +pop_all() List[str]
@@ -539,9 +533,9 @@ classDiagram
     }
 
     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 {
@@ -601,24 +595,19 @@ classDiagram
     BaseSamplingStrategy <|-- TopKStrategy
     BaseSamplingStrategy <|-- TopPStrategy
     SamplingPipeline --> BaseSamplingStrategy : composes
-    Server --> InferenceEngine : uses
-    Server --> ChatMessage : uses
-    Server --> ChatCompletionRequest : uses
-    ParallelSetup --> Trainer : enables
     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
     TrainContextBuilder --> ResumableDistributedSampler : creates
@@ -647,7 +636,7 @@ classDiagram
 | **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, 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
@@ -658,7 +647,7 @@ 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 |
@@ -672,8 +661,8 @@ classDiagram
 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`, uses `PagedCache` for paged KV cache management and `SamplingPipeline` for efficient 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)
@@ -717,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:
 
 $$

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. 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
@@ -222,12 +221,11 @@ curl -X POST http://localhost:8000/v1/chat/completions \
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `messages` | List[dict] | Required | Chat messages with role and content |
-| `temperature` | float | 0.8 | Sampling temperature (0.0-2.0) |
+| `temperature` | float | 1.0 | Sampling temperature (0.0-2.0) |
-| `top_p` | float | 0.95 | Nucleus sampling threshold |
+| `top_p` | float | 1.0 | Nucleus sampling threshold |
 | `top_k` | int | 50 | Top-k sampling parameter |
-| `max_tokens` | int | 2048 | Maximum tokens to generate |
+| `max_tokens` | int | 1024 | Maximum tokens to generate |
 | `stream` | bool | false | Enable streaming response |
-| `system_prompt` | str | None | System prompt override |
 
 **Response (non-streaming):**
 ```json
@@ -242,7 +240,12 @@ curl -X POST http://localhost:8000/v1/chat/completions \
       "message": {"role": "assistant", "content": "Hello! I'm doing well..."},
       "finish_reason": "stop"
     }
-  ]
+  ],
+  "usage": {
+    "prompt_tokens": 20,
+    "completion_tokens": 15,
+    "total_tokens": 35
+  }
 }
 ```
 
@@ -262,25 +265,57 @@ curl -X POST http://localhost:8000/v1/chat/completions \
 
 The server uses Server-Sent Events (SSE) with content type `text/event-stream`.
 
-### Simple Generation Endpoint
+### Anthropic-Compatible Endpoint
 
-For basic text generation without chat format:
+The server also provides an Anthropic-compatible endpoint at `/v1/messages`:
 
 ```bash
-curl -X POST "http://localhost:8000/generate?query=Hello&max_len=1000" \
+curl -X POST http://localhost:8000/v1/messages \
-  -H "Content-Type: application/json"
-```
-
-Or with conversation history:
-
-```bash
-curl -X POST "http://localhost:8000/generate" \
   -H "Content-Type: application/json" \
   -d '{
-    "query": "What is AI?",
+    "model": "astrai",
-    "history": [["Hello", "Hi there!"], ["How are you?", "I'm doing well"]],
+    "system": "You are a helpful assistant.",
-    "temperature": 0.8,
+    "messages": [{"role": "user", "content": "Hello, how are you?"}],
-    "max_len": 2048
+    "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"]
   }'
 ```
 
@@ -290,10 +325,10 @@ Monitor server and model status:
 
 ```bash
 curl http://localhost:8000/health
-# {"status": "ok", "model_loaded": true, "engine_ready": true}
+# {"status": "ok", "model_loaded": true}
 
 curl http://localhost:8000/stats
-# {"requests_total": 10, "tokens_generated": 5000, ...}
+# {"total_tasks": 10, "total_tokens": 5000, "active_tasks": 1, "waiting_queue": 0}
 ```
 
 > Document Update Time: 2026-04-09

assets/docs/params.md

@@ -6,7 +6,7 @@
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--train_type` | Training type (`seq`, `sft`, `dpo`) | required |
+| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`) | required |
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
 | `--n_epoch` | Total training epochs | 1 |
@@ -61,6 +61,10 @@
 |-----------|-------------|---------|---------|
 | `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
 | `--label_smoothing` | Label smoothing for cross-entropy loss | 0.1 | `seq`, `sft` |
+| `--group_size` | GRPO group size | 4 | `grpo` |
+| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
+| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
+| `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` |
 
 ### Usage Example
 

astrai/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.3"
+__version__ = "1.3.4"
 __author__ = "ViperEkura"
 
 from astrai.config import (

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/inference/__init__.py

@@ -3,7 +3,7 @@
 Layers:
   - engine.py:    Facade (InferenceEngine), Value Object (GenerationParams, GenerationRequest)
   - scheduler.py: Continuous-batching loop, Task state machine, TaskStatus enum
-  - cache.py:     Object Pool (SlotAllocator), PrefixCacheManager
+  - 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)
 """

astrai/inference/cache.py

@@ -4,7 +4,7 @@ Provides:
   - PagedCache: paged KV cache combining page pool and tensor storage.
 """
 
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 
 import torch
 from torch import Tensor
@@ -12,12 +12,22 @@ 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.
     """
@@ -45,6 +55,32 @@ class PagedCache:
             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
@@ -68,6 +104,9 @@ class PagedCache:
         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)

astrai/inference/engine.py

@@ -97,7 +97,8 @@ class _Result:
     """Thread-safe token accumulator for streaming and non-streaming modes.
 
     Supports multiple concurrent generation tasks with per-index result tracking.
-    Uses a threading.Event for efficient waiting on completion.
+    Uses a threading.Condition for efficient completion notification
+    and a threading.Event for streaming wakeup.
     """
 
     def __init__(self, count: int = 1):
@@ -106,7 +107,7 @@ class _Result:
         Args:
             count: Number of concurrent generation tasks to track.
         """
-        self._lock = threading.Lock()
+        self._cond = threading.Condition()
         self._event = threading.Event()
         self.tokens: List[str] = []
         self.results: List[str] = [""] * count
@@ -124,7 +125,7 @@ class _Result:
             token: The decoded token string, or STOP sentinel.
             idx: Index of the generation task this token belongs to.
         """
-        with self._lock:
+        with self._cond:
             self.tokens.append(token)
             if token is not STOP:
                 self.results[idx] += token
@@ -132,6 +133,7 @@ class _Result:
                 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]:
@@ -140,7 +142,7 @@ class _Result:
         Returns:
             List of token strings since the last call.
         """
-        with self._lock:
+        with self._cond:
             out = self.tokens.copy()
             self.tokens.clear()
             if not out:
@@ -158,13 +160,22 @@ class _Result:
         """
         return self._event.wait(timeout=timeout)
 
+    def wait_completion(self) -> None:
+        """Blocks until all tasks complete (non-streaming).
+
+        Uses a Condition to sleep efficiently instead of busy-waiting.
+        The calling thread is parked until a STOP signal arrives.
+        """
+        with self._cond:
+            self._cond.wait_for(lambda: self._completed >= self._total)
+
     def get_results(self) -> List[str]:
         """Returns all accumulated results for non-streaming mode.
 
         Returns:
             List of complete generated strings, one per task index.
         """
-        with self._lock:
+        with self._cond:
             return self.results.copy()
 
 
@@ -408,13 +419,14 @@ class InferenceEngine:
             Single string for one prompt, list of strings for batch.
         """
         result = _Result(count=len(prompts))
+        task_ids = []
 
         for i, p in enumerate(prompts):
 
             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,
@@ -422,8 +434,13 @@ class InferenceEngine:
                 top_k=top_k,
                 stream_callback=make_cb(i),
             )
+            task_ids.append(task_id)
+
+        result.wait_completion()
+
+        for task_id in task_ids:
+            self.scheduler.remove_task(task_id)
 
-        result.wait()
         res = result.get_results()
         return res if is_batch else res[0]
 

astrai/inference/scheduler.py

@@ -5,7 +5,7 @@ import threading
 import time
 import uuid
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -53,9 +53,11 @@ class Task:
         self.output_tokens: int = 0
         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:
@@ -87,8 +89,8 @@ class InferenceScheduler:
         max_seq_len: Optional[int] = None,
         max_prompt_len: int = 512,
         page_size: int = 64,
-        device: str = "cuda",
+        device: Optional[str] = None,
-        dtype: torch.dtype = torch.bfloat16,
+        dtype: Optional[torch.dtype] = None,
     ):
         config = model.config
 
@@ -104,7 +106,9 @@ class InferenceScheduler:
         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 - 1) // page_size
+        n_pages = (
+            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
+        ) // page_size
 
         self.page_cache = PagedCache(
             n_layers,
@@ -167,14 +171,21 @@ class InferenceScheduler:
             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:
         finished = []
         for task in self.active_tasks:
@@ -185,9 +196,11 @@ class InferenceScheduler:
                 self._total_tokens += task.output_tokens
 
         for task in finished:
+            if not task._pages_freed:
                 self._free_pages(task.page_table)
                 task.page_table.clear()
                 task.n_pages = 0
+                task._pages_freed = True
 
         self.active_tasks = [
             t for t in self.active_tasks if t.status != TaskStatus.FINISHED
@@ -207,12 +220,25 @@ class InferenceScheduler:
         failed: List[Task] = []
         for task in to_add:
             prompt_len = len(task.prompt_ids)
-            n_pages = self._n_pages_for(prompt_len)
+
-            task.page_table = self.page_cache.alloc_n(n_pages)
+            hit_pages = self.page_cache.lookup_prefix(task.prompt_ids)
-            if not task.page_table:
+            cached_tokens = len(hit_pages) * self.page_size
+            for p in hit_pages:
+                self.page_cache.inc_ref(p)
+
+            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 []
+
+            if remaining > 0 and not new_pages:
+                for p in hit_pages:
+                    self.page_cache.free(p)
                 failed.append(task)
                 continue
+
+            task.page_table = hit_pages + new_pages
             task.n_pages = len(task.page_table)
+            task._prefix_cached_tokens = cached_tokens
             task.status = TaskStatus.RUNNING
             self.active_tasks.append(task)
 
@@ -220,42 +246,20 @@ class InferenceScheduler:
             with self._lock:
                 self.waiting_queue[:0] = failed
 
-    def _execute_prefill(self) -> None:
+    def _execute_prefill(
-        to_prefill = [t for t in self.active_tasks if t.output_tokens == 0]
+        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
-        if not to_prefill:
+    ) -> None:
-            return
-
-        for t in to_prefill:
-            prompt_len = len(t.prompt_ids)
-            t.input_tokens = prompt_len
-            t.output_tokens = 0
-
-        groups: Dict[int, List[Task]] = {}
-        for t in to_prefill:
-            groups.setdefault(len(t.prompt_ids), []).append(t)
-
-        for prompt_len, group in groups.items():
-            self._execute_prefill_batch(group, prompt_len)
-
-    def _execute_prefill_batch(self, tasks: List[Task], prompt_len: int) -> None:
         tasks = sorted(tasks, key=lambda t: t.task_id)
         batch_sz = len(tasks)
 
-        input_ids = torch.zeros(
+        seq_len = prompt_len - start_pos
-            batch_sz,
+        input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
-            prompt_len,
+        input_mask = torch.ones(batch_sz, seq_len, dtype=torch.bool, device=self.device)
-            dtype=torch.long,
-            device=self.device,
-        )
-        input_mask = torch.ones(
-            batch_sz,
-            prompt_len,
-            dtype=torch.bool,
-            device=self.device,
-        )
 
         for i, t in enumerate(tasks):
-            input_ids[i] = torch.tensor(t.prompt_ids, device=self.device)
+            input_ids[i] = torch.tensor(
+                t.prompt_ids[start_pos:prompt_len], device=self.device
+            )
 
         page_tables = self._make_page_table_tensor(tasks)
 
@@ -263,10 +267,14 @@ class InferenceScheduler:
             self.model(
                 input_ids,
                 input_mask=input_mask,
-                start_pos=0,
+                start_pos=start_pos,
                 paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
             )
 
+        start_logical_page = start_pos // self.page_size
+        for t in tasks:
+            self._record_page_hashes(t, start_logical_page=start_logical_page)
+
     def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
         if not tasks:
             return
@@ -274,15 +282,24 @@ class InferenceScheduler:
         tasks = sorted(tasks, key=lambda t: t.task_id)
         batch_sz = len(tasks)
 
-        input_ids = torch.zeros(batch_sz, dtype=torch.long, device=self.device)
+        for t in tasks:
-        for i, t in enumerate(tasks):
+            self._maybe_alloc_page(t, start_pos)
-            input_ids[i] = t.output_ids[-1] if t.output_ids else t.prompt_ids[-1]
+
+        input_ids = torch.tensor(
+            [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_ids.unsqueeze(1),
@@ -294,11 +311,9 @@ class InferenceScheduler:
 
         next_tokens = sample(
             logits,
-            temperature=torch.tensor(
+            temperature=temperatures,
-                [t.temperature for t in tasks], device=logits.device
+            top_k=top_ks,
-            ),
+            top_p=top_ps,
-            top_k=torch.tensor([t.top_k for t in tasks], device=logits.device),
-            top_p=torch.tensor([t.top_p for t in tasks], device=logits.device),
         ).tolist()
 
         for t, ntok in zip(tasks, next_tokens):
@@ -339,7 +354,19 @@ class InferenceScheduler:
                     self._task_event.wait(timeout=1.0)
                     continue
 
-                self._execute_prefill()
+                to_prefill = [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)
+
+                    groups: Dict[Tuple[int, int], List[Task]] = {}
+                    for t in to_prefill:
+                        key = (len(t.prompt_ids), t._prefix_cached_tokens)
+                        groups.setdefault(key, []).append(t)
+
+                    for (prompt_len, start_pos), group in groups.items():
+                        if start_pos < prompt_len:
+                            self._execute_prefill(group, prompt_len, start_pos)
 
                 pos_groups: Dict[int, List[Task]] = {}
                 for t in self.active_tasks:

astrai/inference/server.py

@@ -1,5 +1,5 @@
 """
-OpenAI-compatible chat completion server backed by continuous-batching inference.
+OpenAI / Anthropic-compatible chat completion server backed by continuous-batching inference.
 """
 
 import json
@@ -51,6 +51,7 @@ class ChatCompletionRequest(BaseModel):
     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)
@@ -61,6 +62,25 @@ class ChatCompletionRequest(BaseModel):
     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(
     device: str = "cuda",
     dtype: torch.dtype = torch.bfloat16,
@@ -185,7 +205,7 @@ async def chat_completion(request: ChatCompletionRequest):
             max_tokens=request.max_tokens,
             temperature=request.temperature,
             top_p=request.top_p,
-            top_k=50,
+            top_k=request.top_k,
         )
 
         async def event_stream():
@@ -237,7 +257,7 @@ async def chat_completion(request: ChatCompletionRequest):
         max_tokens=request.max_tokens,
         temperature=request.temperature,
         top_p=request.top_p,
-        top_k=50,
+        top_k=request.top_k,
     )
     async for token in agen:
         chunks.append(token)
@@ -264,55 +284,183 @@ async def chat_completion(request: ChatCompletionRequest):
     }
 
 
-@app.post("/generate")
+def _make_anthropic_sse(event: str, data: Dict[str, Any]) -> str:
-async def generate(
+    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
-    query: str,
+
-    history: Optional[List[List[str]]] = None,
+
-    temperature: float = 0.8,
+def _check_stop_sequence(text: str, stop_sequences: List[str]) -> Optional[str]:
-    top_p: float = 0.95,
+    for seq in stop_sequences:
-    top_k: int = 50,
+        if seq and seq in text:
-    max_len: int = 2048,
+            return seq
-    stream: bool = False,
+    return None
-):
+
-    """Legacy non-OpenAI generation endpoint (kept for backward compat)."""
+
+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
 
-    messages = []
+    chat_messages = _build_anthropic_messages(request.messages, request.system)
-    if history:
+    prompt = engine.tokenizer.apply_chat_template(chat_messages, tokenize=False)
-        for h in history:
+    prompt_tokens = len(engine.tokenizer.encode(prompt))
-            if len(h) >= 2:
-                messages.append({"role": "user", "content": h[0]})
-                messages.append({"role": "assistant", "content": h[1]})
-    messages.append({"role": "user", "content": query})
 
-    prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
+    stop_sequences = request.stop_sequences or []
 
-    if stream:
+    if request.stream:
         agen = engine.generate_async(
             prompt=prompt,
-            max_tokens=max_len,
+            max_tokens=request.max_tokens,
-            temperature=temperature,
+            temperature=request.temperature,
-            top_p=top_p,
+            top_p=request.top_p,
-            top_k=top_k,
+            top_k=request.top_k,
         )
 
-        async def text_stream():
+        async def event_stream():
-            async for token in agen:
+            yield _make_anthropic_sse(
-                yield token + "\n"
+                "message_start",
+                {
+                    "type": "message_start",
+                    "message": {
+                        "id": resp_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "model": model,
+                        "content": [],
+                        "usage": {"input_tokens": prompt_tokens},
+                    },
+                },
+            )
 
-        return StreamingResponse(text_stream(), media_type="text/plain")
+            yield _make_anthropic_sse(
-    else:
+                "content_block_start",
-        chunks = []
+                {
-        for token in engine.generate(
+                    "type": "content_block_start",
+                    "index": 0,
+                    "content_block": {"type": "text", "text": ""},
+                },
+            )
+
+            completion_tokens = 0
+            accumulated = ""
+            stopped_seq: Optional[str] = None
+            async for token in agen:
+                accumulated += token
+                completion_tokens += 1
+
+                matched = _check_stop_sequence(accumulated, stop_sequences)
+                if matched:
+                    text = accumulated[: accumulated.rfind(matched)]
+                    stopped_seq = matched
+                    if text:
+                        yield _make_anthropic_sse(
+                            "content_block_delta",
+                            {
+                                "type": "content_block_delta",
+                                "index": 0,
+                                "delta": {"type": "text_delta", "text": text},
+                            },
+                        )
+                    break
+
+                yield _make_anthropic_sse(
+                    "content_block_delta",
+                    {
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": token},
+                    },
+                )
+
+            yield _make_anthropic_sse(
+                "content_block_stop",
+                {"type": "content_block_stop", "index": 0},
+            )
+
+            stop_reason = "stop_sequence" if stopped_seq else "end_turn"
+            yield _make_anthropic_sse(
+                "message_delta",
+                {
+                    "type": "message_delta",
+                    "delta": {"stop_reason": stop_reason, "stop_sequence": stopped_seq},
+                    "usage": {"output_tokens": completion_tokens},
+                },
+            )
+
+            yield _make_anthropic_sse(
+                "message_stop",
+                {"type": "message_stop"},
+            )
+
+        return StreamingResponse(
+            event_stream(),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
+        )
+
+    completion_tokens = 0
+    chunks: List[str] = []
+    agen = engine.generate_async(
         prompt=prompt,
-            stream=True,
+        max_tokens=request.max_tokens,
-            max_tokens=max_len,
+        temperature=request.temperature,
-            temperature=temperature,
+        top_p=request.top_p,
-            top_p=top_p,
+        top_k=request.top_k,
-            top_k=top_k,
+    )
-        ):
+    stopped_seq: Optional[str] = None
+    accumulated = ""
+    async for token in agen:
         chunks.append(token)
-        return {"response": "".join(chunks)}
+        completion_tokens += 1
+        accumulated += token
+        matched = _check_stop_sequence(accumulated, stop_sequences)
+        if matched:
+            stopped_seq = matched
+            break
+
+    content = "".join(chunks)
+    if stopped_seq:
+        idx = content.rfind(stopped_seq)
+        if idx != -1:
+            content = content[:idx]
+
+    return {
+        "id": resp_id,
+        "type": "message",
+        "role": "assistant",
+        "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(

astrai/model/automodel.py

@@ -84,6 +84,7 @@ class AutoModel(nn.Module):
         cls,
         path: Union[str, Path],
         disable_random_init: bool = True,
+        strict: bool = True,
     ) -> nn.Module:
 
         model_path = Path(path)
@@ -106,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/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/tokenizer.py

@@ -64,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)
 

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,9 +32,14 @@ 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._checkpoint: Optional[Checkpoint] = None
+        self._load_extra_fn = load_extra_fn
 
     def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
         self._checkpoint = checkpoint
@@ -66,6 +71,9 @@ class TrainContextBuilder:
         context.optimizer = self.config.optimizer_fn(context.model)
         context.scheduler = self.config.scheduler_fn(context.optimizer)
 
+        if self._checkpoint and self._checkpoint.extra and self._load_extra_fn:
+            self._load_extra_fn(self._checkpoint.extra, context)
+
         cfg = self.config
         sampler_offset = context.iteration * cfg.batch_size
         sampler = ResumableDistributedSampler(

astrai/trainer/trainer.py

@@ -53,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,
         )
 
@@ -68,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()
@@ -83,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/tools/benchmark.py

@@ -6,8 +6,9 @@ from typing import Any, Dict
 import torch
 from torch import Tensor
 
+from astrai.config import ModelConfig
 from astrai.inference.cache import PagedCache
-from astrai.model.transformer import ModelConfig, Transformer
+from astrai.model.transformer import Transformer
 
 
 @dataclass

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,

tests/inference/test_scheduler_concurrency.py

@@ -5,6 +5,7 @@ import time
 from unittest.mock import MagicMock, patch
 
 import pytest
+import torch
 
 from astrai.inference.scheduler import InferenceScheduler
 
@@ -19,6 +20,9 @@ def mock_model_and_tokenizer():
     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]

tests/inference/test_server.py

@@ -1,7 +1,5 @@
 """Unit tests for the inference HTTP server."""
 
-from unittest.mock import MagicMock
-
 import pytest
 
 
@@ -24,52 +22,6 @@ def test_health_with_model(client, loaded_model):
     assert data["model_loaded"] is True
 
 
-def test_generate_non_stream(client, loaded_model, monkeypatch):
-    """POST /generate with stream=false should return JSON response."""
-    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 "response" in data
-
-
-def test_generate_stream(client, loaded_model, monkeypatch):
-    """POST /generate with stream=true should return plain text stream."""
-
-    async def async_gen():
-        yield "chunk1"
-        yield "chunk2"
-
-    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",
-        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
-    content = response.content.decode("utf-8")
-    assert "chunk1" in content
-    assert "chunk2" in content
-
-
 def test_chat_completions_non_stream(client, loaded_model, monkeypatch):
     """POST /v1/chat/completions with stream=false returns OpenAI-style JSON."""
 
@@ -125,17 +77,87 @@ def test_chat_completions_stream(client, loaded_model, monkeypatch):
     assert any("[DONE]" in line for line in lines)
 
 
-def test_generate_with_history(client, loaded_model, 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."""
+
+    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
+    data = response.json()
+    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