feat: 新增LoRA微调模块

- LoRALinear基于register_parameter托管base weight，state_dict路径不变
- inject_lora/merge_lora/save_lora/load_lora完备封装
- 24个单元测试覆盖注入、合并、存取、边界场景

.github/ISSUE_TEMPLATE/bug_report.md

@@ -2,7 +2,7 @@
 name: Bug report
 about: Create a report to help us improve
 title: "[BUG]"
-labels: enhancement
+labels: bug
 assignees: ''
 
 ---

.github/PULL_REQUEST_TEMPLATE.md

@@ -16,9 +16,9 @@ Please delete options that are not relevant.
 Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce.
 
 ## Checklist:
-- [ ] My code follows the style guidelines of this project (run `ruff format .` and `ruff check --fix .`)
+- [ ] My code follows the style guidelines of this project (run `ruff format .` and `ruff check . --select I`)
 - [ ] I have performed a self-review of my own code
-- [ ] I have commented my code, particularly in hard-to-understand areas
+- [ ] Code is self-documenting (no unnecessary comments)
 - [ ] I have made corresponding changes to the documentation
 - [ ] My changes generate no new warnings
 - [ ] I have added tests that prove my fix is effective or that my feature works

.gitignore

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

CONTRIBUTING.md

@@ -1,68 +1,100 @@
 # Contributing to AstrAI
 
-Thank you for your interest in contributing to AstrAI! This document provides guidelines and steps for contributing.
+Thank you for your interest in contributing! This document provides step-by-step guidelines.
 
-## How to Contribute
+## Quick Start
 
-### Reporting Issues
-If you encounter a bug or have a feature request, please open an issue on GitHub. Include as much detail as possible:
-- A clear description of the problem or request.
-- Steps to reproduce (for bugs).
-- Your environment (Python version, OS, etc.).
+```bash
+git clone https://github.com/your-username/AstrAI.git
+cd AstrAI
+pip install -e ".[dev]"     # install with dev dependencies (pytest, ruff)
+```
 
-### Submitting Changes
-1. **Fork** the repository.
-2. **Clone** your fork:
-   ```bash
-   git clone https://github.com/your-username/AstrAI.git
-   cd AstrAI
-   ```
-3. **Create a feature branch**:
-   ```bash
-   git checkout -b feature/your-feature-name
-   ```
-4. **Make your changes**. Follow the code style guidelines below.
-5. **Commit your changes** with a descriptive commit message:
-   ```bash
-   git commit -m "Add: brief description of the change"
-   ```
-6. **Push** to your fork:
-   ```bash
-   git push origin feature/your-feature-name
-   ```
-7. **Open a Pull Request** (PR) against the `main` branch of the upstream repository.
+## Before You Commit
 
-## Code Style
+Run the following checks **in order** — CI will reject if any fail.
 
-AstrAI uses [Ruff](https://docs.astral.sh/ruff/) for code formatting and linting. Please ensure your code is formatted before submitting.
+### 1. Format
 
-- Run Ruff to format and lint:
-  ```bash
-  ruff format .
-  ruff check --fix .
-  ```
-- The project uses **double quotes** for strings and **4‑space indentation** (as configured in `pyproject.toml`).
+```bash
+ruff format .
+```
 
-## Testing
+> **Note**: `ruff format` may rename parameters (e.g. `mask` → `attn_mask`).  
+> Always review the diff after formatting.
 
-If you add or modify functionality, please include appropriate tests.
+### 2. Import sorting
 
-- Run the test suite with:
-  ```bash
-  pytest
-  ```
-- Ensure all tests pass before submitting your PR.
+```bash
+ruff check . --select I
+```
+
+If this fails, **manually fix** import ordering (ruff does not auto-fix in this project's CI):
+
+```bash
+ruff check . --select I --fix .
+ruff format .    # re-format after fix
+```
+
+### 3. Run tests
+
+```bash
+python -u -m pytest tests/ -v
+```
+
+> Failed tests may leave orphan tempdirs under `%TEMP%`. Clean them manually if needed.
+
+### 4. (Optional) Full pre-commit check
+
+If you have Git Bash available:
+
+```bash
+bash scripts/pre_commit.sh
+```
+
+This runs format check, import sort check, and tests in one go.
+
+## Commit Style
+
+```
+fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description (~50 chars)
+
+- bullet point body (each ~60 chars)
+```
+
+- **Type** must be one of: `fix`, `feat`, `chore`, `docs`, `refactor`, `perf`, `test`, `style`, `ci`, `build`, `revert`.
+- **Subject line** ends with no period.
+- **Body** uses bullet points starting with `-`.
+- No `(scope)` parentheses.
+
+## Common Issues
+
+| Problem | Cause | Fix |
+|---------|-------|-----|
+| `ruff check --select I` fails | Wrong import order | `ruff check . --select I --fix .` then `ruff format .` |
+| `ruff format` changed many files | Not formatted before commit | Review diff carefully before staging |
+| Pre-commit hook rejects | Tests or lint failed | Fix individually, do not `--no-verify` |
+| Tests fail with tempdir left | Test crash | Clean `%TEMP%` manually |
+
+## Submitting Changes
+
+1. Fork the repo.
+2. Create a feature branch: `git checkout -b feat/my-feature`
+3. Make changes following the steps above.
+4. Commit with the commit style above.
+5. Push: `git push origin feat/my-feature`
+6. Open a Pull Request against `main`.
 
 ## Code Review
 
-All submissions will be reviewed. We may request changes or discuss alternatives. Please be responsive to feedback.
+- All PRs are reviewed. We may request changes.
+- CI runs `ruff format --check .` then `ruff check . --select I` (no `--fix` in CI).
+- Ensure all tests pass.
 
 ## License
 
-By contributing, you agree that your contributions will be licensed under the same [GPL-3.0 License](LICENSE) that covers the project.
+By contributing, you agree that your contributions will be licensed under the [GPL-3.0 License](LICENSE).
 
 ---
 
-If you have any questions, feel free to ask in the [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions) or open an issue.
-
-Happy contributing!
+Questions? Ask in [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions) or open an issue.

Dockerfile

@@ -1,7 +1,7 @@
 # AstrAI Dockerfile - Multi-stage Build (Optimized)
 
 # Build stage - use base image with minimal build tools
-FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS builder
+FROM ubuntu:24.04 AS builder
 
 WORKDIR /app
 
@@ -18,7 +18,7 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
 RUN python3.12 -m venv --copies /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
 
-# Copy source code and install dependencies
+# Copy source code and install (deps read from pyproject.toml)
 COPY astrai/ ./astrai/
 COPY pyproject.toml .
 RUN pip install --no-cache-dir --upgrade pip \
@@ -26,13 +26,14 @@ RUN pip install --no-cache-dir --upgrade pip \
     --extra-index-url https://download.pytorch.org/whl/cu126
 
 # Production stage
-FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production
+FROM ubuntu:24.04 AS production
 
 WORKDIR /app
 
-# Install Python 3.12 runtime
+# Install Python 3.12 runtime and healthcheck dependency
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     python3.12 \
+    curl \
     && rm -rf /var/lib/apt/lists/*
 
 # Copy virtual environment from builder

README.md

@@ -27,9 +27,6 @@
 
 ## 📖 Table of Contents
 
-<details open>
-<summary><b>English</b></summary>
-
 - [Features](#features)
 - [Quick Start](#quick-start)
 - [Documentation](#documentation)
@@ -37,8 +34,6 @@
 - [Community](#community)
 - [License](#license)
 
-</details>
-
 ---
 
 <a id="english"></a>
@@ -51,7 +46,8 @@
 - 💡 **Easy to Use**: Simple API with comprehensive examples and demos.
 - 📦 **Lightweight**: Minimal dependencies, easy to deploy.
 - 🔬 **Research‑Friendly**: Modular design, easy to experiment with new ideas.
-- 🤗 **HuggingFace Integration**: Compatible with HuggingFace models and datasets.
+- 🤗 **HuggingFace-Style API**: AutoModel/AutoTokenizer APIs inspired by HuggingFace for easy model and tokenizer loading.
+- 🔌 **Dual API Compatibility**: Supports both OpenAI and Anthropic chat completion APIs out of the box.
 
 ### Quick Start
 
@@ -69,19 +65,51 @@ For development dependencies:
 pip install -e ".[dev]"
 ```
 
+#### Download Pre-trained Model
+
+Download pre-trained model weights (1B bilingual checkpoint) to `params/`:
+
+```bash
+python scripts/demo/download.py
+```
+
+Or download manually from [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) into `params/`.
+
 #### Train a Model
 
 ```bash
-python scripts/tools/train.py \
-  --train_type=seq \
-  --data_root_path=/path/to/dataset \
-  --param_path=/path/to/param_path
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
 ```
 
+Full reference at [Parameter Guide](assets/docs/params.md).
+
 #### Generate Text
 
 ```bash
-python scripts/tools/generate.py --param_path=/path/to/param_path
+python scripts/tools/generate.py \
+    --param_path /path/to/model \
+    --input_json_file /path/to/input.json \
+    --output_json_file /path/to/output.json
 ```
 
 #### Docker
@@ -104,13 +132,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
@@ -119,7 +153,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 '{
@@ -127,7 +161,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 '{
@@ -136,6 +170,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
 ```
@@ -158,16 +213,17 @@ python scripts/demo/generate_batch.py
 python scripts/demo/generate_ar.py
 ```
 
-Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd).
+Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).
 
 ### Documentation
 
 | Document | Description |
 |----------|-------------|
 | [Parameter Guide](./assets/docs/params.md) | Training & inference parameters |
-| [Design Document](./assets/docs/design.md) | Framework architecture & module design |
-| [Data Flow](./assets/docs/dataflow.md) | Data processing pipeline details |
-| [Model Introduction](./assets/docs/introduction.md) | Model architecture & technical details |
+| [Architecture](./assets/docs/architecture.md) | System architecture, class diagram & design patterns |
+| [Training](./assets/docs/training.md) | Training loop, strategies & formulas |
+| [Inference](./assets/docs/inference.md) | KVCache, continuous batching, sampling & HTTP API |
+| [Data Flow](./assets/docs/dataflow.md) | Data pipeline, storage backends & dataset architecture |
 
 ### Contributing
 

assets/docs/README-zh-CN.md

@@ -52,7 +52,8 @@
 - 💡 **易用**: 简洁的 API 与丰富的示例、演示。
 - 📦 **轻量**: 依赖少，部署简单。
 - 🔬 **研究友好**: 模块化设计，便于实验新想法。
-- 🤗 **HuggingFace 集成**: 兼容 HuggingFace 模型与数据集。
+- 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
+- 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
 
 ### 快速开始
 
@@ -70,19 +71,51 @@ pip install -e .
 pip install -e ".[dev]"
 ```
 
+#### 下载预训练模型
+
+下载预训练模型权重（1B 双语检查点）到 `params/` 目录：
+
+```bash
+python scripts/demo/download.py
+```
+
+或从 [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) 手动下载放入 `params/`。
+
 #### 训练模型
 
 ```bash
-python scripts/tools/train.py \
-  --train_type=seq \
-  --data_root_path=/path/to/dataset \
-  --param_path=/path/to/param_path
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
 ```
 
+完整参数列表见[参数说明](./params.md)。
+
 #### 文本生成
 
 ```bash
-python scripts/tools/generate.py --param_path=/path/to/param_path
+python scripts/tools/generate.py \
+    --param_path /path/to/model \
+    --input_json_file /path/to/input.json \
+    --output_json_file /path/to/output.json
 ```
 
 #### Docker
@@ -105,13 +138,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
@@ -120,7 +159,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 '{
@@ -128,7 +167,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 '{
@@ -137,6 +176,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
 ```
@@ -159,16 +219,17 @@ python scripts/demo/generate_batch.py
 python scripts/demo/generate_ar.py
 ```
 
-观看 [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd) 上的视频演示。
+观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
 
 ### 文档
 
 | 文档 | 说明 |
 |------|------|
 | [参数说明](./params.md) | 训练与推理参数配置 |
-| [设计文档](./design.md) | 系统架构与模块设计 |
-| [数据流程](./dataflow.md) | 数据处理管道详解 |
-| [模型介绍](./introduction.md) | 模型架构与技术细节 |
+| [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
+| [训练文档](./training.md) | 训练循环、策略与公式 |
+| [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
+| [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 |
 
 ### 贡献
 

assets/docs/dataflow.md

@@ -1,269 +1,57 @@
-# AstrAI Data Flow Documentation
+# Data Flow
 
-This document describes the data flow of the AstrAI project (a training and inference framework for autoregressive Transformer language models). It covers the complete flow from raw data to model training and inference.
+This document describes the data pipeline: from raw text to model input tensors.
 
 ## Overview
 
-AstrAI adopts a modular design with the following main components:
-- **Dataset Module** (`astrai/dataset/`): Dataset, sampler, serialization tools
-- **Model Module** (`astrai/model/`): AutoModel, Transformer model and its submodules
-- **Training Module** (`astrai/trainer/`): Trainer, training context, strategies, schedulers
-- **Inference Module** (`astrai/inference/`): Inference engine with continuous batching, streaming generation
-- **Config Module** (`astrai/config/`): Model, training, scheduler, and other configurations
-- **Factory Module** (`astrai/factory/`): Registry, BaseFactory for component registration
-- **Parallel Module** (`astrai/parallel/`): Distributed training support
-- **Serialization Module** (`astrai/serialization/`): 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
-flowchart LR
-    subgraph A[Data Preparation]
-        direction TB
-        A1[Raw Text] --> A2[AutoTokenizer]
-        A2 --> A3[Serialize to .h5 files]
-        A3 --> A4[BaseDataset]
-        A4 --> A5[ResumableDistributedSampler]
-        A5 --> A6[PyTorch DataLoader]
-    end
-
-    subgraph B[Training]
-        direction TB
-        B1[Batch Data] --> B2[TrainContextBuilder]
-        B2 --> B3[TrainContext]
-        B3 --> B4[BaseStrategy]
-        B4 --> B5[Transformer]
-        B5 --> B6[Compute Loss]
-        B6 --> B7[Backward]
-        B7 --> B8[Optimizer]
-        B8 --> B9[LRScheduler]
-        B9 --> B10[CheckpointCallback]
-    end
-
-    subgraph C[Inference]
-        direction TB
-        C1[Checkpoint] --> C2[AutoModel]
-        C2 --> C3[Transformer + Tokenizer]
-        C3 --> C4[GenerationRequest + apply_chat_template]
-        C4 --> C5[InferenceEngine]
-        C5 --> C6[InferenceScheduler]
-        C6 --> C7[apply_sampling_strategies]
-        C7 --> C8[Transformer Forward]
-        C8 --> C9[KV Cache + Prefix Cache]
-        C9 --> C10{End Condition?}
-        C10 -->|No| C8
-        C10 -->|Yes| C11[Output Text]
-    end
-
-    A --> B
-    B --> C
+```
+Raw Text → AutoTokenizer → Token IDs → .h5/.json → Dataset → Sampler → DataLoader → Training/Inference
 ```
 
-## Detailed Module Descriptions
+## Data Preparation
 
-### 1. Dataset Module
+Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or JSON (`.json`/`.jsonl`) files with keyed tensor groups.
 
-#### 1.1 Serialization (`serialization.py`)
-- **`save_h5`**: Saves multiple tensors by groups as HDF5 files (`.h5`), each key corresponds to a list of tensors
-- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory (`share_memory=True`)
-- **`Checkpoint` class**: Encapsulates model state dict, training epoch, iteration count; supports safetensors format for saving and loading
+Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
 
-#### 1.2 Dataset (`dataset.py`)
-- **`BaseDataset`**: Abstract base class, defines common logic for window sampling, stride, etc.
-- **`BaseSegmentFetcher`** and **`MultiSegmentFetcher`**: Efficiently fetch data from specified index ranges in multiple segments
-- **`DatasetFactory`**: Factory pattern, supports dynamic registration of dataset types (`seq`, `sft`, `dpo`, `grpo`)
-- After dataset loading, multiple data keys (such as `"sequence"`, `"mask"`) are managed through `MultiSegmentFetcher`
+```
+StorageFactory.create("h5")   → H5Storage
+StorageFactory.create("json") → JSONStorage
+```
 
-#### 1.3 Sampler (`sampler.py`)
-- **`ResumableDistributedSampler`**: Resumable sampler supporting distributed training
-- Records current epoch and iteration position, enabling training resume from breakpoints
-- Supports shuffle and drop_last options
+Both support shared memory via `.share_memory_()`.
 
-### 2. Model Module
+## Data Keys by Training Type
 
-#### 2.1 Transformer / AutoModel (`transformer.py`, `automodel.py`)
-- **`AutoModel`**: Base class for autoregressive language models with `from_pretrained()` and `save_pretrained()` methods
-- **`Transformer`**: Core autoregressive decoder architecture (registered via `@AutoModel.register('transformer')`)
-- Contains embedding layer, multi-layer `DecoderBlock`, RMSNorm, and linear output head
-- Supports weight tying (`tie_weight=True`) to reduce parameter count
-- Uses Rotary Position Embedding (RoPE) to inject position information
-- Supports loading from safetensors format with automatic model type detection from `config.json`
+| Type | Storage Keys |
+|------|-------------|
+| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) |
+| `sft` | `sequence`, `loss_mask` |
+| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` |
+| `grpo` | `prompts`, `responses`, `masks`, `rewards` |
 
-#### 2.2 Submodules (`module.py`)
-- **`RotaryEmbedding`**: Generates RoPE cos/sin cache
-- **`DecoderBlock`**: Contains multi-head attention (supports GQA and MLA), feedforward network (FFN), residual connections
-- **`GQA`**: Grouped Query Attention implementation
-- **`MLA`**: Multi-Latent Attention implementation (like Qwen2-VL)
-- **`MLP`**: Feed-forward network with SiLU activation and gated mechanism
-- **`RMSNorm`**: Layer normalization variant
-- **`Linear`**, **`Embedding`**: Custom linear layer and embedding layer, supporting parallelism wrappers
+## Dataset Architecture
 
-### 3. Training Module
+```
+DatasetFactory.load(train_type, path, window_size, stride)
+  → StorageFactory.create(detect_format(path))
+    → MultiSegmentFetcher(BaseSegmentFetcher per key)
+      → BaseDataset.__getitem__(idx)
+        → sliding window [begin, end) via get_index(idx)
+```
 
-#### 3.1 Training Context (`train_context.py`)
-- **`TrainContext`**: Data class encapsulating all components needed for training (model, optimizer, data loader, strategy, etc.)
-- **`TrainContextBuilder`**: Builder pattern, progressively assembles training context, supports resume from checkpoint
+`window_size` = max input length, `stride` = step between consecutive samples.
 
-#### 3.2 Trainer (`trainer.py`)
-- **`Trainer`**: Main training loop, manages callbacks (progress bar, checkpoint, metric logging, gradient clipping, scheduler)
-- Supports distributed training (launches multi-process via `spawn_parallel_fn`)
-- Training steps include:
-  1. `on_train_begin` → 2. `on_epoch_begin` → 3. `on_batch_begin` → 4. Forward/loss calculation → 5. `on_batch_end` → 6. Gradient accumulation → 7. `on_step_begin` → 8. Optimizer update → 9. `on_step_end` → 10. `on_epoch_end`
+## Sampler
 
-#### 3.3 Strategy (`strategy.py`)
-- **`BaseStrategy`**: Defines training strategy interface
-- **`SEQStrategy`**: Standard next-token prediction training
-- **`SFTStrategy`**: Supervised Fine-tuning with loss masking
-- **`DPOStrategy`**: Direct Preference Optimization
-- **`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
+`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
 
-#### 3.4 Scheduler (`schedule.py`)
-- **`BaseScheduler`**: Abstract base class defining learning rate scheduling interface
-- **`CosineScheduler`**: Cosine decay scheduler with warmup
-- **`SGDRScheduler`**: Stochastic Gradient Descent with Warm Restarts
-- **`SchedulerFactory`**: Factory pattern, supports registration of various schedulers
-- Scheduler is automatically created according to configuration and bound to optimizer
+- Tracks `start_epoch` / `start_iter` for resume
+- Shuffle via `torch.Generator(seed + epoch)`
+- Per-replica index slicing for DDP
 
-#### 3.5 Callbacks (`train_callback.py`)
-- **`TrainCallback`**: Protocol interface for trainer callbacks
-- **`CheckpointCallback`**: Saves model checkpoints at configurable intervals
-- **`ProgressBarCallback`**: Displays training progress
-- **`MetricLoggerCallback`**: Logs training metrics to JSON files
-- **`GradientClippingCallback`**: Clips gradient norms
-- **`SchedulerCallback`**: Steps learning rate scheduler
+## DataLoader
 
-### 4. Factory Module
+Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
 
-#### 4.1 Registry and BaseFactory (`factory.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
-
-### 5. Parallel Module
-
-#### 5.1 Setup (`setup.py`)
-- **`spawn_parallel_fn`**: Spawns multiple processes for distributed training using PyTorch multiprocessing
-- **`setup_parallel`**: Context manager for initializing distributed process group (NCCL/CCL backend)
-- **`only_on_rank`**: Decorator to execute functions only on specific ranks
-- **`get_rank`**: Returns current process rank in distributed group
-- **`get_world_size`**: Returns total number of processes in distributed group
-- **`get_current_device`**: Returns current device from environment
-
-#### 5.2 Parallel Layers (`module.py`)
-- **`ParallelModel`**: Base class for parallel models with process group
-- **`ColumnParallelLinear`**: Column-parallel linear layer with input splitting and output gathering
-- **`RowParallelLinear`**: Row-parallel linear layer with output reduction
-
-### 6. Inference Module
-
-#### 6.1 Inference Engine (`engine.py`)
-- **`InferenceEngine`**: Unified inference interface, supports streaming and non-streaming generation
-- **`InferenceScheduler`**: Continuous batching scheduler with dynamic batch composition
-- **`GenerationRequest`**: Encapsulates generation parameters (top_k, top_p, temperature, max_len, messages, etc.)
-- **`messages` format**: List of message dictionaries with `role` (system/user/assistant) and `content`
-- **`apply_chat_template`** (from `tokenizer.py`): Converts messages into prompt string using ChatML format
-- Provides streaming (`stream=True`) and non-streaming (`stream=False`) generation interfaces
-- Supports continuous batching with `max_batch_size` and `max_seq_len` parameters
-- Uses separate model and tokenizer initialization for flexibility
-
-#### 6.2 Scheduler (`scheduler.py`)
-- **`Task`**: Individual generation task with state management (PENDING, RUNNING, FINISHED, ABORTED)
-- **`TaskStatus`**: Task state enumeration
-- **`apply_sampling_strategies`**: Applies temperature, top-k, top-p sampling to logits
-- **`PrefixCacheManager`**: Radix tree-based prefix cache with LRU eviction for efficient KV cache reuse
-- **`RadixNode`**: Tree node structure for prefix caching
-- Continuous batching: new requests can join at any time, completed requests are released immediately
-
-#### 6.3 Server (`server.py`)
-- FastAPI-based HTTP inference server
-- OpenAI-compatible `/v1/chat/completions` endpoint
-- Health check and statistics endpoints
-- Supports both streaming and non-streaming responses
-
-### 7. Tokenizer Module
-
-#### 7.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
-
-#### 7.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
-   - Token ID sequences (possibly with masks, labels, etc.) are saved by groups as `.h5` files
-   - Files can contain multiple segments, each segment corresponds to a tensor
-
-2. **Dataset Loading**
-   - `BaseDataset`'s `load` method calls `load_h5`, obtaining `segments` dictionary
-   - Create `MultiSegmentFetcher` to manage data for multiple keys
-   - Calculate total sample count, and determine start/end indices for each sample based on window size and stride
-
-3. **Sampling and Batch Loading**
-   - `ResumableDistributedSampler` generates index sequence based on current epoch and iteration position
-   - PyTorch `DataLoader` uses sampler to get indices, calls dataset's `__getitem__` to get actual data
-   - Batch data shape is `[batch_size, window_size]` (or varies according to specific dataset type)
-
-4. **Strategy Forward and Loss Calculation**
-   - Batch data is passed to strategy (such as `SEQStrategy`)
-   - Strategy internally calls `Transformer` model, obtaining logits
-   - Calculate cross-entropy loss (or DPO loss, etc.) according to task type
-   - Return loss tensor
-
-5. **Backpropagation and Optimization**
-   - Loss is normalized by dividing by accumulation steps, then `loss.backward()` is executed
-   - After accumulating `accumulation_steps` batches, optimizer `step()` and `zero_grad()` are executed
-   - Learning rate scheduler updates learning rate after each step
-
-6. **Checkpoint Saving**
-   - `CheckpointCallback` saves checkpoints at set intervals
-   - Checkpoints contain model state dict, current epoch, iteration, and other metadata
-   - Saved in safetensors format, ensuring safety and efficiency
-
-## Inference Data Flow - Detailed Steps
-
-1. **Model Loading**
-   - Load `Transformer` model from checkpoint via `AutoModel.from_pretrained()`
-   - Set model to evaluation mode (`model.eval()`), enable inference mode (`torch.inference_mode`)
-
-2. **Prompt Construction and Encoding**
-   - User messages (list of dict with role and content) are converted to ChatML format string through `apply_chat_template` method in tokenizer
-   - Tokenizer encodes prompt string to token ID sequence `input_ids`
-   - For batch generation, use `pad_sequence` for padding
-
-3. **Autoregressive Generation Loop**
-   - Initialize KV cache (optional) and prefix cache
-   - Loop until generating `max_len` tokens or encountering stop token:
-     - Input current `input_ids` (or cached new token) to model, obtain `logits`
-     - Apply `apply_sampling_strategies` (temperature, top-k, top-p) to `logits`
-     - Sample next token ID from the processed distribution
-     - Append new token to `input_ids`, while updating KV cache
-     - For streaming generation, yield each token to caller immediately
-
-4. **Decoding and Output**
-   - Decode generated token ID sequence to text through tokenizer
-   - Remove special tokens, return plain text response
-
-## Checkpoint and Serialization
-
-- **Training Checkpoint**: Saves model parameters, optimizer state, scheduler state, current epoch and iteration
-- **Model Parameters**: Supports safetensors format, automatically handles special logic like weight tying during loading
-- **Dataset Serialization**: HDF5 format supports efficient random access and shared memory, suitable for large-scale pre-training data
-
-## Summary
-
-The data flow design of AstrAI reflects the characteristics of modularity, extensibility, and resumability. The training data flow supports large-scale distributed training through chunk loading, resumable sampling, gradient accumulation, and other mechanisms; the inference data flow achieves efficient text generation using KV cache, prefix caching, and sampling strategies. Clear interfaces between modules facilitate customization and extension.
-
-> Document Update Time: 2026-04-09
+> Document Update Time: 2026-05-17

assets/docs/design.md

@@ -1,694 +0,0 @@
-## 1. Why I Created This Project
-
-There are many large language models on the market today, such as GPT, LLaMA, and others, with tens of billions or even hundreds of billions of parameters. But honestly, these models have extremely high hardware requirements, making them inaccessible for ordinary developers. I thought: **Can we create a model that is both useful and can run on ordinary computers?** This is also what most people currently hope for - a locally deployable AI project that achieves complete privatization while maintaining some level of intelligence.
-
-Thus, the AstrAI project was born - 1B parameters, Chinese-English bilingual, supporting dialogue, text generation, and the training code is open source!
-
-## 2. System Architecture
-
-```mermaid
-classDiagram
-    namespace config {
-        class ModelConfig {
-            +int vocab_size
-            +int dim
-            +int n_layers
-            +float norm_eps
-            +int dim_ffn
-            +bool tie_weight
-            +int max_len
-            +float rope_theta
-            +int n_heads
-            +int n_kv_heads
-            +bool use_qk_norm
-            +bool use_gated_attention
-            +load(config_path) ModelConfig
-            +save(config_path)
-        }
-
-        class TrainConfig {
-            +nn.Module model
-            +str strategy
-            +Dataset dataset
-            +Callable optimizer_fn
-            +Callable scheduler_fn
-            +int n_epoch
-            +int batch_size
-            +int accumulation_steps
-            +float max_grad_norm
-            +int start_epoch
-            +int start_batch
-            +str ckpt_dir
-            +int ckpt_interval
-            +int random_seed
-            +int num_workers
-            +int prefetch_factor
-            +bool pin_memory
-            +int nprocs
-            +str backend
-            +str master_addr
-            +str master_port
-            +Callable parallel_wrapper
-            +Callable state_dict_fn
-            +List[int] device_ids
-            +str device_type
-            +dict extra_kwargs
-            +validate()
-        }
-
-    }
-
-    namespace dataset {
-        class BaseDataset {
-            +int window_size
-            +int stride
-            +MultiSegmentFetcher fetcher
-            +load(load_path)
-            +__getitem__(index)
-            +__len__()
-        }
-
-        class SEQDataset {
-            +__getitem__(index) Dict
-        }
-
-        class SFTDataset {
-            +__getitem__(index) Dict
-        }
-
-        class DPODataset {
-            +__getitem__(index) Dict
-        }
-
-        class GRPODataset {
-            +__getitem__(index) Dict
-        }
-
-        class BaseSegmentFetcher {
-            +List~Tensor~ segments
-            +List~int~ cum_lengths
-            +int total_length
-            +fetch_data(begin_idx, end_idx) Tensor
-        }
-
-        class MultiSegmentFetcher {
-            +Dict multi_fetchers
-            +List multi_keys
-            +key_fetch(begin_idx, end_idx, keys) Dict
-            +fetch_data(begin_idx, end_idx) Dict
-        }
-
-        class ResumableDistributedSampler {
-            +int start_epoch
-            +int start_iter
-        }
-
-        class DatasetFactory {
-            +Registry _registry
-            +register(name) decorator
-            +create(train_type, window_size, stride) BaseDataset
-            +load(train_type, load_path, window_size, stride) BaseDataset
-        }
-
-        class Checkpoint {
-            +dict state_dict
-            +int epoch
-            +int iteration
-            +save(save_dir)
-            +load(save_dir) Checkpoint
-        }
-    }
-
-    namespace model {
-        class AutoModel {
-            +ModelConfig config
-            +Dict _registry
-            +register(model_type) decorator
-            +get_model_class(model_type) Type
-            +from_pretrained(path, disable_random_init) nn.Module
-            +save_pretrained(save_directory)
-            +to(*args, **kwargs) Self
-        }
-
-        class Transformer {
-            +ModelConfig config
-            +RotaryEmbedding rotary_embedding
-            +Embedding embed_tokens
-            +ModuleList layers
-            +RMSNorm norm
-            +Linear lm_head
-            +forward(input_ids, input_mask, persistent_key_values, start_pos) Dict
-            +load_state_dict(state_dict)
-            +state_dict()
-        }
-
-        class DecoderBlock {
-            +GQA attention
-            +RMSNorm input_norm
-            +MLP mlp
-            +RMSNorm post_attention_norm
-            +forward(x, rotary_emb, attention_mask, kv_cache, start_pos) Tensor
-        }
-
-        class GQA {
-            +int n_heads
-            +int n_kv_heads
-            +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
-        }
-
-        class MLA {
-            +int n_heads
-            +int n_kv_heads
-            +int head_dim
-            +Linear q_a_proj, q_b_proj, q_c_proj
-            +Linear kv_a_proj, kv_b_proj, kv_c_proj
-            +Linear o_proj
-            +RMSNorm q_norm, k_norm
-            +forward(x, rotary_emb, mask, kv_cache, start_pos) Tensor
-        }
-
-        class MLP {
-            +Linear up, gate, down
-            +forward(x) Tensor
-        }
-
-        class RMSNorm {
-            +Parameter weight
-            +float norm_eps
-            +forward(x) Tensor
-        }
-
-        class Linear {
-            +Parameter weight
-            +Parameter bias
-            +forward(x) Tensor
-        }
-
-        class RotaryEmbedding {
-            +int dim
-            +int max_len
-            +float base
-            +forward(x, start_pos) Tuple~Tensor, Tensor~
-        }
-
-        class Embedding {
-            +Parameter weight
-            +forward(x) Tensor
-        }
-    }
-
-    namespace tokenize {
-        class AutoTokenizer {
-            +List~str~ stop_ids
-            +int bos_id
-            +int eos_id
-            +int pad_id
-            +vocab_size int
-            +encode(tokens, out_ids, add_special_tokens) List~int~
-            +decode(tokens, skip_special_tokens) str
-            +apply_chat_template(messages, tokenize) Union~str, List[int]~
-            +set_chat_template(template)
-            +load(path)
-            +from_pretrained(path) AutoTokenizer
-            +save_pretrained(save_path)
-        }
-
-        class ChatTemplate {
-            +String template_str
-            +render(messages, add_generation_prompt) str
-            +from_string(template) ChatTemplate
-        }
-    }
-
-    namespace factory {
-        class Registry {
-            +Dict _entries
-            +register(name, component_cls, category, priority)
-            +get(name) Type
-            +list_names() List~str~
-        }
-
-        class BaseFactory {
-            +Registry _registry
-            +register(name, category, priority) decorator
-            +create(name, *args, **kwargs) T
-            +list_registered() list
-        }
-    }
-
-    namespace trainer {
-        class Trainer {
-            +TrainConfig train_config
-            +List~TrainCallback~ callbacks
-            +train(checkpoint)
-            +_build_context(checkpoint) TrainContext
-            +_get_default_callbacks() List~TrainCallback~
-        }
-
-        class TrainContext {
-            +nn.Module model
-            +BaseStrategy strategy
-            +DataLoader dataloader
-            +Optimizer optimizer
-            +LRScheduler scheduler
-            +Checkpoint checkpoint
-            +int epoch
-            +int iteration
-            +float loss
-            +int world_size
-            +int rank
-        }
-
-        class TrainContextBuilder {
-            +TrainConfig config
-            +with_checkpoint(checkpoint) TrainContextBuilder
-            +with_dataloader() TrainContextBuilder
-            +with_strategy() TrainContextBuilder
-            +build() TrainContext
-        }
-
-        class BaseStrategy {
-            +nn.Module model
-            +str device
-            +compute_loss(batch) Tensor
-        }
-
-        class StrategyFactory {
-            +Registry _registry
-            +register(name) decorator
-            +create(model, train_type, device, **kwargs) BaseStrategy
-        }
-
-        class SEQStrategy {
-            +float label_smoothing
-            +compute_loss(batch) Tensor
-        }
-
-        class SFTStrategy {
-            +float label_smoothing
-            +compute_loss(batch) Tensor
-        }
-
-        class DPOStrategy {
-            +nn.Module ref_model
-            +float beta
-            +str reduction
-            +compute_loss(batch) Tensor
-        }
-
-        class GRPOStrategy {
-            +nn.Module ref_model
-            +float clip_eps
-            +float kl_coef
-            +int group_size
-            +compute_loss(batch) Tensor
-        }
-
-        class BaseScheduler {
-            +get_lr() List~float~
-            +step()
-        }
-
-        class SchedulerFactory {
-            +Registry _registry
-            +register(name) decorator
-            +create(optimizer, schedule_type, **kwargs) BaseScheduler
-        }
-
-        class CosineScheduler {
-            +int warmup_steps
-            +int lr_decay_steps
-            +float min_rate
-        }
-
-        class SGDRScheduler {
-            +int warmup_steps
-            +int cycle_length
-            +float min_rate
-            +int t_mult
-        }
-
-        class TrainCallback {
-            +on_train_begin(context)
-            +on_train_end(context)
-            +on_epoch_begin(context)
-            +on_epoch_end(context)
-            +on_step_begin(context)
-            +on_step_end(context)
-            +on_batch_begin(context)
-            +on_batch_end(context)
-            +on_error(context)
-        }
-
-        class GradientClippingCallback {
-            +float max_grad_norm
-            +on_step_begin(context)
-        }
-
-        class SchedulerCallback {
-            +on_train_begin(context)
-            +on_batch_end(context)
-        }
-
-        class CheckpointCallback {
-            +str save_dir
-            +int interval
-            +_save_checkpoint(context)
-            +on_batch_end(context)
-            +on_train_end(context)
-            +on_error(context)
-        }
-
-        class ProgressBarCallback {
-            +int num_epoch
-            +on_epoch_begin(context)
-            +on_batch_end(context)
-            +on_epoch_end(context)
-        }
-
-        class MetricLoggerCallback {
-            +str log_dir
-            +int save_interval
-            +on_batch_end(context)
-            +on_train_end(context)
-        }
-
-        class CallbackFactory {
-            +Registry _registry
-            +register(name) decorator
-            +create(name, **kwargs) TrainCallback
-        }
-    }
-
-    namespace inference {
-        class InferenceEngine {
-            +nn.Module model
-            +AutoTokenizer tokenizer
-            +InferenceScheduler scheduler
-            +int max_batch_size
-            +Optional int max_seq_len
-            +int max_prefix_len
-            +int cache_capacity
-            +Tensor kv_cache
-            +Tensor seq_mask
-            +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
-            +generate_with_request(request) Union[Generator, str, List[str]]
-            +get_stats() Dict
-            +shutdown()
-        }
-
-        class InferenceScheduler {
-            +nn.Module model
-            +AutoTokenizer tokenizer
-            +ModelConfig config
-            +Tuple kv_cache
-            +Tensor seq_mask
-            +PrefixCacheManager prefix_cache
-            +List waiting_queue
-            +List active_tasks
-            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
-            +remove_task(task_id)
-            +start()
-            +stop()
-            +get_stats() Dict
-        }
-
-        class PrefixCacheManager {
-            +RadixNode root
-            +int max_capacity
-            +List lru
-            +insert(token_ids, slot)
-            +find_longest_prefix(token_ids) Tuple[int, int]
-            +release(token_ids)
-        }
-
-        class RadixNode {
-            +Dict children
-            +int hash
-            +int slot
-            +int ref_count
-            +float last_access
-            +List token_sequence
-        }
-
-        class Task {
-            +str task_id
-            +List prompt_ids
-            +int max_tokens
-            +float temperature
-            +float top_p
-            +int top_k
-            +TaskStatus status
-            +List output_ids
-            +int input_tokens
-            +int output_tokens
-            +int slot
-            +Callable stream_callback
-            +is_finished(stop_ids) bool
-        }
-
-        class TaskStatus {
-            +str PENDING
-            +str RUNNING
-            +str FINISHED
-            +str ABORTED
-        }
-
-        class Server {
-            +start()
-            +predict(request)
-        }
-
-        class GenerationRequest {
-            +int top_k
-            +float top_p
-            +float temperature
-            +int max_len
-            +List~Dict~ messages
-            +stream bool
-        }
-
-        class _Result {
-            +List~str~ tokens
-            +List~str~ results
-            +List~bool~ done_flags
-            +append(token, idx)
-            +get_results() List~str~
-        }
-
-        class ChatMessage {
-            +str role
-            +str content
-        }
-
-        class ChatCompletionRequest {
-            +List~ChatMessage~ messages
-            +float temperature
-            +float top_p
-            +int top_k
-            +int max_tokens
-            +bool stream
-            +Optional~str~ system_prompt
-        }
-
-        class CompletionResponse {
-            +str id
-            +str object
-            +int created
-            +str model
-            +List~Dict~ choices
-        }
-    }
-
-    namespace parallel {
-        class ParallelSetup {
-            +spawn_parallel_fn(fn, nprocs)
-            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type, device_ids)
-        }
-
-        class ParallelModel {
-            +dist.ProcessGroup process_group
-            +int rank
-            +int world_size
-        }
-
-        class ColumnParallelLinear {
-            +forward(x) Tensor
-        }
-
-        class RowParallelLinear {
-            +forward(x) Tensor
-        }
-    }
-
-    %% Relationships
-    TrainConfig --> ModelConfig : uses
-    TrainConfig --> BaseDataset : uses
-    TrainConfig --> StrategyFactory : selects
-    StrategyFactory ..> BaseStrategy : creates
-    BaseStrategy <|-- SEQStrategy
-    BaseStrategy <|-- SFTStrategy
-    BaseStrategy <|-- DPOStrategy
-    BaseStrategy <|-- GRPOStrategy
-    DPOStrategy --> Transformer : uses
-    GRPOStrategy --> Transformer : uses
-    Trainer --> TrainConfig : configures
-    Trainer --> TrainContextBuilder : builds
-    Trainer --> TrainCallback : manages
-    TrainContextBuilder --> TrainContext : creates
-    TrainContext --> Checkpoint : manages
-    TrainContext --> BaseStrategy : uses
-    TrainContext --> BaseScheduler : uses
-    AutoModel --> ModelConfig : contains
-    SchedulerFactory ..> BaseScheduler : creates
-    BaseScheduler <|-- CosineScheduler
-    BaseScheduler <|-- SGDRScheduler
-    CallbackFactory ..> TrainCallback : creates
-    TrainCallback <|-- GradientClippingCallback
-    TrainCallback <|-- SchedulerCallback
-    TrainCallback <|-- CheckpointCallback
-    TrainCallback <|-- ProgressBarCallback
-    TrainCallback <|-- MetricLoggerCallback
-    InferenceEngine --> InferenceScheduler : uses
-    InferenceScheduler --> Task : manages
-    InferenceScheduler --> TaskStatus : uses
-    InferenceScheduler --> Transformer : uses
-    InferenceEngine --> Transformer : uses
-    InferenceEngine --> GenerationRequest : uses
-    Server --> InferenceEngine : uses
-    Server --> ChatMessage : uses
-    Server --> ChatCompletionRequest : uses
-    Server --> CompletionResponse : uses
-    ParallelSetup --> Trainer : enables
-    BaseDataset <|-- SEQDataset
-    BaseDataset <|-- SFTDataset
-    BaseDataset <|-- DPODataset
-    BaseDataset <|-- GRPODataset
-    DatasetFactory ..> BaseDataset : creates
-    BaseSegmentFetcher --> MultiSegmentFetcher : used by
-    MultiSegmentFetcher --> BaseDataset : used by
-    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
-    ResumableDistributedSampler --> BaseDataset : samples
-    ParallelModel <|-- RowParallelLinear
-    ParallelModel <|-- ColumnParallelLinear
-    AutoTokenizer --> ChatTemplate : uses
-    InferenceScheduler --> PrefixCacheManager : uses
-    InferenceScheduler --> RadixNode : uses
-    Checkpoint ..> Checkpoint : saves/loads
-    TrainConfig --> DatasetFactory : selects
-    TrainConfig --> SchedulerFactory : selects
-    TrainConfig --> CallbackFactory : selects
-    AutoModel ..> AutoTokenizer : loads with
-    BaseFactory <|-- DatasetFactory
-    BaseFactory <|-- StrategyFactory
-    BaseFactory <|-- SchedulerFactory
-    BaseFactory <|-- CallbackFactory
-```
-
-### Module Overview
-
-| Module | Components | Description |
-|--------|------------|-------------|
-| **astrai.config** | ModelConfig, TrainConfig | Configuration management |
-| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory, Checkpoint | Dataset loading and management |
-| **astrai.model** | AutoModel, Transformer, DecoderBlock, GQA, MLA, MLP, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
-| **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
-| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy, StrategyFactory, BaseScheduler, SchedulerFactory, TrainCallback, CallbackFactory | Training workflow management |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Task, TaskStatus, Server, GenerationRequest, PrefixCacheManager, ChatMessage, ChatCompletionRequest, CompletionResponse | Inference service with continuous batching |
-| **astrai.parallel** | ParallelSetup, ColumnParallelLinear, RowParallelLinear | Distributed parallel |
-| **astrai.factory** | Registry, BaseFactory | Generic component registration |
-
-### Design Patterns
-
-| Pattern | Classes | Purpose |
-|---------|---------|---------|
-| **Strategy** | `BaseStrategy`, `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy`, `StrategyFactory` | Flexible training strategy switching, supports SEQ/SFT/DPO/GRPO |
-| **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 |
-| **Registry** | `BaseFactory`, `Registry` | Generic component registration with category and priority support |
-| **Producer-Consumer** | `InferenceScheduler`, `Task`, `waiting_queue`, `active_tasks` | Continuous batching with dynamic task queue management |
-| **Event-Driven** | `threading.Event`, `_task_event` | Non-blocking wait mechanism for task scheduling using Python's `threading` module |
-| **AutoModel Registry** | `AutoModel`, `Transformer` | Model type registration and dynamic loading via decorator pattern |
-| **Generator Pattern** | `_Result`, `GenerationRequest` | Event-based result notification for streaming/non-streaming generation |
-
-### Core Relationships
-
-1. **Configuration → Training**: `TrainConfig` contains `ModelConfig`, holds model, dataset, optimizer and other references
-2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` to compute loss
-3. **Strategy Selection**: `StrategyFactory` creates corresponding strategy instance based on `train_type`
-4. **Inference Flow**: `Server` → `InferenceEngine` → `InferenceScheduler` → `Transformer`, supports continuous batching with streaming/non-streaming
-5. **Distributed Support**: `ParallelSetup` provides 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)
-9. **AutoModel Loading**: `AutoModel.from_pretrained()` dynamically loads model based on `config.json` model_type, uses `Registry` pattern for model type registration
-
-## 3. Training Process
-
-The common training process for large language models (LLM) typically includes three stages: **Pre-training (SEQ)**, **Supervised Fine-Tuning (SFT)**, and **Reinforcement Learning from Human Feedback (DPO/GRPO)**. This system is designed to support seamless end-to-end flow, achieving efficient switching and state management of different training stages through modular strategies.
-
-### Core Formulas
-
-**Pre-training (SEQ):**
-
-$$
-L_{\text{PT}} = - \sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
-$$
-
-**SFT:**
-
-$$
-L_{\text{SFT}} = - \sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
-$$
-
-**DPO:**
-
-$$
-L_{\text{DPO}} = -\mathbb{E}_{(x, y_w, y_l) \sim D} \left[ \log \sigma\left( \beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)} \right) \right]
-$$
-
-**GRPO:**
-
-GRPO (Group Relative Policy Optimization) computes advantages from multiple responses to the same prompt, then optimizes using a PPO-style clipped objective:
-
-$$
-\text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
-$$
-
-Where $r_i$ is the reward for the $i$-th response, $\mu$ and $\sigma$ are the mean and standard deviation of group rewards.
-
-$$
-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:
-
-$$
-L_{KL} = \lambda \cdot \mathbb{E} \left[ (\log \pi_\theta - \log \pi_{\text{ref}})^2 \right]
-$$
-
-The final loss is the sum of both: $L = L_{\text{policy}} + L_{KL}$
-
-Through the above three-stage progressive training, the model completes its evolution from a general language foundation to a specialized, highly-aligned dialogue intelligence.
-
-> Document Update Time: 2026-04-09

assets/docs/inference.md

@@ -0,0 +1,140 @@
+# Inference
+
+## KV Cache
+
+At decode time, only the last query token matters. All previous K/V are cached to avoid recomputation:
+
+$$
+o_n = \sum_j \text{softmax}\left(\frac{q_n k_j}{\sqrt{d_k}}\right) v_j
+$$
+
+RoPE is applied **before** KV cache write, not after — otherwise position encoding drift occurs.
+
+## KVCache System
+
+Six classes working together:
+
+```
+KVCache (facade)
+  ├── Allocator        bitmask-based page allocator + ref-count + LRU eviction
+  ├── PrefixCache      hash-based prefix matching (page_hash via rolling hash)
+  ├── PagePool         orchestrates Allocator + PrefixCache
+  ├── TaskTable        maps task_id → page_table + cached token count
+  ├── Storage          k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim)
+  └── KvcacheView      bundles Storage + page_table + total_len for attention layers
+```
+
+`KVCache.bind(page_table, total_len)` returns a `KvcacheView` used by attention layers via `write()` / `gather()`.
+
+## Continuous Batching
+
+`InferenceScheduler` runs a daemon thread with a 4-phase loop:
+
+```
+1. Cleanup → Remove finished tasks, free KV pages
+2. Refill  → Pop from waiting_queue, task_alloc pages, activate
+3. Prefill → Group by (prompt_len, start_pos), run full forward
+4. Decode  → Pick largest same-position group, single-token forward
+```
+
+## Sampling (Strategy Pattern)
+
+```
+BaseSamplingStrategy → TemperatureStrategy → TopKStrategy → TopPStrategy
+```
+
+`SamplingPipeline` composes them: Temperature → Top-K → Top-P → softmax → multinomial.  
+`sample()` is a convenience shortcut for one-shot usage.
+
+## Protocol Handlers (Template Method)
+
+```python
+class ProtocolHandler(ABC):
+    def handle(self):
+        ctx = StreamContext(...)
+        agen = engine.generate_async(prompt, ...)
+        if stream: self._handle_stream(agen, ctx)
+        else:      self._handle_non_stream(agen, ctx)
+```
+
+Subclass hooks: `build_prompt()`, `create_response_id()`, `format_stream_start/token/end()`, `format_non_stream_response()`.
+
+`OpenAIHandler` → `/v1/chat/completions`, `AnthropicHandler` → `/v1/messages`.
+
+## Engine & GenerateResult
+
+```
+InferenceEngine
+  ├── generate(prompt, stream, ...) → str | List[str] | Generator
+  ├── generate_with_request(req)    → same
+  └── generate_async(prompt, ...)   → AsyncGenerator
+```
+
+`GenerateResult` uses `Condition` for non-streaming (`wait_completion()`) and `Event` for streaming (`wait()`). Stream callback is `cb(token)`.
+
+## HTTP API
+
+```
+POST /v1/chat/completions   OpenAI
+POST /v1/messages            Anthropic
+GET  /health                 {"status":"ok","model_loaded":true}
+GET  /stats                  scheduler statistics
+```
+
+### OpenAI
+
+```bash
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+Response:
+```json
+{
+  "id": "chatcmpl-abc123",
+  "object": "chat.completion",
+  "choices": [{"message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"}],
+  "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}
+}
+```
+
+Streaming SSE: `data: {"choices":[{"delta":{"role":"assistant"}}]}` → token chunks → `data: [DONE]`
+
+### Anthropic
+
+```bash
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{"model":"astrai","system":"You are helpful.","messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+Supports `stop_sequences` and streaming via `event: content_block_delta`.
+
+### GenerationRequest Parameters
+
+| Param | Type | Default | Description |
+|-------|------|---------|-------------|
+| `messages` | List[dict] | required | Chat messages (role, content) |
+| `temperature` | float | 1.0 | Sampling temperature (0.0–2.0) |
+| `top_p` | float | 1.0 | Nucleus threshold |
+| `top_k` | int | 50 | Top-k count |
+| `max_tokens` | int | None | Max generation length |
+| `stream` | bool | False | Stream output |
+
+## Engine API
+
+```python
+# Non-streaming
+engine.generate("Hello", stream=False)          # -> str
+engine.generate(["A", "B"], stream=False)       # -> List[str]
+
+# Streaming
+engine.generate("Hello", stream=True)           # -> Generator[str]
+engine.generate(["A", "B"], stream=True)        # -> Generator[Tuple[int, str]]
+
+# Async
+await engine.generate_async("Hello", ...)       # -> AsyncGenerator[str]
+```
+
+> Document Update Time: 2026-05-17

assets/docs/introduction.md

@@ -1,299 +0,0 @@
-## Model Introduction
-
-### 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.
-
-The model now uses the **AutoModel** base class for flexible loading and saving:
-
-```python
-from astrai.model import AutoModel
-
-# Load model from checkpoint
-model = AutoModel.from_pretrained("path/to/model")
-
-# Save model to new directory
-model.save_pretrained("path/to/save")
-```
-
-The Transformer model is registered via `@AutoModel.register('transformer')` decorator, allowing easy extension for new model types.
-
-```mermaid
-flowchart TB
-    subgraph Layers["Transformer Layers"]
-        direction TB
-        A[Input Embedding] --> B[Transformer Block\nLayer 1]
-        B --> C[Transformer Block\nLayer ...]
-        C --> D[Transformer Block\nLayer 32]
-        D --> E[RMSNorm]
-        E --> F[Linear]
-        F --> G[SoftMax]
-    end
-
-    subgraph TransformerBlock["Transformer Block"]
-        direction TB
-        H[x] --> I[RMSNorm]
-        I --> J[Linear → Q/K/V]
-        J --> K[Q]
-        J --> L[K]
-        J --> M[V]
-        K --> N[RoPE]
-        L --> O[RoPE]
-        N --> P["Q @ K^T / sqrt(d)"]
-        O --> P
-        P --> Q[Masked SoftMax]
-        Q --> R[S @ V]
-        M --> R
-        R --> S[Linear]
-        S --> T[+]
-        H --> T
-        T --> U[RMSNorm]
-        U --> V[Linear]
-        V --> W[SiLU]
-        V --> X[×]
-        W --> X
-        X --> Y[Linear]
-        Y --> Z[+]
-        T --> Z
-        Z --> AA[x']
-    end
-
-    classDef main fill:#e6f3ff,stroke:#0066cc;
-    classDef block fill:#fff2e6,stroke:#cc6600;
-    class Layers main;
-    class TransformerBlock block;
-```
-
-What is an autoregressive model? After splitting a sentence into tokens, the model predicts the probability distribution of the next token. This means the model calculates the probability of the next possible token and its corresponding probability based on the given context (the sequence of tokens that have already appeared).
-
-#### 1. Autoregression
-
-In autoregressive modeling, when a sentence is tokenized into a sequence of tokens, the model learns to predict what comes next. Given a sequence of tokens as input, the model calculates a probability distribution over all possible next tokens. This distribution tells us how likely each potential next token is, given the current context.
-
-For instance, if the input sequence contains tokens representing a question, the model might predict that certain response tokens have higher probabilities than others. The sampling process then selects one token from this distribution—controlled by parameters like top_k, top_p, and temperature—to serve as the next token in the sequence.
-
-Once a token is selected, it is appended to the input sequence, and the model repeats this process. The updated sequence is then fed back into the model to predict the next token. This iterative process continues until either a special end-of-sequence token is generated, or the maximum sequence length is reached. These control tokens are essential because without them, the model would continue generating tokens indefinitely, eventually exhausting available memory.
-
-#### 2. Causal Mask
-
-Transformers use attention mechanism. The input shape is generally [bsz, seq_len], and the output is [bsz, seq_len, n_dim]. To predict the next token, the model's input and output must be offset by one position. The target predicted by the model must be offset by one position, and during training we also use the offset-by-one method:
-
-```
-sequence : [[1, 2, 3, 4, 5, 6]]
-input_ids: [[1, 2, 3, 4, 5]]
-target_ids: [[2, 3, 4, 5, 6]]
-```
-
-The attention score calculation formula is:
-
-$$ s_{ij} = softmax(\frac{q_i^Tk_j}{\sqrt{d_k}}) $$
-$$ s_{ij} := s_{ij} + mask_{ij} $$
-
-Here, the attention score represents the degree to which the model attends to the similarity between two tokens.
-
-For decoder-only structure models, to prevent the model from "stealing" information from future positions, a mask needs to be added during attention calculation. We need to apply a mask before attention score calculation. This mask is typically a lower triangular matrix, and for a sequence of length n, its shape is [n, n]. Below is an example of how to create such a causal mask matrix for a sequence of length 5:
-
-```
-[[0, -inf, -inf, -inf, -inf],
- [0,    0, -inf, -inf, -inf],
- [0,    0,    0, -inf, -inf],
- [0,    0,    0,    0, -inf],
- [0,    0,    0,    0,    0]]
-```
-
-In this matrix, 0 represents positions that can be attended to, while -inf represents positions that should be masked (i.e., should not be attended to). Because this matrix ensures that after the softmax, the parts of the attention scores where $j > i$ change from `inf` to 0, meaning the model cannot see future information.
-
-#### 3. Rotary Position Embedding
-
-Rotary Position Embedding (RoPE) is a position encoding method designed to solve the problem of lacking direct modeling of sequence position information in Transformer models. Unlike traditional position encodings (such as sine and cosine function position encodings), RoPE embeds position information directly into the Query (Q) and Key (K) vectors, allowing the model to more naturally handle relative position relationships in sequences.
-
-$$ q_i = R_i W_q x_i $$
-$$ k_j = R_j W_k x_j $$
-$$ q_i^T k_j = (R_i W_q x_i)^T( R_j W_k x_j) = x_i^T W_q^T R_{i-j} W_k x_j $$
-
-The $R_{i-j}$ controls the attenuation of attention for different tokens at different relative distances. When the absolute value of $i - j$ is larger, the degree of attenuation is stronger. This approach allows the model to learn relative position relationships, enabling the model to scale and adapt to longer sequences.
-
-## KV Cache Implementation
-
-According to the attention calculation formula:
-
-$$
-\begin{align*}
-o_i &= \sum_j s_{ij} v_{j} \newline
-s_{ij} &= \text{softmax}\left( \frac{q_{i} k_{j}}{\sqrt{d_k}} \right)
-\end{align*}
-$$
-
-Since the model is an autoregressive model, we only need to calculate for the last part of the sequence, meaning the index $i$ is fixed as the last element of the sequence, and we compute $o_{n}$:
-
-$$
-\begin{align*}
-o_n &= \sum_j s_{j}v_{j} \newline
-s_j &= \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}} \right)
-\end{align*}
-$$
-
-If we expand the expression:
-
-$$
-o_n = \sum_j \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}}\right)v_{j}
-$$
-
-In the above expression, only k and v have length indices, while $q$ does not. Therefore, during the calculation process, the input of $q$ is fixed as the last token from the previous input, while $k$ and $v$ need to be cached for parts of different lengths. Also, when caching, note that position encoding calculation should be performed before KV cache computation, otherwise there will be position encoding calculation errors.
-
-### 4. AutoModel Loading
-
-The project now uses the **AutoModel** base class for flexible model loading and saving:
-
-```python
-from astrai.model import AutoModel
-
-# Load model from checkpoint
-model = AutoModel.from_pretrained("path/to/model")
-
-# Save model to new directory
-model.save_pretrained("path/to/save")
-```
-
-The Transformer model is registered via `@AutoModel.register('transformer')` decorator, allowing easy extension for new model types. The `from_pretrained` method automatically loads the `config.json` to determine the model type and uses safetensors format for weights.
-
-### 5. Continuous Batching Inference
-
-The inference engine supports **continuous batching** for efficient batch processing:
-
-```python
-from astrai.inference import InferenceEngine, GenerationRequest
-
-# Create inference engine with continuous batching
-engine = InferenceEngine(
-    model=model,
-    tokenizer=tokenizer,
-    max_batch_size=8,
-    max_seq_len=4096,
-)
-
-# Use GenerationRequest with messages format
-request = GenerationRequest(
-    messages=[
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": "Hello"},
-    ],
-    temperature=0.8,
-    top_p=0.95,
-    top_k=50,
-    max_len=1024,
-    stream=True,
-)
-
-# Generate with streaming
-for token in engine.generate_with_request(request):
-    print(token, end="", flush=True)
-```
-
-The continuous batching feature allows dynamic batch composition where new requests can join at any time and completed requests are released immediately.
-
-## HTTP API Usage
-
-The inference server provides HTTP endpoints for remote inference. Start the server first:
-
-```bash
-python -m scripts.tools.server --port 8000
-```
-
-### OpenAI-Compatible Endpoint
-
-The server provides an OpenAI-compatible chat completion endpoint at `/v1/chat/completions`:
-
-```bash
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [
-      {"role": "system", "content": "You are a helpful assistant."},
-      {"role": "user", "content": "Hello, how are you?"}
-    ],
-    "temperature": 0.8,
-    "max_tokens": 2048,
-    "stream": false
-  }'
-```
-
-**Request Parameters:**
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `messages` | List[dict] | Required | Chat messages with role and content |
-| `temperature` | float | 0.8 | Sampling temperature (0.0-2.0) |
-| `top_p` | float | 0.95 | Nucleus sampling threshold |
-| `top_k` | int | 50 | Top-k sampling parameter |
-| `max_tokens` | int | 2048 | Maximum tokens to generate |
-| `stream` | bool | false | Enable streaming response |
-| `system_prompt` | str | None | System prompt override |
-
-**Response (non-streaming):**
-```json
-{
-  "id": "chatcmpl-1234567890",
-  "object": "chat.completion",
-  "created": 1234567890,
-  "model": "astrai",
-  "choices": [
-    {
-      "index": 0,
-      "message": {"role": "assistant", "content": "Hello! I'm doing well..."},
-      "finish_reason": "stop"
-    }
-  ]
-}
-```
-
-### Streaming Response
-
-Enable streaming for real-time token-by-token output:
-
-```bash
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "Write a story"}],
-    "stream": true,
-    "max_tokens": 500
-  }'
-```
-
-The server uses Server-Sent Events (SSE) with content type `text/event-stream`.
-
-### Simple Generation Endpoint
-
-For basic text generation without chat format:
-
-```bash
-curl -X POST "http://localhost:8000/generate?query=Hello&max_len=1000" \
-  -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?",
-    "history": [["Hello", "Hi there!"], ["How are you?", "I'm doing well"]],
-    "temperature": 0.8,
-    "max_len": 2048
-  }'
-```
-
-### Health Check
-
-Monitor server and model status:
-
-```bash
-curl http://localhost:8000/health
-# {"status": "ok", "model_loaded": true, "engine_ready": true}
-
-curl http://localhost:8000/stats
-# {"requests_total": 10, "tokens_generated": 5000, ...}
-```
-
-> Document Update Time: 2026-04-09

assets/docs/params.md

@@ -4,138 +4,96 @@
 
 ### Basic Parameters
 
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `--train_type` | Training type (seq, sft, dpo, grpo) | required |
-| `--model_type` | Model type for AutoModel loading (e.g., transformer) | transformer |
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--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 |
-| `--batch_size` | Batch size | 4 |
-| `--accumulation_steps` | Gradient accumulation steps | 1 |
+| `--batch_per_device` | Batch size per device | 1 |
+| `--grad_accum_steps` | Gradient accumulation steps between optimizer steps | 1 |
 
 ### Learning Rate Scheduling
 
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `--warmup_steps` | Warmup steps | 1000 |
-| `--max_lr` | Maximum learning rate (warmup + cosine decay) | 3e-4 |
-| `--max_grad_norm` | Maximum gradient norm | 1.0 |
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
+| `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
+| `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 
-### Checkpoint
+### Optimizer (AdamW)
 
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `--ckpt_interval` | Checkpoint save interval (iterations) | 5000 |
-| `--ckpt_dir` | Checkpoint save directory | checkpoint |
-| `--resume_dir` | Resume training from specified path | - |
-
-### Optimizer Parameters
-
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
+| Parameter | Description | Default |
+|-----------|-------------|---------|
 | `--adamw_beta1` | AdamW beta1 | 0.9 |
 | `--adamw_beta2` | AdamW beta2 | 0.95 |
 | `--adamw_weight_decay` | AdamW weight decay | 0.01 |
 
 ### Data Loading
 
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `--random_seed` | Random seed | 3407 |
-| `--num_workers` | DataLoader workers | 0 |
-| `--prefetch_factor` | Prefetch factor for dataloader | None |
-| `--pin_memory` | Enable pin_memory | False |
-| `--no_pin_memory` | Disable pin_memory | - |
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--window_size` | Max input sequence length | model config `max_len` |
+| `--stride` | Stride for sliding window over sequences | None |
+| `--random_seed` | Random seed for reproducibility | 3407 |
+| `--num_workers` | DataLoader worker processes | 4 |
+| `--no_pin_memory` | Disable pin_memory (enabled by default) | (flag) |
+
+### Checkpoint & Resume
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--ckpt_interval` | Iterations between checkpoints | 5000 |
+| `--ckpt_dir` | Checkpoint save directory | checkpoint |
+| `--start_epoch` | Resume from epoch (0 = from scratch) | 0 |
+| `--start_batch` | Resume from batch iteration | 0 |
 
 ### Distributed Training
 
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `--nprocs` | Number of GPUs | 1 |
-| `--device_type` | Device type (cuda/cpu) | cuda |
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--nprocs` | Number of GPUs / processes | 1 |
+| `--parallel_mode` | Parallel strategy (`none` or `ddp`) | none |
+| `--device_type` | Device type | cuda |
+| `--start_method` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) | spawn |
 
-### Other Parameters
+### Strategy-specific
 
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `--window_size` | Maximum input sequence length | model config max_len |
-| `--stride` | Input sequence stride | - |
-| `--dpo_beta` | DPO beta value | 0.1 |
-| `--grpo_clip_eps` | GRPO clip epsilon | 0.2 |
-| `--grpo_kl_coef` | GRPO KL coefficient | 0.01 |
-| `--grpo_group_size` | GRPO group size | 4 |
-| `--label_smoothing` | Label smoothing parameter | 0.1 |
-| `--start_epoch` | Starting epoch | 0 |
-| `--start_batch` | Starting batch | 0 |
-
----
-
-## Generation Parameters
-
-### GenerationRequest Parameters
-
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `messages` | List of message dictionaries (role, content) | required |
-| `temperature` | Sampling temperature (higher = more random) | 1.0 |
-| `top_p` | Nucleus sampling threshold | 1.0 |
-| `top_k` | Top-k sampling count | 50 |
-| `max_len` | Maximum generation length | 1024 |
-| `stream` | Whether to stream output | False |
+| Parameter | Description | Default | Used by |
+|-----------|-------------|---------|---------|
+| `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.05 | `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
 
-```python
-import torch
-from astrai.model import AutoModel
-from astrai.tokenize import Tokenizer
-from astrai.inference import InferenceEngine, GenerationRequest
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
 
-# Load model using AutoModel
-model = AutoModel.from_pretrained("your_model_dir")
-
-# Load tokenizer
-tokenizer = Tokenizer("your_model_dir")
-
-# Create engine with separate model and tokenizer
-engine = InferenceEngine(
-    model=model,
-    tokenizer=tokenizer,
-)
-
-# Build request with messages format
-request = GenerationRequest(
-    messages=[
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": "Hello"},
-    ],
-    temperature=0.8,
-    top_p=0.95,
-    top_k=50,
-    max_len=1024,
-)
-
-# Generate (streaming)
-for token in engine.generate_with_request(request):
-    print(token, end="", flush=True)
-
-# Or use simple generate interface
-result = engine.generate(
-    prompt="Hello",
-    stream=False,
-    max_tokens=1024,
-    temperature=0.8,
-    top_p=0.95,
-    top_k=50,
-)
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
 ```
 
-### Generation Modes
+---
 
-| Mode | Description |
-|------|-------------|
-| `stream=True` | Streaming output, yields token by token |
-| `stream=False` | Non-streaming output, returns complete result |
-
-> Document Update Time: 2026-04-09
+> Document Update Time: 2026-05-24

assets/docs/training.md

@@ -0,0 +1,226 @@
+# Training
+
+## Model Architecture
+
+The model uses a decoder-only Transformer with **GQA** (Grouped Query Attention) and optional **MLA** (Multi-head Latent Attention). 1.0 billion parameters, Chinese–English bilingual.
+
+```mermaid
+flowchart TB
+    subgraph Layers["Transformer Layers"]
+        direction TB
+        A[Input Embedding] --> B[Transformer Block\nLayer 1]
+        B --> C[Transformer Block\nLayer ...]
+        C --> D[Transformer Block\nLayer ...]
+        D --> E[RMSNorm]
+        E --> F[Linear]
+        F --> G[SoftMax]
+    end
+
+    subgraph TransformerBlock["Transformer Block"]
+        direction TB
+        H[x] --> I[RMSNorm]
+        I --> J[Linear → Q/K/V]
+        J --> K[Q]; J --> L[K]; J --> M[V]
+        K --> N[RoPE]; L --> O[RoPE]
+        N --> P["Q @ K^T / sqrt(d)"]; O --> P
+        P --> Q[Masked SoftMax]; Q --> R[S @ V]; M --> R
+        R --> S[Linear]; S --> T[+]; H --> T
+        T --> U[RMSNorm]
+        U --> V["Linear (gate)"]; U --> W["Linear (up)"]
+        V --> X[SiLU]; X --> Y[×]; W --> Y
+        Y --> Z["Linear (down)"]; Z --> AA[+]; T --> AA
+        AA --> BB[x']
+    end
+```
+
+### Autoregression
+
+Given a token sequence, the model predicts the probability of the next token. Each generated token is appended to the input and fed back, repeating until an end-of-sequence token or max length.
+
+### Causal Mask
+
+```
+sequence : [[1, 2, 3, 4, 5, 6]]
+input_ids: [[1, 2, 3, 4, 5]]
+target_ids: [[2, 3, 4, 5, 6]]
+```
+
+Lower-triangular mask prevents attending to future positions:
+
+```
+[[0, -inf, -inf, -inf, -inf],
+ [0,    0, -inf, -inf, -inf],
+ [0,    0,    0, -inf, -inf],
+ [0,    0,    0,    0, -inf],
+ [0,    0,    0,    0,    0]]
+```
+
+### Rotary Position Embedding (RoPE)
+
+RoPE embeds position into Q/K vectors via complex rotation:
+
+$$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
+
+The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers.
+
+## Training Loop
+
+Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
+
+```
+on_train_begin
+  on_epoch_begin
+    for batch in dataloader:
+      on_batch_begin
+      with executor.accumulate(model):
+        loss = strategy(batch)
+        (loss / grad_accum_steps).backward()
+        iteration += 1
+      on_batch_end
+
+      if executor.sync_gradients:
+        on_optimizer_step
+        optimizer.step()
+        optimizer.zero_grad()
+
+      scheduler.step()  # called every iteration
+    on_epoch_end
+on_train_end
+```
+
+### Callback Lifecycle
+
+| Hook | Fires | Default callback |
+|------|-------|-----------------|
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `ValidationCallback` |
+| `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
+| `on_train_end` | Training ends | `CheckpointCallback`, `MetricLoggerCallback` (final save) |
+
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric_logger` (JSONL, rank-0), `progress_bar` (tqdm), `gradient_clipping`, `validation` (periodic validation on val_dataset).
+
+## Strategies
+
+### SEQ (Pre-training)
+
+Next-token cross-entropy with optional label smoothing:
+
+$$
+L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
+$$
+
+Keys: `input_ids`, `target_ids`
+
+### SFT (Supervised Fine-Tuning)
+
+Masked cross-entropy (`ignore_index=-100`) over response tokens:
+
+$$
+L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
+$$
+
+Keys: `input_ids`, `target_ids`, `loss_mask`
+
+### DPO (Direct Preference Optimization)
+
+Frozen reference model, preference margin via log-ratio:
+
+$$
+L_{\text{DPO}} = -\mathbb{E}\left[\log\sigma\left(\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)} - \beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}\right)\right]
+$$
+
+Parameters: `beta=0.1`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
+
+### GRPO (Group Relative Policy Optimization)
+
+On-policy PPO with group-normalized advantages:
+
+$$
+\text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
+$$
+
+$$
+L_{\text{GRPO}} = -\mathbb{E}\left[\min\left(\frac{\pi_\theta}{\pi_{\text{ref}}}A,\; \text{clip}\left(\frac{\pi_\theta}{\pi_{\text{ref}}}, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}\left[(\log\pi_\theta - \log\pi_{\text{ref}})^2\right]
+$$
+
+Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`, `sync_interval=200`.
+
+Keys: `prompts`, `responses`, `masks`, `rewards`.
+
+## LR Schedulers
+
+| Type | Class | Description |
+|------|-------|-------------|
+| Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
+| SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
+
+Created by `SchedulerFactory.create(optimizer, schedule_type, **kwargs)`.
+
+## Gradient Checkpointing
+
+Trades compute for memory by recomputing activations during backward pass. Specify module types via `gradient_checkpointing_modules`:
+
+```python
+from astrai.model.components.decoder_block import DecoderBlock
+config = TrainConfig(..., gradient_checkpointing_modules=[DecoderBlock])
+```
+
+Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoint(use_reentrant=False)`, compatible with `torch.compile`. Uses `nn.Module.apply()` for traversal — works through DDP wrappers without manual unwrap. Empty list (default) means no-op.
+
+## Checkpoint
+
+```
+Checkpoint(state_dict, epoch, iteration, extra, meta)
+  ├── save(save_dir)    rank-0 only: meta.json (includes training config) + state_dict.safetensors + optional optimizer.pt / scheduler.pt
+  └── load(save_dir)    broadcasts metadata from rank-0
+```
+
+Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
+Training config (`TrainConfig.to_dict()`) saved into `meta.json` during training via `CheckpointCallback`.
+
+## TrainContextBuilder (Builder Pattern)
+
+```python
+context = (
+    TrainContextBuilder(config)
+        .with_checkpoint(checkpoint)
+        .build()
+)
+# Returns TrainContext with model, strategy, optimizer, scheduler, dataloader, checkpoint
+```
+
+- Loads checkpoint weights if provided
+- Creates executor via `ExecutorFactory.create(parallel_mode, **executor_kwargs)`
+- Calls `executor.prepare(model, optimizer, dataloader, scheduler)` for model distribution (e.g. DDP) + gradient accumulation wrappers
+- Creates `ResumableDistributedSampler` for shuffle+resume
+- Builds strategy via `StrategyFactory.create(train_type, ...)`
+
+## Training CLI
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
+```
+
+Full parameter reference at [params.md](params.md).
+
+> Document Update Time: 2026-05-24

astrai/__init__.py

@@ -1,8 +1,9 @@
-__version__ = "1.3.3"
+__version__ = "1.3.6"
 __author__ = "ViperEkura"
 
 from astrai.config import (
-    ModelConfig,
+    AutoRegressiveLMConfig,
+    EncoderConfig,
     TrainConfig,
 )
 from astrai.dataset import DatasetFactory
@@ -11,13 +12,14 @@ from astrai.inference import (
     GenerationRequest,
     InferenceEngine,
 )
-from astrai.model import AutoModel, Transformer
+from astrai.model import AutoModel, AutoRegressiveLM
 from astrai.tokenize import AutoTokenizer
 from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer
 
 __all__ = [
-    "Transformer",
-    "ModelConfig",
+    "AutoRegressiveLM",
+    "AutoRegressiveLMConfig",
+    "EncoderConfig",
     "TrainConfig",
     "DatasetFactory",
     "AutoTokenizer",

astrai/config/__init__.py

@@ -1,8 +1,16 @@
-from astrai.config.model_config import ModelConfig
+from astrai.config.model_config import (
+    AutoRegressiveLMConfig,
+    BaseModelConfig,
+    ConfigFactory,
+    EncoderConfig,
+)
 from astrai.config.train_config import TrainConfig
 
 __all__ = [
     # Model configuration
-    "ModelConfig",
+    "BaseModelConfig",
+    "AutoRegressiveLMConfig",
+    "EncoderConfig",
+    "ConfigFactory",
     "TrainConfig",
 ]

astrai/config/base.py

@@ -0,0 +1,77 @@
+import json
+from dataclasses import MISSING, dataclass, fields
+from typing import Any, Dict, Optional, Self, get_type_hints
+
+
+@dataclass
+class BaseConfig:
+    def to_dict(self) -> Dict[str, Any]:
+        d = {}
+        for fld in fields(self):
+            v = getattr(self, fld.name)
+            if isinstance(v, (str, int, float, bool)):
+                d[fld.name] = v
+            elif v is None:
+                d[fld.name] = None
+            elif isinstance(v, (dict, list)):
+                try:
+                    json.dumps(v)
+                    d[fld.name] = v
+                except (TypeError, ValueError):
+                    pass
+        return d
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Self:
+        hints = get_type_hints(cls)
+        inst = cls.__new__(cls)
+        for fld in fields(cls):
+            if fld.name in d:
+                v = d[fld.name]
+                target = cls._unwrap_optional(hints.get(fld.name))
+                if target is not None:
+                    try:
+                        v = cls._coerce(v, target)
+                    except (TypeError, ValueError):
+                        pass
+                object.__setattr__(inst, fld.name, v)
+            elif fld.default is not MISSING:
+                object.__setattr__(inst, fld.name, fld.default)
+            elif fld.default_factory is not MISSING:
+                object.__setattr__(inst, fld.name, fld.default_factory())
+            else:
+                object.__setattr__(inst, fld.name, None)
+        return inst
+
+    @staticmethod
+    def _unwrap_optional(tp) -> Optional[type]:
+        if tp is None:
+            return None
+        origin = getattr(tp, "__origin__", None)
+        if origin is not None:
+            args = getattr(tp, "__args__", ())
+            non_none = [a for a in args if a is not type(None)]
+            return non_none[0] if non_none else None
+        return tp
+
+    @staticmethod
+    def _coerce(value: Any, target_type: type) -> Any:
+        if target_type is bool and isinstance(value, bool):
+            return value
+        if (
+            target_type is int
+            and isinstance(value, (int, float))
+            and not isinstance(value, bool)
+        ):
+            return int(value)
+        if (
+            target_type is float
+            and isinstance(value, (int, float))
+            and not isinstance(value, bool)
+        ):
+            return float(value)
+        if target_type is str and isinstance(value, str):
+            return value
+        if isinstance(value, target_type):
+            return value
+        raise TypeError

astrai/config/model_config.py

@@ -1,42 +1,90 @@
 import json
-from dataclasses import asdict, dataclass
-from typing import Optional, Self
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Self
+
+from astrai.config.base import BaseConfig
+from astrai.factory import BaseFactory
+
+
+class ConfigFactory(BaseFactory[BaseConfig]):
+    """Factory that dispatches config classes by ``model_type``."""
+
+    @classmethod
+    def load(cls, raw: Dict[str, Any]) -> BaseConfig:
+        model_type = raw.get("model_type") or "autoregressive_lm"
+        config_cls = cls.get_component_class(model_type)
+        return config_cls.from_dict(raw)
 
 
 @dataclass
-class ModelConfig:
-    # basic config
+class BaseModelConfig(BaseConfig):
+    """Base config with ``model_type`` dispatch and file I/O."""
+
     model_type: Optional[str] = None
+
+    @classmethod
+    def from_file(cls, config_path: str) -> Self:
+        with open(config_path, "r") as f:
+            raw: Dict[str, Any] = json.load(f)
+        return cls.from_dict(raw)
+
+    def to_file(self, config_path: str):
+        d = self.to_dict()
+        config_dict = {k: v for k, v in d.items() if v is not None}
+        with open(config_path, "w") as f:
+            json.dump(config_dict, f, indent=4)
+
+
+@dataclass
+@ConfigFactory.register("autoregressive_lm")
+class AutoRegressiveLMConfig(BaseModelConfig):
+    """Configuration for autoregressive language model."""
+
     vocab_size: Optional[int] = None
     dim: Optional[int] = None
-
     n_layers: Optional[int] = None
     norm_eps: Optional[float] = None
     dim_ffn: Optional[int] = None
     tie_weight: Optional[bool] = None
 
-    # RoPE
     max_len: Optional[int] = None
     rope_theta: Optional[float] = None
 
-    # GQA
+    attn_type: str = "gqa"
     n_heads: Optional[int] = None
     n_kv_heads: Optional[int] = None
     use_qk_norm: Optional[bool] = None
     use_gated_attention: Optional[bool] = None
 
-    def load(self, config_path: str) -> Self:
-        config = {}
-        with open(config_path, "r") as f:
-            config.update(json.load(f))
+    kv_lora_rank: Optional[int] = None
+    qk_nope_head_dim: Optional[int] = None
+    qk_rope_head_dim: Optional[int] = None
 
-        for key, value in config.items():
-            if hasattr(self, key):
-                setattr(self, key, value)
+    ffn_type: str = "mlp"
+    n_routed_experts: Optional[int] = None
+    n_shared_experts: Optional[int] = None
+    n_activated_experts: Optional[int] = None
+    topk_method: Optional[str] = None
 
-        return self
 
-    def save(self, config_path: str):
-        config_dict = {k: v for k, v in asdict(self).items() if v is not None}
-        with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
+@dataclass
+@ConfigFactory.register("embedding")
+class EncoderConfig(BaseModelConfig):
+    """Configuration for embedding encoder model."""
+
+    vocab_size: Optional[int] = None
+    dim: Optional[int] = None
+    n_layers: Optional[int] = None
+    norm_eps: Optional[float] = None
+    dim_ffn: Optional[int] = None
+
+    max_len: Optional[int] = None
+    rope_theta: Optional[float] = None
+
+    n_heads: Optional[int] = None
+    n_kv_heads: Optional[int] = None
+    use_qk_norm: Optional[bool] = None
+    use_gated_attention: Optional[bool] = None
+
+    pooling_type: Optional[str] = None
+    normalize_embeddings: Optional[bool] = None

astrai/config/train_config.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
 from typing import Callable, List, Optional
 
 import torch.nn as nn
@@ -6,27 +6,43 @@ from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import Dataset
 
+from astrai.config.base import BaseConfig
+
+
+def required(**kw):
+    return {"required": True, **kw}
+
 
 @dataclass
-class TrainConfig:
+class TrainConfig(BaseConfig):
     # basic setting
-    model: nn.Module = field(default=None, metadata={"help": "Model for training."})
-    strategy: str = field(default=None, metadata={"help": "Training strategy."})
-    dataset: Dataset = field(default=None, metadata={"help": "Dataset for training."})
+    model: nn.Module = field(
+        default=None, metadata=required(help="Model for training.")
+    )
+    strategy: str = field(default=None, metadata=required(help="Training strategy."))
+    dataset: Dataset = field(
+        default=None, metadata=required(help="Dataset for training.")
+    )
     optimizer_fn: Callable[[nn.Module], Optimizer] = field(
-        default=None, metadata={"help": "Optimizer factory for training."}
+        default=None, metadata=required(help="Optimizer factory for training.")
     )
     scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
-        default=None, metadata={"help": "Scheduler factory for training."}
+        default=None, metadata=required(help="Scheduler factory for training.")
     )
     n_epoch: int = field(default=1, metadata={"help": "Number of epochs for training."})
-    batch_size: int = field(default=4, metadata={"help": "Batch size for training."})
-    accumulation_steps: int = field(
+    batch_per_device: int = field(
+        default=4, metadata={"help": "Batch size per device."}
+    )
+    grad_accum_steps: int = field(
         default=1, metadata={"help": "Number of iterations between steps."}
     )
     max_grad_norm: float = field(
         default=1.0, metadata={"help": "Maximum gradient norm."}
     )
+    gradient_checkpointing_modules: list = field(
+        default_factory=list,
+        metadata={"help": "Module types to enable activation checkpointing for."},
+    )
 
     # checkpoint setting
     start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
@@ -40,6 +56,19 @@ class TrainConfig:
         default=5000, metadata={"help": "Number of iterations between checkpoints."}
     )
 
+    # metric setting
+    log_dir: str = field(
+        default="./checkpoint/logs", metadata={"help": "Directory for metric logs."}
+    )
+    log_interval: int = field(
+        default=100,
+        metadata={"help": "Number of batch iterations between metric logs."},
+    )
+    metrics: List[str] = field(
+        default_factory=lambda: ["loss", "lr"],
+        metadata={"help": "Metrics to record during training."},
+    )
+
     # dataloader setting
     random_seed: int = field(default=3407, metadata={"help": "Random seed."})
     num_workers: int = field(
@@ -66,20 +95,31 @@ class TrainConfig:
     master_port: str = field(
         default="29500", metadata={"help": "Master port for distributed training."}
     )
-    parallel_wrapper: Optional[Callable] = field(
-        default=None, metadata={"help": "Parallel function for training."}
+    parallel_mode: str = field(
+        default="none",
+        metadata={"help": "Parallel strategy: none, ddp, fsdp."},
     )
-    state_dict_fn: Optional[Callable] = field(
-        default=None, metadata={"help": "Parallel function for state  dict saving."}
+    start_method: str = field(
+        default="spawn",
+        metadata={"help": "Multiprocessing start method (spawn/fork/forkserver)."},
     )
 
     # 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."}
     )
+    val_dataset: Optional[Dataset] = field(
+        default=None, metadata={"help": "Dataset for validation."}
+    )
+    val_step: int = field(
+        default=1000,
+        metadata={"help": "Number of optimizer steps between validation runs."},
+    )
+
+    executor_kwargs: dict = field(
+        default_factory=dict,
+        metadata={"help": "Extra kwargs passed to ExecutorFactory.create()."},
+    )
     extra_kwargs: dict = field(
         default_factory=dict, metadata={"help": "Other arguments."}
     )
@@ -88,14 +128,6 @@ class TrainConfig:
         self.validate()
 
     def validate(self):
-        required_fields = [
-            "model",
-            "strategy",
-            "dataset",
-            "optimizer_fn",
-            "scheduler_fn",
-        ]
-
-        for field_name in required_fields:
-            if getattr(self, field_name) is None:
-                raise ValueError(f"{field_name} is required.")
+        for fld in fields(self):
+            if fld.metadata.get("required") and getattr(self, fld.name) is None:
+                raise ValueError(f"TrainConfig.{fld.name} is required but got None.")

astrai/dataset/__init__.py

@@ -1,19 +1,35 @@
 from astrai.dataset.dataset import (
     BaseDataset,
-    BaseSegmentFetcher,
     DatasetFactory,
-    MultiSegmentFetcher,
 )
 from astrai.dataset.sampler import ResumableDistributedSampler
+from astrai.dataset.storage import (
+    BaseSegmentFetcher,
+    BaseStorage,
+    H5Storage,
+    JSONStorage,
+    MultiSegmentFetcher,
+    StorageFactory,
+    detect_format,
+    load_h5,
+    load_json,
+    save_h5,
+    save_json,
+)
 
 __all__ = [
-    # Base classes
     "BaseDataset",
-    # Factory
     "DatasetFactory",
-    # Fetchers
     "BaseSegmentFetcher",
     "MultiSegmentFetcher",
-    # Sampler
+    "BaseStorage",
+    "H5Storage",
+    "JSONStorage",
+    "StorageFactory",
+    "detect_format",
+    "save_h5",
+    "load_h5",
+    "save_json",
+    "load_json",
     "ResumableDistributedSampler",
 ]

astrai/dataset/dataset.py

@@ -1,140 +1,97 @@
 """Dataset implementations with factory pattern for training."""
 
-import bisect
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional
 
 import torch
 from torch import Tensor
 from torch.utils.data import Dataset
 
+from astrai.dataset.storage import (
+    BaseStorage,
+    StorageFactory,
+    detect_format,
+)
 from astrai.factory import BaseFactory
-from astrai.serialization import load_h5
-
-
-class BaseSegmentFetcher:
-    """Fetches data segments across multiple tensor segments.
-
-    Maintains cumulative lengths for efficient range queries across
-    multiple discontinuous segments.
-    """
-
-    def __init__(self, segments: List[Tensor]):
-        self.segments = segments
-        self.cum_lengths = []
-
-        total = 0
-        for seg in segments:
-            total += torch.numel(seg)
-            self.cum_lengths.append(total)
-
-        self.total_length = total
-
-    def __len__(self) -> int:
-        return self.total_length
-
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
-        """Fetch data in the range [begin_idx, end_idx).
-
-        Args:
-            begin_idx: Starting index (inclusive)
-            end_idx: Ending index (exclusive)
-
-        Returns:
-            Concatenated tensor of data in the specified range
-        """
-        if not (
-            0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length
-        ):
-            raise ValueError("begin_idx or end_idx out of bounds")
-        if begin_idx >= end_idx:
-            return torch.tensor([], dtype=torch.long)
-
-        # Find segment boundaries for the range
-        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
-        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
-
-        result_segments = []
-
-        for i in range(seg_start_idx, seg_end_idx + 1):
-            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
-            start = max(begin_idx - prev_cum, 0)
-            end = min(end_idx - prev_cum, len(self.segments[i]))
-            data = self.segments[i][start:end]
-            result_segments.append(data)
-
-        return torch.cat(result_segments, dim=0)
-
-
-class MultiSegmentFetcher:
-    """Manages multiple segment fetchers for different data keys.
-
-    Each key corresponds to a different type of data (e.g., "sequence", "mask").
-    """
-
-    def __init__(self, multi_segments: Dict):
-        self.multi_keys = list(multi_segments.keys())
-        self.multi_fetchers = {
-            key: BaseSegmentFetcher(segments)
-            for key, segments in multi_segments.items()
-        }
-
-    def __len__(self) -> int:
-        """Returns the minimum length across all fetchers."""
-        len_list = [len(seg) for seg in self.multi_fetchers.values()]
-        return min(len_list)
-
-    def key_fetch(
-        self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]
-    ) -> Dict:
-        """Fetch data for specific keys.
-
-        Args:
-            begin_idx: Starting index
-            end_idx: Ending index
-            keys: Single key or list of keys to fetch
-
-        Returns:
-            Dictionary of tensors if multiple keys, single tensor if one key
-        """
-        fetch_dict = {}
-        keys = [keys] if isinstance(keys, str) else keys
-
-        for key in keys:
-            fetcher = self.multi_fetchers[key]
-            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
-            fetch_dict[key] = fetch_tensor
-
-        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
-
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
-        """Fetch all keys."""
-        return self.key_fetch(begin_idx, end_idx, self.multi_keys)
 
 
 class BaseDataset(Dataset, ABC):
     """Abstract base class for all dataset types.
 
     Implements common functionality for window-based data fetching.
+    Uses a storage abstraction for format-agnostic data loading.
     """
 
     def __init__(self, window_size: int, stride: int):
         super().__init__()
-        self.segments = {}
         self.window_size = window_size
         self.stride = stride
-        self.total_samples = None
-        self.fetcher: Optional[MultiSegmentFetcher] = None
+        self.storage: Optional[BaseStorage] = None
 
-    def load(self, load_path: str):
-        """Load dataset from HDF5 file.
+    @property
+    def required_keys(self) -> List[str]:
+        """Return required storage keys for this dataset type.
+
+        Subclasses should override to specify expected keys.
+        """
+        return []
+
+    def _validate_keys(self):
+        if not self.required_keys:
+            return
+        actual_keys = set(self.storage.keys)
+        missing = [k for k in self.required_keys if k not in actual_keys]
+        if missing:
+            raise KeyError(
+                f"Dataset {type(self).__name__} requires keys {self.required_keys}, "
+                f"but storage at {self._load_path} only has {sorted(actual_keys)}. "
+                f"Missing: {missing}"
+            )
+
+    def load(self, load_path: str, storage_type: Optional[str] = None, tokenizer=None):
+        """Load dataset from the given path.
+
+        Auto-detects the storage format if not specified.
 
         Args:
-            load_path: Path to the HDF5 data file
+            load_path: Path to the data directory or file
+            storage_type: Force a specific storage type ("h5", "json"),
+                          or None for auto-detection
+            tokenizer: Callable str -> List[int], used to tokenize raw text
+                       in JSON files. Ignored for HDF5.
+
+        Raises:
+            KeyError: If the loaded storage is missing required keys.
         """
-        self.segments = load_h5(load_path)
-        self.fetcher = MultiSegmentFetcher(self.segments)
-        self.total_samples = len(self.fetcher)
+        if storage_type is None:
+            storage_type = detect_format(load_path)
+        self.storage = StorageFactory.create(storage_type)
+        self._load_path = load_path
+        self.storage.load(load_path, tokenizer=tokenizer)
+        self._validate_keys()
+
+    def load_json(self, load_path: str, tokenizer=None):
+        """Load dataset from JSON files explicitly.
+
+        Args:
+            load_path: Path to the JSON data file or directory
+            tokenizer: Optional tokenizer callable for raw text JSON.
+        """
+        self.load(load_path, storage_type="json", tokenizer=tokenizer)
+
+    @property
+    def count(self) -> int:
+        """Return the total number of raw elements (tokens) in the dataset."""
+        if self.storage is None:
+            return 0
+        return len(self.storage)
+
+    @property
+    def keys(self) -> List[str]:
+        """Return the available data keys."""
+        if self.storage is None:
+            return []
+        return self.storage.keys
 
     def get_index(self, index: int) -> tuple:
         """Calculate begin and end indices for a sample.
@@ -145,10 +102,16 @@ class BaseDataset(Dataset, ABC):
         Returns:
             Tuple of (begin_idx, end_idx)
         """
-        assert self.total_samples > self.window_size
+        if self.storage is None:
+            raise RuntimeError("Dataset not loaded, call load() first")
+        total = len(self.storage)
+        if total <= self.window_size:
+            raise ValueError(
+                f"Data too short: {total} tokens <= window_size {self.window_size}"
+            )
 
-        begin_idx = min(index * self.stride, self.total_samples - 1 - self.window_size)
-        end_idx = min(begin_idx + self.window_size, self.total_samples - 1)
+        begin_idx = min(index * self.stride, total - 1 - self.window_size)
+        end_idx = min(begin_idx + self.window_size, total - 1)
 
         return begin_idx, end_idx
 
@@ -161,10 +124,12 @@ class BaseDataset(Dataset, ABC):
         raise NotImplementedError
 
     def __len__(self) -> int:
-        assert self.total_samples is not None
-        if self.total_samples <= self.window_size:
+        if self.storage is None:
             return 0
-        return (self.total_samples - 1 - self.window_size) // self.stride + 1
+        total = len(self.storage)
+        if total <= self.window_size:
+            return 0
+        return (total - 1 - self.window_size) // self.stride + 1
 
 
 class DatasetFactory(BaseFactory["BaseDataset"]):
@@ -209,6 +174,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         load_path: str,
         window_size: int,
         stride: Optional[int] = None,
+        storage_type: Optional[str] = None,
+        tokenizer=None,
     ) -> "BaseDataset":
         """Create and load a dataset in one step.
 
@@ -217,6 +184,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
             load_path: Path to the data file
             window_size: Window size for data sampling
             stride: Stride between consecutive samples (default: same as window_size)
+            storage_type: Storage type ("h5", "json") or None for auto-detection
+            tokenizer: Callable str -> List[int] for raw text JSON tokenization
 
         Returns:
             Loaded dataset instance
@@ -225,7 +194,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
             stride = window_size
 
         dataset = cls.create(train_type, window_size, stride)
-        dataset.load(load_path)
+        dataset.load(load_path, storage_type=storage_type, tokenizer=tokenizer)
 
         return dataset
 
@@ -235,10 +204,6 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         return cls.list_registered()
 
 
-# ============== Dataset Classes ==============
-# All dataset classes are registered at class definition time using the decorator
-
-
 @DatasetFactory.register("seq")
 class SEQDataset(BaseDataset):
     """Dataset for sequential next-token prediction training."""
@@ -246,8 +211,12 @@ class SEQDataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["sequence"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
+        return self.storage.fetch(begin_idx, end_idx, "sequence")
 
     def __getitem__(self, index):
         begin_idx, end_idx = self.get_index(index)
@@ -265,8 +234,12 @@ class SFTDataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["sequence", "loss_mask"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+        return self.storage.fetch(begin_idx, end_idx, key)
 
     def __getitem__(self, index):
         begin_idx, end_idx = self.get_index(index)
@@ -289,8 +262,12 @@ class DPODataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+        return self.storage.fetch(begin_idx, end_idx, key)
 
     def __getitem__(self, index: int):
         begin_idx, end_idx = self.get_index(index)
@@ -319,8 +296,12 @@ class GRPODataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["prompts", "responses", "masks", "rewards"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+        return self.storage.fetch(begin_idx, end_idx, key)
 
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
         begin_idx, end_idx = self.get_index(index)

astrai/dataset/storage.py

@@ -0,0 +1,301 @@
+"""Storage backends for different data formats.
+
+Each storage handles format-specific loading (HDF5, JSON, etc.) and provides
+a uniform interface for data access and length observation via fetchers.
+"""
+
+import bisect
+import json
+import os
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Union
+
+import h5py
+import torch
+from torch import Tensor
+
+from astrai.factory import BaseFactory
+
+
+def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
+    os.makedirs(file_path, exist_ok=True)
+    full_file_path = os.path.join(file_path, f"{file_name}.h5")
+    with h5py.File(full_file_path, "w") as f:
+        for key, tensors in tensor_group.items():
+            grp = f.create_group(key)
+            for idx, tensor in enumerate(tensors):
+                arr = tensor.cpu().numpy()
+                grp.create_dataset(f"data_{idx}", data=arr)
+
+
+def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
+    tensor_group: Dict[str, List[Tensor]] = {}
+
+    root_path = Path(file_path)
+    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
+
+    for h5_file in h5_files:
+        with h5py.File(h5_file, "r") as f:
+            for key in f.keys():
+                grp = f[key]
+                dsets = []
+                for dset_name in grp.keys():
+                    dset = grp[dset_name]
+                    tensor = torch.from_numpy(dset[:])
+                    if share_memory:
+                        tensor = tensor.share_memory_()
+                    dsets.append(tensor)
+
+                if tensor_group.get(key) is None:
+                    tensor_group[key] = []
+                tensor_group[key].extend(dsets)
+
+    return tensor_group
+
+
+def save_json(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
+    os.makedirs(file_path, exist_ok=True)
+    full_file_path = os.path.join(file_path, f"{file_name}.json")
+    json_data = {}
+    for key, tensors in tensor_group.items():
+        json_data[key] = [tensor.tolist() for tensor in tensors]
+    with open(full_file_path, "w", encoding="utf-8") as f:
+        json.dump(json_data, f, ensure_ascii=False)
+
+
+def load_json(
+    file_path: str,
+    share_memory: bool = True,
+    tokenizer: Optional[Callable[[str], List[int]]] = None,
+) -> Dict[str, List[Tensor]]:
+    """Load tensor data from JSON files.
+
+    Supports two modes:
+    - Pre-tokenized: JSON values are List[List[int]] (token IDs), loaded as-is.
+    - Raw text: JSON values are List[str], tokenized via ``tokenizer`` callable
+      at load time. A ``tokenizer`` receives a str and returns List[int].
+
+    Non-data JSON files (e.g. config.json) with scalar/object values are
+    silently skipped.
+    """
+    tensor_group: Dict[str, List[Tensor]] = {}
+    root_path = Path(file_path)
+    json_files = list(root_path.rglob("*.json")) + list(root_path.rglob("*.jsonl"))
+    for json_file in json_files:
+        with open(json_file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        if not isinstance(data, dict):
+            continue
+        for key, sequences in data.items():
+            if not isinstance(sequences, list):
+                continue
+            tensors = []
+            for seq in sequences:
+                if tokenizer is not None and isinstance(seq, str):
+                    seq = tokenizer(seq)
+                tensor = torch.tensor(seq, dtype=torch.long)
+                if share_memory:
+                    tensor = tensor.share_memory_()
+                tensors.append(tensor)
+            if tensor_group.get(key) is None:
+                tensor_group[key] = []
+            tensor_group[key].extend(tensors)
+    return tensor_group
+
+
+def detect_format(load_path: str) -> str:
+    """Auto-detect storage format from files in the directory.
+
+    Args:
+        load_path: Directory or file path
+
+    Returns:
+        Format string ("h5" or "json")
+
+    Raises:
+        FileNotFoundError: If no supported data files are found
+    """
+    root = Path(load_path)
+    if root.is_file():
+        suffix = root.suffix.lower()
+        if suffix in (".h5", ".hdf5"):
+            return "h5"
+        if suffix in (".json", ".jsonl"):
+            return "json"
+        raise ValueError(f"Unsupported file format: {suffix}")
+
+    h5_files = list(root.rglob("*.h5")) + list(root.rglob("*.hdf5"))
+    if h5_files:
+        return "h5"
+    json_files = list(root.rglob("*.json")) + list(root.rglob("*.jsonl"))
+    if json_files:
+        return "json"
+    raise FileNotFoundError(f"No supported data files found at {load_path}")
+
+
+class BaseSegmentFetcher:
+    """Fetches data segments across multiple tensor segments.
+
+    Maintains cumulative lengths for efficient range queries across
+    multiple discontinuous segments.
+    """
+
+    def __init__(self, segments: List[Tensor]):
+        self.segments = segments
+        self.cum_lengths = []
+
+        total = 0
+        for seg in segments:
+            total += torch.numel(seg)
+            self.cum_lengths.append(total)
+
+        self.total_length = total
+
+    def __len__(self) -> int:
+        return self.total_length
+
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        """Fetch data in the range [begin_idx, end_idx)."""
+        if not (
+            0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length
+        ):
+            raise ValueError("begin_idx or end_idx out of bounds")
+        if begin_idx >= end_idx:
+            return torch.tensor([], dtype=torch.long)
+
+        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
+        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
+
+        result_segments = []
+
+        for i in range(seg_start_idx, seg_end_idx + 1):
+            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
+            start = max(begin_idx - prev_cum, 0)
+            end = min(end_idx - prev_cum, len(self.segments[i]))
+            result_segments.append(self.segments[i][start:end])
+
+        return torch.cat(result_segments, dim=0)
+
+
+class MultiSegmentFetcher:
+    """Manages multiple segment fetchers for different data keys."""
+
+    def __init__(self, multi_segments: Dict):
+        self.multi_keys = list(multi_segments.keys())
+        self.multi_fetchers = {
+            key: BaseSegmentFetcher(segments)
+            for key, segments in multi_segments.items()
+        }
+
+    def __len__(self) -> int:
+        """Returns the minimum length across all fetchers."""
+        if not self.multi_fetchers:
+            return 0
+        len_list = [len(seg) for seg in self.multi_fetchers.values()]
+        return min(len_list)
+
+    def key_fetch(
+        self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]
+    ) -> Dict:
+        """Fetch data for specific keys."""
+        fetch_dict = {}
+        keys = [keys] if isinstance(keys, str) else keys
+
+        for key in keys:
+            fetcher = self.multi_fetchers[key]
+            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
+            fetch_dict[key] = fetch_tensor
+
+        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
+
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
+        """Fetch all keys."""
+        return self.key_fetch(begin_idx, end_idx, self.multi_keys)
+
+
+class BaseStorage(ABC):
+    """Abstract storage backend for loading and dispatching data.
+
+    Storage encapsulates format-specific loading and provides a uniform
+    interface for data access and length observation. Subclasses handle
+    different data formats (HDF5, JSON, etc.) while exposing the same
+    fetch interface.
+    """
+
+    def __init__(self):
+        self._fetcher: Optional[MultiSegmentFetcher] = None
+
+    @abstractmethod
+    def load(self, load_path: str, tokenizer=None) -> None:
+        """Load data from the given path into internal fetcher."""
+        raise NotImplementedError
+
+    def __len__(self) -> int:
+        """Total number of raw elements (tokens) in storage."""
+        if self._fetcher is None:
+            return 0
+        return len(self._fetcher)
+
+    def fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]):
+        """Fetch data for the given keys and index range.
+
+        Args:
+            begin_idx: Starting index (inclusive)
+            end_idx: Ending index (exclusive)
+            keys: Single key or list of keys to fetch
+
+        Returns:
+            Tensor if single key, Dict[str, Tensor] if multiple keys
+        """
+        if self._fetcher is None:
+            raise RuntimeError("Storage not loaded")
+        return self._fetcher.key_fetch(begin_idx, end_idx, keys)
+
+    @property
+    def keys(self) -> List[str]:
+        """Return the data keys available in this storage."""
+        if self._fetcher is None:
+            return []
+        return self._fetcher.multi_keys
+
+
+class StorageFactory(BaseFactory["BaseStorage"]):
+    """Factory for creating storage backends by type name.
+
+    Example:
+        @StorageFactory.register("custom")
+        class CustomStorage(BaseStorage):
+            ...
+
+        storage = StorageFactory.create("custom")
+    """
+
+    @classmethod
+    def _validate_component(cls, storage_cls: type) -> None:
+        if not issubclass(storage_cls, BaseStorage):
+            raise TypeError(f"{storage_cls.__name__} must inherit from BaseStorage")
+
+
+@StorageFactory.register("h5")
+class H5Storage(BaseStorage):
+    """HDF5-based storage backend (pre-tokenized data)."""
+
+    def load(self, load_path: str, tokenizer=None) -> None:
+        segments = load_h5(load_path)
+        self._fetcher = MultiSegmentFetcher(segments)
+
+
+@StorageFactory.register("json")
+class JSONStorage(BaseStorage):
+    """JSON-based storage backend.
+
+    Supports two modes:
+    - Pre-tokenized: JSON values are List[List[int]], loaded as-is.
+    - Raw text: JSON values are List[str], tokenized via ``tokenizer``
+      callable (str -> List[int]) at load time.
+    """
+
+    def load(self, load_path: str, tokenizer=None) -> None:
+        segments = load_json(load_path, tokenizer=tokenizer)
+        self._fetcher = MultiSegmentFetcher(segments)

astrai/factory.py

@@ -1,5 +1,6 @@
 """Base factory class for extensible component registration."""
 
+import inspect
 from abc import ABC
 from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar
 
@@ -122,6 +123,10 @@ class BaseFactory(ABC, Generic[T]):
     def create(cls, name: str, *args, **kwargs) -> T:
         """Create a component instance by name.
 
+        Filters kwargs to match the component's __init__ signature,
+        so components don't need to declare **kwargs just to absorb
+        parameters meant for other components.
+
         Args:
             name: Registered name of the component
             *args: Positional arguments passed to component constructor
@@ -139,6 +144,17 @@ class BaseFactory(ABC, Generic[T]):
                 f"Supported types: {sorted(cls._registry.list_names())}"
             )
         component_cls = cls._registry.get(name)
+        sig = inspect.signature(component_cls.__init__)
+        has_var_kwargs = any(
+            p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
+        )
+        if not has_var_kwargs:
+            valid = {
+                p.name
+                for p in sig.parameters.values()
+                if p.name != "self" and p.kind != inspect.Parameter.VAR_KEYWORD
+            }
+            kwargs = {k: v for k, v in kwargs.items() if k in valid}
         return component_cls(*args, **kwargs)
 
     @classmethod
@@ -155,6 +171,26 @@ class BaseFactory(ABC, Generic[T]):
         """
         pass
 
+    @classmethod
+    def get_component_class(cls, name: str) -> Type[T]:
+        """Get the registered component class by name without instantiating it.
+
+        Args:
+            name: Registered name of the component
+
+        Returns:
+            The component class itself
+
+        Raises:
+            ValueError: If the component name is not registered
+        """
+        if not cls._registry.contains(name):
+            raise ValueError(
+                f"Unknown component: '{name}'. "
+                f"Supported types: {sorted(cls._registry.list_names())}"
+            )
+        return cls._registry.get(name)
+
     @classmethod
     def list_registered(cls) -> list:
         """List all registered component names.

astrai/inference/__init__.py

@@ -1,25 +1,92 @@
-"""Inference module for continuous batching."""
+"""Inference module for continuous batching.
 
+Layers:
+  - core/:           Core inference loop (cache, executor, scheduler, task)
+  - api/:            HTTP protocol handlers (OpenAI, Anthropic)
+  - engine.py:       Facade (InferenceEngine), Value Object (GenerationRequest)
+  - sample.py:       Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
+"""
+
+from astrai.inference.api import (
+    AnthropicHandler,
+    AnthropicMessage,
+    ChatCompletionRequest,
+    ChatMessage,
+    MessagesRequest,
+    OpenAIHandler,
+    ProtocolHandler,
+    StopChecker,
+    StreamContext,
+    app,
+    run_server,
+)
+from astrai.inference.core import (
+    STOP,
+    Allocator,
+    Executor,
+    InferenceScheduler,
+    KVCache,
+    KvcacheView,
+    PagePool,
+    PrefixCache,
+    Storage,
+    Task,
+    TaskManager,
+    TaskStatus,
+    TaskTable,
+    page_hash,
+)
 from astrai.inference.engine import (
     GenerationRequest,
     InferenceEngine,
 )
-from astrai.inference.scheduler import (
-    InferenceScheduler,
-    Task,
-    TaskStatus,
-    apply_sampling_strategies,
+from astrai.inference.sample import (
+    BaseSamplingStrategy,
+    SamplingPipeline,
+    TemperatureStrategy,
+    TopKStrategy,
+    TopPStrategy,
+    sample,
 )
 
 __all__ = [
-    # Engine
+    # Engine / Requests
     "InferenceEngine",
-    # Scheduler
-    "InferenceScheduler",
-    "Task",
-    "TaskStatus",
-    # Request
     "GenerationRequest",
-    # Sampling
-    "apply_sampling_strategies",
+    # Core scheduler
+    "InferenceScheduler",
+    "Executor",
+    "STOP",
+    "Task",
+    "TaskManager",
+    "TaskStatus",
+    # Core cache
+    "Allocator",
+    "KVCache",
+    "KvcacheView",
+    "PagePool",
+    "PrefixCache",
+    "Storage",
+    "TaskTable",
+    "page_hash",
+    # Sampling (Strategy pattern)
+    "sample",
+    "BaseSamplingStrategy",
+    "TemperatureStrategy",
+    "TopKStrategy",
+    "TopPStrategy",
+    "SamplingPipeline",
+    # Protocol
+    "ProtocolHandler",
+    "StopChecker",
+    "StreamContext",
+    "AnthropicHandler",
+    "OpenAIHandler",
+    # Server
+    "ChatMessage",
+    "ChatCompletionRequest",
+    "AnthropicMessage",
+    "MessagesRequest",
+    "app",
+    "run_server",
 ]

astrai/inference/api/__init__.py

@@ -0,0 +1,31 @@
+"""Inference API: protocol handlers and FastAPI server."""
+
+from astrai.inference.api.protocol import (
+    AnthropicHandler,
+    OpenAIHandler,
+    ProtocolHandler,
+    StopChecker,
+    StreamContext,
+)
+from astrai.inference.api.server import (
+    AnthropicMessage,
+    ChatCompletionRequest,
+    ChatMessage,
+    MessagesRequest,
+    app,
+    run_server,
+)
+
+__all__ = [
+    "AnthropicHandler",
+    "OpenAIHandler",
+    "ProtocolHandler",
+    "StopChecker",
+    "StreamContext",
+    "AnthropicMessage",
+    "ChatCompletionRequest",
+    "ChatMessage",
+    "MessagesRequest",
+    "app",
+    "run_server",
+]

astrai/inference/api/protocol.py

@@ -0,0 +1,445 @@
+"""Protocol handlers for OpenAI and Anthropic chat completion APIs.
+
+Template Method + Builder patterns eliminate the 45% code duplication between
+stream/non-stream branches and across protocol adapters.
+"""
+
+import json
+import time
+import uuid
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+from astrai.inference.engine import InferenceEngine
+
+
+def _sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
+    lines: List[str] = []
+    if event:
+        lines.append(f"event: {event}")
+    lines.append(f"data: {json.dumps(data, ensure_ascii=False)}")
+    lines.append("")
+    return "\n".join(lines)
+
+
+def _sse_done() -> str:
+    return "data: [DONE]\n\n"
+
+
+@dataclass
+class StreamContext:
+    """Shared state across the streaming generation lifecycle."""
+
+    resp_id: str
+    created: int
+    model: str
+    prompt_tokens: int
+    completion_tokens: int = 0
+    accumulated: str = ""
+    stop_matched: Optional[str] = None
+    last_yield_trimmed: str = ""
+
+
+class StopChecker:
+    """Scans accumulated text for stop sequence matches."""
+
+    def __init__(self, sequences: List[str]):
+        self._sequences = [s for s in sequences if s]
+
+    def check(self, text: str) -> Optional[str]:
+        for seq in self._sequences:
+            if seq in text:
+                return seq
+        return None
+
+    def trim(self, text: str, matched: str) -> str:
+        idx = text.rfind(matched)
+        return text[:idx] if idx != -1 else text
+
+    @property
+    def has_sequences(self) -> bool:
+        return len(self._sequences) > 0
+
+
+class ProtocolHandler(ABC):
+    """Template-method base for API protocol handlers.
+
+    Subclasses implement format hooks; the base class orchestrates the
+    generate-async loop and SSE/JSON response construction.
+
+    Lifecycle::
+
+        handle()
+          ├─ build_prompt()          # protocol-specific prompt assembly
+          ├─ create_response_id()    # unique response identifier
+          ├─ [stream]
+          │   ├─ format_stream_start()
+          │   ├─ format_stream_token()  × N
+          │   │   └─ on_token() hook for stop-sequence interception
+          │   └─ format_stream_end()
+          └─ [non-stream]
+              ├─ (accumulate tokens)
+              └─ format_non_stream_response()
+    """
+
+    request_model: type[BaseModel]
+
+    def __init__(self, request: BaseModel, engine: InferenceEngine):
+        self.request = request
+        self.engine = engine
+
+    @abstractmethod
+    def build_prompt(self) -> str:
+        """Build the full prompt string from the request messages."""
+
+    @abstractmethod
+    def create_response_id(self) -> str:
+        """Generate a unique response ID following the protocol convention."""
+
+    @abstractmethod
+    def format_stream_start(self, ctx: StreamContext) -> List[str]:
+        """Yield SSE events that open the stream (role marker, metadata)."""
+
+    @abstractmethod
+    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
+        """Yield an SSE event for a single generated token."""
+
+    @abstractmethod
+    def format_stream_end(self, ctx: StreamContext) -> List[str]:
+        """Yield SSE events that close the stream (finish reason, usage stats)."""
+
+    @abstractmethod
+    def format_non_stream_response(
+        self, ctx: StreamContext, content: str
+    ) -> Dict[str, Any]:
+        """Build the JSON response body for non-streaming mode."""
+
+    def get_stop_sequences(self) -> List[str]:
+        return []
+
+    def create_stop_checker(self) -> StopChecker:
+        return StopChecker(self.get_stop_sequences())
+
+    def on_token(
+        self, ctx: StreamContext, token: str, stop_checker: StopChecker
+    ) -> Optional[str]:
+        """Hook after each token is appended to accumulated.
+
+        Return a matched stop-sequence string to break the loop,
+        or None to continue.
+
+        """
+        return None
+
+    async def handle(self) -> Union[StreamingResponse, Dict[str, Any]]:
+        ctx = StreamContext(
+            resp_id=self.create_response_id(),
+            created=int(time.time()),
+            model=self.request.model,
+            prompt_tokens=self._count_prompt_tokens(),
+        )
+
+        agen = self.engine.generate_async(
+            prompt=self.build_prompt(),
+            max_tokens=self.request.max_tokens,
+            temperature=self.request.temperature,
+            top_p=self.request.top_p,
+            top_k=self.request.top_k,
+        )
+
+        if self.request.stream:
+            return self._handle_stream(agen, ctx)
+        else:
+            return await self._handle_non_stream(agen, ctx)
+
+    def _count_prompt_tokens(self) -> int:
+        return len(self.engine.tokenizer.encode(self.build_prompt()))
+
+    def _handle_stream(self, agen, ctx: StreamContext) -> StreamingResponse:
+        stop_checker = self.create_stop_checker()
+
+        async def event_stream():
+            for event in self.format_stream_start(ctx):
+                yield event
+
+            async for token in agen:
+                ctx.completion_tokens += 1
+                ctx.accumulated += token
+
+                matched = self.on_token(ctx, token, stop_checker)
+                if matched:
+                    break
+
+                yield self.format_stream_token(ctx, token)
+
+            for event in self.format_stream_end(ctx):
+                yield event
+            yield _sse_done()
+
+        return StreamingResponse(
+            event_stream(),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
+        )
+
+    async def _handle_non_stream(self, agen, ctx: StreamContext) -> Dict[str, Any]:
+        stop_checker = self.create_stop_checker()
+        chunks: List[str] = []
+
+        async for token in agen:
+            ctx.completion_tokens += 1
+            ctx.accumulated += token
+            chunks.append(token)
+
+            matched = self.on_token(ctx, token, stop_checker)
+            if matched:
+                break
+
+        content = "".join(chunks)
+        return self.format_non_stream_response(ctx, content)
+
+
+def _extract_text_content(content: Union[str, List[Dict[str, Any]]]) -> str:
+    """Extract plain text from an Anthropic content block (string or list)."""
+    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 ""
+
+
+class OpenAIHandler(ProtocolHandler):
+    """OpenAI-compatible /v1/chat/completions handler."""
+
+    def build_prompt(self) -> str:
+        messages = [
+            {"role": m.role, "content": m.content} for m in self.request.messages
+        ]
+        return self.engine.tokenizer.apply_chat_template(messages, tokenize=False)
+
+    def create_response_id(self) -> str:
+        return f"chatcmpl-{uuid.uuid4().hex[:12]}"
+
+    def get_stop_sequences(self) -> List[str]:
+        stop = self.request.stop
+        if stop is None:
+            return []
+        return [stop] if isinstance(stop, str) else stop
+
+    def on_token(
+        self, ctx: StreamContext, token: str, stop_checker: StopChecker
+    ) -> Optional[str]:
+        return stop_checker.check(ctx.accumulated)
+
+    def format_stream_start(self, ctx: StreamContext) -> List[str]:
+        return [
+            _sse_event(
+                {
+                    "id": ctx.resp_id,
+                    "object": "chat.completion.chunk",
+                    "created": ctx.created,
+                    "model": ctx.model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"role": "assistant"},
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+            )
+        ]
+
+    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
+        return _sse_event(
+            {
+                "id": ctx.resp_id,
+                "object": "chat.completion.chunk",
+                "created": ctx.created,
+                "model": ctx.model,
+                "choices": [
+                    {"index": 0, "delta": {"content": token}, "finish_reason": None}
+                ],
+            }
+        )
+
+    def format_stream_end(self, ctx: StreamContext) -> List[str]:
+        return [
+            _sse_event(
+                {
+                    "id": ctx.resp_id,
+                    "object": "chat.completion.chunk",
+                    "created": ctx.created,
+                    "model": ctx.model,
+                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+                }
+            ),
+            _sse_event(
+                {
+                    "prompt_tokens": ctx.prompt_tokens,
+                    "completion_tokens": ctx.completion_tokens,
+                    "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
+                }
+            ),
+        ]
+
+    def format_non_stream_response(
+        self, ctx: StreamContext, content: str
+    ) -> Dict[str, Any]:
+        return {
+            "id": ctx.resp_id,
+            "object": "chat.completion",
+            "created": ctx.created,
+            "model": ctx.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": ctx.prompt_tokens,
+                "completion_tokens": ctx.completion_tokens,
+                "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
+            },
+        }
+
+
+class AnthropicHandler(ProtocolHandler):
+    """Anthropic-compatible /v1/messages handler."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._yielded = ""
+
+    def build_prompt(self) -> str:
+        messages: List[Dict[str, str]] = []
+        system = getattr(self.request, "system", None)
+        if system:
+            messages.append({"role": "system", "content": system})
+        for m in self.request.messages:
+            content = _extract_text_content(m.content)
+            if content:
+                messages.append({"role": m.role, "content": content})
+        return self.engine.tokenizer.apply_chat_template(messages, tokenize=False)
+
+    def create_response_id(self) -> str:
+        return f"msg_{uuid.uuid4().hex[:24]}"
+
+    def get_stop_sequences(self) -> List[str]:
+        return getattr(self.request, "stop_sequences", None) or []
+
+    def on_token(
+        self, ctx: StreamContext, token: str, stop_checker: StopChecker
+    ) -> Optional[str]:
+        matched = stop_checker.check(ctx.accumulated)
+        if not matched:
+            return None
+
+        ctx.stop_matched = matched
+        trimmed = ctx.accumulated[: ctx.accumulated.rfind(matched)]
+        unyielded = trimmed[len(self._yielded) :]
+        if unyielded:
+            ctx.last_yield_trimmed = unyielded
+        return matched
+
+    def format_stream_start(self, ctx: StreamContext) -> List[str]:
+        return [
+            _sse_event(
+                {
+                    "type": "message_start",
+                    "message": {
+                        "id": ctx.resp_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "model": ctx.model,
+                        "content": [],
+                        "usage": {"input_tokens": ctx.prompt_tokens},
+                    },
+                },
+                event="message_start",
+            ),
+            _sse_event(
+                {
+                    "type": "content_block_start",
+                    "index": 0,
+                    "content_block": {"type": "text", "text": ""},
+                },
+                event="content_block_start",
+            ),
+        ]
+
+    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
+        self._yielded += token
+        return _sse_event(
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "text_delta", "text": token},
+            },
+            event="content_block_delta",
+        )
+
+    def format_stream_end(self, ctx: StreamContext) -> List[str]:
+        matched = ctx.stop_matched
+        events: List[str] = []
+        last_yielded = ctx.last_yield_trimmed
+        if last_yielded:
+            events.append(
+                _sse_event(
+                    {
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": last_yielded},
+                    },
+                    event="content_block_delta",
+                )
+            )
+        events.append(
+            _sse_event(
+                {"type": "content_block_stop", "index": 0},
+                event="content_block_stop",
+            )
+        )
+        events.append(
+            _sse_event(
+                {
+                    "type": "message_delta",
+                    "delta": {
+                        "stop_reason": "stop_sequence" if matched else "end_turn",
+                        "stop_sequence": matched,
+                    },
+                    "usage": {"output_tokens": ctx.completion_tokens},
+                },
+                event="message_delta",
+            )
+        )
+        events.append(_sse_event({"type": "message_stop"}, event="message_stop"))
+        return events
+
+    def format_non_stream_response(
+        self, ctx: StreamContext, content: str
+    ) -> Dict[str, Any]:
+        matched = ctx.stop_matched
+        if matched:
+            content = content[: content.rfind(matched)]
+        return {
+            "id": ctx.resp_id,
+            "type": "message",
+            "role": "assistant",
+            "model": ctx.model,
+            "content": [{"type": "text", "text": content}],
+            "stop_reason": "stop_sequence" if matched else "end_turn",
+            "stop_sequence": matched,
+            "usage": {
+                "input_tokens": ctx.prompt_tokens,
+                "output_tokens": ctx.completion_tokens,
+            },
+        }

astrai/inference/api/server.py

@@ -0,0 +1,167 @@
+"""
+OpenAI / Anthropic-compatible chat completion server backed by continuous-batching inference.
+
+Protocol-specific formatting is delegated to ``astrai.inference.protocol``.
+This module owns the FastAPI app, request/response schemas, and dependency wiring.
+"""
+
+import logging
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+from astrai.inference.api.protocol import AnthropicHandler, OpenAIHandler
+from astrai.inference.engine import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+_project_root = Path(__file__).parent.parent.parent
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+
+class ChatCompletionRequest(BaseModel):
+    """OpenAI Chat Completion API request body."""
+
+    model: str = "astrai"
+    messages: List[ChatMessage]
+    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
+    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
+    top_k: Optional[int] = Field(default=50, ge=1)
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
+    max_tokens: Optional[int] = Field(default=2048, ge=1)
+    n: Optional[int] = Field(default=1, ge=1)
+    presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
+    frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
+    logit_bias: Optional[Dict[int, float]] = None
+    user: Optional[str] = None
+
+
+class AnthropicMessage(BaseModel):
+    role: str
+    content: Union[str, List[Dict[str, Any]]]
+
+
+class MessagesRequest(BaseModel):
+    """Anthropic Messages API request body."""
+
+    model: str = "astrai"
+    max_tokens: int = Field(default=1024, ge=1)
+    messages: List[AnthropicMessage]
+    system: Optional[str] = None
+    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
+    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
+    top_k: Optional[int] = Field(default=50, ge=1)
+    stream: Optional[bool] = False
+    stop_sequences: Optional[List[str]] = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    config = app.state.server_config
+    if not config.get("_test", False):
+        try:
+            app.state.engine = _create_engine(**config)
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            raise
+    yield
+    if app.state.engine:
+        app.state.engine.shutdown()
+        logger.info("Inference engine shutdown complete")
+
+
+app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
+
+
+def _create_engine(
+    param_path: Optional[Path] = None,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    max_batch_size: int = 16,
+) -> InferenceEngine:
+    if param_path is None:
+        param_path = _project_root / "params"
+    if not param_path.exists():
+        raise FileNotFoundError(f"Parameter directory not found: {param_path}")
+
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
+    model = AutoModel.from_pretrained(param_path)
+    model.to(device=device, dtype=dtype)
+    logger.info(f"Model loaded on {device} with dtype {dtype}")
+
+    engine = InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=max_batch_size,
+    )
+    logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
+    return engine
+
+
+def _get_engine() -> InferenceEngine:
+    engine = app.state.engine
+    if engine is None:
+        raise HTTPException(status_code=503, detail="Engine not initialized")
+    return engine
+
+
+@app.get("/health")
+async def health():
+    return {
+        "status": "ok",
+        "model_loaded": app.state.engine is not None,
+    }
+
+
+@app.get("/stats")
+async def get_stats():
+    return _get_engine().get_stats()
+
+
+@app.post("/v1/chat/completions")
+async def chat_completion(request: ChatCompletionRequest):
+    engine = _get_engine()
+    handler = OpenAIHandler(request, engine)
+    return await handler.handle()
+
+
+@app.post("/v1/messages")
+async def create_message(request: MessagesRequest):
+    engine = _get_engine()
+    handler = AnthropicHandler(request, engine)
+    return await handler.handle()
+
+
+def run_server(
+    host: str = "0.0.0.0",
+    port: int = 8000,
+    reload: bool = False,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    param_path: Optional[Path] = None,
+    max_batch_size: int = 16,
+):
+    app.state.server_config = {
+        "device": device,
+        "dtype": dtype,
+        "param_path": param_path,
+        "max_batch_size": max_batch_size,
+    }
+    uvicorn.run(
+        app,
+        host=host,
+        port=port,
+        reload=reload,
+    )

astrai/inference/core/__init__.py

@@ -0,0 +1,32 @@
+"""Inference core: cache, executor, scheduler, task management."""
+
+from astrai.inference.core.cache import (
+    Allocator,
+    KVCache,
+    KvcacheView,
+    PagePool,
+    PrefixCache,
+    Storage,
+    TaskTable,
+    page_hash,
+)
+from astrai.inference.core.executor import Executor
+from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
+
+__all__ = [
+    "Allocator",
+    "KVCache",
+    "KvcacheView",
+    "PagePool",
+    "PrefixCache",
+    "Storage",
+    "TaskTable",
+    "page_hash",
+    "Executor",
+    "InferenceScheduler",
+    "STOP",
+    "Task",
+    "TaskManager",
+    "TaskStatus",
+]

astrai/inference/core/cache.py

@@ -0,0 +1,372 @@
+import threading
+from collections import OrderedDict
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+
+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 Allocator:
+    """Bitmask-based page allocator with ref-counting and LRU eviction."""
+
+    def __init__(self, n_pages: int):
+        self._free_mask = (1 << n_pages) - 1
+        self._refs: List[int] = [0] * n_pages
+        self._lru: OrderedDict[int, None] = OrderedDict()
+        self.on_evict: Optional[Callable[[int], None]] = None
+        self._lock = threading.Lock()
+
+    def alloc(self) -> int:
+        with self._lock:
+            if self._free_mask:
+                lsb = self._free_mask & -self._free_mask
+                idx = lsb.bit_length() - 1
+                self._free_mask ^= lsb
+                self._refs[idx] = 1
+                return idx
+            if self._lru:
+                idx, _ = self._lru.popitem(last=False)
+                if self.on_evict:
+                    self.on_evict(idx)
+                self._refs[idx] = 1
+                self._free_mask &= ~(1 << idx)
+                return idx
+            return -1
+
+    def free(self, idx: int, keep_cached: bool = False) -> None:
+        with self._lock:
+            self._refs[idx] -= 1
+            if self._refs[idx] == 0:
+                if keep_cached:
+                    self._lru[idx] = None
+                else:
+                    self._free_mask |= 1 << idx
+
+    def inc_ref(self, idx: int) -> None:
+        with self._lock:
+            self._refs[idx] += 1
+            self._lru.pop(idx, None)
+
+    def ref_count(self, idx: int) -> int:
+        with self._lock:
+            return self._refs[idx]
+
+    def touch(self, idx: int) -> None:
+        with self._lock:
+            self._lru.move_to_end(idx)
+
+
+class PrefixCache:
+    """Hash-based prefix matching: maps page hashes to physical page indices."""
+
+    def __init__(self, page_size: int):
+        self._page_size = page_size
+        self._page_to_hash: Dict[int, int] = {}
+        self._hash_to_page: Dict[int, int] = {}
+        self._lock = threading.Lock()
+
+    def evict(self, idx: int) -> None:
+        with self._lock:
+            h = self._page_to_hash.pop(idx, None)
+            if h is not None:
+                self._hash_to_page.pop(h, None)
+
+    def has_page(self, idx: int) -> bool:
+        with self._lock:
+            return idx in self._page_to_hash
+
+    def lookup(self, token_ids: List[int]) -> List[int]:
+        with self._lock:
+            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 record(
+        self, page_idx: int, token_ids: List[int], logical_page_idx: int
+    ) -> None:
+        with self._lock:
+            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
+
+
+class PagePool:
+    """Orchestrates allocator (page management) and PrefixCache (content addressing)."""
+
+    def __init__(self, allocator: Allocator, prefix: PrefixCache):
+        self._alloc = allocator
+        self._prefix = prefix
+        self._alloc.on_evict = prefix.evict
+
+    @property
+    def allocator(self) -> Allocator:
+        return self._alloc
+
+    @property
+    def prefix(self) -> PrefixCache:
+        return self._prefix
+
+    def alloc(self) -> int:
+        return self._alloc.alloc()
+
+    def free(self, idx: int) -> None:
+        keep = self._prefix.has_page(idx)
+        self._alloc.free(idx, keep_cached=keep)
+        if not keep:
+            self._prefix.evict(idx)
+
+    def inc_ref(self, idx: int) -> None:
+        self._alloc.inc_ref(idx)
+
+    def lookup(self, token_ids: List[int]) -> List[int]:
+        hits = self._prefix.lookup(token_ids)
+        for p in hits:
+            self._alloc.touch(p)
+        return hits
+
+    def record(
+        self, page_idx: int, token_ids: List[int], logical_page_idx: int
+    ) -> None:
+        self._prefix.record(page_idx, token_ids, logical_page_idx)
+
+
+class TaskTable:
+    """Maps task_ids to page tables and cached token counts."""
+
+    def __init__(self, page_size: int):
+        self._page_size = page_size
+        self._pages: Dict[str, List[int]] = {}
+        self._cached: Dict[str, int] = {}
+        self._lock = threading.Lock()
+
+    def set(self, task_id: str, page_table: List[int], cached: int) -> None:
+        with self._lock:
+            self._pages[task_id] = page_table
+            self._cached[task_id] = cached
+
+    def get(self, task_id: str) -> List[int]:
+        with self._lock:
+            return self._pages.get(task_id, [])
+
+    def get_cached(self, task_id: str) -> int:
+        with self._lock:
+            return self._cached.get(task_id, 0)
+
+    def pop(self, task_id: str) -> Tuple[List[int], int]:
+        with self._lock:
+            pages = self._pages.pop(task_id, [])
+            cached = self._cached.pop(task_id, 0)
+            return pages, cached
+
+    def get_ref(self, task_id: str) -> List[int]:
+        with self._lock:
+            return self._pages.setdefault(task_id, [])
+
+    def table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor:
+        with self._lock:
+            states = [self._pages.get(tid, []) for tid in task_ids]
+            max_pages = max((len(s) for s in states), default=0)
+            rows = [s + [-1] * (max_pages - len(s)) for s in states]
+            return torch.tensor(rows, dtype=torch.long, device=device)
+
+
+class Storage:
+    """KV-cache tensor storage with paged write/gather."""
+
+    def __init__(
+        self,
+        n_layers: int,
+        n_pages: int,
+        page_size: int,
+        n_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        self.page_size = page_size
+        self.k_cache = torch.empty(
+            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
+            device=device,
+            dtype=dtype,
+        )
+        self.v_cache = torch.empty(
+            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
+            device=device,
+            dtype=dtype,
+        )
+
+    def write(
+        self,
+        layer_id: int,
+        page_table: Tensor,
+        start_pos: int,
+        k: Tensor,
+        v: Tensor,
+    ) -> None:
+        seq_len = k.size(1)
+        if seq_len == 0:
+            return
+        page_size = self.page_size
+        written = 0
+        first_page = start_pos // page_size
+        last_page = (start_pos + seq_len - 1) // page_size
+        for pi in range(first_page, last_page + 1):
+            phys_pages = page_table[:, pi]
+            page_start = pi * page_size
+            write_start = max(page_start, start_pos)
+            write_end = min(page_start + page_size, start_pos + seq_len)
+            offset = write_start - page_start
+            chunk = write_end - write_start
+            valid = phys_pages >= 0
+            if not valid.all():
+                if valid.any():
+                    valid_pages = phys_pages[valid]
+                    self.k_cache[layer_id, valid_pages, offset : offset + chunk] = k[
+                        valid, written : written + chunk
+                    ]
+                    self.v_cache[layer_id, valid_pages, offset : offset + chunk] = v[
+                        valid, written : written + chunk
+                    ]
+                written += chunk
+                continue
+            self.k_cache[layer_id, phys_pages, offset : offset + chunk] = k[
+                :, written : written + chunk
+            ]
+            self.v_cache[layer_id, phys_pages, offset : offset + chunk] = v[
+                :, written : written + chunk
+            ]
+            written += chunk
+
+    def gather(
+        self, layer_id: int, page_table: Tensor, total_len: int
+    ) -> Tuple[Tensor, Tensor]:
+        safe = page_table.clamp(min=0)
+        k = self.k_cache[layer_id, safe]
+        v = self.v_cache[layer_id, safe]
+        k = k.flatten(1, 2)
+        v = v.flatten(1, 2)
+        if (page_table < 0).any():
+            invalid = (
+                (page_table < 0)
+                .unsqueeze(-1)
+                .expand(-1, -1, self.page_size)
+                .flatten(1, 2)
+            )
+            invalid = invalid[:, :, None, None].expand_as(k)
+            k = k.masked_fill(invalid, 0.0)
+            v = v.masked_fill(invalid, 0.0)
+        k = k[:, :total_len]
+        v = v[:, :total_len]
+        return k, v
+
+
+class KvcacheView:
+    """Bundles Storage + page_table + total_len for attention layers."""
+
+    def __init__(self, storage: Storage, page_table: Tensor, total_len: int = 0):
+        self._storage = storage
+        self._page_table = page_table
+        self._total_len = total_len
+
+    def write(self, layer_id: int, k: Tensor, v: Tensor) -> None:
+        start_pos = self._total_len - k.size(1)
+        self._storage.write(layer_id, self._page_table, start_pos, k, v)
+
+    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
+        return self._storage.gather(layer_id, self._page_table, self._total_len)
+
+
+class KVCache:
+    """Facade: page management + KV-cache I/O for continuous batching."""
+
+    def __init__(
+        self,
+        n_layers: int,
+        n_pages: int,
+        page_size: int,
+        n_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        self.page_size = page_size
+        self._pool = PagePool(Allocator(n_pages), PrefixCache(page_size))
+        self._table = TaskTable(page_size)
+        self._storage = Storage(
+            n_layers, n_pages, page_size, n_kv_heads, head_dim, device, dtype
+        )
+
+    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
+        hits = self._pool.lookup(prompt_ids)
+        cached = len(hits) * self.page_size
+        for p in hits:
+            self._pool.inc_ref(p)
+
+        remaining = len(prompt_ids) - cached
+        n_new = (
+            (remaining + self.page_size - 1) // self.page_size if remaining > 0 else 0
+        )
+        new_pages: List[int] = []
+        if n_new > 0:
+            for _ in range(n_new):
+                p = self._pool.alloc()
+                if p < 0:
+                    for hp in hits:
+                        self._pool.free(hp)
+                    for np in new_pages:
+                        self._pool.free(np)
+                    return False
+                new_pages.append(p)
+
+        self._table.set(task_id, hits + new_pages, cached)
+        return True
+
+    def task_free(self, task_id: str) -> None:
+        page_table, _ = self._table.pop(task_id)
+        for idx in page_table:
+            self._pool.free(idx)
+
+    def task_extend(self, task_id: str, pos: int) -> bool:
+        page_table = self._table.get(task_id)
+        needed = (pos + 1 + self.page_size - 1) // self.page_size
+        while len(page_table) < needed:
+            p = self._pool.alloc()
+            if p < 0:
+                return False
+            page_table.append(p)
+        return True
+
+    def task_cached(self, task_id: str) -> int:
+        return self._table.get_cached(task_id)
+
+    def task_record_hashes(
+        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
+    ) -> None:
+        page_table = self._table.get(task_id)
+        full_pages = len(prompt_ids) // self.page_size
+        for i in range(start_logical_page, full_pages):
+            self._pool.record(page_table[i], prompt_ids, i)
+
+    def make_table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor:
+        return self._table.table_tensor(task_ids, device)
+
+    def bind(self, page_table: Tensor, total_len: int = 0) -> KvcacheView:
+        return KvcacheView(self._storage, page_table, total_len)

astrai/inference/core/executor.py

@@ -0,0 +1,96 @@
+import logging
+from typing import List, Optional
+
+import torch
+
+from astrai.inference.core.cache import KVCache
+from astrai.inference.core.task import Task
+from astrai.inference.sample import sample
+from astrai.model.automodel import AutoModel
+from astrai.tokenize.tokenizer import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class Executor:
+    """Model forward passes for prefill and decode phases."""
+
+    def __init__(
+        self,
+        model: AutoModel,
+        tokenizer: AutoTokenizer,
+        page_cache: KVCache,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.page_cache = page_cache
+        self.device = device or next(model.parameters()).device
+        self.dtype = dtype or next(model.parameters()).dtype
+
+    def execute_prefill(
+        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
+    ) -> None:
+        if start_pos >= prompt_len:
+            return
+
+        tasks = sorted(tasks, key=lambda t: t.task_id)
+        batch_sz = len(tasks)
+
+        input_ids = torch.tensor(
+            [t.prompt_ids[start_pos:prompt_len] for t in tasks],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        task_ids = [t.task_id for t in tasks]
+        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
+
+        with torch.inference_mode():
+            self.model(
+                input_ids,
+                position_ids=torch.arange(
+                    start_pos, prompt_len, dtype=torch.long, device=self.device
+                )
+                .unsqueeze(0)
+                .expand(batch_sz, -1),
+                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
+            )
+
+    def execute_decode(self, tasks: List[Task]) -> List[int]:
+        if not tasks:
+            return []
+
+        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,
+        )
+
+        position_ids = torch.tensor(
+            [t.next_pos for t in tasks], dtype=torch.long, device=self.device
+        )
+        total_len = position_ids.max().item() + 1
+
+        task_ids = [t.task_id for t in tasks]
+        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
+
+        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),
+                paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
+                position_ids=position_ids.unsqueeze(1),
+            )
+            logits = outputs["logits"][:, -1, :]
+
+        return sample(
+            logits,
+            temperature=temperatures,
+            top_k=top_ks,
+            top_p=top_ps,
+        ).tolist()

astrai/inference/core/scheduler.py

@@ -0,0 +1,202 @@
+import logging
+import threading
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from astrai.inference.core.cache import KVCache
+from astrai.inference.core.executor import Executor
+from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
+from astrai.model.automodel import AutoModel
+from astrai.tokenize.tokenizer import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class InferenceScheduler:
+    """Four-phase continuous batching loop: cleanup -> refill -> prefill -> decode."""
+
+    def __init__(
+        self,
+        model: AutoModel,
+        tokenizer: AutoTokenizer,
+        max_batch_size: int = 16,
+        max_seq_len: Optional[int] = None,
+        max_prompt_len: int = 2048,
+        page_size: int = 64,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        config = model.config
+
+        if max_seq_len is not None:
+            self.max_seq_len = max_seq_len
+        elif config.max_len is not None:
+            self.max_seq_len = config.max_len
+        else:
+            raise ValueError(
+                "max_seq_len must be provided either as argument "
+                "or in model config (config.max_len)"
+            )
+        self.device = device or next(model.parameters()).device
+        self.dtype = dtype or next(model.parameters()).dtype
+
+        n_pages = (
+            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
+        ) // page_size
+
+        self._page_cache = KVCache(
+            config.n_layers,
+            n_pages,
+            page_size,
+            config.n_kv_heads,
+            config.dim // config.n_heads,
+            self.device,
+            self.dtype,
+        )
+
+        self._task_mgr = TaskManager(
+            tokenizer=tokenizer,
+            max_batch_size=max_batch_size,
+            max_seq_len=self.max_seq_len,
+            max_prompt_len=max_prompt_len,
+        )
+
+        self._executor = Executor(
+            model=model,
+            tokenizer=tokenizer,
+            page_cache=self._page_cache,
+            device=self.device,
+            dtype=self.dtype,
+        )
+
+        self._running = False
+
+    def add_task(self, prompt: str, **kwargs) -> str:
+        return self._task_mgr.add_task(prompt, **kwargs)
+
+    def remove_task(self, task_id: str) -> None:
+        for task in self._task_mgr.remove_task(task_id):
+            self._page_cache.task_free(task.task_id)
+
+    def get_stats(self) -> Dict[str, Any]:
+        return self._task_mgr.get_stats()
+
+    def _run_generation_loop(self) -> None:
+        stop_ids = self._task_mgr.tokenizer.stop_ids
+        try:
+            while self._running:
+                finished = self._task_mgr.remove_finished_tasks(stop_ids)
+                for task in finished:
+                    self._page_cache.task_free(task.task_id)
+
+                active = self._task_mgr.get_active_tasks()
+                available = self._task_mgr.max_batch_size - len(active)
+                if available > 0:
+                    candidates = self._task_mgr.pull_candidates(available)
+                    failed = []
+                    for task in candidates:
+                        if self._page_cache.task_alloc(task.task_id, task.prompt_ids):
+                            self._task_mgr.activate(task)
+                        else:
+                            failed.append(task)
+                    if failed:
+                        self._task_mgr.return_to_waiting(failed)
+
+                if not self._task_mgr.has_work():
+                    self._task_mgr.wait_for_tasks(timeout=1.0)
+                    continue
+
+                to_prefill = [
+                    t for t in self._task_mgr.get_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),
+                            self._page_cache.task_cached(t.task_id),
+                        )
+                        groups.setdefault(key, []).append(t)
+
+                    for (prompt_len, start_pos), group in groups.items():
+                        self._executor.execute_prefill(group, prompt_len, start_pos)
+                        start_logical_page = start_pos // self._page_cache.page_size
+                        for t in group:
+                            self._page_cache.task_record_hashes(
+                                t.task_id,
+                                t.prompt_ids,
+                                start_logical_page=start_logical_page,
+                            )
+
+                pos_groups: Dict[int, List[Task]] = {}
+                for t in self._task_mgr.get_active_tasks():
+                    pos_groups.setdefault(t.next_pos, []).append(t)
+
+                if pos_groups:
+                    best_key = max(pos_groups, key=lambda k: len(pos_groups[k]))
+                    group = sorted(pos_groups[best_key], key=lambda t: t.task_id)
+
+                    valid: List[Task] = []
+                    for t in group:
+                        if self._page_cache.task_extend(t.task_id, t.next_pos):
+                            valid.append(t)
+                        else:
+                            t.status = TaskStatus.ABORTED
+                            if t.stream_callback:
+                                t.stream_callback(STOP)
+
+                    if valid:
+                        next_tokens = self._executor.execute_decode(valid)
+
+                        for t, ntok in zip(valid, next_tokens):
+                            t.output_ids.append(ntok)
+                            t.output_tokens += 1
+                            pos = t.input_tokens + t.output_tokens
+                            extend_ok = self._page_cache.task_extend(t.task_id, pos)
+                            if t.stream_callback:
+                                t.stream_callback(
+                                    self._task_mgr.tokenizer.decode([ntok])
+                                )
+                            if not extend_ok:
+                                t.status = TaskStatus.ABORTED
+                                if t.stream_callback:
+                                    t.stream_callback(STOP)
+
+                        for t in valid:
+                            if t.is_finished(stop_ids):
+                                if t.stream_callback:
+                                    t.stream_callback(STOP)
+
+        except Exception as e:
+            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
+            for task in self._task_mgr.get_active_tasks():
+                if task.stream_callback:
+                    task.stream_callback(STOP)
+                self._page_cache.task_free(task.task_id)
+            for task in self._task_mgr.get_waiting_tasks():
+                if task.stream_callback:
+                    task.stream_callback(STOP)
+            self._task_mgr.clear_queues()
+            raise
+
+    def start(self) -> None:
+        if not self._running:
+            self._running = True
+            t = threading.Thread(target=self._run_generation_loop, daemon=True)
+            t.start()
+            self._loop_thread = t
+
+    def stop(self) -> None:
+        self._running = False
+        self._task_mgr.wake()
+        if hasattr(self, "_loop_thread"):
+            self._loop_thread.join(timeout=2.0)
+        for task in self._task_mgr.get_active_tasks():
+            self._page_cache.task_free(task.task_id)
+        self._task_mgr.clear_queues()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()

astrai/inference/core/task.py

@@ -0,0 +1,206 @@
+import logging
+import threading
+import time
+import uuid
+from collections import deque
+from enum import Enum
+from typing import Any, Callable, Deque, Dict, List, Optional
+
+from astrai.tokenize.tokenizer import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+STOP = object()
+
+
+class TaskStatus(Enum):
+    """Task lifecycle states."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    FINISHED = "finished"
+    ABORTED = "aborted"
+
+
+class Task:
+    """Single generation request: prompt, sampling params, output state."""
+
+    def __init__(
+        self,
+        task_id: str,
+        prompt_ids: List[int],
+        max_tokens: Optional[int] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ):
+        self.task_id = task_id
+        self.prompt_ids = prompt_ids
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
+
+        self.status = TaskStatus.PENDING
+        self.output_ids: List[int] = []
+        self.input_tokens: int = 0
+        self.output_tokens: int = 0
+        self.arrival_time = time.time()
+        self.finish_time: Optional[float] = None
+        self.stream_callback = stream_callback
+
+    @property
+    def next_pos(self) -> int:
+        return self.input_tokens + len(self.output_ids)
+
+    def is_finished(self, stop_ids: List[int]) -> bool:
+        if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
+            return True
+        if self.output_ids and self.output_ids[-1] in stop_ids:
+            return True
+        return False
+
+
+class TaskManager:
+    """Thread-safe task queues and lifecycle transitions (no page ops)."""
+
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        max_batch_size: int = 16,
+        max_seq_len: int = 8192,
+        max_prompt_len: int = 512,
+    ):
+        self.tokenizer = tokenizer
+        self.max_batch_size = max_batch_size
+        self.max_seq_len = max_seq_len
+        self.max_prompt_len = max_prompt_len
+
+        self.waiting_queue: Deque[Task] = deque()
+        self.active_tasks: List[Task] = []
+
+        self._task_event = threading.Event()
+        self._lock = threading.Lock()
+
+        self._total_tasks = 0
+        self._total_tokens = 0
+
+    def add_task(
+        self,
+        prompt: str,
+        max_tokens: Optional[int] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
+        prompt_ids = self.tokenizer.encode(prompt)
+        if len(prompt_ids) > self.max_prompt_len:
+            prompt_ids = prompt_ids[-self.max_prompt_len :]
+
+        if len(prompt_ids) >= self.max_seq_len:
+            if stream_callback:
+                stream_callback(STOP)
+            return task_id
+
+        if max_tokens is None:
+            max_tokens = self.max_seq_len - len(prompt_ids)
+        else:
+            max_tokens = min(max_tokens, self.max_seq_len - len(prompt_ids))
+
+        task = Task(
+            task_id=task_id,
+            prompt_ids=prompt_ids,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            stream_callback=stream_callback,
+        )
+
+        with self._lock:
+            self.waiting_queue.append(task)
+            self._total_tasks += 1
+
+        self._task_event.set()
+        return task_id
+
+    def remove_task(self, task_id: str) -> List[Task]:
+        with self._lock:
+            removed_active = [t for t in self.active_tasks if t.task_id == task_id]
+            self.waiting_queue = deque(
+                t for t in self.waiting_queue if t.task_id != task_id
+            )
+            self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
+        return removed_active
+
+    def get_stats(self) -> Dict[str, Any]:
+        return {
+            "total_tasks": self._total_tasks,
+            "total_tokens": self._total_tokens,
+            "active_tasks": len(self.active_tasks),
+            "waiting_queue": len(self.waiting_queue),
+        }
+
+    def remove_finished_tasks(self, stop_ids: List[int]) -> List[Task]:
+        with self._lock:
+            finished = []
+            for task in self.active_tasks:
+                if task.status == TaskStatus.ABORTED:
+                    task.finish_time = time.time()
+                    finished.append(task)
+                elif task.is_finished(stop_ids):
+                    task.status = TaskStatus.FINISHED
+                    task.finish_time = time.time()
+                    finished.append(task)
+                    self._total_tokens += task.output_tokens
+
+            self.active_tasks = [
+                t
+                for t in self.active_tasks
+                if t.status not in (TaskStatus.FINISHED, TaskStatus.ABORTED)
+            ]
+            return finished
+
+    def pull_candidates(self, n: int) -> List[Task]:
+        to_add: List[Task] = []
+        with self._lock:
+            take = min(n, len(self.waiting_queue))
+            for _ in range(take):
+                to_add.append(self.waiting_queue.popleft())
+        return to_add
+
+    def activate(self, task: Task) -> None:
+        task.status = TaskStatus.RUNNING
+        with self._lock:
+            self.active_tasks.append(task)
+
+    def return_to_waiting(self, tasks: List[Task]) -> None:
+        with self._lock:
+            for task in reversed(tasks):
+                self.waiting_queue.appendleft(task)
+
+    def has_work(self) -> bool:
+        return bool(self.active_tasks or self.waiting_queue)
+
+    def wait_for_tasks(self, timeout: float = 1.0) -> None:
+        self._task_event.clear()
+        self._task_event.wait(timeout=timeout)
+
+    def get_active_tasks(self) -> List[Task]:
+        with self._lock:
+            return list(self.active_tasks)
+
+    def get_waiting_tasks(self) -> List[Task]:
+        with self._lock:
+            return list(self.waiting_queue)
+
+    def clear_queues(self) -> None:
+        with self._lock:
+            self.waiting_queue.clear()
+            self.active_tasks.clear()
+
+    def wake(self) -> None:
+        self._task_event.set()

astrai/inference/engine.py

@@ -1,17 +1,77 @@
-"""Unified inference engine."""
+"""Unified inference engine for continuous batching."""
 
+import asyncio
 import gc
-import logging
 import threading
-from typing import Any, Dict, Generator, List, Optional, Union
+from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 
-from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.core.task import STOP
 from astrai.tokenize import AutoTokenizer
 
-logger = logging.getLogger(__name__)
+
+def _validate_sampling_params(
+    top_k: int, top_p: float, temperature: float, max_tokens: Optional[int] = None
+):
+    if not (isinstance(top_k, int) and top_k >= 0):
+        raise ValueError("top_k must be a non-negative integer")
+    if not (0.0 <= top_p <= 1.0):
+        raise ValueError("top_p must be a float between 0.0 and 1.0")
+    if not (isinstance(temperature, (int, float)) and temperature >= 0):
+        raise ValueError("temperature must be a non-negative number")
+
+
+class GenerateResult:
+    """Thread-safe token accumulator for streaming and non-streaming modes."""
+
+    def __init__(self, count: int = 1):
+        self._cond = threading.Condition()
+        self._event = threading.Event()
+        self.tokens: List[Tuple[int, str]] = []
+        self.results: List[str] = [""] * count
+        self._done: List[bool] = [False] * count
+        self._completed = 0
+        self._total = count
+
+    def append(self, token: str, idx: int = 0):
+        with self._cond:
+            self.tokens.append((idx, token))
+            if token is not STOP:
+                self.results[idx] += token
+            else:
+                if not self._done[idx]:
+                    self._done[idx] = True
+                    self._completed += 1
+                    self._cond.notify_all()
+            self._event.set()
+
+    def pop_all(self) -> List[Tuple[int, str]]:
+        with self._cond:
+            out = self.tokens.copy()
+            self.tokens.clear()
+            if not out:
+                self._event.clear()
+            return out
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        return self._event.wait(timeout=timeout)
+
+    def wait_completion(self, timeout: float = 300.0) -> None:
+        with self._cond:
+            if not self._cond.wait_for(
+                lambda: self._completed >= self._total, timeout=timeout
+            ):
+                raise TimeoutError(
+                    f"Generation timeout after {timeout}s "
+                    f"({self._completed}/{self._total} completed)"
+                )
+
+    def get_results(self) -> List[str]:
+        with self._cond:
+            return self.results.copy()
 
 
 class GenerationRequest:
@@ -23,73 +83,21 @@ class GenerationRequest:
         top_k: int = 50,
         top_p: float = 1.0,
         temperature: float = 1.0,
-        max_len: int = 1024,
+        max_tokens: Optional[int] = None,
         stream: bool = False,
     ):
+        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
+
         self.messages = messages
         self.top_k = top_k
         self.top_p = top_p
         self.temperature = temperature
-        self.max_len = max_len
+        self.max_tokens = max_tokens
         self.stream = stream
 
-        self._validate()
-
-    def _validate(self):
-        """Validate request parameters."""
-        if not (isinstance(self.top_k, int) and self.top_k >= 0):
-            raise ValueError("top_k must be a non-negative integer")
-        if not (0.0 <= self.top_p <= 1.0):
-            raise ValueError("top_p must be a float between 0.0 and 1.0")
-        if not (isinstance(self.temperature, (int, float)) and self.temperature >= 0):
-            raise ValueError("temperature must be a non-negative number")
-
-
-class _Result:
-    """Unified result holder for streaming/non-streaming modes."""
-
-    def __init__(self, count: int = 1, stream: bool = False):
-        self._stream = stream
-        self._lock = threading.Lock()
-        self._event = threading.Event()
-        self.tokens: List[str] = []
-        self.results: List[str] = [""] * count if count > 1 else [""]
-        self.done_flags: List[bool] = [False] * count
-        self._completed_count = 0
-
-    def append(self, token: str, idx: int = 0):
-        with self._lock:
-            if self._stream:
-                self.tokens.append(token)
-            else:
-                if token == "[DONE]":
-                    if not self.done_flags[idx]:
-                        self.done_flags[idx] = True
-                        self._completed_count += 1
-                        if self._completed_count == len(self.results):
-                            self._event.set()
-                else:
-                    self.results[idx] += token
-        self._event.set()
-
-    def pop_all(self) -> List[str]:
-        with self._lock:
-            tokens = self.tokens.copy()
-            self.tokens.clear()
-            if not tokens:
-                self._event.clear()
-            return tokens
-
-    def wait(self, timeout: float = None) -> bool:
-        return self._event.wait(timeout=timeout)
-
-    def get_results(self) -> List[str]:
-        with self._lock:
-            return self.results.copy()
-
 
 class InferenceEngine:
-    """Unified inference engine for continuous batching."""
+    """Unified inference engine backed by continuous-batching scheduler."""
 
     def __init__(
         self,
@@ -97,55 +105,26 @@ class InferenceEngine:
         tokenizer: AutoTokenizer,
         max_batch_size: int = 1,
         max_seq_len: Optional[int] = None,
-        max_prefix_len: int = 512,
-        cache_capacity: int = 1000,
+        max_prompt_len: int = 2048,
+        page_size: int = 128,
     ):
-        """
-        Initialize inference engine with separate model and tokenizer.
-
-        Args:
-            model: The language model for inference (nn.Module, e.g., Transformer)
-            tokenizer: The tokenizer for encoding/decoding text
-            config: Model configuration
-            max_batch_size: Maximum batch size for continuous batching
-            max_seq_len: Maximum sequence length (defaults to config.max_len)
-            max_prefix_len: Maximum prefix length for cache (default: 512)
-            cache_capacity: Maximum number of cached prefixes (default: 1000)
-        """
         self.model = model
         self.tokenizer = tokenizer
-
-        # Get device and dtype from model parameters
-        try:
-            first_param = next(model.parameters())
-            device = first_param.device
-            dtype = first_param.dtype
-        except StopIteration:
-            # Model has no parameters, use default device/dtype
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            dtype = torch.float32
-
         self.scheduler = InferenceScheduler(
             model=self.model,
             tokenizer=self.tokenizer,
             max_batch_size=max_batch_size,
             max_seq_len=max_seq_len,
-            max_prefix_len=max_prefix_len,
-            cache_capacity=cache_capacity,
-            device=device,
-            dtype=dtype,
+            max_prompt_len=max_prompt_len,
+            page_size=page_size,
         )
 
-        self.kv_cache = self.scheduler.kv_cache
-        self.seq_mask = self.scheduler.seq_mask
-
         self.scheduler.start()
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        """Handle exceptions on exit."""
         self.shutdown()
         return False
 
@@ -153,139 +132,164 @@ class InferenceEngine:
         self,
         prompt: Union[str, List[str]],
         stream: bool = False,
-        max_tokens: int = 1024,
+        max_tokens: Optional[int] = None,
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = 50,
-        abort_on_exception: bool = True,
-    ) -> Union[Generator[str, None, None], str, List[str]]:
-        """Unified generation interface.
-
-        Args:
-            abort_on_exception: If True, abort the generation when consumer
-                stops iterating (GeneratorExit/StopIteration). Default: True.
-        """
+    ) -> Union[Generator, str, List[str]]:
+        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
         is_batch = isinstance(prompt, list)
         prompts = prompt if is_batch else [prompt]
 
         if stream:
             return self._generate_streaming(
-                prompts,
-                is_batch,
-                max_tokens,
-                temperature,
-                top_p,
-                top_k,
-                abort_on_exception,
+                prompts, is_batch, max_tokens, temperature, top_p, top_k
             )
         else:
             return self._generate_non_streaming(
                 prompts, is_batch, max_tokens, temperature, top_p, top_k
             )
 
+    def generate_async(
+        self,
+        prompt: str,
+        max_tokens: Optional[int] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+    ) -> AsyncGenerator[str, None]:
+        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
+        sync_gen = self._generate_streaming(
+            [prompt], False, max_tokens, temperature, top_p, top_k
+        )
+
+        async def _agen():
+            loop = asyncio.get_event_loop()
+            while True:
+                token = await loop.run_in_executor(None, self._next_token, sync_gen)
+                if token is None:
+                    break
+                yield token
+
+        return _agen()
+
+    @staticmethod
+    def _next_token(gen: Generator) -> Optional[str]:
+        try:
+            return next(gen)
+        except StopIteration:
+            return None
+
     def generate_with_request(
         self, request: GenerationRequest
     ) -> Union[Generator[str, None, None], str, List[str]]:
-        """Generate with GenerationRequest object."""
-        # Use tokenizer's chat template with messages
         prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
-
         return self.generate(
             prompt=prompt,
             stream=request.stream,
-            max_tokens=request.max_len,
+            max_tokens=request.max_tokens,
             temperature=request.temperature,
             top_p=request.top_p,
             top_k=request.top_k,
         )
 
+    def _submit_tasks(
+        self,
+        prompts: List[str],
+        max_tokens: Optional[int],
+        temperature: float,
+        top_p: float,
+        top_k: int,
+    ) -> Tuple[GenerateResult, List[str]]:
+        n = len(prompts)
+        result = GenerateResult(count=n)
+        task_ids = []
+        for i, p in enumerate(prompts):
+            cb = self._make_callback(result, i)
+            task_id = self.scheduler.add_task(
+                prompt=p,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                stream_callback=cb,
+            )
+            task_ids.append(task_id)
+        return result, task_ids
+
+    @staticmethod
+    def _make_callback(result: GenerateResult, idx: int):
+        def cb(token):
+            result.append(token, idx)
+
+        return cb
+
     def _generate_streaming(
         self,
         prompts: List[str],
         is_batch: bool,
-        max_tokens: int,
+        max_tokens: Optional[int],
         temperature: float,
         top_p: float,
         top_k: int,
-        abort_on_exception: bool = True,
-    ) -> Union[Generator[str, None, None], List[Generator[str, None, None]]]:
-        """Generate with streaming output.
-
-        Args:
-            abort_on_exception: If True, abort the task when generator is
-                stopped early by consumer (GeneratorExit/StopIteration).
-        """
-        if is_batch:
-            raise NotImplementedError("Batch streaming is not implemented yet")
-
-        result = _Result(stream=True)
-
-        task_id = self.scheduler.add_task(
-            prompt=prompts[0],
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            stream_callback=result.append,
+    ) -> Generator:
+        result, task_ids = self._submit_tasks(
+            prompts, max_tokens, temperature, top_p, top_k
         )
+        n = len(prompts)
+        remaining = n
+        finished = [False] * n
 
         def gen():
+            nonlocal remaining
             try:
-                while True:
-                    tokens = result.pop_all()
-                    for token in tokens:
-                        if token == "[DONE]":
-                            return
-                        yield token
-                    result.wait(timeout=0.05)
-            except Exception:
-                # Consumer stopped iterating - abort the task
-                if abort_on_exception:
-                    self.scheduler.remove_task(task_id)
-                raise
+                while remaining > 0:
+                    items = result.pop_all()
+                    for idx, token in items:
+                        if token is STOP:
+                            if not finished[idx]:
+                                finished[idx] = True
+                                remaining -= 1
+                        else:
+                            yield (idx, token) if is_batch else token
+                    if remaining > 0:
+                        result.wait(timeout=0.05)
+            finally:
+                for tid in task_ids:
+                    self.scheduler.remove_task(tid)
 
-        gen.task_id = task_id
         return gen()
 
     def _generate_non_streaming(
         self,
         prompts: List[str],
         is_batch: bool,
-        max_tokens: int,
+        max_tokens: Optional[int],
         temperature: float,
         top_p: float,
         top_k: int,
     ) -> Union[str, List[str]]:
-        """Generate without streaming."""
-        result = _Result(count=len(prompts))
+        result, task_ids = self._submit_tasks(
+            prompts, max_tokens, temperature, top_p, top_k
+        )
 
-        for i, p in enumerate(prompts):
-            # Create closure to capture current index value using factory function
-            def make_callback(idx):
-                def callback(token):
-                    result.append(idx, token)
+        try:
+            result.wait_completion()
+        except TimeoutError:
+            for tid in task_ids:
+                self.scheduler.remove_task(tid)
+            raise
 
-                return callback
+        for tid in task_ids:
+            self.scheduler.remove_task(tid)
 
-            self.scheduler.add_task(
-                prompt=p,
-                max_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                stream_callback=make_callback(i),
-            )
-
-        result.wait()
-        results = result.get_results()
-        return results if is_batch else results[0]
+        res = result.get_results()
+        return res if is_batch else res[0]
 
     def get_stats(self) -> Dict[str, Any]:
-        """Get engine statistics."""
         return self.scheduler.get_stats()
 
     def shutdown(self) -> None:
-        """Shutdown the engine and release all resources."""
         self.scheduler.stop()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

astrai/inference/sample.py

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

astrai/inference/scheduler.py

@@ -1,637 +0,0 @@
-"""Inference scheduler for continuous batching."""
-
-import threading
-import time
-import uuid
-from typing import Any, Callable, Dict, List, Optional, Tuple
-
-import torch
-from torch import Tensor
-
-from astrai.model.automodel import AutoModel
-from astrai.tokenize import AutoTokenizer
-
-
-class RadixNode:
-    """Radix tree node for prefix cache."""
-
-    def __init__(self):
-        self.children: Dict[int, "RadixNode"] = {}  # token_id -> child node
-        self.hash: Optional[int] = None  # 64-bit hash of the prefix
-        self.slot: int = -1  # KV Cache slot, valid only for leaf nodes
-        self.ref_count: int = 0  # number of tasks referencing this prefix
-        self.last_access: float = 0.0  # timestamp for LRU
-        self.token_sequence: list = []  # full token sequence from root to this node
-
-
-class PrefixCacheManager:
-    """Prefix cache manager using Radix tree with LRU eviction."""
-
-    def __init__(self, max_capacity: int = 1000, base: int = 131, mod: int = 10**9 + 7):
-        self.root = RadixNode()
-        self.base = base
-        self.mod = mod
-        self.max_capacity = max_capacity
-        self.lru: List[Tuple[float, RadixNode]] = []  # (timestamp, node) for LRU
-
-    def insert(self, token_ids: Tuple[int, ...], slot: int) -> None:
-        """Insert a prefix, increase ref_count if already exists, otherwise create new node."""
-        node = self.root
-        path = []
-        h = 0
-        for i, token_id in enumerate(token_ids):
-            if token_id not in node.children:
-                node.children[token_id] = RadixNode()
-            node = node.children[token_id]
-            h = (h * self.base + token_id) % self.mod
-            node.hash = h
-            path.append(token_id)
-            node.token_sequence = list(
-                path
-            )  # store full sequence for exact verification
-
-        # Leaf node: set slot and increase ref_count
-        if node.slot == -1:
-            node.slot = slot
-        node.ref_count += 1
-        node.last_access = time.time()
-        self._update_lru(node)
-        self._evict_if_needed()
-
-    def find_longest_prefix(self, token_ids: List[int]) -> Optional[Tuple[int, int]]:
-        """Find longest matching prefix, return (prefix_len, slot).
-
-        During traversal, compute hash per token and compare with node hash.
-        If hash matches, perform full token sequence verification to avoid
-        hash collision errors.
-        """
-        node = self.root
-        best_len = 0
-        best_slot = -1
-        h = 0
-
-        for i, token_id in enumerate(token_ids):
-            if token_id not in node.children:
-                break
-            node = node.children[token_id]
-            h = (h * self.base + token_id) % self.mod
-            if node.hash == h:  # hash matches
-                # Exact verification: compare full token sequence
-                if node.token_sequence == token_ids[: i + 1]:
-                    best_len = i + 1
-                    best_slot = node.slot
-                    node.last_access = time.time()
-                    self._update_lru(node)
-
-        if best_len > 0:
-            return (best_len, best_slot)
-        return None
-
-    def release(self, token_ids: Tuple[int, ...]) -> None:
-        """Release reference to a prefix, decrease ref_count. If zero, mark as evictable."""
-        node = self.root
-        for token_id in token_ids:
-            if token_id not in node.children:
-                return
-            node = node.children[token_id]
-        if node.ref_count > 0:
-            node.ref_count -= 1
-            if node.ref_count == 0:
-                node.slot = -1  # slot can be reused
-
-    def _update_lru(self, node: RadixNode) -> None:
-        """Update LRU list, move node to most recently used position."""
-        self.lru = [(ts, n) for (ts, n) in self.lru if n is not node]
-        self.lru.append((node.last_access, node))
-
-    def _evict_if_needed(self) -> None:
-        """If cache entries exceed capacity, evict least recently used leaf nodes (ref_count must be 0)."""
-        if len(self.lru) <= self.max_capacity:
-            return
-        # Sort by timestamp
-        self.lru.sort(key=lambda x: x[0])
-        for ts, node in self.lru:
-            if node.ref_count == 0:
-                # Remove leaf node from tree (need to recursively delete empty branches)
-                self._remove_node(node)
-                self.lru.remove((ts, node))
-                if len(self.lru) <= self.max_capacity:
-                    break
-
-    def _remove_node(
-        self,
-        node: RadixNode,
-        parent: Optional[RadixNode] = None,
-        child_key: Optional[int] = None,
-    ) -> None:
-        """Remove node from tree, including empty parent nodes."""
-        # First, recursively remove all children
-        for child_key, child_node in list(node.children.items()):
-            self._remove_node(child_node, node, child_key)
-
-        # Clear the node's leaf properties
-        node.slot = -1
-        node.hash = None
-        node.token_sequence = []
-        node.children.clear()
-
-        # If this node has no children and has a parent, remove the reference from parent
-        if parent is not None and child_key is not None and len(node.children) == 0:
-            if child_key in parent.children:
-                del parent.children[child_key]
-
-
-class TaskStatus:
-    """Task state for continuous batching."""
-
-    PENDING = "pending"
-    RUNNING = "running"
-    FINISHED = "finished"
-    ABORTED = "aborted"
-
-
-class Task:
-    """Individual task for continuous batching."""
-
-    def __init__(
-        self,
-        task_id: str,
-        prompt_ids: List[int],
-        max_tokens: int = 1024,
-        temperature: float = 1.0,
-        top_p: float = 1.0,
-        top_k: int = 50,
-        stream_callback: Optional[Callable[[str], None]] = None,
-    ):
-        self.task_id = task_id
-        self.prompt_ids = prompt_ids
-        self.max_tokens = max_tokens
-        self.temperature = temperature
-        self.top_p = top_p
-        self.top_k = top_k
-
-        self.status = TaskStatus.PENDING
-        self.output_ids: List[int] = []
-        self.input_tokens: int = 0
-        self.output_tokens: int = 0
-        self.slot: int = -1
-        self.prefix_len: int = 0  # prefix cache matched length
-        self.arrival_time = time.time()
-        self.finish_time: Optional[float] = None
-
-        self.stream_callback = stream_callback
-
-    def is_finished(self, stop_ids: List[int]) -> bool:
-        """Check if task is finished."""
-        return (
-            bool(self.output_ids and self.output_ids[-1] in stop_ids)
-            or self.output_tokens >= self.max_tokens
-        )
-
-
-def apply_sampling_strategies(
-    logits: Tensor,
-    temperature: float,
-    top_k: int,
-    top_p: float,
-    filter_value: float = -float("inf"),
-) -> Tensor:
-    """Apply sampling strategies to the logits tensor."""
-    # Clone logits to avoid inplace updates on inference tensor
-    logits = logits.clone()
-
-    if temperature != 1.0:
-        logits = logits / temperature
-
-    if top_k > 0:
-        top_k = min(top_k, logits.size(-1))
-        indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
-        logits[indices_to_remove] = filter_value
-
-    if top_p < 1.0:
-        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
-        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-
-        sorted_indices_to_remove = cumulative_probs > top_p
-        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-        sorted_indices_to_remove[..., 0] = 0
-
-        indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
-        indices_to_remove.scatter_(
-            dim=1, index=sorted_indices, src=sorted_indices_to_remove
-        )
-
-        logits[indices_to_remove] = filter_value
-
-    return logits
-
-
-class InferenceScheduler:
-    """Inference scheduler with continuous batching support."""
-
-    def __init__(
-        self,
-        model: AutoModel,
-        tokenizer: AutoTokenizer,
-        max_batch_size: int = 16,
-        max_seq_len: Optional[int] = None,
-        max_prefix_len: int = 512,
-        cache_capacity: int = 1000,
-        device: str = "cuda",
-        dtype: torch.dtype = torch.bfloat16,
-    ):
-        config = model.config
-
-        self.model = model
-        self.tokenizer = tokenizer
-        self.max_batch_size = max_batch_size
-        self.max_seq_len = max_seq_len or config.max_len
-        self.max_prefix_len = max_prefix_len
-        self.device = device or next(model.parameters()).device
-        self.dtype = dtype or next(model.parameters()).dtype
-
-        # Initialize prefix cache
-        self.prefix_cache = PrefixCacheManager(max_capacity=cache_capacity)
-
-        num_kv_heads = config.n_kv_heads
-        head_dim = config.dim // config.n_heads
-        n_layers = config.n_layers
-
-        k_cache = torch.empty(
-            (
-                max_batch_size,
-                self.max_seq_len,
-                n_layers,
-                num_kv_heads,
-                head_dim,
-            ),
-            device=self.device,
-            dtype=self.dtype,
-        )
-        v_cache = torch.empty(
-            (
-                max_batch_size,
-                self.max_seq_len,
-                n_layers,
-                num_kv_heads,
-                head_dim,
-            ),
-            device=self.device,
-            dtype=self.dtype,
-        )
-        self.kv_cache = (k_cache, v_cache)
-        self.seq_mask = torch.ones(
-            (max_batch_size, self.max_seq_len), device=self.device, dtype=torch.bool
-        )
-
-        self.waiting_queue: List[Task] = []
-        self.active_tasks: List[Task] = []
-
-        self._running = False
-        self._task_event = threading.Event()
-        self._lock = threading.Lock()
-
-        self._total_tasks = 0
-        self._total_tokens = 0
-
-    def add_task(
-        self,
-        prompt: str,
-        max_tokens: int = 1024,
-        temperature: float = 1.0,
-        top_p: float = 1.0,
-        top_k: int = 50,
-        stream_callback: Optional[Callable[[str], None]] = None,
-    ) -> str:
-        """Add a new task to the waiting queue."""
-        task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
-        prompt_ids = self.tokenizer.encode(prompt)
-
-        # Truncate if exceeds max_prefix_len
-        if len(prompt_ids) > self.max_prefix_len:
-            prompt_ids = prompt_ids[: self.max_prefix_len]
-
-        task = Task(
-            task_id=task_id,
-            prompt_ids=prompt_ids,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            stream_callback=stream_callback,
-        )
-
-        # Find longest matching prefix from cache
-        match = self.prefix_cache.find_longest_prefix(prompt_ids)
-        if match:
-            prefix_len, slot = match
-            task.prefix_len = prefix_len
-            task.slot = slot
-        else:
-            task.prefix_len = 0
-            task.slot = -1
-
-        with self._lock:
-            self.waiting_queue.append(task)
-            self._total_tasks += 1
-
-        self._task_event.set()
-        return task_id
-
-    def remove_task(self, task_id: str) -> None:
-        """Remove a task from the scheduler."""
-        with self._lock:
-            self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
-            self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
-
-    def _remove_finished_tasks(self) -> None:
-        """Remove finished tasks from active batch."""
-        finished = []
-        for task in self.active_tasks:
-            if task.is_finished(self.tokenizer.stop_ids):
-                task.status = TaskStatus.FINISHED
-                task.finish_time = time.time()
-                finished.append(task)
-                self._total_tokens += task.output_tokens
-
-        for task in finished:
-            slot = task.slot
-            if slot >= 0 and slot < len(self.active_tasks):
-                self.seq_mask[slot, :] = False
-
-            # Release prefix cache reference
-            if task.prefix_len > 0:
-                self.prefix_cache.release(tuple(task.prompt_ids[: task.prefix_len]))
-
-            task.slot = -1
-
-        self.active_tasks = [
-            t for t in self.active_tasks if t.status != TaskStatus.FINISHED
-        ]
-
-    def _refill_active_batch(self) -> None:
-        """Refill active batch with waiting tasks."""
-        available_slots = self.max_batch_size - len(self.active_tasks)
-        if available_slots <= 0:
-            return
-
-        with self._lock:
-            to_add = [
-                self.waiting_queue.pop(0)
-                for _ in range(min(available_slots, len(self.waiting_queue)))
-            ]
-            for task in to_add:
-                task.slot = self._allocate_slot()
-                task.status = TaskStatus.RUNNING
-                self.active_tasks.append(task)
-
-    def _allocate_slot(self) -> int:
-        """Allocate an available slot for a task."""
-        for i in range(self.max_batch_size):
-            if not any(t.slot == i for t in self.active_tasks):
-                return i
-        return -1
-
-    def _execute_prefill(self, tasks: List[Task]) -> None:
-        """Execute Prefill phase with incremental prefill support."""
-        if not tasks:
-            return
-
-        # Group tasks by prefix cache status
-        fully_cached, partial, full = [], [], []
-        for task in tasks:
-            total_len, prefix_len = len(task.prompt_ids), task.prefix_len
-            if prefix_len == total_len:
-                fully_cached.append(task)
-            elif prefix_len > 0:
-                partial.append(task)
-            else:
-                full.append(task)
-
-        # Handle fully cached tasks
-        for t in fully_cached:
-            t.input_tokens, t.output_tokens = len(t.prompt_ids), 0
-            if t.slot >= 0:
-                self.seq_mask[t.slot, : t.input_tokens] = True
-
-        if full:
-            self._execute_full_prefill(full)
-        if partial:
-            self._execute_partial_prefill(partial)
-
-    def _execute_full_prefill(self, tasks: List[Task]) -> None:
-        """Execute full prefill for tasks without prefix cache."""
-        if not tasks:
-            return
-
-        tasks = sorted(tasks, key=lambda t: t.slot)
-
-        prompt_lens = [len(task.prompt_ids) for task in tasks]
-        max_len = max(prompt_lens)
-
-        input_ids = torch.zeros(
-            len(tasks), max_len, dtype=torch.long, device=self.device
-        )
-        for i, task in enumerate(tasks):
-            if len(task.prompt_ids) > 0:
-                input_ids[i, : len(task.prompt_ids)] = torch.tensor(
-                    task.prompt_ids, device=self.device
-                )
-
-        if self.tokenizer.pad_id is not None:
-            input_mask = torch.ne(input_ids, self.tokenizer.pad_id)
-        else:
-            input_mask = torch.ones(
-                input_ids.shape, dtype=torch.bool, device=self.device
-            )
-
-        with torch.inference_mode():
-            self.model(
-                input_ids,
-                input_mask=input_mask,
-                start_pos=0,
-                persistent_key_values=self.kv_cache,
-            )
-
-        for i, task in enumerate(tasks):
-            task.input_tokens = prompt_lens[i]
-            task.output_tokens = 0
-            # Insert new prefix into cache
-            self.prefix_cache.insert(tuple(task.prompt_ids), task.slot)
-
-        for task in tasks:
-            if task.slot >= 0:
-                self.seq_mask[task.slot, : task.input_tokens] = True
-
-    def _execute_partial_prefill(self, tasks: List[Task]) -> None:
-        """Execute incremental prefill for tasks with partial prefix cache match."""
-        for task in tasks:
-            total_len = len(task.prompt_ids)
-            prefix_len = task.prefix_len
-
-            if prefix_len >= total_len:
-                task.input_tokens = total_len
-                task.output_tokens = 0
-                continue
-
-            # Get new tokens that need prefill
-            new_ids = task.prompt_ids[prefix_len:]
-            new_len = len(new_ids)
-
-            if new_len == 0:
-                task.input_tokens = total_len
-                task.output_tokens = 0
-                continue
-
-            # Build input for incremental prefill
-            input_ids = torch.tensor([new_ids], dtype=torch.long, device=self.device)
-
-            # Input mask should cover from position 0 to prefix_len + new_len
-            # The prefix part uses cached KV, new part needs computation
-            input_mask = torch.ones(
-                (1, prefix_len + new_len), dtype=torch.bool, device=self.device
-            )
-
-            with torch.inference_mode():
-                self.model(
-                    input_ids,
-                    input_mask=input_mask,
-                    start_pos=prefix_len,
-                    persistent_key_values=self.kv_cache,
-                )
-
-            task.input_tokens = total_len
-            task.output_tokens = 0
-
-            # Insert full prefix into cache (ref_count already increased in add_task)
-            self.prefix_cache.insert(tuple(task.prompt_ids), task.slot)
-
-            if task.slot >= 0:
-                self.seq_mask[task.slot, : task.input_tokens] = True
-
-    def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
-        """Execute Decode phase."""
-        if not tasks:
-            return
-
-        tasks = sorted(tasks, key=lambda t: t.slot)
-
-        input_ids = torch.zeros(len(tasks), dtype=torch.long, device=self.device)
-        for i, task in enumerate(tasks):
-            if task.output_ids:
-                input_ids[i] = task.output_ids[-1]
-            else:
-                input_ids[i] = task.prompt_ids[-1]
-
-        input_tensor = input_ids.unsqueeze(1)
-        active_mask = torch.ones((len(tasks), 1), dtype=torch.bool, device=self.device)
-
-        with torch.inference_mode():
-            outputs = self.model(
-                input_tensor,
-                input_mask=active_mask,
-                persistent_key_values=self.kv_cache,
-                start_pos=start_pos,
-            )
-            logits = outputs["logits"][:, -1, :]
-
-        next_token_ids = []
-        for i, task in enumerate(tasks):
-            logit = logits[i : i + 1]
-            logit = apply_sampling_strategies(
-                logit,
-                task.temperature,
-                task.top_k,
-                task.top_p,
-            )
-            probs = torch.softmax(logit, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            next_token_ids.append(next_token.item())
-
-        for task, next_token in zip(tasks, next_token_ids):
-            task.output_ids.append(next_token)
-            task.output_tokens += 1
-
-            pos = task.input_tokens + task.output_tokens
-            if task.slot >= 0 and pos < self.max_seq_len:
-                self.seq_mask[task.slot, pos] = True
-
-            if task.stream_callback:
-                token_str = self.tokenizer.decode([next_token])
-                task.stream_callback(token_str)
-
-        for task in tasks:
-            if task.output_tokens >= task.max_tokens or (
-                task.output_ids and task.output_ids[-1] in self.tokenizer.stop_ids
-            ):
-                if task.stream_callback:
-                    task.stream_callback("[DONE]")
-
-    def _run_generation_loop(self) -> None:
-        """Main generation loop."""
-        while self._running:
-            self._remove_finished_tasks()
-            self._refill_active_batch()
-
-            if not self.active_tasks:
-                self._task_event.wait(timeout=0.01)
-                self._task_event.clear()
-                continue
-
-            new_tasks = [t for t in self.active_tasks if t.output_tokens == 0]
-            decode_tasks = [t for t in self.active_tasks if t.output_tokens > 0]
-
-            if decode_tasks:
-                start_pos = max(t.input_tokens + t.output_tokens for t in decode_tasks)
-            else:
-                start_pos = 0
-
-            if new_tasks:
-                self._execute_prefill(new_tasks)
-                decode_tasks = new_tasks
-                start_pos = max(t.input_tokens for t in decode_tasks)
-
-            if decode_tasks:
-                self._execute_decode(decode_tasks, start_pos)
-
-            if not self.active_tasks and not self.waiting_queue:
-                self._task_event.wait(timeout=0.05)
-                self._task_event.clear()
-
-    def start(self) -> None:
-        """Start the generation loop."""
-        if not self._running:
-            self._running = True
-            self._loop_thread = threading.Thread(target=self._run_generation_loop)
-            self._loop_thread.daemon = True
-            self._loop_thread.start()
-
-    def stop(self) -> None:
-        """Stop the generation loop."""
-        self._running = False
-        if hasattr(self, "_loop_thread"):
-            self._loop_thread.join(timeout=1.0)
-
-        # Clear KV cache to free GPU memory
-        if self.kv_cache is not None:
-            k_cache, v_cache = self.kv_cache
-            if k_cache is not None:
-                k_cache.detach()
-            if v_cache is not None:
-                v_cache.detach()
-
-        # Clear seq mask
-        self.seq_mask.detach()
-
-        # Clear task lists
-        self.waiting_queue.clear()
-        self.active_tasks.clear()
-
-    def get_stats(self) -> Dict[str, Any]:
-        """Get scheduler statistics."""
-        return {
-            "total_tasks": self._total_tasks,
-            "total_tokens": self._total_tokens,
-            "active_tasks": len(self.active_tasks),
-            "waiting_queue": len(self.waiting_queue),
-        }

astrai/inference/server.py

@@ -1,321 +0,0 @@
-"""
-Inference Server with Continuous Batching Support
-
-FastAPI server for inference with continuous batching.
-Provides OpenAI-compatible chat completion endpoints.
-"""
-
-import json
-import logging
-from contextlib import asynccontextmanager
-from pathlib import Path
-from typing import Any, Dict, List, Optional
-
-import torch
-import uvicorn
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, Field
-
-from astrai.inference.engine import InferenceEngine
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
-
-logger = logging.getLogger(__name__)
-
-# Global model parameter and engine (loaded once)
-_engine: Optional[InferenceEngine] = None
-_model_param: Optional[Any] = None
-_project_root = Path(__file__).parent.parent.parent
-
-# Server configuration (set before running server)
-_server_config: Dict[str, Any] = {
-    "device": "cuda",
-    "dtype": torch.bfloat16,
-    "param_path": None,
-    "max_batch_size": 16,
-}
-
-
-def configure_server(
-    device: str = "cuda",
-    dtype: torch.dtype = torch.bfloat16,
-    param_path: Optional[Path] = None,
-    max_batch_size: int = 16,
-):
-    """Configure server settings before starting.
-
-    Args:
-        device: Device to load model on (e.g., "cuda", "cpu", "cuda:0")
-        dtype: Data type for model weights (e.g., torch.bfloat16, torch.float16)
-        param_path: Path to model parameters directory
-        max_batch_size: Maximum batch size for continuous batching
-    """
-    _server_config["device"] = device
-    _server_config["dtype"] = dtype
-    _server_config["param_path"] = param_path
-    _server_config["max_batch_size"] = max_batch_size
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    """Lifespan context manager for startup and shutdown events."""
-    global _model_param, _engine
-    # Startup: Load model with configured settings
-    try:
-        load_model(
-            param_path=_server_config["param_path"],
-            device=_server_config["device"],
-            dtype=_server_config["dtype"],
-            max_batch_size=_server_config["max_batch_size"],
-        )
-    except Exception as e:
-        logger.error(f"Failed to load model: {e}")
-        raise
-    yield
-    # Shutdown: Cleanup engine
-    if _engine:
-        _engine.shutdown()
-        logger.info("Inference engine shutdown complete")
-
-
-app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
-
-
-def load_model(
-    param_path: Optional[Path] = None,
-    device: str = "cuda",
-    dtype: torch.dtype = torch.bfloat16,
-    max_batch_size: int = 16,
-):
-    """Load model parameters and initialize inference engine."""
-    global _model_param, _engine
-    if param_path is None:
-        param_path = _project_root / "params"
-    if not param_path.exists():
-        raise FileNotFoundError(f"Parameter directory not found: {param_path}")
-
-    # Load tokenizer separately
-    tokenizer = AutoTokenizer.from_pretrained(param_path)
-    _model_param = AutoModel.from_pretrained(param_path)
-    _model_param.to(device=device, dtype=dtype)
-    logger.info(f"Model loaded on {device} with dtype {dtype}")
-
-    # Initialize inference engine with separate model and tokenizer
-    _engine = InferenceEngine(
-        model=_model_param,
-        tokenizer=tokenizer,
-        max_batch_size=max_batch_size,
-    )
-    logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
-
-
-# Pydantic models for API request/response
-class ChatMessage(BaseModel):
-    role: str  # "user", "assistant", "system"
-    content: str
-
-
-class ChatCompletionRequest(BaseModel):
-    messages: List[ChatMessage]
-    temperature: float = Field(0.8, ge=0.0, le=2.0)
-    top_p: float = Field(0.95, ge=0.0, le=1.0)
-    top_k: int = Field(50, ge=0)
-    max_tokens: int = Field(2048, ge=1)
-    stream: bool = False
-    system_prompt: Optional[str] = None
-
-
-class CompletionResponse(BaseModel):
-    id: str = "chatcmpl-default"
-    object: str = "chat.completion"
-    created: int = 0
-    model: str = "astrai"
-    choices: List[Dict[str, Any]]
-
-
-@app.get("/health")
-async def health():
-    return {
-        "status": "ok",
-        "model_loaded": _model_param is not None,
-        "engine_ready": _engine is not None,
-    }
-
-
-@app.get("/stats")
-async def get_stats():
-    """Get inference engine statistics."""
-    if _engine is None:
-        raise HTTPException(status_code=503, detail="Engine not initialized")
-    return _engine.get_stats()
-
-
-@app.post("/v1/chat/completions", response_model=CompletionResponse)
-async def chat_completion(request: ChatCompletionRequest):
-    """OpenAI-compatible chat completion endpoint.
-
-    Supports both streaming and non-streaming modes with continuous batching.
-    """
-    if _engine is None:
-        raise HTTPException(status_code=503, detail="Engine not initialized")
-
-    # Convert messages to prompt using engine's tokenizer
-    # Extract system prompt if present, then apply chat template
-    # Apply chat template directly with messages
-    prompt = _engine.tokenizer.apply_chat_template(
-        [{"role": m.role, "content": m.content} for m in request.messages],
-        tokenize=False,
-    )
-
-    if request.stream:
-        # Streaming response (use synchronous generator)
-        generator = _engine.generate(
-            prompt=prompt,
-            stream=True,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            top_k=request.top_k,
-        )
-
-        def generate_stream():
-            for token in generator:
-                if token == "[DONE]":
-                    break
-                yield f"data: {json.dumps({'choices': [{'delta': {'content': token}}]})}\n\n"
-            yield "data: [DONE]\n\n"
-
-        return StreamingResponse(
-            generate_stream(),
-            media_type="text/event-stream",
-            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
-        )
-    else:
-        # Non-streaming response
-        result = _engine.generate(
-            prompt=prompt,
-            stream=False,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            top_k=request.top_k,
-        )
-
-        # Build OpenAI-style response
-        import time
-
-        resp = CompletionResponse(
-            id=f"chatcmpl-{int(time.time())}",
-            created=int(time.time()),
-            choices=[
-                {
-                    "index": 0,
-                    "message": {"role": "assistant", "content": result},
-                    "finish_reason": "stop",
-                }
-            ],
-        )
-        return resp
-
-
-@app.post("/generate")
-async def generate(
-    query: str,
-    history: Optional[List[List[str]]] = None,
-    temperature: float = 0.8,
-    top_p: float = 0.95,
-    top_k: int = 50,
-    max_len: int = 2048,
-    stream: bool = False,
-):
-    """Simple generation endpoint.
-
-    Args:
-        query: Input query string
-        history: Conversation history as list of [user, assistant] pairs
-        temperature: Sampling temperature
-        top_p: Top-p sampling parameter
-        top_k: Top-k sampling parameter
-        max_len: Maximum tokens to generate
-        stream: Enable streaming output
-
-    Returns:
-        dict: Generation result with response field
-    """
-    if _engine is None:
-        raise HTTPException(status_code=503, detail="Engine not initialized")
-
-    # Build messages for chat template
-    messages = []
-    if history:
-        # Convert history format: List[List[str]] -> List[Dict]
-        for h in history:
-            if len(h) >= 2:
-                messages.append({"role": "user", "content": h[0]})
-                messages.append({"role": "assistant", "content": h[1]})
-    messages.append({"role": "user", "content": query})
-
-    # Use tokenizer's chat template
-    prompt = _engine.tokenizer.apply_chat_template(messages, tokenize=False)
-
-    if stream:
-        # Synchronous streaming
-        result = _engine.generate(
-            prompt=prompt,
-            stream=True,
-            max_tokens=max_len,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-        )
-
-        def stream_generator():
-            for token in result:
-                yield token + "\n"
-
-        return StreamingResponse(stream_generator(), media_type="text/plain")
-    else:
-        result = _engine.generate(
-            prompt=prompt,
-            stream=False,
-            max_tokens=max_len,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-        )
-        return {"response": result}
-
-
-def run_server(
-    host: str = "0.0.0.0",
-    port: int = 8000,
-    reload: bool = False,
-    device: str = "cuda",
-    dtype: torch.dtype = torch.bfloat16,
-    param_path: Optional[Path] = None,
-    max_batch_size: int = 16,
-):
-    """Run the FastAPI server with uvicorn.
-
-    Args:
-        host: Server host address
-        port: Server port number
-        reload: Enable auto-reload for development
-        device: Device to load model on (e.g., "cuda", "cpu", "cuda:0")
-        dtype: Data type for model weights (e.g., torch.bfloat16, torch.float16)
-        param_path: Path to model parameters directory
-        max_batch_size: Maximum batch size for continuous batching
-    """
-    configure_server(
-        device=device,
-        dtype=dtype,
-        param_path=param_path,
-        max_batch_size=max_batch_size,
-    )
-    uvicorn.run(
-        "astrai.inference.server:app",
-        host=host,
-        port=port,
-        reload=reload,
-    )

astrai/model/__init__.py

@@ -1,12 +1,18 @@
 from astrai.model.automodel import AutoModel
-from astrai.model.module import (
-    GQA,
-    MLP,
-    DecoderBlock,
-    Linear,
-    RMSNorm,
+from astrai.model.components.attention import GQA
+from astrai.model.components.decoder_block import DecoderBlock
+from astrai.model.components.linear import Linear
+from astrai.model.components.lora import (
+    LoRAConfig,
+    inject_lora,
+    load_lora,
+    merge_lora,
+    save_lora,
 )
-from astrai.model.transformer import Transformer
+from astrai.model.components.mlp import MLP
+from astrai.model.components.norm import RMSNorm
+from astrai.model.encoder import EmbeddingEncoder
+from astrai.model.transformer import AutoRegressiveLM
 
 __all__ = [
     # Modules
@@ -16,6 +22,13 @@ __all__ = [
     "GQA",
     "DecoderBlock",
     # Models
-    "Transformer",
+    "AutoRegressiveLM",
+    "EmbeddingEncoder",
     "AutoModel",
+    # LoRA
+    "LoRAConfig",
+    "inject_lora",
+    "merge_lora",
+    "save_lora",
+    "load_lora",
 ]

astrai/model/automodel.py

@@ -2,14 +2,16 @@
 AutoModel base class for model loading and saving.
 """
 
+import json
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Dict, Self, Type, Union
+from typing import Self, Union
 
 import safetensors.torch as st
 import torch.nn as nn
 
-from astrai.config import ModelConfig
+from astrai.config.model_config import BaseModelConfig, ConfigFactory
+from astrai.factory import BaseFactory
 
 
 @contextmanager
@@ -38,72 +40,37 @@ def _disable_random_init(enable: bool = True):
                 setattr(nn.init, name, orig_func)
 
 
-class AutoModel(nn.Module):
+class AutoModel(BaseFactory["AutoModel"], nn.Module):
     """
     Autoregressive language model base class.
-    Provides model loading/saving and generation capabilities.
+    Provides model loading/saving, registration, and generation.
     """
 
-    # Model registry - stored as class attribute
-    _registry: Dict[str, Type["AutoModel"]] = {}
-
-    def __init__(self, config: ModelConfig):
+    def __init__(self, config: BaseModelConfig):
         super().__init__()
         self.config = config
 
-    @classmethod
-    def register(cls, model_type: str):
-        """
-        Class method decorator to register model type.
-
-        Usage:
-            @AutoModel.register('transformer')
-            class Transformer(AutoModel):
-                ...
-        """
-
-        def decorator(sub_cls: Type["AutoModel"]) -> Type["AutoModel"]:
-            cls._registry[model_type.lower()] = sub_cls
-            return sub_cls
-
-        return decorator
-
-    @classmethod
-    def get_model_class(cls, model_type: str) -> Type["AutoModel"]:
-        """Get model class by model_type string."""
-        model_type = model_type.lower()
-        if model_type not in cls._registry:
-            available = list(cls._registry.keys())
-            raise ValueError(
-                f"Unknown model_type: {model_type}. Available: {available}"
-            )
-        return cls._registry[model_type]
-
     @classmethod
     def from_pretrained(
         cls,
         path: Union[str, Path],
         disable_random_init: bool = True,
+        strict: bool = True,
     ) -> nn.Module:
 
         model_path = Path(path)
 
         # Load config
-        config = ModelConfig()
         config_path = model_path / "config.json"
         if config_path.exists():
-            config.load(str(config_path))
+            with open(config_path, "r") as f:
+                raw = json.load(f)
+            config = ConfigFactory.load(raw)
+            model_type = config.model_type or "autoregressive_lm"
         else:
             raise FileNotFoundError(f"Config file not found: {config_path}")
 
-        # If called from base class, use model_type to determine actual model class
-        if cls is AutoModel:
-            model_type = config.model_type or "transformer"
-            actual_cls = cls.get_model_class(model_type)
-        else:
-            raise ValueError(
-                f"Cannot call from_pretrained() on subclass {cls.__name__}"
-            )
+        actual_cls = AutoModel.get_component_class(model_type)
 
         with _disable_random_init(enable=disable_random_init):
             model = actual_cls(config)
@@ -112,7 +79,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
 
@@ -124,7 +91,7 @@ class AutoModel(nn.Module):
         save_path.mkdir(parents=True, exist_ok=True)
 
         # Save config
-        self.config.save(str(save_path / "config.json"))
+        self.config.to_file(str(save_path / "config.json"))
 
         # Save weights
         st.save_file(self.state_dict(), str(save_path / "model.safetensors"))

astrai/model/components/__init__.py

@@ -0,0 +1,25 @@
+from astrai.model.components.attention import GQA, MLA, repeat_kv
+from astrai.model.components.decoder_block import DecoderBlock
+from astrai.model.components.embedding import Embedding
+from astrai.model.components.linear import Linear
+from astrai.model.components.mlp import MLP
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import (
+    RotaryEmbedding,
+    apply_rotary_emb,
+    get_rotary_emb,
+)
+
+__all__ = [
+    "Linear",
+    "RMSNorm",
+    "MLP",
+    "Embedding",
+    "GQA",
+    "MLA",
+    "DecoderBlock",
+    "RotaryEmbedding",
+    "apply_rotary_emb",
+    "get_rotary_emb",
+    "repeat_kv",
+]

astrai/model/components/attention.py

@@ -0,0 +1,212 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from astrai.factory import BaseFactory
+from astrai.inference.core.cache import KvcacheView
+from astrai.model.components.linear import Linear
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import apply_rotary_emb
+
+
+def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
+    bs, slen, n_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, slen, n_heads, n_rep, head_dim)
+        .reshape(bs, slen, n_heads * n_rep, head_dim)
+    )
+
+
+class AttnFactory(BaseFactory[nn.Module]):
+    @classmethod
+    def create(cls, attn_type: str, **kwargs) -> nn.Module:
+        return super().create(attn_type, **kwargs)
+
+
+@AttnFactory.register("gqa")
+class GQA(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        use_qk_norm: bool,
+        norm_eps: float,
+        use_gated_attention: bool,
+        layer_id: int,
+    ):
+        super().__init__()
+        assert dim % n_heads == 0
+        assert n_heads % n_kv_heads == 0
+
+        self.head_dim = dim // n_heads
+        self.layer_id = layer_id
+        self.dim = dim
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.n_rep = n_heads // n_kv_heads
+        self.use_qk_norm = use_qk_norm
+        self.use_gated_attention = use_gated_attention
+
+        self.q_proj = Linear(dim, n_heads * self.head_dim)
+        self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
+        self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
+        self.o_proj = Linear(dim, dim)
+
+        if self.use_qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, norm_eps)
+
+        if self.use_gated_attention:
+            self.gate = Linear(dim, dim)
+
+    def _split_heads(self, x: Tensor, n_heads) -> Tensor:
+        batch_size, seq_len, _ = x.shape
+        x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
+        return x
+
+    def forward(
+        self,
+        x: Tensor,
+        rotary_emb: Tensor,
+        attn_mask: Tensor = None,
+        paged_cache: Optional[KvcacheView] = None,
+    ) -> Tensor:
+        is_causal = attn_mask is None
+
+        q = self._split_heads(self.q_proj(x), self.n_heads)
+        k = self._split_heads(self.k_proj(x), self.n_kv_heads)
+        v = self._split_heads(self.v_proj(x), self.n_kv_heads)
+        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
+
+        if self.use_qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        if paged_cache is not None:
+            paged_cache.write(self.layer_id, k, v)
+            k, v = paged_cache.gather(self.layer_id)
+
+        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
+
+        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
+        sdqa_out = (
+            F.scaled_dot_product_attention(q, k, v, attn_mask, is_causal=is_causal)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+            .flatten(2)
+        )
+
+        if self.use_gated_attention:
+            sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
+
+        out = self.o_proj(sdqa_out)
+        return out
+
+
+@AttnFactory.register("mla")
+class MLA(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        norm_eps: float,
+        use_qk_norm: bool,
+        use_gated_attention: bool,
+        layer_id: int,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.layer_id = layer_id
+        self.n_rep = n_heads // n_kv_heads
+        self.use_qk_norm = use_qk_norm
+        self.use_gated_attention = use_gated_attention
+
+        self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
+
+        if self.use_qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, norm_eps)
+        self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
+        self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
+
+        self.kv_b_proj = Linear(
+            kv_lora_rank,
+            n_kv_heads * (2 * self.head_dim),
+        )
+
+        self.o_proj = Linear(dim, dim, bias=False)
+
+        if use_gated_attention:
+            self.gate = Linear(dim, dim, bias=False)
+
+    def forward(
+        self,
+        x: Tensor,
+        rotary_emb: Tensor,
+        attn_mask: Tensor = None,
+        paged_cache: Optional[KvcacheView] = None,
+    ) -> Tensor:
+        bsz, seq_len, _ = x.size()
+        is_causal = attn_mask is None
+
+        q = self.q_proj(x)
+        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
+
+        kv_compressed = self.kv_a_proj(x)
+        kv_compressed = self.kv_norm(kv_compressed)
+
+        kv = self.kv_b_proj(kv_compressed)
+        kv = kv.view(bsz, seq_len, self.n_kv_heads, -1)
+
+        k_nope, k_rope, v = torch.split(
+            kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
+        )
+
+        q_nope, q_rope = (
+            q[..., : self.qk_nope_head_dim],
+            q[..., self.qk_nope_head_dim :],
+        )
+        q_rope = apply_rotary_emb(q_rope, rotary_emb)
+        k_rope = apply_rotary_emb(k_rope, rotary_emb)
+
+        q = torch.cat([q_nope, q_rope], dim=-1)
+        k = torch.cat([k_nope, k_rope], dim=-1)
+
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if paged_cache is not None:
+            paged_cache.write(self.layer_id, k, v)
+            k, v = paged_cache.gather(self.layer_id)
+
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v, attn_mask, is_causal=is_causal
+        )
+        attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
+
+        if self.use_gated_attention:
+            attn_out = attn_out * F.sigmoid(self.gate(x))
+
+        out = self.o_proj(attn_out)
+        return out

astrai/model/components/decoder_block.py

@@ -0,0 +1,59 @@
+from typing import Optional
+
+import torch.nn as nn
+from torch import Tensor
+
+from astrai.inference.core.cache import KvcacheView
+from astrai.model.components.attention import AttnFactory
+from astrai.model.components.mlp import FFNFactory
+from astrai.model.components.norm import RMSNorm
+
+
+class DecoderBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        dim_ffn: int,
+        n_kv_heads: int,
+        norm_eps: float,
+        use_qk_norm: bool,
+        use_gated_attention: bool,
+        layer_id: int,
+        attn_type: str = "gqa",
+        ffn_type: str = "mlp",
+        **kwargs,
+    ):
+        super().__init__()
+        self.attention = AttnFactory.create(
+            attn_type,
+            dim=dim,
+            n_heads=n_heads,
+            n_kv_heads=n_kv_heads,
+            use_qk_norm=use_qk_norm,
+            norm_eps=norm_eps,
+            use_gated_attention=use_gated_attention,
+            layer_id=layer_id,
+            **kwargs,
+        )
+        self.input_norm = RMSNorm(dim, norm_eps)
+        self.post_attention_norm = RMSNorm(dim, norm_eps)
+        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **kwargs)
+
+    def forward(
+        self,
+        x: Tensor,
+        rotary_emb: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        paged_cache: Optional[KvcacheView] = None,
+    ) -> Tensor:
+        attn_output = self.attention(
+            self.input_norm(x),
+            rotary_emb,
+            attention_mask,
+            paged_cache,
+        )
+        x = attn_output + x
+        x = self.mlp(self.post_attention_norm(x)) + x
+
+        return x

astrai/model/components/embedding.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class Embedding(nn.Module):
+    def __init__(self, vocab_size: int, embedding_dim: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
+
+    def reset_parameters(self):
+        nn.init.normal_(self.weight, mean=0.0, std=0.02)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return F.embedding(x, self.weight)

astrai/model/components/linear.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class Linear(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
+        self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / (fan_in**0.5)
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return F.linear(x, self.weight, self.bias)

astrai/model/components/lora.py

@@ -0,0 +1,192 @@
+import json
+import logging
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Optional, Set
+
+import safetensors.torch as st
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from astrai.model.components.linear import Linear
+
+logger = logging.getLogger(__name__)
+
+TARGET_MODULES_ATTN = {"q_proj", "k_proj", "v_proj", "o_proj"}
+TARGET_MODULES_FFN = {"up", "gate", "down"}
+
+
+@dataclass
+class LoRAConfig:
+    r: int = 16
+    alpha: int = 32
+    target_modules: tuple = ("q_proj", "v_proj")
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: Linear, r: int = 16, alpha: int = 32):
+        super().__init__()
+        self.register_parameter("weight", base.weight)
+        self.weight.requires_grad_(False)
+        self.bias = base.bias
+        if self.bias is not None:
+            self.bias.requires_grad_(False)
+
+        self.r = r
+        self.scaling = alpha / r
+        self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1]) / r)
+        self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r))
+        self._merged = False
+
+    def forward(self, x):
+        out = F.linear(x, self.weight, self.bias)
+        if not self._merged:
+            out += (F.linear(x, self.lora_A) @ self.lora_B.T) * self.scaling
+        return out
+
+    def merge(self):
+        if self._merged:
+            return
+        self.weight.data += (self.lora_B @ self.lora_A) * self.scaling
+        self._merged = True
+        del self.lora_A
+        del self.lora_B
+
+
+def _collect_lora_info(model: nn.Module) -> dict:
+    names = {}
+    for n, m in model.named_modules():
+        if isinstance(m, Linear):
+            _, _, child = n.rpartition(".")
+            names.setdefault(child, []).append(n)
+    return names
+
+
+def _get_lora_count(model: nn.Module) -> int:
+    return sum(1 for m in model.modules() if isinstance(m, LoRALinear))
+
+
+def inject_lora(
+    model: nn.Module,
+    r: int = 16,
+    alpha: int = 32,
+    target_modules: Optional[Set[str]] = None,
+) -> LoRAConfig:
+    if target_modules is None:
+        target_modules = TARGET_MODULES_ATTN
+
+    available = _collect_lora_info(model)
+    injected = 0
+
+    for name, module in list(model.named_modules()):
+        if not isinstance(module, Linear):
+            continue
+        parent_name, _, child_name = name.rpartition(".")
+        if child_name not in target_modules:
+            continue
+        parent = model.get_submodule(parent_name) if parent_name else model
+        setattr(parent, child_name, LoRALinear(module, r=r, alpha=alpha))
+        injected += 1
+
+    if injected == 0:
+        logger.warning(
+            "No LoRA layers injected. Available Linear child names: %s. "
+            "target_modules: %s. Check model type and target_modules.",
+            sorted(available),
+            sorted(target_modules),
+        )
+    else:
+        logger.info("LoRA injected: %d layers (r=%d, alpha=%d)", injected, r, alpha)
+
+    return LoRAConfig(r=r, alpha=alpha, target_modules=tuple(target_modules))
+
+
+def merge_lora(model: nn.Module):
+    n = 0
+    for module in model.modules():
+        if isinstance(module, LoRALinear):
+            module.merge()
+            n += 1
+    if n == 0:
+        logger.warning("No LoRA layers to merge.")
+    else:
+        logger.info("Merged %d LoRA layers", n)
+
+
+def save_lora(model: nn.Module, save_dir: str, config: LoRAConfig):
+    lora_sd = {
+        k: v
+        for k, v in model.state_dict().items()
+        if k.endswith((".lora_A", ".lora_B"))
+    }
+    if not lora_sd:
+        raise RuntimeError(
+            "No LoRA parameters found in model. "
+            "The model may not have been injected or was already merged."
+        )
+
+    path = Path(save_dir)
+    path.mkdir(parents=True, exist_ok=True)
+    st.save_file(lora_sd, str(path / "adapter_model.safetensors"))
+    with open(path / "adapter_config.json", "w") as f:
+        json.dump(asdict(config), f, indent=2)
+    logger.info("LoRA adapter saved to %s (%d keys)", save_dir, len(lora_sd))
+
+
+def load_lora(model: nn.Module, load_dir: str) -> LoRAConfig:
+    path = Path(load_dir)
+    with open(path / "adapter_config.json") as f:
+        raw = json.load(f)
+    config = LoRAConfig(
+        r=raw["r"], alpha=raw["alpha"], target_modules=tuple(raw["target_modules"])
+    )
+
+    existing = _get_lora_count(model)
+    if existing > 0:
+        logger.warning(
+            "Model already has %d LoRA layers. Skipping injection, "
+            "loading weights onto existing layers only.",
+            existing,
+        )
+    else:
+        inject_lora(
+            model,
+            r=config.r,
+            alpha=config.alpha,
+            target_modules=set(config.target_modules),
+        )
+
+    weights = st.load_file(str(path / "adapter_model.safetensors"))
+    try:
+        missing, unexpected = model.load_state_dict(weights, strict=False)
+    except RuntimeError as e:
+        msg = str(e)
+        if "size mismatch" in msg:
+            raise RuntimeError(
+                f"LoRA weight shapes do not match the model. "
+                f"The adapter config (r={config.r}) may not match the injected layers. "
+                f"Original error: {msg}"
+            ) from e
+        raise
+
+    injected = _get_lora_count(model)
+    if injected == 0:
+        raise RuntimeError(
+            "No LoRA layers found after loading. "
+            "Inject LoRA before calling load_lora, or check the adapter config."
+        )
+
+    if missing:
+        lora_missing = [k for k in missing if "lora" in k]
+        if lora_missing:
+            raise RuntimeError(
+                f"LoRA weight keys not found in model: {lora_missing}. "
+                f"The adapter config (r={config.r}) may not match the model."
+            )
+        logger.debug("LoRA load: %d missing base-weight keys (expected)", len(missing))
+    if unexpected:
+        logger.warning("LoRA load: %d unexpected keys", len(unexpected))
+
+    logger.info("LoRA adapter loaded from %s", load_dir)
+    return config

astrai/model/components/mlp.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from astrai.factory import BaseFactory
+from astrai.model.components.linear import Linear
+
+
+class FFNFactory(BaseFactory[nn.Module]):
+    @classmethod
+    def create(cls, ffn_type: str, dim: int, dim_ffn: int, **kwargs) -> nn.Module:
+        return super().create(ffn_type, dim, dim_ffn, **kwargs)
+
+
+@FFNFactory.register("mlp")
+class MLP(nn.Module):
+    def __init__(self, dim: int, dim_ffn: int):
+        super().__init__()
+        self.up = Linear(dim, dim_ffn)
+        self.gate = Linear(dim, dim_ffn)
+        self.down = Linear(dim_ffn, dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        gated = self.up(x) * F.silu(self.gate(x))
+        out = self.down(gated)
+        return out
+
+
+@FFNFactory.register("moe")
+class DeepSeekMoE(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_ffn: int,
+        n_routed_experts: int,
+        n_shared_experts: int = 1,
+        n_activated_experts: int = 2,
+        topk_method: str = "greedy",
+    ):
+        super().__init__()
+        self.dim = dim
+        self.n_routed_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.n_activated_experts = n_activated_experts
+        self.topk_method = topk_method
+
+        self.router = Linear(dim, n_routed_experts, bias=False)
+
+        self.shared_experts = nn.ModuleList(
+            [MLP(dim, dim_ffn) for _ in range(n_shared_experts)]
+        )
+        self.routed_experts = nn.ModuleList(
+            [MLP(dim, dim_ffn) for _ in range(n_routed_experts)]
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        bsz, seq_len, dim = x.shape
+        x_flat = x.view(-1, dim)
+
+        shared_out = self._shared_forward(x_flat)
+        routed_out = self._routed_forward(x_flat)
+
+        out = (shared_out + routed_out).view(bsz, seq_len, dim)
+        return out
+
+    def _shared_forward(self, x: Tensor) -> Tensor:
+        if self.n_shared_experts == 0:
+            return torch.zeros_like(x)
+        return sum(e(x) for e in self.shared_experts) / self.n_shared_experts
+
+    def _routed_forward(self, x: Tensor) -> Tensor:
+        N, D = x.shape
+        K = self.n_activated_experts
+
+        router_logits = self.router(x)
+        router_probs = torch.softmax(router_logits.float(), dim=-1).to(x.dtype)
+
+        topk_weights, topk_indices = torch.topk(router_probs, K, dim=-1)
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+        output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
+        for expert_idx in range(self.n_routed_experts):
+            expert_mask = topk_indices == expert_idx
+            token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
+            if token_idx.numel() == 0:
+                continue
+            expert_input = x[token_idx]
+            expert_output = self.routed_experts[expert_idx](expert_input)
+            weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
+            output.index_add_(0, token_idx, expert_output * weights)
+
+        return output

astrai/model/components/norm.py

@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, norm_eps):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.normalized_shape = (dim,)
+        self.norm_eps = norm_eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)

astrai/model/components/rope.py

@@ -0,0 +1,53 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+def get_rotary_emb(
+    dim: int,
+    max_len: int,
+    base: float = 10000,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
+    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
+    freqs = torch.outer(t, theta).float()
+    cos = torch.cos(freqs)
+    sin = torch.sin(freqs)
+    return torch.complex(cos, sin)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
+    dtype = x.dtype
+    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
+    x_complex = torch.view_as_complex(x_)
+    freqs_cis = freqs_cis.unsqueeze(2)
+    x_rotated = x_complex * freqs_cis
+    x_out = torch.view_as_real(x_rotated).flatten(-2)
+    return x_out.to(dtype)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, max_len: int, base: float = 10000):
+        super().__init__()
+        self.dim = dim
+        self.max_len = max_len
+        self.base = base
+        self._set_rotary_buffer(self.max_len)
+
+    def _set_rotary_buffer(self, max_len: int):
+        rotary_emb = get_rotary_emb(self.dim, max_len, self.base)
+        freqs_cis = torch.view_as_real(rotary_emb)
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
+        if position_ids is None:
+            position_ids = (
+                torch.arange(x.size(1), device=x.device)
+                .unsqueeze(0)
+                .expand(x.size(0), -1)
+            )
+        position_freq_cis = self.freqs_cis[position_ids].float()
+        return torch.view_as_complex(position_freq_cis)

astrai/model/encoder.py

@@ -0,0 +1,100 @@
+from typing import Any, Mapping, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from astrai.config.model_config import EncoderConfig
+from astrai.model.automodel import AutoModel
+from astrai.model.components.decoder_block import DecoderBlock
+from astrai.model.components.embedding import Embedding
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import RotaryEmbedding
+from astrai.model.transformer import process_attention_mask
+
+
+@AutoModel.register("embedding")
+class EmbeddingEncoder(AutoModel):
+    def __init__(self, config: EncoderConfig):
+        super().__init__(config)
+        self.config = config
+        rope_dim = config.dim // config.n_heads
+        rope_base = config.rope_theta if config.rope_theta is not None else 10000
+        self.rotary_embedding = RotaryEmbedding(rope_dim, config.max_len, rope_base)
+        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+
+        self.layers = nn.ModuleList(
+            [
+                DecoderBlock(
+                    config.dim,
+                    config.n_heads,
+                    config.dim_ffn,
+                    config.n_kv_heads,
+                    config.norm_eps,
+                    config.use_qk_norm,
+                    config.use_gated_attention,
+                    layer_id,
+                )
+                for layer_id in range(config.n_layers)
+            ]
+        )
+
+        self.norm = RMSNorm(config.dim, config.norm_eps)
+
+        self.pooling_type = config.pooling_type or "mean"
+        self.normalize_embeddings = config.normalize_embeddings or False
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if hasattr(module, "reset_parameters"):
+            module.reset_parameters()
+
+    def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
+        state_dict = dict(state_dict)
+        state_dict.pop("lm_head.weight", None)
+        return super().load_state_dict(state_dict, strict=strict, assign=assign)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        input_mask: Optional[Tensor] = None,
+        position_ids: Optional[Tensor] = None,
+    ) -> Tensor:
+        assert input_ids.ndim == 2
+        B, S = input_ids.shape
+
+        x = self.embed_tokens(input_ids)
+
+        if position_ids is None:
+            position_ids = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1)
+
+        rotary_emb = self.rotary_embedding(x, position_ids)
+        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=False)
+
+        for layer in self.layers:
+            x = layer(x, rotary_emb, attn_mask, paged_cache=None)
+
+        hidden_states = self.norm(x)
+
+        if self.pooling_type == "cls":
+            pooled = hidden_states[:, 0]
+        elif self.pooling_type == "last":
+            if input_mask is not None:
+                lengths = input_mask.sum(dim=1) - 1
+                pooled = hidden_states[torch.arange(B, device=x.device), lengths]
+            else:
+                pooled = hidden_states[:, -1]
+        else:
+            if input_mask is not None:
+                mask = input_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
+                pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(
+                    min=1.0
+                )
+            else:
+                pooled = hidden_states.mean(dim=1)
+
+        if self.normalize_embeddings:
+            pooled = torch.nn.functional.normalize(pooled, p=2, dim=-1)
+
+        return pooled

astrai/model/module.py

@@ -1,382 +0,0 @@
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-
-
-def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
-    """
-    Repeat k times along the dimension for attention heads.
-    Args:
-        x (Tensor): The input tensor.
-        n_rep (int): The number of repetitions.
-    Returns:
-        Tensor: The repeated tensor.
-    """
-
-    bs, slen, n_heads, head_dim = x.shape
-    if n_rep == 1:
-        return x
-    return (
-        x[:, :, :, None, :]
-        .expand(bs, slen, n_heads, n_rep, head_dim)
-        .reshape(bs, slen, n_heads * n_rep, head_dim)
-    )
-
-
-def get_rotary_emb(
-    dim: int,
-    max_len: int,
-    base: float = 10000,
-    device: Optional[torch.device] = None,
-) -> Tuple[Tensor, Tensor]:
-    """
-    Get the rotary embedding for the given dimension and maximum length.
-    Args:
-        dim (int): The dimension of the input.
-        max_len (int): The maximum length of the input.
-        base (float, optional): The base for the frequency. Defaults to 10000.
-        device (optional): The device to create tensors on. Defaults to None.
-    Returns:
-        Tensor: The rotary embedding tensor.
-    """
-
-    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
-    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
-    freqs = torch.outer(t, theta)
-
-    return torch.cos(freqs).float(), torch.sin(freqs).float()
-
-
-def apply_rotary_emb(x: torch.Tensor, rotary_emb: Tuple[Tensor, Tensor]) -> Tensor:
-    """
-    Apply rotary embedding to the input tensor using cos/sin form.
-    Args:
-        x (Tensor): The input tensor (shape [..., seq_len, dim]).
-        rotary_emb (Tuple[Tensor, Tensor]): The rotary embedding (shape [seq_len, dim//2]).
-    Returns:
-        Tensor: The output tensor (rotated, same shape as input).
-    """
-
-    dtype = x.dtype
-    cos, sin = rotary_emb
-
-    cos = cos.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, dim//2]
-    sin = sin.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, dim//2]
-
-    x_real = x[..., 0::2]  # [batch, seq_len, dim//2]
-    x_imag = x[..., 1::2]  # [batch, seq_len, dim//2]
-
-    x_real_rot = x_real * cos - x_imag * sin
-    x_imag_rot = x_real * sin + x_imag * cos
-
-    x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1)  # [batch, seq_len, dim//2, 2]
-    x_out = x_out.view(*x_out.shape[:-2], -1)  # [batch, seq_len, dim]
-
-    return x_out.to(dtype)
-
-
-class RotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, max_len: int, base: int = 10000):
-        super().__init__()
-        self.dim = dim
-        self.max_len = max_len
-        self.base = base
-        self.max_len_cached = None
-        self._set_rotary_buffer(self.max_len, None)
-
-    def _set_rotary_buffer(self, max_len: int, device: Optional[torch.device] = None):
-        cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base, device)
-        self.register_buffer("cos_cached", cos_cached, persistent=False)
-        self.register_buffer("sin_cached", sin_cached, persistent=False)
-        self.max_len_cached = max_len
-
-    def forward(self, x: Tensor, start_pos: int = 0) -> Tuple[Tensor, Tensor]:
-        seq_len = x.size(1)
-
-        if self.max_len_cached < seq_len + start_pos:
-            self._set_rotary_buffer(self.max_len_cached * 2, x.device)
-
-        cos = self.cos_cached[start_pos : start_pos + seq_len]
-        sin = self.sin_cached[start_pos : start_pos + seq_len]
-
-        return (cos, sin)
-
-
-class Linear(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
-        super().__init__()
-        self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
-        self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
-
-    def forward(self, x: Tensor) -> Tensor:
-        return F.linear(x, self.weight, self.bias)
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim, norm_eps):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.normalized_shape = (dim,)
-        self.norm_eps = norm_eps
-
-    def forward(self, x: Tensor) -> Tensor:
-        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)
-
-
-class MLP(nn.Module):
-    def __init__(self, dim: int, dim_feed_forward: int):
-        super().__init__()
-        self.up = Linear(dim, dim_feed_forward)
-        self.gate = Linear(dim, dim_feed_forward)
-        self.down = Linear(dim_feed_forward, dim)
-
-    def forward(self, x: Tensor) -> Tensor:
-        gated = self.up(x) * F.silu(self.gate(x))
-        out = self.down(gated)
-        return out
-
-
-class GQA(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        n_heads: int,
-        n_kv_heads: int,
-        use_qk_norm: bool,
-        norm_eps: float,
-        use_gated_attention: bool,
-        layer_id: int,
-    ):
-        super().__init__()
-        assert dim % n_heads == 0
-        assert n_heads % n_kv_heads == 0
-
-        self.head_dim = dim // n_heads
-        self.layer_id = layer_id
-        self.dim = dim
-        self.n_heads = n_heads
-        self.n_kv_heads = n_kv_heads
-        self.n_rep = n_heads // n_kv_heads
-        self.use_qk_norm = use_qk_norm
-        self.use_gated_attention = use_gated_attention
-
-        self.q_proj = Linear(dim, n_heads * self.head_dim)
-        self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
-        self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
-        self.o_proj = Linear(dim, dim)
-
-        if self.use_qk_norm:
-            self.q_norm = RMSNorm(self.head_dim, norm_eps)
-            self.k_norm = RMSNorm(self.head_dim, norm_eps)
-
-        if self.use_gated_attention:
-            self.gate = Linear(dim, dim)
-
-    def _split_heads(self, x: Tensor, n_heads) -> Tensor:
-        batch_size, seq_len, _ = x.shape
-        x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
-        return x
-
-    def forward(
-        self,
-        x: Tensor,
-        rotary_emb: Tuple[Tensor, Tensor],
-        mask: Tensor = None,
-        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
-        start_pos: int = 0,
-    ) -> Tensor:
-        bsz, seq_len, _ = x.size()
-        is_causal = mask is None
-
-        # x(bsz, seq_len, n_heads * head_dim) -> (bsz, seq_len, n_heads, head_dim)
-        q = self._split_heads(self.q_proj(x), self.n_heads)
-        k = self._split_heads(self.k_proj(x), self.n_kv_heads)
-        v = self._split_heads(self.v_proj(x), self.n_kv_heads)
-        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
-
-        if self.use_qk_norm:
-            q, k = self.q_norm(q), self.k_norm(k)
-
-        if kv_cache is not None:
-            k_cache, v_cache = kv_cache
-
-            # copy to cache
-            k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
-            v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
-
-            # get cache
-            k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
-            v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
-
-        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
-
-        # (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
-        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
-        # (bsz, n_heads, seq_len, head_dim) - > (bsz, seq_len, n_heads*head_dim)
-        sdqa_out = (
-            F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
-            .permute(0, 2, 1, 3)
-            .contiguous()
-            .flatten(2)
-        )
-
-        if self.use_gated_attention:
-            sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
-
-        out = self.o_proj(sdqa_out)
-
-        return out
-
-
-class MLA(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        n_heads: int,
-        n_kv_heads: int,
-        kv_lora_rank: int,
-        qk_nope_head_dim: int,
-        qk_rope_head_dim: int,
-        norm_eps: float,
-        use_gated_attention: bool,
-        layer_id: int,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.n_heads = n_heads
-        self.n_kv_heads = n_kv_heads
-        self.kv_lora_rank = kv_lora_rank
-        self.qk_nope_head_dim = qk_nope_head_dim
-        self.qk_rope_head_dim = qk_rope_head_dim
-        self.head_dim = qk_nope_head_dim + qk_rope_head_dim
-        self.layer_id = layer_id
-        self.n_rep = n_heads // n_kv_heads
-        self.use_gated_attention = use_gated_attention
-
-        self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
-        self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
-        self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
-
-        # KV (k_nope, k_rope, v)
-        self.kv_b_proj = Linear(
-            kv_lora_rank,
-            n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
-        )
-
-        self.o_proj = Linear(dim, dim, bias=False)
-
-        if use_gated_attention:
-            self.gate = Linear(dim, dim, bias=False)
-
-    def forward(
-        self,
-        x: Tensor,
-        rotary_emb: Tuple[Tensor, Tensor],
-        mask: Tensor = None,
-        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
-        start_pos: int = 0,
-    ) -> Tensor:
-        bsz, seq_len, _ = x.size()
-        is_causal = mask is None
-
-        q = self.q_proj(x)
-        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
-
-        kv_compressed = self.kv_a_proj(x)
-        kv_compressed = self.kv_norm(kv_compressed)
-
-        kv = self.kv_b_proj(kv_compressed)
-        kv = kv.view(bsz, seq_len, self.n_kv_heads, -1)
-
-        k_nope, k_rope, v = torch.split(
-            kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
-        )
-
-        q_nope, q_rope = (
-            q[..., : self.qk_nope_head_dim],
-            q[..., self.qk_rope_head_dim :],
-        )
-        q_rope = apply_rotary_emb(q_rope, rotary_emb)
-        k_rope = apply_rotary_emb(k_rope, rotary_emb)
-
-        q = torch.cat([q_nope, q_rope], dim=-1)
-        k = torch.cat([k_nope, k_rope], dim=-1)
-
-        if kv_cache is not None:
-            k_cache, v_cache = kv_cache
-            k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
-            v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
-            k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
-            v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
-
-        q = q.permute(0, 2, 1, 3)
-        k = k.permute(0, 2, 1, 3)
-        v = v.permute(0, 2, 1, 3)
-
-        attn_out = F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
-        attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
-
-        if self.use_gated_attention:
-            attn_out = attn_out * F.sigmoid(self.gate(x))
-
-        out = self.o_proj(attn_out)
-
-        return out
-
-
-class DecoderBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        n_heads: int,
-        dim_ffn: int,
-        n_kv_heads: int,
-        norm_eps: int,
-        use_qk_norm: bool,
-        use_gated_attention: bool,
-        layer_id: int,
-    ):
-        super().__init__()
-        self.attention = GQA(
-            dim,
-            n_heads,
-            n_kv_heads,
-            use_qk_norm,
-            norm_eps,
-            use_gated_attention,
-            layer_id,
-        )
-        self.input_norm = RMSNorm(dim, norm_eps)
-        self.mlp = MLP(dim, dim_ffn)
-        self.post_attention_norm = RMSNorm(dim, norm_eps)
-
-    def forward(
-        self,
-        x: Tensor,
-        rotary_emb: Tuple[Tensor, Tensor],
-        attention_mask: Optional[Tensor] = None,
-        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
-        start_pos: int = 0,
-    ) -> Tensor:
-        # attention
-        attn_output = self.attention(
-            self.input_norm(x), rotary_emb, attention_mask, kv_cache, start_pos
-        )
-        x = attn_output + x
-
-        # feed forward
-        x = self.mlp(self.post_attention_norm(x)) + x
-
-        return x
-
-
-class Embedding(nn.Module):
-    def __init__(self, vocab_size: int, embedding_dim: int):
-        super().__init__()
-        self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
-
-    def forward(self, x: Tensor) -> Tensor:
-        return F.embedding(x, self.weight)

astrai/model/transformer.py

@@ -1,84 +1,65 @@
-from typing import Any, Mapping, Optional, Tuple
+from typing import Any, Mapping, Optional
 
 import torch
 import torch.nn as nn
 from torch import Tensor
 
-from astrai.config.model_config import ModelConfig
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.inference.core.cache import KvcacheView
 from astrai.model.automodel import AutoModel
-from astrai.model.module import (
-    DecoderBlock,
-    Embedding,
-    Linear,
-    RMSNorm,
-    RotaryEmbedding,
-)
+from astrai.model.components.decoder_block import DecoderBlock
+from astrai.model.components.embedding import Embedding
+from astrai.model.components.linear import Linear
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import RotaryEmbedding
 
 
 def process_attention_mask(
-    seq_mask: Tensor,
     input_tensor: Tensor,
-    start_pos: int = 0,
+    position_ids: Optional[Tensor],
+    input_mask: Optional[Tensor] = None,
     is_causal: bool = False,
-) -> Tensor:
-    """
-    Create attention mask for GQA
-    Args:
-        seq_mask (Tensor): A tensor indicating whether each position is valid or not.
-        input_tensor (Tensor): The input tensor.
-        start_pos (int): The starting position of the sequence.
-        is_causal (bool): Whether the attention is causal or not.
-    Returns:
-        Tensor: The attention mask tensor.
-    """
+) -> Optional[Tensor]:
+    if position_ids is None:
+        return None
+    if input_mask is not None and input_mask.dim() > 2:
+        return input_mask
+
     device = input_tensor.device
     dtype = input_tensor.dtype
-    seq_len = input_tensor.size(1)
+    B, S = input_tensor.size()[:2]
+    T = position_ids.max().item() + 1
 
-    if seq_mask is None:
-        if start_pos != 0:
-            # for single prompt chat
-            seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
-        else:
+    if input_mask is None:
+        if position_ids.min().item() == 0 and is_causal:
             return None
+        pad = torch.ones(B, T, dtype=torch.bool, device=device)
+    else:
+        pad = input_mask[:, :T].to(device=device, dtype=torch.bool)
 
-    if seq_mask.dim() > 2:
-        # shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos)
-        # if ndim > 2, it's 4D tensor
-        return seq_mask
-
-    batch_size = seq_mask.size(0)
-    seq_mask = seq_mask[:, : start_pos + seq_len].to(device=device, dtype=torch.bool)
-    # (bsz, start_pos + seq_len)
-    expanded_mask = seq_mask.unsqueeze(1).expand(
-        batch_size, seq_len, start_pos + seq_len
-    )
-    # (bsz, seq_len, start_pos + seq_len)
-
+    attend = pad.view(B, 1, T).expand(B, S, T).clone()
     if is_causal:
-        expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)
+        attend &= position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
 
-    attention_mask = torch.zeros_like(expanded_mask, dtype=dtype, device=device)
-    attention_mask = attention_mask.masked_fill_(
-        ~expanded_mask, -torch.finfo(dtype).max / 2
-    ).unsqueeze(1)
-    # (bsz, 1, seq_len, seq_len + start_pos)
-
-    return attention_mask
+    return torch.full(
+        (B, 1, S, T), -torch.finfo(dtype).max / 2, dtype=dtype, device=device
+    ).masked_fill_(attend.unsqueeze(1), 0.0)
 
 
-@AutoModel.register("transformer")
-class Transformer(AutoModel):
-    """
-    Transformer language model.
-    """
+@AutoModel.register("autoregressive_lm")
+class AutoRegressiveLM(AutoModel):
+    """Autoregressive language model with paged KV cache."""
 
-    def __init__(self, config: ModelConfig):
+    def __init__(self, config: AutoRegressiveLMConfig):
         super().__init__(config)
         self.config = config
-        self.rotary_embedding = RotaryEmbedding(
-            config.dim // config.n_heads, config.max_len
+        rope_dim = (
+            config.qk_rope_head_dim
+            if config.attn_type == "mla"
+            else config.dim // config.n_heads
         )
+        rope_base = config.rope_theta if config.rope_theta is not None else 10000
+        self.rotary_embedding = RotaryEmbedding(rope_dim, config.max_len, rope_base)
         self.embed_tokens = Embedding(config.vocab_size, config.dim)
 
         self.layers = nn.ModuleList(
@@ -92,6 +73,15 @@ class Transformer(AutoModel):
                     config.use_qk_norm,
                     config.use_gated_attention,
                     layer_id,
+                    attn_type=config.attn_type,
+                    ffn_type=config.ffn_type,
+                    n_routed_experts=config.n_routed_experts,
+                    n_shared_experts=config.n_shared_experts,
+                    n_activated_experts=config.n_activated_experts,
+                    topk_method=config.topk_method,
+                    kv_lora_rank=config.kv_lora_rank,
+                    qk_nope_head_dim=config.qk_nope_head_dim,
+                    qk_rope_head_dim=config.qk_rope_head_dim,
                 )
                 for layer_id in range(config.n_layers)
             ]
@@ -100,32 +90,28 @@ class Transformer(AutoModel):
         self.norm = RMSNorm(config.dim, config.norm_eps)
         self.lm_head = Linear(config.dim, config.vocab_size)
 
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
             self.lm_head.weight = self.embed_tokens.weight
 
-        self._init_weights()
+        self.apply(self._init_weights)
 
-    def _init_weights(self):
-        for param in self.parameters():
-            if param.dim() > 1:
-                nn.init.normal_(param, mean=0.0, std=0.006)
+    def _init_weights(self, module):
+        if hasattr(module, "reset_parameters"):
+            module.reset_parameters()
 
     def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
         lm_head_key = "lm_head.weight"
         embed_key = "embed_tokens.weight"
 
-        # Make a copy to avoid modifying the original state_dict
         state_dict = dict(state_dict)
 
-        if self.config.tie_weight:
-            # same tensor
+        if self.config.tie_weight is True:
+            # same tensor for embed and lm_head
             if embed_key in state_dict:
                 state_dict[lm_head_key] = state_dict[embed_key]
         else:
-            # If lm_head.weight exists in checkpoint, use it directly
-            # If not, copy from embed_tokens.weight
             if lm_head_key not in state_dict and embed_key in state_dict:
-                # use clone to avoid sharing the same tensor
+                # clone to avoid sharing gradients
                 state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
 
         return super().load_state_dict(state_dict, strict, assign)
@@ -135,7 +121,7 @@ class Transformer(AutoModel):
             destination=destination, prefix=prefix, keep_vars=keep_vars
         )
 
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
             lm_head_key = prefix + "lm_head.weight"
             if lm_head_key in state_dict:
                 del state_dict[lm_head_key]
@@ -146,18 +132,17 @@ class Transformer(AutoModel):
         self,
         input_ids: Tensor,
         input_mask: Optional[Tensor] = None,
-        persistent_key_values: Optional[Tuple[Tensor, Tensor]] = None,
-        start_pos: int = 0,
+        paged_cache: Optional[KvcacheView] = None,
+        position_ids: Optional[Tensor] = None,
     ) -> Tensor:
         assert input_ids.ndim == 2
 
         x = self.embed_tokens(input_ids)
-        rotary_emb = self.rotary_embedding(x, start_pos)
-
-        attn_mask = process_attention_mask(input_mask, x, start_pos, is_causal=True)
+        rotary_emb = self.rotary_embedding(x, position_ids)
+        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=True)
 
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, persistent_key_values, start_pos)
+            x = layer(x, rotary_emb, attn_mask, paged_cache)
 
         hidden_states = self.norm(x)
         logits = self.lm_head(hidden_states)

astrai/parallel/__init__.py

@@ -1,3 +1,13 @@
+from astrai.parallel.executor import (
+    AccumOptimizer,
+    AccumScheduler,
+    BaseExecutor,
+    DDPExecutor,
+    ExecutorFactory,
+    FSDPExecutor,
+    GradientState,
+    NoneExecutor,
+)
 from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
 from astrai.parallel.setup import (
     get_current_device,
@@ -17,4 +27,12 @@ __all__ = [
     "spawn_parallel_fn",
     "RowParallelLinear",
     "ColumnParallelLinear",
+    "ExecutorFactory",
+    "BaseExecutor",
+    "GradientState",
+    "AccumOptimizer",
+    "AccumScheduler",
+    "NoneExecutor",
+    "DDPExecutor",
+    "FSDPExecutor",
 ]

astrai/parallel/executor.py

@@ -0,0 +1,231 @@
+"""Unified training executor — parallel strategy + gradient accumulation."""
+
+import contextlib
+import logging
+from contextlib import contextmanager
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from torch.utils.data import DataLoader
+
+from astrai.factory import BaseFactory
+from astrai.parallel.setup import get_rank, get_world_size
+
+logger = logging.getLogger(__name__)
+
+
+class GradientState:
+    def __init__(self, grad_accum_steps: int = 1):
+        self.num_steps = max(grad_accum_steps, 1)
+        self._step: int = 0
+        self._sync_gradients: bool = True
+
+    @property
+    def sync_gradients(self) -> bool:
+        return self._sync_gradients
+
+    def _do_sync(self):
+        self._step += 1
+        self._sync_gradients = self._step % self.num_steps == 0
+
+
+class AccumOptimizer:
+    def __init__(self, optimizer: Optimizer, gradient_state: GradientState):
+        self.optimizer = optimizer
+        self.gradient_state = gradient_state
+
+    def step(self, closure=None):
+        if self.gradient_state.sync_gradients:
+            self.optimizer.step(closure)
+
+    def zero_grad(self):
+        if self.gradient_state.sync_gradients:
+            self.optimizer.zero_grad()
+
+    @property
+    def param_groups(self):
+        return self.optimizer.param_groups
+
+    def state_dict(self):
+        return self.optimizer.state_dict()
+
+    def load_state_dict(self, d):
+        self.optimizer.load_state_dict(d)
+
+
+class AccumScheduler:
+    def __init__(self, scheduler: LRScheduler, gradient_state: GradientState):
+        self.scheduler = scheduler
+        self.gradient_state = gradient_state
+
+    def step(self):
+        if self.gradient_state.sync_gradients:
+            self.scheduler.step()
+
+    def state_dict(self):
+        return self.scheduler.state_dict()
+
+    def load_state_dict(self, d):
+        self.scheduler.load_state_dict(d)
+
+    def get_last_lr(self):
+        return self.scheduler.get_last_lr()
+
+
+class BaseExecutor:
+    def __init__(self, grad_accum_steps: int = 1):
+        self.gradient_state = GradientState(grad_accum_steps)
+
+    def prepare(
+        self,
+        model: nn.Module,
+        optimizer: Optional[Optimizer] = None,
+        dataloader: Optional[DataLoader] = None,
+        scheduler: Optional[LRScheduler] = None,
+    ) -> Tuple[
+        nn.Module, Optional[Optimizer], Optional[DataLoader], Optional[LRScheduler]
+    ]:
+        model = self._prepare_model(model)
+        if optimizer is not None:
+            optimizer = AccumOptimizer(optimizer, self.gradient_state)
+        if scheduler is not None:
+            scheduler = AccumScheduler(scheduler, self.gradient_state)
+        return model, optimizer, dataloader, scheduler
+
+    def _prepare_model(self, model: nn.Module) -> nn.Module:
+        return model
+
+    def _no_sync(self, model: nn.Module):
+        return contextlib.nullcontext()
+
+    @contextmanager
+    def accumulate(self, model: nn.Module):
+        self.gradient_state._do_sync()
+        if not self.gradient_state.sync_gradients:
+            with self._no_sync(model):
+                yield
+        else:
+            yield
+
+    def backward(self, loss: torch.Tensor):
+        loss.backward()
+
+    def unwrap_model(self, model: nn.Module) -> nn.Module:
+        return model
+
+    @property
+    def use_distributed(self) -> bool:
+        return get_world_size() > 1
+
+    @property
+    def sync_gradients(self) -> bool:
+        return self.gradient_state.sync_gradients
+
+    @property
+    def grad_accum_steps(self) -> int:
+        return self.gradient_state.num_steps
+
+
+class ExecutorFactory(BaseFactory[BaseExecutor]):
+    pass
+
+
+@ExecutorFactory.register("none")
+class NoneExecutor(BaseExecutor):
+    pass
+
+
+@ExecutorFactory.register("ddp")
+class DDPExecutor(BaseExecutor):
+    def __init__(
+        self,
+        grad_accum_steps: int = 1,
+        dim: int = 0,
+        broadcast_buffers: bool = True,
+        init_sync: bool = True,
+        process_group=None,
+        bucket_cap_mb: int = 25,
+        find_unused_parameters: bool = False,
+        check_reduction: bool = False,
+        gradient_as_bucket_view: bool = False,
+        static_graph: bool = False,
+        delay_all_reduce_named_params=None,
+        param_to_hook_all_reduce=None,
+        mixed_precision=None,
+        device_mesh=None,
+    ):
+        super().__init__(grad_accum_steps=grad_accum_steps)
+        self._ddp_kwargs = dict(
+            dim=dim,
+            broadcast_buffers=broadcast_buffers,
+            init_sync=init_sync,
+            process_group=process_group,
+            bucket_cap_mb=bucket_cap_mb,
+            find_unused_parameters=find_unused_parameters,
+            check_reduction=check_reduction,
+            gradient_as_bucket_view=gradient_as_bucket_view,
+            static_graph=static_graph,
+            delay_all_reduce_named_params=delay_all_reduce_named_params,
+            param_to_hook_all_reduce=param_to_hook_all_reduce,
+            mixed_precision=mixed_precision,
+            device_mesh=device_mesh,
+        )
+
+    def _prepare_model(self, model: nn.Module) -> nn.Module:
+        if not self.use_distributed:
+            logger.warning("DDP backend selected but world_size=1, model not wrapped")
+            return model
+        local_rank = get_rank()
+        model = DDP(
+            model,
+            device_ids=[local_rank],
+            output_device=local_rank,
+            **self._ddp_kwargs,
+        )
+        logger.info("Model wrapped with DDP (world_size=%d)", get_world_size())
+        return model
+
+    def _no_sync(self, model: nn.Module):
+        if isinstance(model, DDP):
+            return model.no_sync()
+        return contextlib.nullcontext()
+
+    def unwrap_model(self, model: nn.Module) -> nn.Module:
+        if isinstance(model, DDP):
+            return model.module
+        return model
+
+
+@ExecutorFactory.register("fsdp")
+class FSDPExecutor(BaseExecutor):
+    def __init__(self, grad_accum_steps: int = 1, **fsdp_kwargs):
+        super().__init__(grad_accum_steps=grad_accum_steps)
+        self._fsdp_kwargs = fsdp_kwargs
+        self._original_model: Optional[nn.Module] = None
+
+    def _prepare_model(self, model: nn.Module) -> nn.Module:
+        if not self.use_distributed:
+            logger.warning("FSDP backend selected but world_size=1, model not wrapped")
+            return model
+        self._original_model = model
+        device_id = torch.device("cuda", get_rank())
+        model = FSDP(model, device_id=device_id, **self._fsdp_kwargs)
+        logger.info("Model wrapped with FSDP (world_size=%d)", get_world_size())
+        return model
+
+    def _no_sync(self, model: nn.Module):
+        if isinstance(model, FSDP):
+            return model.no_sync()
+        return contextlib.nullcontext()
+
+    def unwrap_model(self, model: nn.Module) -> nn.Module:
+        if self._original_model is not None:
+            return self._original_model
+        if isinstance(model, FSDP):
+            return model._fsdp_wrapped_module
+        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_ids = [i for i in range(world_size)]
-
-    rank = device_ids[rank % len(device_ids)]
-    device_id = torch.device(device_type, device_ids[rank])
+    device_id = torch.device(device_type, 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,7 @@ def spawn_parallel_fn(
     master_addr: str = "localhost",
     master_port: str = "29500",
     device_type: str = "cuda",
-    device_ids: Optional[List[int]] = None,
+    start_method: str = "spawn",
     **kwargs,
 ):
     # clear environment variables
@@ -147,8 +139,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, device_ids[0])
+        device_id = torch.device(device_type, 0)
+        os.environ["LOCAL_RANK"] = "0"
+        os.environ["WORLD_SIZE"] = "1"
         os.environ["LOCAL_DEVICE"] = str(device_id)
 
         func(**kwargs)
@@ -160,11 +153,14 @@ def spawn_parallel_fn(
         master_addr,
         master_port,
         device_type,
-        device_ids,
         func,
         kwargs,
     )
 
-    mp.spawn(
-        wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
+    mp.start_processes(
+        wrapper_spawn_func,
+        args=wrapper_spawn_func_args,
+        nprocs=world_size,
+        start_method=start_method,
+        join=True,
     )

astrai/protocols.py

@@ -0,0 +1,21 @@
+"""Training component protocols — structural subtyping for optimizer/scheduler wrappers."""
+
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class OptimizerProtocol(Protocol):
+    def step(self, closure=None): ...
+    def zero_grad(self): ...
+    @property
+    def param_groups(self) -> Any: ...
+    def state_dict(self) -> dict: ...
+    def load_state_dict(self, d: dict): ...
+
+
+@runtime_checkable
+class SchedulerProtocol(Protocol):
+    def step(self): ...
+    def state_dict(self) -> dict: ...
+    def load_state_dict(self, d: dict): ...
+    def get_last_lr(self): ...

astrai/serialization.py

@@ -1,63 +1,29 @@
 import json
-import os
+import time
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, Optional
 
-import h5py
 import safetensors.torch as st
 import torch
 import torch.distributed as dist
-from torch import Tensor
 
 from astrai.parallel.setup import get_rank
 
 
-def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
-    os.makedirs(file_path, exist_ok=True)
-    full_file_path = os.path.join(file_path, f"{file_name}.h5")
-    with h5py.File(full_file_path, "w") as f:
-        for key, tensors in tensor_group.items():
-            grp = f.create_group(key)
-            for idx, tensor in enumerate(tensors):
-                arr = tensor.cpu().numpy()
-                grp.create_dataset(f"data_{idx}", data=arr)
-
-
-def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
-    tensor_group: Dict[str, List[Tensor]] = {}
-
-    root_path = Path(file_path)
-    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
-
-    for h5_file in h5_files:
-        with h5py.File(h5_file, "r") as f:
-            for key in f.keys():
-                grp = f[key]
-                dsets = []
-                for dset_name in grp.keys():
-                    dset = grp[dset_name]
-                    tensor = torch.from_numpy(dset[:])
-                    if share_memory:
-                        tensor = tensor.share_memory_()
-                    dsets.append(tensor)
-
-                if tensor_group.get(key) is None:
-                    tensor_group[key] = []
-                tensor_group[key].extend(dsets)
-
-    return tensor_group
-
-
 class Checkpoint:
     def __init__(
         self,
         state_dict: Dict[str, Any],
         epoch: int = 0,
         iteration: int = 0,
+        extra: Optional[Dict[str, Any]] = None,
+        meta: Optional[Dict[str, Any]] = None,
     ):
         self.state_dict = state_dict
         self.epoch = epoch
         self.iteration = iteration
+        self.extra = extra or {}
+        self.meta = meta or {}
 
     def save(
         self,
@@ -72,11 +38,16 @@ class Checkpoint:
             meta = {
                 "epoch": self.epoch,
                 "iteration": self.iteration,
+                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             }
+            meta.update(self.meta)
             with open(save_path / "meta.json", "w") as f:
                 json.dump(meta, f, indent=2)
 
             st.save_file(self.state_dict, save_path / "state_dict.safetensors")
+            if self.extra:
+                for key, value in self.extra.items():
+                    torch.save(value, save_path / f"{key}.pt")
 
     @classmethod
     def load(
@@ -99,8 +70,14 @@ class Checkpoint:
 
         state_dict = st.load_file(save_path / "state_dict.safetensors")
 
+        extra = {}
+        for f in save_path.iterdir():
+            if f.suffix == ".pt" and f.stem not in ("meta",):
+                extra[f.stem] = torch.load(f, map_location="cpu", weights_only=False)
+
         return cls(
             state_dict=state_dict,
             epoch=meta["epoch"],
             iteration=meta["iteration"],
+            extra=extra or None,
         )

astrai/tokenize/tokenizer.py

@@ -51,9 +51,26 @@ class AutoTokenizer:
                 self.set_chat_template(config["chat_template"])
 
     @classmethod
-    def from_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoTokenizer":
-        """Load tokenizer from pretrained directory."""
+    def from_pretrained(cls, path: Union[str, Path]) -> "AutoTokenizer":
+        """Load tokenizer from pretrained directory.
+
+        Raises:
+            FileNotFoundError: If tokenizer.json is missing.
+            RuntimeError: If tokenizer failed to initialize.
+        """
+        path = Path(path)
+        tokenizer_file = path / "tokenizer.json"
+        if not tokenizer_file.exists():
+            raise FileNotFoundError(
+                f"Tokenizer file not found: {tokenizer_file}. "
+                "A valid tokenizer.json is required."
+            )
         instance = cls(path)
+        if instance._tokenizer is None:
+            raise RuntimeError(
+                f"Failed to load tokenizer from {path}. "
+                "The tokenizer.json may be corrupted or incompatible."
+            )
         return instance
 
     def save_pretrained(self, save_path: str):
@@ -64,6 +81,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/__init__.py

@@ -1,3 +1,4 @@
+from astrai.trainer.optim import Muon
 from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 from astrai.trainer.train_callback import (
@@ -9,6 +10,8 @@ from astrai.trainer.trainer import Trainer
 __all__ = [
     # Main trainer
     "Trainer",
+    # Optimizer
+    "Muon",
     # Strategy factory
     "StrategyFactory",
     "BaseStrategy",

astrai/trainer/metric_util.py

@@ -1,75 +1,42 @@
-from typing import Dict
+from typing import Any, Callable, Dict
 
+import torch
 import torch.nn as nn
 
 
-def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
-    """Compute gradient norm for each parameter in the model."""
-    norms = {}
+def _grad_stat(
+    model: nn.Module, fn: Callable[[torch.Tensor], Any], default: Any
+) -> dict:
+    results = {}
     for name, param in model.named_parameters():
-        norms[name] = 0.0
-        if param.grad:
-            norm = param.grad.data.norm(norm_type).item()
-            norms[name] = norm
-    return norms
+        results[name] = default
+        if param.grad is not None:
+            results[name] = fn(param.grad.data)
+    return results
+
+
+def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+    return _grad_stat(model, lambda g: g.norm(norm_type).item(), 0.0)
 
 
 def grad_std(model: nn.Module) -> Dict[str, float]:
-    """Compute standard deviation of gradients for each parameter."""
-    stds = {}
-    for name, param in model.named_parameters():
-        stds[name] = 0.0
-        if param.grad:
-            std = param.grad.data.std().item()
-            stds[name] = std
-    return stds
+    return _grad_stat(model, lambda g: g.std().item(), 0.0)
 
 
 def grad_max(model: nn.Module) -> Dict[str, float]:
-    """Find the maximum absolute gradient value for each parameter."""
-    max_vals = {}
-    for name, param in model.named_parameters():
-        max_vals[name] = -float("inf")
-        if param.grad:
-            max_val = param.grad.data.max().item()
-            max_vals[name] = max_val
-
-    return max_vals
+    return _grad_stat(model, lambda g: g.max().item(), -float("inf"))
 
 
 def grad_min(model: nn.Module) -> Dict[str, float]:
-    """Find the minimum absolute gradient value for each parameter."""
-    min_vals = {}
-    for name, param in model.named_parameters():
-        min_vals[name] = float("inf")
-        if param.grad:
-            min_val = param.grad.data.min().item()
-            min_vals[name] = min_val
-
-    return min_vals
+    return _grad_stat(model, lambda g: g.min().item(), float("inf"))
 
 
 def grad_mean(model: nn.Module) -> Dict[str, float]:
-    """Compute mean of gradients for each parameter."""
-    means = {}
-    for name, param in model.named_parameters():
-        means[name] = 0.0
-        if param.grad:
-            mean = param.grad.data.mean().item()
-            means[name] = mean
-
-    return means
+    return _grad_stat(model, lambda g: g.mean().item(), 0.0)
 
 
 def grad_nan_num(model: nn.Module) -> Dict[str, int]:
-    """Count the number of NaNs in gradients for each parameter."""
-    nan_nums = {}
-    for name, param in model.named_parameters():
-        nan_nums[name] = 0
-        if param.grad:
-            nan_num = param.grad.isnan().sum().item()
-            nan_nums[name] = nan_num
-    return nan_nums
+    return _grad_stat(model, lambda g: g.isnan().sum().item(), 0)
 
 
 def ctx_get_loss(ctx):
@@ -80,6 +47,10 @@ def ctx_get_lr(ctx):
     return ctx.optimizer.param_groups[-1]["lr"]
 
 
+def ctx_get_val_loss(ctx):
+    return ctx.val_loss
+
+
 def ctx_get_grad_norm(ctx):
     return grad_norm(ctx.model)
 

astrai/trainer/optim.py

@@ -0,0 +1,143 @@
+import torch
+from torch.optim import Optimizer
+
+
+def _zeropower_via_newtonschulz(G: torch.Tensor, steps: int = 5):
+    assert G.ndim == 2
+    X = G
+    scale = max(1, G.size(0) / G.size(1)) ** 0.5
+    X = X / (X.norm() + 1e-7) * scale
+    if steps == 0:
+        return X
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    for _ in range(steps):
+        A = X @ X.T
+        B = A @ X
+        X = a * X + b * B + c * (A @ B)
+    return X
+
+
+class Muon(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr: float = 2e-3,
+        momentum: float = 0.95,
+        weight_decay: float = 0.0,
+        nesterov: bool = True,
+        ns_steps: int = 5,
+        adamw_lr: float = None,
+        adamw_betas: tuple = (0.9, 0.95),
+        adamw_eps: float = 1e-8,
+        adamw_wd: float = 0.0,
+    ):
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            ns_steps=ns_steps,
+            adamw_lr=adamw_lr if adamw_lr is not None else lr * 0.1,
+            adamw_betas=adamw_betas,
+            adamw_eps=adamw_eps,
+            adamw_wd=adamw_wd,
+        )
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_2d, params_1d = [], []
+            grads_2d, grads_1d = [], []
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                if p.grad.is_sparse:
+                    raise RuntimeError("Muon does not support sparse gradients")
+                if p.ndim >= 2:
+                    params_2d.append(p)
+                    grads_2d.append(p.grad)
+                else:
+                    params_1d.append(p)
+                    grads_1d.append(p.grad)
+
+            if params_2d:
+                self._muon_update_foreach(params_2d, grads_2d, group)
+            if params_1d:
+                self._adamw_update_foreach(params_1d, grads_1d, group)
+
+        return loss
+
+    def _muon_update_foreach(self, params_2d, grads_2d, group):
+        lr = group["lr"]
+        momentum = group["momentum"]
+        wd = group["weight_decay"]
+        nesterov = group["nesterov"]
+        ns_steps = group["ns_steps"]
+
+        if wd != 0:
+            torch._foreach_mul_(params_2d, 1 - lr * wd)
+
+        if nesterov:
+            grads_2d = torch._foreach_add(grads_2d, params_2d, alpha=wd)
+
+        bufs = []
+        for p, grad in zip(params_2d, grads_2d):
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(grad)
+            bufs.append(state["momentum_buffer"])
+
+        torch._foreach_lerp_(bufs, grads_2d, 1 - momentum)
+
+        for p, buf in zip(params_2d, bufs):
+            update = _zeropower_via_newtonschulz(buf, steps=ns_steps)
+            scale = max(1, p.size(0) / p.size(1)) ** 0.5
+            p.add_(update, alpha=-lr * scale)
+
+    def _adamw_update_foreach(self, params_1d, grads_1d, group):
+        lr = group["adamw_lr"]
+        betas = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        wd = group["adamw_wd"]
+
+        steps: list[int] = []
+        exp_avgs, exp_avg_sqs = [], []
+        has_state = []
+        for p in params_1d:
+            state = self.state[p]
+            if not state:
+                state["step"] = 0
+                state["exp_avg"] = torch.zeros_like(p)
+                state["exp_avg_sq"] = torch.zeros_like(p)
+                has_state.append(False)
+            else:
+                has_state.append(True)
+            state["step"] += 1
+            steps.append(state["step"])
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+
+        beta1, beta2 = betas
+
+        torch._foreach_lerp_(exp_avgs, grads_1d, 1 - beta1)
+        grads_sq = torch._foreach_mul(grads_1d, grads_1d)
+        torch._foreach_lerp_(exp_avg_sqs, grads_sq, 1 - beta2)
+
+        bias_correction1 = [1 - beta1**s for s in steps]
+        bias_correction2 = [1 - beta2**s for s in steps]
+
+        if wd != 0:
+            torch._foreach_mul_(params_1d, 1 - lr * wd)
+
+        exp_avg_corrected = torch._foreach_div(exp_avgs, bias_correction1)
+        denom = torch._foreach_div(exp_avg_sqs, bias_correction2)
+        denom = torch._foreach_sqrt(denom)
+        torch._foreach_add_(denom, eps)
+        torch._foreach_addcdiv_(params_1d, exp_avg_corrected, denom, value=-lr)

astrai/trainer/strategy.py

@@ -8,15 +8,17 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from astrai.factory import BaseFactory
 
 
 def unwrap_model(model: nn.Module) -> nn.Module:
-    """Unwrap DDP wrapper if present to get the original model."""
     if isinstance(model, DDP):
         return model.module
+    if isinstance(model, FSDP):
+        return model._fsdp_wrapped_module
     return model
 
 
@@ -265,7 +267,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 +280,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 +289,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 +313,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 +327,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(0)  # Off-policy: policy_model = old_model
+        ratio = torch.exp(log_probs_policy - log_probs_ref)
+
         surr1 = ratio * advantages
         surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
 

astrai/trainer/train_callback.py

@@ -1,15 +1,21 @@
 import json
+import logging
 import os
+import sys
 import time
 from pathlib import Path
-from typing import Callable, List, Optional, Protocol, runtime_checkable
+from typing import IO, Callable, List, Optional, Protocol, runtime_checkable
 
+import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch.nn.utils import clip_grad_norm_
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
 from tqdm import tqdm
 
 from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
+from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
     ctx_get_grad_max,
@@ -20,9 +26,12 @@ from astrai.trainer.metric_util import (
     ctx_get_grad_std,
     ctx_get_loss,
     ctx_get_lr,
+    ctx_get_val_loss,
 )
 from astrai.trainer.train_context import TrainContext
 
+logger = logging.getLogger(__name__)
+
 
 @runtime_checkable
 class TrainCallback(Protocol):
@@ -42,18 +51,15 @@ class TrainCallback(Protocol):
     def on_epoch_end(self, context: TrainContext):
         """Called at the end of each epoch."""
 
-    def on_step_begin(self, context: TrainContext):
-        """Called at the beginning of each step."""
-
-    def on_step_end(self, context: TrainContext):
-        """Called at the end of each step."""
-
     def on_batch_begin(self, context: TrainContext):
         """Called at the beginning of each batch."""
 
     def on_batch_end(self, context: TrainContext):
         """Called at the end of each batch."""
 
+    def on_optimizer_step(self, context: TrainContext):
+        """Called on every optimizer step (sync step only)."""
+
     def on_error(self, context: TrainContext):
         """Called when an error occurs during training."""
 
@@ -69,12 +75,6 @@ class CallbackFactory(BaseFactory[TrainCallback]):
         callback = CallbackFactory.create("my_callback", **kwargs)
     """
 
-    @classmethod
-    def _validate_component(cls, callback_cls: type) -> None:
-        """Validate that the callback class inherits from TrainCallback."""
-        if not issubclass(callback_cls, TrainCallback):
-            raise TypeError(f"{callback_cls.__name__} must inherit from TrainCallback")
-
 
 @CallbackFactory.register("gradient_clipping")
 class GradientClippingCallback(TrainCallback):
@@ -85,28 +85,43 @@ class GradientClippingCallback(TrainCallback):
     def __init__(self, max_grad_norm: float):
         self.max_grad_norm = max_grad_norm
 
-    def on_step_begin(self, context: TrainContext):
-        _ = context
+    def on_optimizer_step(self, context: TrainContext):
         clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
 
 
-@CallbackFactory.register("scheduler")
-class SchedulerCallback(TrainCallback):
+@CallbackFactory.register("gradient_checkpointing")
+class GradientCheckpointingCallback(TrainCallback):
     """
-    Scheduler callback for trainer.
+    Activation checkpointing callback — trades compute for memory
+    by recomputing specified module activations during the backward pass.
+
+    Args:
+        modules: Module types to apply checkpointing to.
     """
 
-    def __init__(self):
-        pass
+    def __init__(self, modules: Optional[List[type]] = None):
+        self.modules = tuple(modules) if modules else ()
+
+    def _enable(self, module: nn.Module):
+        if self.modules and isinstance(module, self.modules):
+            fn = module.forward
+            module._original_forward = fn
+            module.forward = lambda *a, **kw: torch_checkpoint(
+                fn, *a, use_reentrant=False, **kw
+            )
+
+    @staticmethod
+    def _disable(module: nn.Module):
+        if hasattr(module, "_original_forward"):
+            module.forward = module._original_forward
+            del module._original_forward
 
     def on_train_begin(self, context: TrainContext):
-        for group in context.optimizer.param_groups:
-            if "initial_lr" not in group:
-                group["initial_lr"] = group["lr"]
+        context.model.apply(self._enable)
+        logger.info("Gradient checkpointing enabled")
 
-    def on_batch_end(self, context: TrainContext):
-        if context.scheduler:
-            context.scheduler.step()
+    def on_train_end(self, context: TrainContext):
+        context.model.apply(self._disable)
 
 
 @CallbackFactory.register("checkpoint")
@@ -115,17 +130,23 @@ class CheckpointCallback(TrainCallback):
     Checkpoint callback for trainer.
     """
 
+    extra_keys = ("optimizer", "scheduler")
+
     def __init__(
         self,
         save_dir: str,
         interval: int,
         weight_only: bool = False,
         state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
+        save_extra_fn: Optional[Callable[["TrainContext"], dict]] = None,
+        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = 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 or CheckpointCallback.save_extra
+        self.load_extra_fn = load_extra_fn or CheckpointCallback.load_extra
         self.last_ckpt_iter = 0
 
     @only_on_rank(0)
@@ -139,13 +160,22 @@ class CheckpointCallback(TrainCallback):
             else context.model.state_dict()
         )
 
+        extra = self.save_extra_fn(context)
         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,
+            meta=context.config.to_dict(),
         )
 
         context.checkpoint.save(save_path)
         self.last_ckpt_iter = context.iteration
 
+    def on_train_begin(self, context: TrainContext):
+        if context.checkpoint and context.checkpoint.extra:
+            self.load_extra_fn(context.checkpoint.extra, context)
+
     def on_batch_end(self, context: TrainContext):
         if context.iteration - self.last_ckpt_iter >= self.interval:
             self._save_checkpoint(context)
@@ -157,6 +187,21 @@ class CheckpointCallback(TrainCallback):
     def on_error(self, context: TrainContext):
         self._save_checkpoint(context)
 
+    @staticmethod
+    def save_extra(context: TrainContext) -> dict:
+        extra = {}
+        for name in CheckpointCallback.extra_keys:
+            obj = getattr(context, name, None)
+            if obj:
+                extra[name] = obj.state_dict()
+        return extra
+
+    @staticmethod
+    def load_extra(extra: dict, context: TrainContext):
+        for name in CheckpointCallback.extra_keys:
+            if name in extra:
+                getattr(context, name).load_state_dict(extra[name])
+
 
 @CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):
@@ -164,8 +209,12 @@ class ProgressBarCallback(TrainCallback):
     Progress bar callback for trainer.
     """
 
-    def __init__(self, num_epoch: int):
+    def __init__(
+        self, num_epoch: int, log_interval: int = 100, file: IO[str] = sys.stdout
+    ):
         self.num_epoch = num_epoch
+        self.log_interval = log_interval
+        self.file = file
         self.progress_bar: tqdm = None
 
     @only_on_rank(0)
@@ -174,16 +223,18 @@ class ProgressBarCallback(TrainCallback):
             context.dataloader,
             desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
             dynamic_ncols=True,
+            file=self.file,
         )
 
     @only_on_rank(0)
     def on_batch_end(self, context: TrainContext):
-        self.progress_bar.set_postfix(
-            {
-                "loss": f"{context.loss:.4f}",
-                "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
-            }
-        )
+        postfix = {
+            "loss": f"{context.loss:.4f}",
+            "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
+        }
+        if context.val_loss > 0:
+            postfix["val_loss"] = f"{context.val_loss:.4f}"
+        self.progress_bar.set_postfix(postfix)
         self.progress_bar.update(1)
 
     @only_on_rank(0)
@@ -215,6 +266,7 @@ class MetricLoggerCallback(TrainCallback):
         self._metric_funcs = {
             "loss": ctx_get_loss,
             "lr": ctx_get_lr,
+            "val_loss": ctx_get_val_loss,
             "grad_norm": ctx_get_grad_norm,
             "grad_std": ctx_get_grad_std,
             "grad_max": ctx_get_grad_max,
@@ -225,7 +277,7 @@ class MetricLoggerCallback(TrainCallback):
 
     def _get_log_data(self, context: TrainContext):
         return {
-            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             "epoch": context.epoch,
             "iter": context.iteration,
             **{m: self._metric_funcs[m](context) for m in self.metrics},
@@ -258,3 +310,43 @@ class MetricLoggerCallback(TrainCallback):
 
     def on_error(self, context):
         self._save_log(context.epoch, context.iteration)
+
+
+@CallbackFactory.register("validation")
+class ValidationCallback(TrainCallback):
+    def _run_validation(self, context: TrainContext):
+        context.model.eval()
+
+        total_loss = 0.0
+        num_batches = 0
+
+        with torch.no_grad():
+            for batch in context.val_dataloader:
+                loss = context.strategy(batch)
+                total_loss += loss.item()
+                num_batches += 1
+
+        avg_loss = total_loss / max(num_batches, 1)
+
+        if context.world_size > 1 and dist.is_initialized():
+            loss_tensor = torch.tensor([avg_loss], device=get_current_device())
+            dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
+            avg_loss = loss_tensor.item()
+
+        context.val_loss = avg_loss
+        context.model.train()
+
+        step_count = context.iteration // context.config.grad_accum_steps
+        logger.info(
+            f"Epoch {context.epoch + 1}, Step {step_count}, Val Loss: {avg_loss:.4f}"
+        )
+
+    def on_optimizer_step(self, context: TrainContext):
+        if context.val_dataloader is None:
+            return
+        cfg = context.config
+        if cfg.val_step <= 0:
+            return
+        step_count = context.iteration // cfg.grad_accum_steps
+        if step_count % cfg.val_step == 0:
+            self._run_validation(context)

astrai/trainer/train_context.py

@@ -2,13 +2,13 @@ from dataclasses import dataclass, field
 from typing import Optional, Self
 
 import torch.nn as nn
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import ResumableDistributedSampler
+from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
+from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import Checkpoint
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
@@ -18,13 +18,17 @@ class TrainContext:
     model: nn.Module = field(default=None)
     strategy: BaseStrategy = field(default=None)
     dataloader: DataLoader = field(default=None)
-    optimizer: Optimizer = field(default=None)
-    scheduler: LRScheduler = field(default=None)
+    optimizer: OptimizerProtocol = field(default=None)
+    scheduler: SchedulerProtocol = field(default=None)
     checkpoint: Checkpoint = field(default=None)
+    config: TrainConfig = field(default=None)
+    executor: BaseExecutor = field(default=None)
 
     epoch: int = field(default=0)
     iteration: int = field(default=0)
     loss: float = field(default=0.0)
+    val_dataloader: DataLoader = field(default=None)
+    val_loss: float = field(default=0.0)
 
     world_size: int = field(default=1)
     rank: int = field(default=0)
@@ -32,68 +36,97 @@ class TrainContext:
 
 
 class TrainContextBuilder:
-    def __init__(self, config: TrainConfig):
+    def __init__(
+        self,
+        config: TrainConfig,
+    ):
         self.config = config
-        self._context = TrainContext(
-            model=config.model,
-            world_size=get_world_size(),
-            rank=get_rank(),
-        )
-
-        device = get_current_device()
-        self._context.model = self._context.model.to(device=device)
-
-        if self.config.nprocs > 1:
-            fn = self.config.parallel_wrapper
-            self._context.model = fn(self._context.model)
-
-        self._context.optimizer = self.config.optimizer_fn(self._context.model)
-        self._context.scheduler = self.config.scheduler_fn(self._context.optimizer)
+        self._checkpoint: Optional[Checkpoint] = None
 
     def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
-        if checkpoint is None:
-            checkpoint = Checkpoint(
-                state_dict=self._context.model.state_dict(),
-            )
-        else:
-            # resume from the assigned checkpoint or assigned iteration
-            self._context.epoch = max(checkpoint.epoch, self.config.start_epoch)
-            self._context.iteration = max(checkpoint.iteration, self.config.start_batch)
-            self._context.model.load_state_dict(checkpoint.state_dict)
-
-        self._context.checkpoint = checkpoint
-        return self
-
-    def with_dataloader(self) -> Self:
-        # fix: change batch level iteration to sample level offset
-        config = self.config
-        sampler_offset = self._context.iteration * config.batch_size
-        resumeable_sampler = ResumableDistributedSampler(
-            data_source=config.dataset,
-            start_epoch=self._context.epoch,
-            start_iter=sampler_offset,
-            seed=config.random_seed,
-        )
-
-        dataloader = DataLoader(
-            config.dataset,
-            batch_size=config.batch_size,
-            sampler=resumeable_sampler,
-            num_workers=config.num_workers,
-            pin_memory=config.pin_memory,
-            prefetch_factor=config.prefetch_factor,
-        )
-        self._context.dataloader = dataloader
-        return self
-
-    def with_strategy(self) -> Self:
-        self._context.strategy = StrategyFactory.create(
-            model=self._context.model,
-            train_type=self.config.strategy,
-            device=get_current_device(),
-            **self.config.extra_kwargs,
-        )
+        self._checkpoint = checkpoint
         return self
 
     def build(self) -> TrainContext:
-        return self._context
+        cfg = self.config
+        device = get_current_device()
+
+        executor = ExecutorFactory.create(
+            cfg.parallel_mode,
+            grad_accum_steps=cfg.grad_accum_steps,
+            **cfg.executor_kwargs,
+        )
+
+        context = TrainContext(
+            model=cfg.model,
+            world_size=get_world_size(),
+            rank=get_rank(),
+            config=cfg,
+            executor=executor,
+        )
+
+        context.model = context.model.to(device=device)
+
+        if self._checkpoint is not None:
+            context.epoch = max(self._checkpoint.epoch, cfg.start_epoch)
+            context.iteration = max(self._checkpoint.iteration, cfg.start_batch)
+            context.model.load_state_dict(self._checkpoint.state_dict)
+            context.checkpoint = self._checkpoint
+        else:
+            context.checkpoint = Checkpoint(
+                state_dict=context.model.state_dict(),
+            )
+
+        context.optimizer = cfg.optimizer_fn(context.model)
+        context.scheduler = cfg.scheduler_fn(context.optimizer)
+
+        sampler_offset = context.iteration * cfg.batch_per_device
+        sampler = ResumableDistributedSampler(
+            data_source=cfg.dataset,
+            start_epoch=context.epoch,
+            start_iter=sampler_offset,
+            seed=cfg.random_seed,
+        )
+        context.dataloader = DataLoader(
+            cfg.dataset,
+            batch_size=cfg.batch_per_device,
+            sampler=sampler,
+            num_workers=cfg.num_workers,
+            pin_memory=cfg.pin_memory,
+            prefetch_factor=cfg.prefetch_factor,
+        )
+
+        if cfg.val_dataset is not None:
+            val_sampler = ResumableDistributedSampler(
+                data_source=cfg.val_dataset,
+                start_epoch=0,
+                start_iter=0,
+                seed=cfg.random_seed,
+                shuffle=False,
+            )
+            context.val_dataloader = DataLoader(
+                cfg.val_dataset,
+                batch_size=cfg.batch_per_device,
+                sampler=val_sampler,
+                num_workers=cfg.num_workers,
+                pin_memory=cfg.pin_memory,
+                prefetch_factor=cfg.prefetch_factor,
+            )
+
+        context.model, context.optimizer, context.dataloader, context.scheduler = (
+            executor.prepare(
+                context.model,
+                context.optimizer,
+                context.dataloader,
+                context.scheduler,
+            )
+        )
+
+        context.strategy = StrategyFactory.create(
+            model=context.model,
+            train_type=cfg.strategy,
+            device=device,
+            **cfg.extra_kwargs,
+        )
+
+        return context

astrai/trainer/trainer.py

@@ -25,22 +25,28 @@ class Trainer:
 
     def _get_default_callbacks(self) -> List[TrainCallback]:
         cfg = self.train_config
-        return [
+        callbacks = [
+            CallbackFactory.create(
+                "gradient_checkpointing",
+                modules=cfg.gradient_checkpointing_modules,
+            ),
+            CallbackFactory.create(
+                "checkpoint",
+                cfg.ckpt_dir,
+                cfg.ckpt_interval,
+            ),
+            CallbackFactory.create(
+                "metric_logger",
+                log_dir=cfg.log_dir,
+                save_interval=cfg.ckpt_interval,
+                log_interval=cfg.log_interval,
+                metrics=cfg.metrics,
+            ),
             CallbackFactory.create("progress_bar", cfg.n_epoch),
-            CallbackFactory.create("checkpoint", cfg.ckpt_dir, cfg.ckpt_interval),
-            CallbackFactory.create("metric_logger", cfg.ckpt_dir, cfg.ckpt_interval),
             CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
-            CallbackFactory.create("scheduler"),
+            CallbackFactory.create("validation"),
         ]
-
-    def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
-        return (
-            TrainContextBuilder(self.train_config)
-            .with_checkpoint(checkpoint)
-            .with_dataloader()
-            .with_strategy()
-            .build()
-        )
+        return callbacks
 
     def _call_callbacks(self, method_name: str, context: TrainContext):
         for callback in self.callbacks:
@@ -48,49 +54,38 @@ class Trainer:
             if method:
                 method(context)
 
-    def train(self, checkpoint: Optional[Checkpoint] = None):
-        config = self.train_config
-        spawn_parallel_fn(
-            self._train_impl,
-            backend=config.backend,
-            world_size=config.nprocs,
-            master_addr=config.master_addr,
-            master_port=config.master_port,
-            device_type=config.device_type,
-            device_ids=config.device_ids,
-            checkpoint=checkpoint,
+    def _trainer_loop(self, checkpoint: Optional[Checkpoint] = None):
+        context = (
+            TrainContextBuilder(self.train_config).with_checkpoint(checkpoint).build()
         )
-
-    def _train_impl(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
-        context = self._build_context(checkpoint)
+        executor = context.executor
         self._call_callbacks("on_train_begin", context)
 
         try:
             context.model.train()
-            # 1.epoch
-            for epoch in range(context.epoch, self.train_config.n_epoch):
+
+            for epoch in range(context.epoch, context.config.n_epoch):
                 context.epoch = epoch
                 self._call_callbacks("on_epoch_begin", context)
 
                 for batch in context.dataloader:
-                    if context.iteration % self.train_config.accumulation_steps == 0:
-                        # 2. step
-                        self._call_callbacks("on_step_begin", context)
-                        context.optimizer.step()
-                        context.optimizer.zero_grad()
-                        self._call_callbacks("on_step_end", context)
-
-                    # 3. batch
                     self._call_callbacks("on_batch_begin", context)
-                    loss = context.strategy(batch)
-                    context.loss = loss.item()
-                    context.iteration += 1
 
-                    # to make the loss normalized by accumulation steps
-                    stand_loss = loss / self.train_config.accumulation_steps
-                    stand_loss.backward()
+                    with executor.accumulate(context.model):
+                        loss = context.strategy(batch)
+                        context.loss = loss.item()
+                        stand_loss = loss / executor.grad_accum_steps
+                        executor.backward(stand_loss)
+                        context.iteration += 1
+                        self._call_callbacks("on_batch_end", context)
 
-                    self._call_callbacks("on_batch_end", context)
+                        if executor.sync_gradients:
+                            self._call_callbacks("on_optimizer_step", context)
+                            context.optimizer.step()
+                            context.optimizer.zero_grad()
+
+                        if context.scheduler:
+                            context.scheduler.step()
 
                 self._call_callbacks("on_epoch_end", context)
 
@@ -100,3 +95,16 @@ class Trainer:
             raise
         finally:
             self._call_callbacks("on_train_end", context)
+
+    def train(self, checkpoint: Optional[Checkpoint] = None):
+        cfg = self.train_config
+        spawn_parallel_fn(
+            self._trainer_loop,
+            backend=cfg.backend,
+            world_size=cfg.nprocs,
+            master_addr=cfg.master_addr,
+            master_port=cfg.master_port,
+            device_type=cfg.device_type,
+            start_method=cfg.start_method,
+            checkpoint=checkpoint,
+        )

docker-compose.yml

@@ -0,0 +1,44 @@
+services:
+  server:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
+    ports:
+      - "8000:8000"
+    volumes:
+      - ./params:/app/params:ro
+    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:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
+    ports:
+      - "8000:8000"
+    volumes:
+      - ./params:/app/params:ro
+    command: python -m scripts.tools.server --port 8000 --device cpu
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+      interval: 30s
+      timeout: 10s
+      retries: 3
+      start_period: 120s
+    restart: unless-stopped

scripts/demo/generate_ar.py

@@ -11,7 +11,6 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def generate_text():
-    # Load model from pretrained
     model = AutoModel.from_pretrained(PARAMETER_ROOT)
     tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
     model.to(device="cuda", dtype=torch.bfloat16)
@@ -22,16 +21,15 @@ def generate_text():
         model=model,
         tokenizer=tokenizer,
     )
-    response = engine.generate(
+    for token in engine.generate(
         prompt=query,
-        stream=False,
+        stream=True,
         max_tokens=2048,
         temperature=0.8,
         top_p=0.95,
         top_k=50,
-    )
-
-    print(response)
+    ):
+        print(token, end="", flush=True)
 
 
 if __name__ == "__main__":

scripts/demo/generate_batch.py

@@ -24,12 +24,23 @@ def batch_generate():
         "请问什么是显卡",
     ]
 
+    prompts = [
+        tokenizer.apply_chat_template(
+            [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": q},
+            ],
+            tokenize=False,
+        )
+        for q in inputs
+    ]
+
     engine = InferenceEngine(
         model=model,
         tokenizer=tokenizer,
     )
     responses = engine.generate(
-        prompt=inputs,
+        prompt=prompts,
         stream=False,
         max_tokens=2048,
         temperature=0.8,

scripts/demo/stream_chat.py

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

scripts/docker.sh

@@ -16,6 +16,7 @@ NC='\033[0m' # No Color
 IMAGE_NAME="astrai"
 IMAGE_TAG="latest"
 REGISTRY=""
+CONTAINER_ID=""
 
 # Print colored messages
 print_info() {
@@ -175,6 +176,10 @@ main() {
                 PORT="$2"
                 shift 2
                 ;;
+            --container)
+                CONTAINER_ID="$2"
+                shift 2
+                ;;
             --gpu)
                 GPU=true
                 shift
@@ -197,6 +202,7 @@ main() {
                 echo "  --dockerfile FILE  Dockerfile path (default: Dockerfile)"
                 echo "  --context PATH     Build context (default: .)"
                 echo "  --port PORT        Port for run (default: 8000)"
+                echo "  --container ID     Container ID for logs"
                 echo "  --gpu              Enable GPU support"
                 echo "  --help             Show this help message"
                 echo ""
@@ -205,6 +211,7 @@ main() {
                 echo "  $0 build --tag v1.0.0"
                 echo "  $0 run --port 8080"
                 echo "  $0 run --gpu"
+                echo "  $0 logs --container abc123"
                 echo "  $0 push --registry ghcr.io/username"
                 exit 0
                 ;;
@@ -237,7 +244,7 @@ main() {
             show_info
             ;;
         logs)
-            show_logs "$2"
+            show_logs "$CONTAINER_ID"
             ;;
         "")
             print_error "No command specified. Use --help for usage"

scripts/tools/benchmark.py

@@ -1,9 +1,13 @@
+"""Benchmark AutoRegressiveLM with KVCache"""
+
 from dataclasses import dataclass
 from typing import Any, Dict
 
 import torch
 
-from astrai.model.transformer import ModelConfig, Transformer
+from astrai.config import AutoRegressiveLMConfig
+from astrai.inference import KVCache
+from astrai.model.transformer import AutoRegressiveLM
 
 
 @dataclass
@@ -17,29 +21,27 @@ class BenchmarkResult:
 class GenerationBenchmark:
     def __init__(
         self,
-        config: ModelConfig,
+        config: AutoRegressiveLMConfig,
         device: str = "cuda",
-        dtype: torch.dtype = torch.float16,
+        dtype: torch.dtype = torch.bfloat16,
+        page_size: int = 128,
     ):
         self.config = config
         self.device = device
         self.dtype = dtype
-        self.model = Transformer(config).to(device=device, dtype=dtype)
+        self.model = AutoRegressiveLM(config).to(device=device, dtype=dtype)
         self.model.eval()
-
-    def _initialize_kv_cache(self, batch_size: int) -> list:
-        """初始化KV缓存"""
-        config = self.config
-        shape = (
-            batch_size,
-            config.max_len,
+        head_dim = config.dim // config.n_heads
+        n_pages = (config.max_len * 4 + page_size - 1) // page_size
+        self._page_cache = KVCache(
             config.n_layers,
+            n_pages,
+            page_size,
             config.n_kv_heads,
-            config.dim // config.n_heads,
+            head_dim,
+            device,
+            dtype,
         )
-        k_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
-        v_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
-        return (k_cache, v_cache)
 
     def _prepare_inputs(self, batch_size: int, prompt_length: int, total_length: int):
         prompt_ids = torch.randint(
@@ -49,7 +51,6 @@ class GenerationBenchmark:
             device=self.device,
             dtype=torch.long,
         )
-
         gen_ids = torch.randint(
             low=0,
             high=self.config.vocab_size,
@@ -57,7 +58,6 @@ class GenerationBenchmark:
             device=self.device,
             dtype=torch.long,
         )
-
         return prompt_ids, gen_ids
 
     @torch.inference_mode()
@@ -67,13 +67,11 @@ class GenerationBenchmark:
         prompt_length: int = 512,
         num_trials: int = 10,
     ) -> BenchmarkResult:
-
         for _ in range(3):
             prompt_ids, _ = self._prepare_inputs(
                 batch_size, prompt_length, prompt_length
             )
             _ = self.model(prompt_ids)
-
         torch.cuda.synchronize()
 
         total_time = 0.0
@@ -83,20 +81,20 @@ class GenerationBenchmark:
             prompt_ids, _ = self._prepare_inputs(
                 batch_size, prompt_length, prompt_length
             )
-            start_event = torch.cuda.Event(enable_timing=True)
-            end_event = torch.cuda.Event(enable_timing=True)
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
 
-            start_event.record()
+            start.record()
             _ = self.model(prompt_ids)
-            end_event.record()
+            end.record()
             torch.cuda.synchronize()
 
-            trial_time = start_event.elapsed_time(end_event) / 1000
+            trial_time = start.elapsed_time(end) / 1000
             total_time += trial_time
 
             print(
-                f"Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
-                f"({prompt_length / trial_time:.1f} tokens/s)"
+                f"  Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
+                f"({prompt_length / trial_time:.1f} tok/s)"
             )
 
         return BenchmarkResult(
@@ -107,7 +105,7 @@ class GenerationBenchmark:
                 "benchmark_type": "prefill",
                 "batch_size": batch_size,
                 "prompt_length": prompt_length,
-                "dtype": self.dtype,
+                "dtype": str(self.dtype),
                 "device": self.device,
             },
         )
@@ -120,41 +118,74 @@ class GenerationBenchmark:
         gen_length: int = 128,
         num_trials: int = 5,
     ) -> BenchmarkResult:
-
         total_time = 0.0
         total_tokens = batch_size * gen_length * num_trials
+        page_size = self._page_cache.page_size
 
         for trial in range(num_trials):
             prompt_ids, gen_ids = self._prepare_inputs(
-                batch_size, prompt_length, prompt_length + gen_length
+                batch_size,
+                prompt_length,
+                prompt_length + gen_length,
+            )
+
+            n_pages = (prompt_length + gen_length + page_size - 1) // page_size
+            total = n_pages * batch_size
+            pages = []
+            for _ in range(total):
+                p = self._page_cache._pool.alloc()
+                assert p >= 0, "OOM"
+                pages.append(p)
+            page_table = torch.tensor(
+                [pages[i * n_pages : (i + 1) * n_pages] for i in range(batch_size)],
+                dtype=torch.long,
+                device=self.device,
+            )
+
+            cv = self._page_cache.bind(page_table, total_len=prompt_length)
+            _ = self.model(
+                prompt_ids,
+                paged_cache=cv,
+                position_ids=torch.arange(
+                    prompt_length, dtype=torch.long, device=self.device
+                )
+                .unsqueeze(0)
+                .expand(batch_size, -1),
             )
-            kv_cache = self._initialize_kv_cache(batch_size)
-            _ = self.model(prompt_ids, persistent_key_values=kv_cache, start_pos=0)
 
             torch.cuda.synchronize()
 
-            start_event = torch.cuda.Event(enable_timing=True)
-            end_event = torch.cuda.Event(enable_timing=True)
-
-            start_event.record()
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
 
+            start.record()
             current_pos = prompt_length
             for i in range(gen_length):
                 input_token = gen_ids[:, i : i + 1]
+                cv = self._page_cache.bind(page_table, total_len=current_pos + 1)
                 _ = self.model(
-                    input_token, persistent_key_values=kv_cache, start_pos=current_pos
+                    input_token,
+                    paged_cache=cv,
+                    position_ids=torch.full(
+                        (batch_size, 1),
+                        current_pos,
+                        dtype=torch.long,
+                        device=self.device,
+                    ),
                 )
                 current_pos += 1
-
-            end_event.record()
+            end.record()
             torch.cuda.synchronize()
 
-            trial_time = start_event.elapsed_time(end_event) / 1000
+            trial_time = start.elapsed_time(end) / 1000
             total_time += trial_time
 
+            for idx in pages:
+                self._page_cache._pool.free(idx)
+
             print(
-                f"Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
-                f"({gen_length / trial_time:.1f} tokens/s)"
+                f"  Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
+                f"({gen_length / trial_time:.1f} tok/s)"
             )
 
         return BenchmarkResult(
@@ -166,36 +197,26 @@ class GenerationBenchmark:
                 "batch_size": batch_size,
                 "prompt_length": prompt_length,
                 "gen_length": gen_length,
-                "dtype": self.dtype,
+                "dtype": str(self.dtype),
                 "device": self.device,
             },
         )
 
 
 def print_benchmark_result(result: BenchmarkResult):
-    """打印基准测试结果"""
-    benchmark_type = result.metadata["benchmark_type"]
-
-    print(f"\n{' ' + benchmark_type.upper().replace('_', ' ') + ' Benchmark ':-^80}")
+    btype = result.metadata["benchmark_type"]
+    print(f"\n{' ' + btype.upper() + ' Benchmark ':-^80}")
     print(f"Total Tokens Processed: {result.total_tokens:,}")
     print(f"Time Consumed: {result.total_time:.3f}s")
-    print(f"Throughput: {result.tokens_per_second:,.1f} tokens/s")
-
-    if benchmark_type == "prefill":
-        print(
-            f"Batch Size: {result.metadata['batch_size']} | Prompt Length: {result.metadata['prompt_length']}"
-        )
-    elif benchmark_type == "decoding":
-        print(
-            f"Batch Size: {result.metadata['batch_size']} | Gen Length: {result.metadata['gen_length']}"
-        )
-
-    print(f"Device: {result.metadata['device']} | Dtype: {result.metadata['dtype']}")
+    print(f"Throughput: {result.tokens_per_second:,.1f} tok/s")
+    for k, v in result.metadata.items():
+        if k != "benchmark_type":
+            print(f"{k.replace('_', ' ').title()}: {v}")
     print("-" * 80)
 
 
 if __name__ == "__main__":
-    config = ModelConfig(
+    config = AutoRegressiveLMConfig(
         vocab_size=10000,
         dim=1536,
         n_heads=24,
@@ -209,15 +230,20 @@ if __name__ == "__main__":
     benchmark = GenerationBenchmark(config)
 
     print("=" * 80)
-    print("Running Transformer Generation Benchmark")
+    print("Running AutoRegressiveLM Generation Benchmark (KVCache)")
     print("=" * 80)
 
     prefill_result = benchmark.run_prefill_benchmark(
-        batch_size=4, prompt_length=512, num_trials=5
+        batch_size=4,
+        prompt_length=512,
+        num_trials=5,
     )
     print_benchmark_result(prefill_result)
 
     gen_result = benchmark.run_decoding_benchmark(
-        batch_size=4, prompt_length=512, gen_length=128, num_trials=5
+        batch_size=4,
+        prompt_length=512,
+        gen_length=128,
+        num_trials=5,
     )
     print_benchmark_result(gen_result)

scripts/tools/generate.py

@@ -9,7 +9,7 @@ from astrai.tokenize import AutoTokenizer
 
 
 def processor(
-    model_dir: str,
+    param_path: str,
     input_json_file: str,
     output_json_file: str,
     temperature: float,
@@ -18,14 +18,17 @@ def processor(
     question_key: str,
     response_key: str,
     max_tokens: int,
+    batch_size: int,
 ):
     # Load model and tokenizer
-    model = AutoModel.from_pretrained(model_dir)
-    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModel.from_pretrained(param_path)
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
     model.to(device="cuda", dtype=torch.bfloat16)
 
     # Create inference engine
-    engine = InferenceEngine(model=model, tokenizer=tokenizer)
+    engine = InferenceEngine(
+        model=model, tokenizer=tokenizer, max_batch_size=batch_size
+    )
 
     with open(input_json_file, "r", encoding="utf-8") as f:
         input_data = [json.loads(line) for line in f]
@@ -72,7 +75,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/server.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torch
 
-from astrai.inference.server import run_server
+from astrai.inference import run_server
 
 
 def main():

scripts/tools/train.py

@@ -4,26 +4,23 @@ from functools import partial
 
 import safetensors.torch as st
 import torch
-import torch.nn as nn
 import torch.optim as optim
-from torch.nn.parallel import DistributedDataParallel as DDP
 
-from astrai.config import ModelConfig, TrainConfig
+from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory
-from astrai.model import Transformer
-from astrai.parallel import get_rank
+from astrai.model import AutoRegressiveLM
 from astrai.trainer import SchedulerFactory, Trainer
 
 
 def parse_args() -> argparse.Namespace:
 
-    parser = argparse.ArgumentParser(description="Train the Transformer model.")
+    parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.")
 
     parser.add_argument(
         "--train_type",
         type=str,
         required=True,
-        choices=["seq", "sft", "dpo"],
+        choices=["seq", "sft", "dpo", "grpo"],
         help="Train type.",
     )
     parser.add_argument(
@@ -43,19 +40,19 @@ def parse_args() -> argparse.Namespace:
         "--n_epoch", type=int, default=1, help="Number of epochs to train."
     )
     parser.add_argument(
-        "--batch_size", type=int, default=1, help="Batch size for training."
+        "--batch_per_device", type=int, default=1, help="Batch size per GPU."
     )
     parser.add_argument(
-        "--accumulation_steps",
+        "--grad_accum_steps",
         type=int,
         default=1,
         help="Number of iterations between each optimizer step.",
     )
     parser.add_argument(
-        "--warmup_steps",
-        type=int,
-        default=1000,
-        help="Number of iters between warnings.",
+        "--warmup_ratio",
+        type=float,
+        default=0.05,
+        help="Fraction of total steps used for LR warmup.",
     )
     parser.add_argument(
         "--max_lr", type=float, default=3e-4, help="Max learning rate for training."
@@ -70,13 +67,13 @@ def parse_args() -> argparse.Namespace:
         "--adamw_beta1",
         type=float,
         default=0.9,
-        help="Beta values for AdamW optimizer.",
+        help="Beta1 for AdamW optimizer.",
     )
     parser.add_argument(
         "--adamw_beta2",
         type=float,
         default=0.95,
-        help="Beta values for AdamW optimizer.",
+        help="Beta2 for AdamW optimizer.",
     )
     parser.add_argument(
         "--adamw_weight_decay",
@@ -100,16 +97,23 @@ 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,
-        default=0.1,
+        default=0.05,
         help="cross_entropy function label smoothing parameter",
     )
 
@@ -125,6 +129,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."
     )
@@ -133,29 +143,31 @@ def parse_args() -> argparse.Namespace:
     )
 
     parser.add_argument("--nprocs", type=int, default=1, help="Number of GPUs to use.")
+    parser.add_argument(
+        "--parallel_mode",
+        type=str,
+        default="none",
+        choices=["none", "ddp"],
+        help="Parallel training strategy.",
+    )
     parser.add_argument(
         "--device_type", type=str, default="cuda", help="Device type to use."
     )
+    parser.add_argument(
+        "--start_method",
+        type=str,
+        default="spawn",
+        choices=["spawn", "fork", "forkserver"],
+        help="Multiprocessing start method.",
+    )
 
     args = parser.parse_args()
 
     return args
 
 
-def ddp_wrap(model: nn.Module):
-    local_rank = get_rank()
-    model = model.to(device=f"cuda:{local_rank}", dtype=torch.bfloat16)
-    ddp_model = DDP(
-        model,
-        device_ids=[local_rank],
-        output_device=local_rank,
-        find_unused_parameters=False,
-    )
-    return ddp_model
-
-
-def create_optimizer(model: nn.Module, **kwargs) -> optim.Optimizer:
-    return optim.AdamW(model.parameters(), **kwargs)
+def create_optimizer(model, **kwargs) -> optim.Optimizer:
+    return optim.AdamW(model.parameters(), fused=True, **kwargs)
 
 
 def create_scheduler(
@@ -164,8 +176,21 @@ def create_scheduler(
     return SchedulerFactory.create(optimizer, **kwargs)
 
 
-def prepare_checkpoint(model: nn.Module) -> dict:
-    return model.module.state_dict()
+def compute_total_steps(
+    dataset_len: int,
+    n_epoch: int,
+    batch_per_device: int,
+    nprocs: int,
+    grad_accum_steps: int,
+) -> int:
+
+    def ceil_div(a: int, b: int) -> int:
+        return (a + b - 1) // b
+
+    samples_per_replica = ceil_div(dataset_len, nprocs)
+    batches_per_replica = ceil_div(samples_per_replica, batch_per_device)
+    total_steps = (batches_per_replica // grad_accum_steps) * n_epoch
+    return total_steps
 
 
 def train(
@@ -174,14 +199,18 @@ def train(
     data_root_path: str,
     max_lr: float,
     n_epoch: int,
-    batch_size: int,
+    batch_per_device: int,
     start_epoch: int,
     start_batch: int,
-    accumulation_steps: int,
-    warmup_steps: int,
+    grad_accum_steps: int,
+    warmup_ratio: float,
     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,
@@ -193,22 +222,22 @@ def train(
     window_size: int,
     stride: int,
     nprocs: int,
+    parallel_mode: str,
     device_type: str,
+    start_method: str,
 ):
-    assert train_type in ["seq", "sft", "dpo"]
+    assert train_type in ["seq", "sft", "dpo", "grpo"]
     assert os.path.exists(param_path)
 
     # Load config
-    config = ModelConfig()
     config_path = os.path.join(param_path, "config.json")
-    if os.path.exists(config_path):
-        config.load(config_path)
+    config = AutoRegressiveLMConfig.from_file(config_path)
 
     if window_size is None:
         window_size = config.max_len
 
-    # Create bare Transformer (for training, no tokenizer needed)
-    model = Transformer(config)
+    # Create bare AutoRegressiveLM (for training, no tokenizer needed)
+    model = AutoRegressiveLM(config)
 
     # Load weights if available
     weights_path = os.path.join(param_path, "model.safetensors")
@@ -216,7 +245,23 @@ 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}
+    model = model.to(dtype=torch.bfloat16)
+
+    strategy_kwargs = {
+        "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,
+    }
+
+    executor_kwargs = {
+        "static_graph": True,
+        "find_unused_parameters": False,
+        "gradient_as_bucket_view": True,
+        "broadcast_buffers": False,
+    }
 
     dataset = DatasetFactory.load(
         train_type=train_type,
@@ -234,13 +279,17 @@ def train(
         },
     )
 
-    total_steps = len(dataset) * n_epoch // (batch_size * nprocs)
+    total_steps = compute_total_steps(
+        len(dataset), n_epoch, batch_per_device, nprocs, grad_accum_steps
+    )
+    warmup_steps = int(warmup_ratio * total_steps)
+
     scheduler_fn = partial(
         create_scheduler,
         **{
             "schedule_type": "cosine",
-            "warmup_steps": warmup_steps,
-            "lr_decay_steps": total_steps - warmup_steps,
+            "warmup_steps": min(warmup_steps, total_steps),
+            "lr_decay_steps": total_steps - min(warmup_steps, total_steps),
         },
     )
 
@@ -252,19 +301,20 @@ def train(
         scheduler_fn=scheduler_fn,
         ckpt_dir=ckpt_dir,
         n_epoch=n_epoch,
-        batch_size=batch_size,
+        batch_per_device=batch_per_device,
         start_epoch=start_epoch,
         start_batch=start_batch,
         ckpt_interval=ckpt_interval,
-        accumulation_steps=accumulation_steps,
+        grad_accum_steps=grad_accum_steps,
         max_grad_norm=max_grad_norm,
         random_seed=random_seed,
         num_workers=num_workers,
         pin_memory=pin_memory,
         nprocs=nprocs,
-        parallel_wrapper=ddp_wrap,
-        state_dict_fn=prepare_checkpoint,
+        parallel_mode=parallel_mode,
         device_type=device_type,
+        start_method=start_method,
+        executor_kwargs=executor_kwargs,
         extra_kwargs=strategy_kwargs,
     )
 

tests/conftest.py

@@ -3,18 +3,22 @@ import os
 import shutil
 import tempfile
 
-import numpy as np
 import pytest
-import safetensors.torch as st
 import torch
 from tokenizers import Tokenizer, models, pre_tokenizers, trainers
 from torch.utils.data import Dataset
 
-from astrai.config.model_config import ModelConfig
-from astrai.model.transformer import Transformer
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.transformer import AutoRegressiveLM
 from astrai.tokenize import AutoTokenizer
 
 
+def pytest_configure(config):
+    config.addinivalue_line("markers", "slow: marks tests as slow")
+    config.addinivalue_line("markers", "integration: integration tests")
+    config.addinivalue_line("markers", "unit: fast unit tests")
+
+
 def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
     """Create a simple tokenizer for testing purposes."""
     tokenizer = Tokenizer(models.BPE())
@@ -22,7 +26,6 @@ def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
     trainer = trainers.BpeTrainer(
         vocab_size=vocab_size, min_frequency=1, special_tokens=["<unk>", "<pad>"]
     )
-    # Train on empty iterator with single character
     tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer)
     auto_tokenizer = AutoTokenizer()
     auto_tokenizer._tokenizer = tokenizer
@@ -34,7 +37,7 @@ class RandomDataset(Dataset):
     """Random dataset for testing purposes."""
 
     def __init__(self, length=None, max_length=64, vocab_size=1000):
-        self.length = length or int(np.random.randint(100, 200))
+        self.length = length or int(torch.randint(100, 200, (1,)).item())
         self.max_length = max_length
         self.vocab_size = vocab_size
 
@@ -52,7 +55,7 @@ class MultiTurnDataset(Dataset):
     """Multi-turn dataset with loss mask for SFT training tests."""
 
     def __init__(self, length=None, max_length=64, vocab_size=1000):
-        self.length = length or int(np.random.randint(100, 200))
+        self.length = length or int(torch.randint(100, 200, (1,)).item())
         self.max_length = max_length
         self.vocab_size = vocab_size
 
@@ -93,46 +96,65 @@ class EarlyStoppingDataset(Dataset):
         }
 
 
-@pytest.fixture
-def base_test_env(request: pytest.FixtureRequest):
-    """Create base test environment with randomly configured model and tokenizer"""
-    func_name = request.function.__name__
-    test_dir = tempfile.mkdtemp(prefix=f"{func_name}_")
-    config_path = os.path.join(test_dir, "config.json")
+@pytest.fixture(scope="session")
+def test_tokenizer():
+    """Session-scoped tokenizer, created once for the entire test run."""
+    return create_test_tokenizer()
 
-    n_dim_choices = [8, 16, 32]
-    n_head_choices = [2, 4]
 
-    dim = int(np.random.choice(n_dim_choices))
-    n_heads = int(np.random.choice(n_head_choices))
-    n_kv_heads = n_heads // 2
-    dim_ffn = dim * 2
+@pytest.fixture(scope="session")
+def test_model():
+    """Session-scoped small AutoRegressiveLM model, created once."""
+    config = AutoRegressiveLMConfig(
+        vocab_size=1000,
+        dim=8,
+        n_heads=2,
+        n_kv_heads=1,
+        dim_ffn=16,
+        max_len=64,
+        n_layers=2,
+        norm_eps=1e-5,
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoRegressiveLM(config).to(device=device)
 
-    config = {
-        "vocab_size": 1000,
-        "dim": dim,
-        "n_heads": n_heads,
-        "n_kv_heads": n_kv_heads,
-        "dim_ffn": dim_ffn,
-        "max_len": 1024,
-        "n_layers": 4,
-        "norm_eps": 1e-5,
+    return {
+        "model": model,
+        "device": device,
+        "config": config,
     }
 
+
+@pytest.fixture
+def base_test_env(test_model, test_tokenizer):
+    """Function-scoped test environment with isolated temp directory.
+
+    Composes session-scoped model and tokenizer with a per-test temp dir.
+    """
+    test_dir = tempfile.mkdtemp()
+    config_path = os.path.join(test_dir, "config.json")
     with open(config_path, "w") as f:
-        json.dump(config, f)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    transformer_config = ModelConfig().load(config_path)
-    model = Transformer(transformer_config).to(device=device)
-    tokenizer = create_test_tokenizer()
+        json.dump(
+            {
+                "vocab_size": 1000,
+                "dim": 8,
+                "n_heads": 2,
+                "n_kv_heads": 1,
+                "dim_ffn": 16,
+                "max_len": 64,
+                "n_layers": 2,
+                "norm_eps": 1e-5,
+            },
+            f,
+        )
 
     yield {
-        "device": device,
+        "device": test_model["device"],
         "test_dir": str(test_dir),
         "config_path": config_path,
-        "transformer_config": transformer_config,
-        "model": model,
-        "tokenizer": tokenizer,
+        "transformer_config": test_model["config"],
+        "model": test_model["model"],
+        "tokenizer": test_tokenizer,
     }
 
     shutil.rmtree(test_dir)
@@ -154,43 +176,3 @@ def multi_turn_dataset():
 def early_stopping_dataset():
     dataset = EarlyStoppingDataset()
     yield dataset
-
-
-@pytest.fixture
-def test_env(request: pytest.FixtureRequest):
-    """Create a test environment with saved model and tokenizer files."""
-
-    func_name = request.function.__name__
-    test_dir = tempfile.mkdtemp(prefix=f"{func_name}_")
-    config_path = os.path.join(test_dir, "config.json")
-    tokenizer_path = os.path.join(test_dir, "tokenizer.json")
-    model_path = os.path.join(test_dir, "model.safetensors")
-
-    config = {
-        "vocab_size": 1000,
-        "dim": 128,
-        "n_heads": 4,
-        "n_kv_heads": 2,
-        "dim_ffn": 256,
-        "max_len": 64,
-        "n_layers": 2,
-        "norm_eps": 1e-5,
-    }
-    with open(config_path, "w") as f:
-        json.dump(config, f)
-
-    tokenizer = create_test_tokenizer(vocab_size=config["vocab_size"])
-    tokenizer.save(tokenizer_path)
-
-    transformer_config = ModelConfig().load(config_path)
-    model = Transformer(transformer_config)
-    st.save_file(model.state_dict(), model_path)
-
-    yield {
-        "test_dir": test_dir,
-        "model": model,
-        "tokenizer": tokenizer,
-        "transformer_config": transformer_config,
-    }
-
-    shutil.rmtree(test_dir)

tests/data/test_checkpoint.py

@@ -35,6 +35,33 @@ def test_single_process():
         assert loaded_checkpoint.iteration == 30
 
 
+def test_checkpoint_with_extra():
+    """Verify extra keys are saved as individual .pt files and loaded back."""
+    model = torch.nn.Linear(10, 5)
+    optimizer = AdamW(model.parameters(), lr=1e-3)
+    optimizer.step()
+
+    extra = {
+        "optimizer": optimizer.state_dict(),
+        "scheduler": {"last_epoch": 5},
+    }
+    checkpoint = Checkpoint(
+        state_dict=model.state_dict(), epoch=1, iteration=10, extra=extra
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        checkpoint.save(tmpdir)
+
+        import os
+
+        assert os.path.exists(os.path.join(tmpdir, "optimizer.pt"))
+        assert os.path.exists(os.path.join(tmpdir, "scheduler.pt"))
+
+        loaded = Checkpoint.load(tmpdir)
+        assert loaded.extra["scheduler"]["last_epoch"] == 5
+        assert "state" in loaded.extra["optimizer"]
+
+
 def simple_training():
     model = torch.nn.Linear(10, 5)
     optimizer = AdamW(model.parameters(), lr=1e-3)

tests/data/test_dataset.py

@@ -1,8 +1,20 @@
+import json
+import os
+
 import numpy as np
+import pytest
 import torch
 
-from astrai.dataset.dataset import DatasetFactory
-from astrai.serialization import save_h5
+from astrai.dataset.dataset import DatasetFactory, SEQDataset
+from astrai.dataset.storage import (
+    BaseSegmentFetcher,
+    H5Storage,
+    MultiSegmentFetcher,
+    StorageFactory,
+    detect_format,
+    load_json,
+    save_h5,
+)
 
 
 def test_dataset_loader_random_paths(base_test_env):
@@ -64,7 +76,7 @@ def test_dpo_strategy_with_random_data(base_test_env):
     )
 
     assert dpo_dataset is not None
-    assert hasattr(dpo_dataset, "fetcher")
+    assert dpo_dataset.storage is not None
     assert len(dpo_dataset) > 0
 
     # Test that we can get DPO items without errors
@@ -100,7 +112,7 @@ def test_sft_dataset_with_random_data(base_test_env):
     )
 
     assert sft_dataset is not None
-    assert hasattr(sft_dataset, "fetcher")
+    assert sft_dataset.storage is not None
     assert len(sft_dataset) > 0
 
     # Test that we can get SFT items without errors
@@ -143,3 +155,266 @@ def test_dataset_with_custom_stride(base_test_env):
     )
 
     assert len(dataset) > len(default_stride_dataset)
+
+
+# ============== JSON Storage Tests (raw text + tokenizer) ==============
+
+
+def _make_tokenizer_fn(tokenizer):
+    """Wrap tokenizer.encode() as a str -> List[int] callable."""
+    return lambda text: tokenizer.encode(text, add_special_tokens=False)
+
+
+def test_seq_dataset_from_json_text(base_test_env):
+    """Test loading SEQ dataset from raw-text JSON with tokenizer"""
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_fn = _make_tokenizer_fn(tokenizer)
+    test_dir = base_test_env["test_dir"]
+    data_dir = os.path.join(test_dir, "json_text")
+    os.makedirs(data_dir, exist_ok=True)
+
+    texts = [
+        "hello world this is a test sentence for tokenizer",
+        "another sentence with different words and tokens",
+        "machine learning is fascinating and powerful",
+    ]
+
+    json_path = os.path.join(data_dir, "seq_data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"sequence": texts}, f, ensure_ascii=False)
+
+    dataset = DatasetFactory.load(
+        train_type="seq",
+        load_path=data_dir,
+        window_size=16,
+        tokenizer=tokenizer_fn,
+    )
+    assert dataset is not None
+    assert len(dataset) > 0
+    assert dataset.count > 0
+    assert "sequence" in dataset.keys
+
+    item = dataset[0]
+    assert "input_ids" in item
+    assert "target_ids" in item
+    assert item["input_ids"].shape[0] == 16
+
+
+def test_sft_dataset_from_json_text(base_test_env):
+    """Test loading SFT dataset from raw-text JSON with tokenizer"""
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_fn = _make_tokenizer_fn(tokenizer)
+    test_dir = base_test_env["test_dir"]
+    data_dir = os.path.join(test_dir, "json_sft")
+    os.makedirs(data_dir, exist_ok=True)
+
+    texts = [
+        "user asks a question about the weather",
+        "assistant provides a helpful response to the user",
+    ]
+
+    json_path = os.path.join(data_dir, "sft_data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(
+            {"sequence": texts, "loss_mask": texts},
+            f,
+            ensure_ascii=False,
+        )
+
+    dataset = DatasetFactory.load(
+        train_type="sft",
+        load_path=data_dir,
+        window_size=16,
+        tokenizer=tokenizer_fn,
+    )
+    assert dataset is not None
+    assert len(dataset) > 0
+
+    item = dataset[0]
+    assert "loss_mask" in item
+
+
+def test_json_storage_explicit_tokenizer(base_test_env):
+    """Test explicit JSON storage with tokenizer"""
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_fn = _make_tokenizer_fn(tokenizer)
+    test_dir = base_test_env["test_dir"]
+    data_dir = os.path.join(test_dir, "json_explicit")
+    os.makedirs(data_dir, exist_ok=True)
+
+    texts = ["abcdefghijklmnopqrstuvwxyz" * 10]
+
+    json_path = os.path.join(data_dir, "data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"sequence": texts}, f, ensure_ascii=False)
+
+    token_count = len(tokenizer_fn(texts[0]))
+
+    dataset = DatasetFactory.load(
+        train_type="seq",
+        load_path=data_dir,
+        window_size=32,
+        storage_type="json",
+        tokenizer=tokenizer_fn,
+    )
+    assert dataset is not None
+    assert len(dataset) > 0
+    assert dataset.count == token_count
+
+
+def test_dataset_count_property(base_test_env):
+    """Test the count property returns correct raw token count"""
+    test_dir = base_test_env["test_dir"]
+
+    seq_length = 200
+    dummy_data = {
+        "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
+    }
+
+    save_h5(test_dir, "count_test_data", dummy_data)
+
+    dataset = DatasetFactory.load(
+        train_type="seq",
+        load_path=test_dir,
+        window_size=64,
+    )
+
+    assert dataset.count == seq_length
+    assert dataset.count > len(dataset)  # raw tokens > windows
+    assert len(dataset) == (seq_length - 1 - 64) // 64 + 1
+
+
+def test_empty_dataset_count():
+    """Test count returns 0 when no data is loaded"""
+    dataset = SEQDataset(window_size=64, stride=32)
+    assert dataset.count == 0
+    assert dataset.keys == []
+
+
+def test_dataset_too_short_for_window(base_test_env):
+    """Dataset shorter than window_size returns __len__ == 0"""
+    test_dir = base_test_env["test_dir"]
+    seq_length = 30
+    save_h5(
+        test_dir,
+        "short",
+        {"sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)]},
+    )
+    dataset = DatasetFactory.load("seq", test_dir, window_size=64)
+    assert len(dataset) == 0
+    assert dataset.count == seq_length
+
+
+def test_unloaded_dataset_getitem_raises():
+    """__getitem__ without load() should fail clearly"""
+    dataset = SEQDataset(window_size=64, stride=32)
+    with pytest.raises(RuntimeError, match="not loaded"):
+        dataset.get_index(0)
+
+
+def test_unloaded_dataset_len():
+    """__len__ without load() returns 0"""
+    dataset = SEQDataset(window_size=64, stride=32)
+    assert len(dataset) == 0
+
+
+def test_base_segment_fetcher_empty():
+    """BaseSegmentFetcher with empty segments list"""
+    fetcher = BaseSegmentFetcher([])
+    assert len(fetcher) == 0
+    with pytest.raises(ValueError, match="out of bounds"):
+        fetcher.fetch_data(0, 1)
+
+
+def test_base_segment_fetcher_begin_equals_end(base_test_env):
+    """fetch_data with begin == end returns empty tensor"""
+    test_dir = base_test_env["test_dir"]
+    dummy = {"sequence": [torch.randint(0, 1000, (100,), dtype=torch.int64)]}
+    save_h5(test_dir, "empty_fetch", dummy)
+
+    dataset = DatasetFactory.load("seq", test_dir, window_size=32)
+    fetcher = dataset.storage._fetcher.multi_fetchers["sequence"]
+    result = fetcher.fetch_data(10, 10)
+    assert result.numel() == 0
+
+
+def test_multi_segment_fetcher_empty_dict():
+    """MultiSegmentFetcher with empty dict has __len__ == 0"""
+    fetcher = MultiSegmentFetcher({})
+    assert len(fetcher) == 0
+
+
+def test_storage_fetch_before_load():
+    """BaseStorage.fetch before load raises RuntimeError"""
+    storage = H5Storage()
+    with pytest.raises(RuntimeError, match="not loaded"):
+        storage.fetch(0, 10, "sequence")
+
+
+def test_detect_format_nonexistent_path():
+    """detect_format raises FileNotFoundError for bad path"""
+    with pytest.raises(FileNotFoundError, match="No supported"):
+        detect_format("/nonexistent/path/xyz")
+
+
+def test_detect_format_unsupported_file(base_test_env):
+    """detect_format raises ValueError for unsupported file extension"""
+    test_dir = base_test_env["test_dir"]
+    path = os.path.join(test_dir, "data.txt")
+    with open(path, "w") as f:
+        f.write("hello")
+    with pytest.raises(ValueError, match="Unsupported"):
+        detect_format(path)
+
+
+def test_create_storage_invalid_type():
+    """StorageFactory.create raises ValueError for unknown type"""
+    with pytest.raises(ValueError, match="Unknown component"):
+        StorageFactory.create("parquet")
+
+
+def test_json_pretokenized_without_tokenizer(base_test_env):
+    """Pre-tokenized JSON (List[List[int]]) loads without tokenizer"""
+    test_dir = base_test_env["test_dir"]
+    data_dir = os.path.join(test_dir, "json_pretok")
+    os.makedirs(data_dir, exist_ok=True)
+
+    json_path = os.path.join(data_dir, "data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"sequence": [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]}, f)
+
+    dataset = DatasetFactory.load("seq", data_dir, window_size=4, storage_type="json")
+    assert len(dataset) > 0
+    assert dataset.count == 10
+
+    item = dataset[0]
+    assert item["input_ids"].tolist() == [1, 2, 3, 4]
+    assert item["target_ids"].tolist() == [2, 3, 4, 5]
+
+
+def test_load_json_skips_config_file(base_test_env):
+    """load_json skips scalar-value config files"""
+    test_dir = base_test_env["test_dir"]
+    with open(os.path.join(test_dir, "config.json"), "w") as f:
+        json.dump({"vocab_size": 1000, "dim": 16}, f)
+
+    with open(os.path.join(test_dir, "data.json"), "w") as f:
+        json.dump({"sequence": [[1, 2, 3, 4, 5]]}, f)
+
+    result = load_json(test_dir)
+    assert "sequence" in result
+    assert "vocab_size" not in result
+    assert len(result["sequence"]) == 1
+
+
+def test_base_segment_fetcher_multi_segment():
+    """fetch_data across multiple segment boundaries"""
+    segs = [
+        torch.tensor([1, 2, 3]),
+        torch.tensor([4, 5, 6, 7]),
+        torch.tensor([8, 9]),
+    ]
+    fetcher = BaseSegmentFetcher(segs)
+    assert len(fetcher) == 9
+    result = fetcher.fetch_data(2, 7)
+    assert result.tolist() == [3, 4, 5, 6, 7]

tests/inference/conftest.py

@@ -5,46 +5,49 @@ from unittest.mock import MagicMock
 import pytest
 from fastapi.testclient import TestClient
 
-from astrai.inference.server import app
+from astrai.inference import app
 
 
 @pytest.fixture
 def client():
     """Provide a test client for the FastAPI app."""
+    app.state.server_config = {
+        "device": "cpu",
+        "dtype": "bfloat16",
+        "param_path": None,
+        "max_batch_size": 1,
+        "_test": True,
+    }
+    app.state.engine = None
     return TestClient(app)
 
 
-@pytest.fixture
-def mock_model_param():
-    """Create a mock ModelParameter."""
-    mock_param = MagicMock()
-    mock_param.model = MagicMock()
-    mock_param.tokenizer = MagicMock()
-    mock_param.config = MagicMock()
-    mock_param.config.max_len = 100
-    mock_param.tokenizer.encode = MagicMock(return_value=[1, 2, 3])
-    mock_param.tokenizer.decode = MagicMock(return_value="mock response")
-    mock_param.tokenizer.stop_ids = []
-    mock_param.tokenizer.pad_id = 0
-    return mock_param
-
-
 @pytest.fixture
 def mock_engine():
     """Create a mock InferenceEngine."""
+
+    async def _async_gen():
+        yield "chunk1"
+        yield "chunk2"
+        yield "[DONE]"
+
     mock = MagicMock()
     mock.generate.return_value = "mock response"
+    mock.generate_async.return_value = _async_gen()
     mock.get_stats.return_value = {
         "total_tasks": 0,
         "total_tokens": 0,
         "active_tasks": 0,
         "waiting_queue": 0,
     }
+    mock.tokenizer.encode.return_value = [1, 2, 3]
+    mock.tokenizer.decode.return_value = "mock response"
+    mock.tokenizer.apply_chat_template.return_value = "mock prompt"
     return mock
 
 
 @pytest.fixture
-def loaded_model(mock_model_param, monkeypatch):
-    """Simulate that the model is loaded."""
-    monkeypatch.setattr("astrai.inference.server._model_param", mock_model_param)
-    return mock_model_param
+def loaded_model(client, mock_engine):
+    """Simulate that the engine is loaded."""
+    app.state.engine = mock_engine
+    return mock_engine

tests/inference/test_cache.py

@@ -0,0 +1,279 @@
+"""Unit tests for inference cache components."""
+
+import torch
+
+from astrai.inference import (
+    Allocator,
+    KVCache,
+    PagePool,
+    PrefixCache,
+    Storage,
+    TaskTable,
+    page_hash,
+)
+
+
+def make_pool(n_pages: int, page_size: int) -> PagePool:
+    return PagePool(Allocator(n_pages), PrefixCache(page_size))
+
+
+def test_page_hash_full_page():
+    token_ids = list(range(256))
+    h = page_hash(token_ids, 0, 64)
+    assert isinstance(h, int)
+    assert h >= 0
+
+
+def test_page_hash_different_page_differs():
+    token_ids = list(range(256))
+    assert page_hash(token_ids, 0, 64) != page_hash(token_ids, 1, 64)
+
+
+def test_page_pool_alloc_free_cycle():
+    pool = make_pool(4, 64)
+    a = pool.alloc()
+    b = pool.alloc()
+    assert a != b
+    pool.free(a)
+    pool.free(b)
+    c = pool.alloc()
+    assert c in (a, b)
+
+
+def test_page_pool_alloc_when_full():
+    pool = make_pool(2, 64)
+    pool.alloc()
+    pool.alloc()
+    assert pool.alloc() == -1
+
+
+def test_page_pool_lru_eviction():
+    pool = make_pool(2, 64)
+    p0 = pool.alloc()
+    p1 = pool.alloc()
+    pool.record(p0, list(range(64)), 0)
+    pool.record(p1, list(range(64, 128)), 0)
+    pool.free(p0)
+    pool.free(p1)
+    pool.alloc()
+    assert p0 in pool._alloc._lru or p1 in pool._alloc._lru
+
+
+def test_page_pool_inc_ref_and_free():
+    pool = make_pool(2, 64)
+    p = pool.alloc()
+    pool.inc_ref(p)
+    assert pool._alloc._refs[p] == 2
+    pool.free(p)
+    assert pool._alloc._refs[p] == 1
+    pool.free(p)
+    assert pool._alloc._refs[p] == 0
+
+
+def test_page_pool_keep_cached_realloc():
+    """Free mask has priority over LRU; cached page returned only when no free pages."""
+    pool = make_pool(3, 64)
+    p0 = pool.alloc()
+    p1 = pool.alloc()
+    p2 = pool.alloc()
+    for p in (p0, p1, p2):
+        pool.record(p, [p] * 64, 0)
+    pool.free(p0)
+    pool.free(p1)
+    pool.free(p2)
+    assert pool.alloc() == p0
+
+
+def test_prefix_cache_lookup_returns_hits():
+    token_ids = list(range(256))
+    pool = make_pool(16, 64)
+    pages = [pool.alloc() for _ in range(4)]
+    for i, p in enumerate(pages):
+        pool.record(p, token_ids, i)
+        pool.free(p)
+    hits = pool.lookup(token_ids)
+    assert hits == pages
+
+
+def test_prefix_cache_lookup_stops_at_first_miss():
+    token_ids = list(range(256))
+    pool = make_pool(16, 64)
+    p0 = pool.alloc()
+    pool.record(p0, token_ids, 0)
+    pool.free(p0)
+    p1 = pool.alloc()
+    pool.record(p1, [99] * 64, 1)
+    pool.free(p1)
+    hits = pool.lookup(token_ids)
+    assert len(hits) == 1
+    assert hits[0] == p0
+
+
+def test_prefix_cache_ignores_partial_last_page():
+    token_ids = list(range(100))
+    pool = make_pool(16, 64)
+    p = pool.alloc()
+    pool.record(p, token_ids, 0)
+    pool.free(p)
+    hits = pool.lookup(token_ids)
+    assert len(hits) == 1
+
+
+def test_prefix_cache_on_evict_clears_mappings():
+    pool = make_pool(4, 64)
+    p = pool.alloc()
+    pool.record(p, list(range(64)), 0)
+    pool.free(p)
+    assert p in pool._prefix._page_to_hash
+    pool._prefix.evict(p)
+    assert p not in pool._prefix._page_to_hash
+
+
+def test_prefix_cache_has_page():
+    pool = make_pool(4, 64)
+    p = pool.alloc()
+    assert p not in pool._prefix._page_to_hash
+    pool.record(p, list(range(64)), 0)
+    pool.free(p)
+    assert p in pool._prefix._page_to_hash
+
+
+def test_task_table_set_get():
+    table = TaskTable(page_size=64)
+    table.set("task1", [0, 1, 2], 128)
+    assert table.get("task1") == [0, 1, 2]
+    assert table.get_cached("task1") == 128
+
+
+def test_task_table_get_missing():
+    table = TaskTable(page_size=64)
+    assert table.get("nonexistent") == []
+    assert table.get_cached("nonexistent") == 0
+
+
+def test_task_table_pop():
+    table = TaskTable(page_size=64)
+    table.set("task1", [0, 1], 64)
+    pages, cached = table.pop("task1")
+    assert pages == [0, 1]
+    assert cached == 64
+    assert table.get("task1") == []
+
+
+def test_kv_cache_task_extend_allocates():
+    cache = KVCache(
+        n_layers=1,
+        n_pages=8,
+        page_size=64,
+        n_kv_heads=2,
+        head_dim=8,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    cache._table.set("task1", [], 0)
+    ok = cache.task_extend("task1", 200)
+    assert ok
+    assert len(cache._table.get("task1")) == 4
+
+
+def test_kv_cache_task_extend_fails_when_pool_full():
+    cache = KVCache(
+        n_layers=1,
+        n_pages=2,
+        page_size=64,
+        n_kv_heads=2,
+        head_dim=8,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    cache._table.set("task1", [0, 1], 0)
+    ok = cache.task_extend("task1", 300)
+    assert not ok
+
+
+def test_task_table_table_tensor():
+    table = TaskTable(page_size=64)
+    table.set("a", [0, 1], 0)
+    table.set("b", [2, 3, 4], 0)
+    t = table.table_tensor(["a", "b"], torch.device("cpu"))
+    assert t.shape == (2, 3)
+    assert t[0].tolist() == [0, 1, -1]
+    assert t[1].tolist() == [2, 3, 4]
+
+
+def test_task_table_table_tensor_empty_input():
+    table = TaskTable(page_size=64)
+    t = table.table_tensor([], torch.device("cpu"))
+    assert t.numel() == 0
+
+
+def test_storage_write_gather_single_page():
+    storage = Storage(
+        n_layers=2,
+        n_pages=8,
+        page_size=4,
+        n_kv_heads=2,
+        head_dim=8,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    page_table = torch.tensor([[0]], dtype=torch.long)
+    k = torch.randn(1, 2, 2, 8)
+    v = torch.randn(1, 2, 2, 8)
+
+    storage.write(0, page_table, 0, k, v)
+    gk, gv = storage.gather(0, page_table, 2)
+    assert torch.allclose(gk, k)
+
+
+def test_storage_write_cross_page():
+    storage = Storage(
+        n_layers=1,
+        n_pages=8,
+        page_size=4,
+        n_kv_heads=2,
+        head_dim=8,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    page_table = torch.tensor([[0, 1]], dtype=torch.long)
+    k = torch.randn(1, 8, 2, 8)
+    v = torch.randn(1, 8, 2, 8)
+
+    storage.write(0, page_table, 0, k, v)
+    gk, gv = storage.gather(0, page_table, 8)
+    assert torch.allclose(gk, k)
+
+
+def test_storage_gather_truncates_to_total_len():
+    storage = Storage(
+        n_layers=1,
+        n_pages=8,
+        page_size=4,
+        n_kv_heads=2,
+        head_dim=8,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    page_table = torch.tensor([[0, 1]], dtype=torch.long)
+    k = torch.randn(1, 6, 2, 8)
+    v = torch.randn(1, 6, 2, 8)
+    storage.write(0, page_table, 0, k, v)
+
+    gk, gv = storage.gather(0, page_table, 5)
+    assert gk.shape == (1, 5, 2, 8)
+
+
+def test_storage_gather_clamps_negative_padding():
+    storage = Storage(
+        n_layers=1,
+        n_pages=8,
+        page_size=4,
+        n_kv_heads=2,
+        head_dim=8,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    page_table = torch.tensor([[0, -1]], dtype=torch.long)
+    gk, gv = storage.gather(0, page_table, 4)
+    assert gk.shape == (1, 4, 2, 8)

tests/inference/test_engine.py

@@ -0,0 +1,181 @@
+"""Unit tests for GenerateResult accumulator and InferenceEngine.generate()."""
+
+import threading
+from unittest.mock import MagicMock, patch
+
+from astrai.inference import STOP
+from astrai.inference.engine import GenerateResult
+
+
+def test_result_append_single():
+    r = GenerateResult(count=1)
+    r.append("hello", 0)
+    assert r.results[0] == "hello"
+
+
+def test_result_append_multiple_tasks():
+    r = GenerateResult(count=3)
+    r.append("a", 0)
+    r.append("b", 1)
+    r.append("c", 2)
+    assert r.results[0] == "a"
+    assert r.results[1] == "b"
+    assert r.results[2] == "c"
+
+
+def test_result_stop_marks_complete():
+    r = GenerateResult(count=2)
+    r.append("text", 0)
+    r.append(STOP, 0)
+    r.append("more", 1)
+    assert r._done[0] is True
+    assert r._done[1] is False
+    assert r._completed == 1
+
+
+def test_result_stop_does_not_double_count():
+    r = GenerateResult(count=1)
+    r.append(STOP, 0)
+    r.append(STOP, 0)
+    assert r._completed == 1
+
+
+def test_result_pop_all_returns_and_clears():
+    r = GenerateResult(count=2)
+    r.append("a", 0)
+    r.append("b", 1)
+    out = r.pop_all()
+    assert len(out) == 2
+    assert out[0] == (0, "a")
+    assert out[1] == (1, "b")
+    assert r.pop_all() == []
+
+
+def test_result_wait_blocks_until_data():
+    r = GenerateResult(count=1)
+
+    def delayed_append():
+        import time
+
+        time.sleep(0.05)
+        r.append("delayed", 0)
+
+    t = threading.Thread(target=delayed_append)
+    t.start()
+    ok = r.wait(timeout=5.0)
+    t.join()
+    assert ok
+    assert r.results[0] == "delayed"
+
+
+def test_result_wait_timeout():
+    r = GenerateResult(count=1)
+    ok = r.wait(timeout=0.01)
+    assert not ok
+
+
+def test_result_wait_completion_non_streaming():
+    r = GenerateResult(count=2)
+
+    def finish_later():
+        import time
+
+        time.sleep(0.05)
+        r.append(STOP, 0)
+        time.sleep(0.05)
+        r.append(STOP, 1)
+
+    t = threading.Thread(target=finish_later)
+    t.start()
+    r.wait_completion()
+    t.join()
+    assert r._completed == 2
+
+
+def test_result_get_results():
+    r = GenerateResult(count=2)
+    r.append("hello", 0)
+    r.append("world", 1)
+    results = r.get_results()
+    assert results == ["hello", "world"]
+
+
+def test_engine_generate_non_streaming_single():
+    from astrai.inference.engine import InferenceEngine
+
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.encode.return_value = [1, 2, 3]
+    mock_tokenizer.decode.return_value = "response"
+    mock_tokenizer.stop_ids = [0]
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        instance = MockSched.return_value
+
+        def fake_add(prompt, **kw):
+            cb = kw["stream_callback"]
+            cb("response")
+            cb(STOP)
+
+        instance.add_task.side_effect = fake_add
+        instance.remove_task.return_value = []
+
+        eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
+        result = eng.generate("hello")
+        assert result == "response"
+
+
+def test_engine_generate_streaming_yields_tokens():
+    from astrai.inference.engine import InferenceEngine
+
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.encode.return_value = [1, 2, 3]
+    mock_tokenizer.decode.return_value = "tok"
+    mock_tokenizer.stop_ids = [0]
+
+    callbacks_saved = []
+
+    def capture_cb(prompt, **kw):
+        callbacks_saved.append(kw.get("stream_callback"))
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        instance = MockSched.return_value
+        instance.add_task.side_effect = capture_cb
+        instance.remove_task.return_value = []
+
+        eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
+        gen = eng.generate("hello", stream=True)
+
+        cb = callbacks_saved[0]
+        cb("t1")
+        cb("t2")
+        cb(STOP)
+
+        tokens = list(gen)
+        assert tokens == ["t1", "t2"]
+
+
+def test_engine_generate_non_streaming_batch():
+    from astrai.inference.engine import InferenceEngine
+
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.encode.return_value = [1, 2, 3]
+    mock_tokenizer.decode.return_value = "r"
+    mock_tokenizer.stop_ids = [0]
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        instance = MockSched.return_value
+
+        def fake_add(prompt, **kw):
+            cb = kw["stream_callback"]
+            cb("r")
+            cb(STOP)
+
+        instance.add_task.side_effect = fake_add
+        instance.remove_task.return_value = []
+
+        eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=2)
+        results = eng.generate(["hello", "world"])
+        assert results == ["r", "r"]

tests/inference/test_sample.py

@@ -0,0 +1,127 @@
+"""Unit tests for inference sampling strategies."""
+
+import torch
+
+from astrai.inference.sample import (
+    SamplingPipeline,
+    TemperatureStrategy,
+    TopKStrategy,
+    TopPStrategy,
+    sample,
+)
+
+
+def test_temperature_scalar():
+    logits = torch.tensor([[1.0, 2.0, 3.0]])
+    s = TemperatureStrategy(0.5)
+    result = s.apply(logits.clone())
+    assert torch.allclose(result, logits / 0.5)
+
+
+def test_temperature_skip_when_one():
+    logits = torch.tensor([[1.0, 2.0, 3.0]])
+    s = TemperatureStrategy(1.0)
+    result = s.apply(logits.clone())
+    assert torch.equal(result, logits)
+
+
+def test_temperature_per_sample_tensor():
+    logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    s = TemperatureStrategy(torch.tensor([0.5, 0.5]))
+    result = s.apply(logits.clone())
+    assert torch.allclose(result, logits / 0.5)
+
+
+def test_top_k_keeps_top():
+    logits = torch.tensor([[0.1, 0.5, 0.3, 0.9, 0.2]])
+    s = TopKStrategy(top_k=2)
+    result = s.apply(logits.clone(), filter_value=-1e9)
+    kept = (result > -1e9).sum().item()
+    assert kept == 2
+
+
+def test_top_k_skip_when_zero():
+    logits = torch.tensor([[1.0, 2.0, 3.0]])
+    s = TopKStrategy(top_k=0)
+    result = s.apply(logits.clone())
+    assert torch.equal(result, logits)
+
+
+def test_top_k_batch_tensor():
+    """Each row respects its own top_k."""
+    logits = torch.tensor([[0.1, 0.5, 0.3], [0.9, 0.2, 0.1]])
+    s = TopKStrategy(top_k=torch.tensor([2, 1]))
+    result = s.apply(logits.clone(), filter_value=-1e9)
+    assert (result[0] > -1e9).sum() == 2
+    assert (result[1] > -1e9).sum() == 1
+
+
+def test_top_p_nucleus_filtering():
+    logits = torch.tensor([[10.0, 1.0, 1.0, 1.0, 1.0]])
+    s = TopPStrategy(top_p=0.5)
+    result = s.apply(logits.clone(), filter_value=-1e9)
+    kept = (result > -1e9).sum().item()
+    assert kept >= 1
+
+
+def test_top_p_skip_when_one():
+    logits = torch.tensor([[1.0, 2.0, 3.0]])
+    s = TopPStrategy(top_p=1.0)
+    result = s.apply(logits.clone())
+    assert torch.equal(result, logits)
+
+
+def test_top_p_filter_all_except_max_when_zero():
+    logits = torch.tensor([[0.1, 0.5, 0.3, 0.9, 0.2]])
+    s = TopPStrategy(top_p=0.0)
+    result = s.apply(logits.clone(), filter_value=-1e9)
+    kept = (result > -1e9).sum().item()
+    assert kept == 1
+
+
+def test_sampling_pipeline_composes_strategies():
+    logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
+    pipeline = SamplingPipeline(
+        [
+            TemperatureStrategy(0.8),
+            TopKStrategy(3),
+            TopPStrategy(0.95),
+        ]
+    )
+    result = pipeline.apply(logits.clone(), filter_value=-1e9)
+    kept = (result > -1e9).sum().item()
+    assert 1 <= kept <= 3
+
+
+def test_sampling_pipeline_sample_returns_valid_token():
+    logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
+    pipeline = SamplingPipeline(
+        [
+            TemperatureStrategy(0.8),
+            TopKStrategy(3),
+            TopPStrategy(0.95),
+        ]
+    )
+    tokens = pipeline.sample(logits)
+    assert tokens.shape == (1,)
+    assert 0 <= tokens[0] < logits.size(-1)
+
+
+def test_module_sample_shortcut():
+    logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
+    tokens = sample(logits, temperature=0.8, top_k=3, top_p=0.95)
+    assert tokens.shape == (1,)
+    assert 0 <= tokens[0] < logits.size(-1)
+
+
+def test_module_sample_batch():
+    logits = torch.tensor(
+        [
+            [1.0, 2.0, 3.0, 4.0, 5.0],
+            [5.0, 4.0, 3.0, 2.0, 1.0],
+        ]
+    )
+    tokens = sample(logits, temperature=0.8, top_k=3, top_p=0.95)
+    assert tokens.shape == (2,)
+    for t in tokens:
+        assert 0 <= t < logits.size(-1)

tests/inference/test_scheduler.py

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

tests/inference/test_scheduler_concurrency.py

@@ -1,320 +0,0 @@
-"""Tests for scheduler concurrency."""
-
-import threading
-import time
-from unittest.mock import MagicMock, patch
-
-import pytest
-
-from astrai.inference.scheduler import (
-    InferenceScheduler,
-    PrefixCacheManager,
-)
-
-
-def test_prefix_cache_concurrent_insert_find():
-    """Test concurrent insert and find operations."""
-    cache = PrefixCacheManager(max_capacity=100)
-
-    results = {"errors": [], "inserts": 0, "finds": 0}
-
-    def insert_worker():
-        try:
-            for i in range(50):
-                cache.insert((i,), slot=i % 10)
-                results["inserts"] += 1
-        except Exception as e:
-            results["errors"].append(str(e))
-
-    def find_worker():
-        try:
-            for i in range(50):
-                cache.find_longest_prefix([i])
-                results["finds"] += 1
-        except Exception as e:
-            results["errors"].append(str(e))
-
-    threads = [threading.Thread(target=insert_worker) for _ in range(3)]
-    threads += [threading.Thread(target=find_worker) for _ in range(3)]
-
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-
-    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
-    assert results["inserts"] == 150
-    assert results["finds"] == 150
-
-
-def test_prefix_cache_concurrent_release():
-    """Test concurrent release operations."""
-    cache = PrefixCacheManager(max_capacity=100)
-
-    # Insert some prefixes
-    for i in range(10):
-        cache.insert((i,), slot=i)
-
-    results = {"errors": []}
-
-    def release_worker():
-        try:
-            for i in range(10):
-                cache.release((i,))
-        except Exception as e:
-            results["errors"].append(str(e))
-
-    threads = [threading.Thread(target=release_worker) for _ in range(3)]
-
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-
-    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
-
-
-def test_prefix_cache_concurrent_insert_release_find():
-    """Test mixed concurrent operations."""
-    cache = PrefixCacheManager(max_capacity=50)
-
-    results = {"errors": []}
-
-    def worker(worker_id):
-        try:
-            for i in range(20):
-                token_ids = (worker_id * 100 + i,)
-                cache.insert(token_ids, slot=worker_id)
-
-                # Find after insert
-                cache.find_longest_prefix(list(token_ids))
-
-                # Release
-                cache.release(token_ids)
-        except Exception as e:
-            results["errors"].append(f"Worker {worker_id}: {str(e)}")
-
-    threads = [threading.Thread(target=worker, args=(i,)) for i in range(5)]
-
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-
-    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
-
-
-@pytest.fixture
-def mock_model_and_tokenizer():
-    """Create mock model and tokenizer."""
-    mock_model = MagicMock()
-    mock_model.config = MagicMock()
-    mock_model.config.n_kv_heads = 8
-    mock_model.config.n_heads = 8
-    mock_model.config.dim = 128
-    mock_model.config.n_layers = 2
-    mock_model.config.max_len = 100
-
-    mock_tokenizer = MagicMock()
-    mock_tokenizer.encode.return_value = [1, 2, 3, 4, 5]
-    mock_tokenizer.decode.return_value = "token"
-    mock_tokenizer.stop_ids = [0]
-    mock_tokenizer.pad_id = None
-
-    return mock_model, mock_tokenizer
-
-
-def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
-    """Test concurrent add_task operations."""
-    mock_model, mock_tokenizer = mock_model_and_tokenizer
-
-    with patch("astrai.inference.scheduler.AutoModel"):
-        with patch("astrai.inference.scheduler.AutoTokenizer"):
-            scheduler = InferenceScheduler(
-                model=mock_model,
-                tokenizer=mock_tokenizer,
-                max_batch_size=4,
-                device="cpu",
-            )
-
-    results = {"task_ids": [], "errors": []}
-    lock = threading.Lock()
-
-    def add_task_worker(worker_id):
-        try:
-            for i in range(10):
-                task_id = scheduler.add_task(f"prompt from worker {worker_id}-{i}")
-                with lock:
-                    results["task_ids"].append(task_id)
-        except Exception as e:
-            results["errors"].append(str(e))
-
-    threads = [threading.Thread(target=add_task_worker, args=(i,)) for i in range(5)]
-
-    for t in threads:
-        t.start()
-
-    # Let some tasks be processed
-    time.sleep(0.1)
-
-    scheduler.stop()
-
-    for t in threads:
-        t.join()
-
-    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
-    assert len(results["task_ids"]) == 50
-
-
-def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
-    """Test concurrent add and remove task operations."""
-    mock_model, mock_tokenizer = mock_model_and_tokenizer
-
-    with patch("astrai.inference.scheduler.AutoModel"):
-        with patch("astrai.inference.scheduler.AutoTokenizer"):
-            scheduler = InferenceScheduler(
-                model=mock_model,
-                tokenizer=mock_tokenizer,
-                max_batch_size=4,
-                device="cpu",
-            )
-
-    results = {"added": [], "removed": [], "errors": []}
-
-    def add_worker():
-        try:
-            for i in range(20):
-                task_id = scheduler.add_task(f"prompt {i}")
-                results["added"].append(task_id)
-                time.sleep(0.001)
-        except Exception as e:
-            results["errors"].append(f"Add: {str(e)}")
-
-    def remove_worker():
-        try:
-            time.sleep(0.05)  # Wait for some tasks to be added
-            for task_id in results["added"][:10]:
-                scheduler.remove_task(task_id)
-                results["removed"].append(task_id)
-        except Exception as e:
-            results["errors"].append(f"Remove: {str(e)}")
-
-    add_thread = threading.Thread(target=add_worker)
-    remove_thread = threading.Thread(target=remove_worker)
-
-    add_thread.start()
-    remove_thread.start()
-
-    time.sleep(0.2)
-    scheduler.stop()
-
-    add_thread.join()
-    remove_thread.join()
-
-    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
-    assert len(results["added"]) == 20
-
-
-def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
-    """Test concurrent get_stats operations."""
-    mock_model, mock_tokenizer = mock_model_and_tokenizer
-
-    with patch("astrai.inference.scheduler.AutoModel"):
-        with patch("astrai.inference.scheduler.AutoTokenizer"):
-            scheduler = InferenceScheduler(
-                model=mock_model,
-                tokenizer=mock_tokenizer,
-                max_batch_size=4,
-                device="cpu",
-            )
-
-    results = {"stats": [], "errors": []}
-
-    def add_tasks():
-        try:
-            for i in range(20):
-                scheduler.add_task(f"prompt {i}")
-                time.sleep(0.001)
-        except Exception as e:
-            results["errors"].append(f"Add: {str(e)}")
-
-    def get_stats():
-        try:
-            for _ in range(50):
-                stats = scheduler.get_stats()
-                results["stats"].append(stats)
-                time.sleep(0.001)
-        except Exception as e:
-            results["errors"].append(f"Get stats: {str(e)}")
-
-    add_thread = threading.Thread(target=add_tasks)
-    stats_thread = threading.Thread(target=get_stats)
-
-    add_thread.start()
-    stats_thread.start()
-
-    time.sleep(0.3)
-    scheduler.stop()
-
-    add_thread.join()
-    stats_thread.join()
-
-    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
-    assert len(results["stats"]) == 50
-
-    # Verify stats are consistent
-    for stats in results["stats"]:
-        assert "total_tasks" in stats
-        assert stats["total_tasks"] >= 0
-
-
-def test_prefix_cache_insert_same_prefix_concurrently():
-    """Test inserting the same prefix concurrently."""
-    cache = PrefixCacheManager(max_capacity=100)
-
-    results = {"slot_values": [], "errors": []}
-
-    def insert_worker():
-        try:
-            # All workers try to insert the same prefix
-            cache.insert((1, 2, 3), slot=threading.current_thread().name)
-            node = cache.root.children.get(1)
-            if node:
-                node = node.children.get(2)
-                if node:
-                    node = node.children.get(3)
-                    if node:
-                        results["slot_values"].append(node.slot)
-        except Exception as e:
-            results["errors"].append(str(e))
-
-    threads = [threading.Thread(target=insert_worker) for _ in range(10)]
-
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-
-    # All inserts should succeed, final slot should be one of the values
-    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
-    # Check ref_count is correct (should be 10)
-    node = cache.root.children.get(1).children.get(2).children.get(3)
-    assert node.ref_count == 10, f"Expected ref_count=10, got {node.ref_count}"
-
-
-def test_prefix_cache_ref_count_underflow_prevention():
-    """Test that ref_count doesn't go negative."""
-    cache = PrefixCacheManager(max_capacity=100)
-
-    # Insert a prefix
-    cache.insert((1, 2, 3), slot=0)
-
-    # Release multiple times
-    for _ in range(5):
-        cache.release((1, 2, 3))
-
-    # Try to find it - should return None since ref_count would be negative
-    # or handle it gracefully
-    node = cache.root.children.get(1).children.get(2).children.get(3)
-    # The ref_count should be 0, not negative
-    assert node.ref_count >= 0, f"ref_count went negative: {node.ref_count}"

tests/inference/test_server.py

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

tests/inference/test_task.py

@@ -0,0 +1,170 @@
+"""Unit tests for Task and TaskManager."""
+
+from unittest.mock import MagicMock
+
+from astrai.inference import STOP, Task, TaskManager, TaskStatus
+
+
+def _make_mock_tokenizer():
+    t = MagicMock()
+    t.encode.return_value = [1, 2, 3, 4, 5]
+    t.stop_ids = [0]
+    return t
+
+
+def test_task_default_status_is_pending():
+    task = Task("id1", [1, 2, 3])
+    assert task.status == TaskStatus.PENDING
+
+
+def test_task_next_pos():
+    task = Task("id1", [1, 2, 3])
+    task.input_tokens = 5
+    assert task.next_pos == 5
+    task.output_ids.append(4)
+    assert task.next_pos == 6
+
+
+def test_task_is_finished_max_tokens():
+    task = Task("id1", [1, 2, 3], max_tokens=2)
+    task.output_tokens = 2
+    assert task.is_finished([])
+
+
+def test_task_is_finished_stop_id():
+    task = Task("id1", [1, 2, 3])
+    task.output_ids = [5, 0]
+    assert task.is_finished([0])
+
+
+def test_task_is_finished_not_yet():
+    task = Task("id1", [1, 2, 3], max_tokens=10)
+    task.output_ids = [1, 2]
+    assert not task.is_finished([0])
+
+
+def test_task_manager_add_task():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tid = tm.add_task("hello")
+    assert tid.startswith("task_")
+    assert tm._total_tasks == 1
+    assert len(tm.waiting_queue) == 1
+
+
+def test_task_manager_add_task_too_long_immediate_stop():
+    t = _make_mock_tokenizer()
+    t.encode.return_value = list(range(9000))
+    cb_calls = []
+
+    tm = TaskManager(tokenizer=t, max_seq_len=16)
+    tm.add_task("long", stream_callback=lambda tok: cb_calls.append(tok))
+    assert cb_calls[0] is STOP
+    assert len(tm.waiting_queue) == 0
+
+
+def test_task_manager_remove_task():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tid = tm.add_task("test")
+    tm.remove_task(tid)
+    assert len(tm.waiting_queue) == 0
+
+
+def test_task_manager_remove_active_task():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tid = tm.add_task("test")
+    tasks = tm.pull_candidates(1)
+    tm.activate(tasks[0])
+    assert len(tm.active_tasks) == 1
+    removed = tm.remove_task(tid)
+    assert len(removed) == 1
+    assert len(tm.active_tasks) == 0
+
+
+def test_task_manager_pull_candidates_fifo():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tm.add_task("a")
+    tm.add_task("b")
+    tm.add_task("c")
+    pulled = tm.pull_candidates(2)
+    assert len(pulled) == 2
+    assert pulled[0].prompt_ids == [1, 2, 3, 4, 5]
+    assert len(tm.waiting_queue) == 1
+
+
+def test_task_manager_activate():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tm.add_task("test")
+    task = tm.pull_candidates(1)[0]
+    tm.activate(task)
+    assert task.status == TaskStatus.RUNNING
+    assert task in tm.active_tasks
+
+
+def test_task_manager_return_to_waiting():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tm.add_task("a")
+    tm.add_task("b")
+    t1 = tm.pull_candidates(1)[0]
+    tm.return_to_waiting([t1])
+    assert len(tm.waiting_queue) == 2
+    assert tm.waiting_queue[0] == t1
+
+
+def test_task_manager_remove_finished_aborted():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tm.add_task("test")
+    task = tm.pull_candidates(1)[0]
+    tm.activate(task)
+    task.status = TaskStatus.ABORTED
+    finished = tm.remove_finished_tasks([0])
+    assert len(finished) == 1
+    assert len(tm.active_tasks) == 0
+
+
+def test_task_manager_remove_finished_stop_id():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tm.add_task("test")
+    task = tm.pull_candidates(1)[0]
+    tm.activate(task)
+    task.output_ids = [0]
+    task.output_tokens = 1
+    finished = tm.remove_finished_tasks([0])
+    assert len(finished) == 1
+    assert task.status == TaskStatus.FINISHED
+    assert len(tm.active_tasks) == 0
+
+
+def test_task_manager_has_work():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    assert not tm.has_work()
+    tm.add_task("test")
+    assert tm.has_work()
+
+
+def test_task_manager_wake():
+    import threading
+
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    called = threading.Event()
+
+    def waiter():
+        tm.wait_for_tasks(timeout=5.0)
+        called.set()
+
+    t = threading.Thread(target=waiter)
+    t.start()
+    import time
+
+    time.sleep(0.05)
+    tm.wake()
+    t.join(timeout=2.0)
+    assert called.is_set()
+
+
+def test_task_manager_get_stats():
+    tm = TaskManager(tokenizer=_make_mock_tokenizer())
+    tm.add_task("test")
+    stats = tm.get_stats()
+    assert stats["total_tasks"] == 1
+    assert stats["waiting_queue"] == 1
+    assert stats["active_tasks"] == 0

tests/module/test_encoder.py

@@ -0,0 +1,166 @@
+import torch
+
+from astrai.config.model_config import EncoderConfig
+from astrai.model.encoder import EmbeddingEncoder
+
+TINY_CONFIG = dict(
+    vocab_size=128,
+    dim=8,
+    n_heads=2,
+    n_kv_heads=1,
+    dim_ffn=16,
+    max_len=64,
+    n_layers=2,
+    norm_eps=1e-5,
+)
+
+
+def test_encoder_forward_mean():
+    config = EncoderConfig(**TINY_CONFIG)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_forward_cls():
+    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "cls"})
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_forward_last():
+    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "last"})
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_forward_with_padding():
+    config = EncoderConfig(**TINY_CONFIG)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+    input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
+    input_mask[:, 4:] = False
+
+    with torch.no_grad():
+        output = model(input_ids, input_mask=input_mask)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_normalize():
+    config = EncoderConfig(
+        **{**TINY_CONFIG, "pooling_type": "mean", "normalize_embeddings": True}
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    norms = output.norm(p=2, dim=-1)
+    assert torch.allclose(norms, torch.ones_like(norms), atol=1e-4)
+
+
+def test_encoder_register():
+    from astrai.model.automodel import AutoModel
+
+    assert AutoModel.is_registered("embedding")
+    cls = AutoModel.get_component_class("embedding")
+    assert cls is EmbeddingEncoder
+
+
+def test_encoder_from_transformer_checkpoint():
+    config = EncoderConfig(**TINY_CONFIG)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+
+    state_dict = model.state_dict()
+    state_dict["lm_head.weight"] = torch.randn(
+        config.vocab_size, config.dim, device=device
+    )
+
+    new_model = EmbeddingEncoder(config).to(device=device)
+    new_model.load_state_dict(state_dict, strict=True)
+
+    for key in model.state_dict():
+        assert torch.equal(new_model.state_dict()[key], model.state_dict()[key])
+
+
+def test_encoder_save_load():
+    import json
+    import os
+    import tempfile
+
+    import safetensors.torch as st
+
+    test_dir = tempfile.mkdtemp(prefix="encoder_test_")
+    config_path = os.path.join(test_dir, "config.json")
+    weights_path = os.path.join(test_dir, "model.safetensors")
+
+    try:
+        config_data = {**TINY_CONFIG, "pooling_type": "mean"}
+        with open(config_path, "w") as f:
+            json.dump(config_data, f)
+
+        config = EncoderConfig.from_file(config_path)
+        original = EmbeddingEncoder(config)
+        st.save_file(original.state_dict(), weights_path)
+
+        loaded = EmbeddingEncoder(config)
+        loaded.load_state_dict(st.load_file(weights_path))
+
+        for key in original.state_dict():
+            assert torch.equal(original.state_dict()[key], loaded.state_dict()[key])
+    finally:
+        if os.path.exists(test_dir):
+            for f in os.listdir(test_dir):
+                os.remove(os.path.join(test_dir, f))
+            os.rmdir(test_dir)

tests/module/test_forward_configs.py

@@ -0,0 +1,108 @@
+import pytest
+import torch
+
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.transformer import AutoRegressiveLM
+
+TINY_CONFIG = dict(
+    vocab_size=128,
+    dim=8,
+    n_heads=2,
+    n_kv_heads=1,
+    dim_ffn=16,
+    max_len=64,
+    n_layers=2,
+    norm_eps=1e-5,
+)
+
+
+CONFIGS = [
+    pytest.param(
+        {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp"},
+        id="gqa_mlp",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "mla",
+            "ffn_type": "mlp",
+            "kv_lora_rank": 4,
+            "qk_nope_head_dim": 2,
+            "qk_rope_head_dim": 2,
+        },
+        id="mla_mlp",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+        },
+        id="gqa_moe",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "mlp",
+            "rope_theta": 100000.0,
+        },
+        id="gqa_rope_theta",
+    ),
+    pytest.param(
+        {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp", "use_qk_norm": True},
+        id="gqa_qk_norm",
+    ),
+    pytest.param(
+        {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp", "tie_weight": True},
+        id="gqa_tie_weight",
+    ),
+]
+
+
+@pytest.mark.parametrize("config_kwargs", CONFIGS)
+def test_model_forward(config_kwargs):
+    config = AutoRegressiveLMConfig(**config_kwargs)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoRegressiveLM(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert "logits" in output
+    assert "hidden_states" in output
+    assert output["logits"].shape == (batch_size, seq_len, config.vocab_size)
+    assert output["hidden_states"].shape == (batch_size, seq_len, config.dim)
+    assert not torch.isnan(output["logits"]).any()
+    assert not torch.isnan(output["hidden_states"]).any()
+
+
+@pytest.mark.parametrize("config_kwargs", CONFIGS)
+def test_model_forward_with_padding(config_kwargs):
+    config = AutoRegressiveLMConfig(**config_kwargs)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoRegressiveLM(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+    input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
+    input_mask[:, 4:] = False
+
+    with torch.no_grad():
+        output = model(input_ids, input_mask=input_mask)
+
+    assert output["logits"].shape == (batch_size, seq_len, config.vocab_size)
+    assert not torch.isnan(output["logits"]).any()

tests/module/test_lora.py

@@ -0,0 +1,355 @@
+import tempfile
+
+import pytest
+import torch
+
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model import AutoRegressiveLM
+from astrai.model.components.lora import (
+    LoRAConfig,
+    LoRALinear,
+    _get_lora_count,
+    _collect_lora_info,
+    inject_lora,
+    load_lora,
+    merge_lora,
+    save_lora,
+)
+from astrai.model.components.linear import Linear
+
+MODEL_KWARGS = dict(
+    vocab_size=1000,
+    dim=64,
+    n_heads=4,
+    n_kv_heads=2,
+    dim_ffn=128,
+    n_layers=2,
+    max_len=32,
+    norm_eps=1e-5,
+)
+
+
+def _make_model(**kwargs):
+    kw = {**MODEL_KWARGS, **kwargs}
+    config = AutoRegressiveLMConfig(**kw)
+    model = AutoRegressiveLM(config)
+    model.eval()
+    return model
+
+
+def test_loralinear_init():
+    base = Linear(64, 128)
+    lora = LoRALinear(base, r=8, alpha=16)
+
+    assert lora.weight is base.weight
+    assert not lora.weight.requires_grad
+    assert lora.lora_A.shape == (8, 64)
+    assert lora.lora_B.shape == (128, 8)
+    assert lora.scaling == 2.0
+    assert not lora._merged
+    assert lora.lora_A.requires_grad
+    assert lora.lora_B.requires_grad
+
+
+def test_loralinear_forward_init_zero_delta():
+    base = Linear(4, 4)
+    with torch.no_grad():
+        base.weight.zero_()
+
+    x = torch.randn(2, 4)
+    lora = LoRALinear(base, r=2, alpha=2)
+    base_out = base(x)
+    lora_out = lora(x)
+
+    assert torch.allclose(base_out, lora_out)
+
+
+def test_loralinear_forward_with_delta():
+    base = Linear(4, 4)
+    with torch.no_grad():
+        base.weight.zero_()
+
+    x = torch.randn(2, 4)
+    lora = LoRALinear(base, r=2, alpha=2)
+    base_out = base(x)
+
+    with torch.no_grad():
+        lora.lora_B.fill_(1.0)
+
+    lora_out = lora(x)
+    assert not torch.allclose(base_out, lora_out)
+
+
+def test_loralinear_merge():
+    base = Linear(4, 4)
+    with torch.no_grad():
+        base.weight.zero_()
+
+    x = torch.randn(2, 4)
+    lora = LoRALinear(base, r=2, alpha=2)
+    with torch.no_grad():
+        lora.lora_B.fill_(1.0)
+
+    out_before = lora(x).clone()
+    lora.merge()
+    out_after = lora(x)
+
+    torch.testing.assert_close(out_before, out_after)
+    assert lora._merged
+    assert not hasattr(lora, "lora_A")
+
+
+def test_loralinear_merge_is_idempotent():
+    base = Linear(4, 4)
+    with torch.no_grad():
+        base.weight.zero_()
+
+    lora = LoRALinear(base, r=2, alpha=2)
+    with torch.no_grad():
+        lora.lora_B.fill_(1.0)
+
+    lora.merge()
+    lora.merge()
+
+
+def test_inject_lora_default_target():
+    model = _make_model()
+    n_before = sum(1 for m in model.modules() if isinstance(m, Linear))
+
+    inject_lora(model, r=4, alpha=8)
+
+    lora_count = _get_lora_count(model)
+    assert lora_count > 0
+    assert lora_count < n_before
+
+
+def test_inject_lora_ffn():
+    model = _make_model()
+    from astrai.model.components.lora import TARGET_MODULES_FFN
+
+    inject_lora(model, r=4, alpha=8, target_modules=TARGET_MODULES_FFN)
+    assert _get_lora_count(model) > 0
+
+
+def test_inject_lora_returns_config():
+    model = _make_model()
+    cfg = inject_lora(model, r=8, alpha=32)
+    assert isinstance(cfg, LoRAConfig)
+    assert cfg.r == 8
+    assert cfg.alpha == 32
+
+
+def test_inject_lora_no_matching_targets_warns(caplog):
+    model = _make_model()
+    inject_lora(model, r=4, alpha=8, target_modules={"nonexistent"})
+    assert "No LoRA layers injected" in caplog.text
+
+
+def test_inject_lora_preserves_base_output():
+    model = _make_model()
+    x = torch.randint(0, 1000, (2, 16))
+
+    with torch.no_grad():
+        out_before = model(x)["logits"].clone()
+
+    inject_lora(model, r=4, alpha=8)
+
+    with torch.no_grad():
+        out_after = model(x)["logits"]
+
+    torch.testing.assert_close(out_before, out_after)
+
+
+def test_inject_lora_does_not_reinject():
+    model = _make_model()
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+    first_count = _get_lora_count(model)
+
+    inject_lora(model, r=2, alpha=4, target_modules={"q_proj"})
+    assert _get_lora_count(model) == first_count
+
+
+def test_inject_lora_adds_new_modules():
+    model = _make_model()
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+    first = _get_lora_count(model)
+
+    inject_lora(model, r=4, alpha=8, target_modules={"v_proj"})
+    assert _get_lora_count(model) > first
+
+
+def test_inject_lora_on_mla_model():
+    model = _make_model(
+        attn_type="mla", kv_lora_rank=16, qk_nope_head_dim=16, qk_rope_head_dim=16
+    )
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj", "o_proj"})
+    assert _get_lora_count(model) > 0
+
+
+def test_inject_lora_on_moe_model():
+    model = _make_model(
+        ffn_type="moe",
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+        dim_ffn=32,
+    )
+    inject_lora(model, r=4, alpha=8, target_modules={"up", "gate", "down"})
+    assert _get_lora_count(model) > 0
+
+
+def test_state_dict_key_format():
+    model = _make_model()
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+
+    sd = model.state_dict()
+    assert "layers.0.attention.q_proj.weight" in sd
+    assert "layers.0.attention.q_proj.lora_A" in sd
+    assert "layers.0.attention.q_proj.lora_B" in sd
+
+
+def test_only_lora_params_trainable():
+    model = _make_model()
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj", "v_proj"})
+
+    for name, param in model.named_parameters():
+        if isinstance(name.split(".")[-1], str) and "lora" in name:
+            assert param.requires_grad, f"lora param should be trainable: {name}"
+        elif any(name.endswith(f".{t}.weight") for t in ("q_proj", "v_proj")):
+            assert not param.requires_grad, f"injected weight should be frozen: {name}"
+
+
+def test_state_dict_after_inject_consistent_with_original():
+    model = _make_model()
+    sd_before = {k: v for k, v in model.state_dict().items()}
+
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+    sd_after = model.state_dict()
+
+    # original keys unchanged
+    for k in sd_before:
+        assert k in sd_after
+        assert sd_before[k].shape == sd_after[k].shape
+
+    # new lora keys present
+    lora_keys = [k for k in sd_after if "lora" in k]
+    assert len(lora_keys) > 0
+
+
+def test_save_load_roundtrip():
+    model = _make_model()
+    cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+
+    with torch.no_grad():
+        for m in model.modules():
+            if isinstance(m, LoRALinear):
+                m.lora_B.fill_(0.5)
+
+    x = torch.randint(0, 1000, (2, 16))
+    with torch.no_grad():
+        out_src = model(x)["logits"].clone()
+
+    tmpdir = tempfile.mkdtemp()
+    save_lora(model, tmpdir, cfg)
+
+    model2 = _make_model()
+    model2.load_state_dict(model.state_dict(), strict=False)
+    load_lora(model2, tmpdir)
+
+    with torch.no_grad():
+        out_dst = model2(x)["logits"]
+
+    torch.testing.assert_close(out_src, out_dst)
+
+
+def test_save_after_merge_raises():
+    model = _make_model()
+    cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+
+    with torch.no_grad():
+        for m in model.modules():
+            if isinstance(m, LoRALinear):
+                m.lora_B.fill_(0.5)
+
+    tmpdir = tempfile.mkdtemp()
+    save_lora(model, tmpdir, cfg)
+    merge_lora(model)
+
+    tmpdir2 = tempfile.mkdtemp()
+    with pytest.raises(RuntimeError, match="No LoRA parameters"):
+        save_lora(model, tmpdir2, cfg)
+
+
+def test_load_lora_on_already_injected():
+    model = _make_model()
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+
+    with torch.no_grad():
+        for m in model.modules():
+            if isinstance(m, LoRALinear):
+                m.lora_B.fill_(0.5)
+
+    tmpdir = tempfile.mkdtemp()
+    save_lora(model, tmpdir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",)))
+
+    model2 = _make_model()
+    model2.load_state_dict(model.state_dict(), strict=False)
+    inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
+
+    # load onto already-injected model
+    load_lora(model2, tmpdir)
+    assert _get_lora_count(model2) > 0
+
+
+def test_load_lora_mismatched_r_raises():
+    model = _make_model()
+    cfg = inject_lora(model, r=8, alpha=16, target_modules={"q_proj"})
+
+    with torch.no_grad():
+        for m in model.modules():
+            if isinstance(m, LoRALinear):
+                m.lora_B.fill_(0.5)
+
+    tmpdir = tempfile.mkdtemp()
+    save_lora(model, tmpdir, cfg)
+
+    model2 = _make_model()
+    model2.load_state_dict(model.state_dict(), strict=False)
+    inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
+
+    with pytest.raises(RuntimeError, match="size mismatch"):
+        load_lora(model2, tmpdir)  # strict=False, only lora keys
+
+
+def test_merge_preserves_output():
+    model = _make_model()
+    inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
+
+    with torch.no_grad():
+        for m in model.modules():
+            if isinstance(m, LoRALinear):
+                m.lora_B.fill_(0.5)
+
+    x = torch.randint(0, 1000, (2, 16))
+    with torch.no_grad():
+        out_before = model(x)["logits"].clone()
+
+    merge_lora(model)
+
+    with torch.no_grad():
+        out_after = model(x)["logits"]
+    torch.testing.assert_close(out_before, out_after)
+
+
+def test_merge_no_lora_warns(caplog):
+    model = _make_model()
+    merge_lora(model)
+    assert "No LoRA layers to merge" in caplog.text
+
+
+def test_collect_lora_info():
+    model = _make_model()
+    info = _collect_lora_info(model)
+    assert "q_proj" in info
+    assert "o_proj" in info
+    assert "q_proj" in info  # each layer has one

tests/module/test_tie_weight.py

@@ -6,8 +6,8 @@ import pytest
 import safetensors.torch as st
 import torch
 
-from astrai.config.model_config import ModelConfig
-from astrai.model.transformer import Transformer
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.transformer import AutoRegressiveLM
 
 
 @pytest.fixture
@@ -17,10 +17,10 @@ def transformer_test_env():
 
     config = {
         "vocab_size": 1000,
-        "dim": 128,
-        "n_heads": 4,
-        "n_kv_heads": 2,
-        "dim_ffn": 256,
+        "dim": 8,
+        "n_heads": 2,
+        "n_kv_heads": 1,
+        "dim_ffn": 16,
         "max_len": 64,
         "n_layers": 2,
         "norm_eps": 1e-5,
@@ -50,8 +50,8 @@ def test_tie_weight_init(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    config = ModelConfig().load(config_path)
-    model = Transformer(config)
+    config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(config)
 
     assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
     assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
@@ -68,8 +68,8 @@ def test_tie_weight_init(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    config = ModelConfig().load(config_path)
-    model = Transformer(config)
+    config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(config)
 
     assert not torch.equal(model.lm_head.weight, model.embed_tokens.weight)
     assert model.lm_head.weight.data_ptr() != model.embed_tokens.weight.data_ptr()
@@ -94,13 +94,13 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    config = ModelConfig().load(config_path)
-    original_model = Transformer(config)
+    config = AutoRegressiveLMConfig.from_file(config_path)
+    original_model = AutoRegressiveLM(config)
 
     st.save_file(original_model.state_dict(), model_path)
 
-    loaded_config = ModelConfig().load(config_path)
-    model = Transformer(loaded_config)
+    loaded_config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(loaded_config)
     model.load_state_dict(st.load_file(model_path))
 
     assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
@@ -112,8 +112,8 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    loaded_config = ModelConfig().load(config_path)
-    model = Transformer(loaded_config)
+    loaded_config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(loaded_config)
     model.load_state_dict(st.load_file(model_path))
 
     assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)

tests/trainer/conftest.py

@@ -1,3 +1,5 @@
+import os
+
 import pytest
 import torch
 from torch.utils.data import Dataset
@@ -31,8 +33,8 @@ def create_train_config(
     device: str,
     strategy: str = "seq",
     n_epoch: int = 1,
-    batch_size: int = 2,
-    accumulation_steps: int = 1,
+    batch_per_device: int = 2,
+    grad_accum_steps: int = 1,
     max_grad_norm: float = 1.0,
     ckpt_interval: int = 5,
     random_seed: int = 42,
@@ -47,8 +49,8 @@ def create_train_config(
         device: Device type ("cuda" or "cpu")
         strategy: Training strategy type (default: "seq")
         n_epoch: Number of epochs (default: 1)
-        batch_size: Batch size (default: 2)
-        accumulation_steps: Gradient accumulation steps (default: 1)
+        batch_per_device: Batch size per device (default: 2)
+        grad_accum_steps: Gradient accumulation steps (default: 1)
         max_grad_norm: Maximum gradient norm for clipping (default: 1.0)
         ckpt_interval: Checkpoint save interval in iterations (default: 5)
         random_seed: Random seed for reproducibility (default: 42)
@@ -73,10 +75,11 @@ def create_train_config(
         optimizer_fn=optimizer_fn,
         scheduler_fn=scheduler_fn,
         ckpt_dir=test_dir,
+        log_dir=os.path.join(test_dir, "logs"),
         n_epoch=n_epoch,
-        batch_size=batch_size,
+        batch_per_device=batch_per_device,
         ckpt_interval=ckpt_interval,
-        accumulation_steps=accumulation_steps,
+        grad_accum_steps=grad_accum_steps,
         max_grad_norm=max_grad_norm,
         random_seed=random_seed,
         device_type=device,

tests/trainer/test_callbacks.py

@@ -1,11 +1,133 @@
+import os
+
 import torch
 
 from astrai.config.train_config import TrainConfig
+from astrai.model.components.decoder_block import DecoderBlock
 from astrai.trainer.schedule import SchedulerFactory
-from astrai.trainer.train_callback import TrainCallback
+from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
 from astrai.trainer.trainer import Trainer
 
 
+def test_gradient_checkpointing_enable_disable(test_model):
+    """Enable wraps forward, _disable restores it."""
+    model = test_model["model"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    originals = [layer.forward for layer in model.layers]
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    for layer in model.layers:
+        assert hasattr(layer, "_original_forward")
+        assert layer.forward is not originals[0]
+
+    for layer in model.layers:
+        callback._disable(layer)
+
+    for layer in model.layers:
+        assert not hasattr(layer, "_original_forward")
+
+
+def test_gradient_checkpointing_empty_modules_noop(test_model):
+    """modules=None should leave forwards untouched."""
+    model = test_model["model"]
+    callback = GradientCheckpointingCallback()
+
+    originals = [layer.forward for layer in model.layers]
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    for layer, orig in zip(model.layers, originals):
+        assert layer.forward is orig
+
+
+def test_gradient_checkpointing_forward_unchanged(test_model):
+    """Forward output unchanged after patching (no_grad)."""
+    model = test_model["model"]
+    device = test_model["device"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    input_ids = torch.randint(0, 1000, (2, 32)).to(device)
+
+    with torch.no_grad():
+        ref = model(input_ids)["logits"].clone()
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    with torch.no_grad():
+        out = model(input_ids)["logits"]
+
+    assert torch.equal(ref, out)
+
+
+def test_gradient_checkpointing_backward(test_model):
+    """backward passes gradients through checkpointed layers."""
+    model = test_model["model"]
+    device = test_model["device"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    input_ids = torch.randint(0, 1000, (2, 32)).to(device)
+    target_ids = torch.randint(0, 1000, (2, 32)).to(device)
+
+    logits = model(input_ids)["logits"]
+    loss = torch.nn.functional.cross_entropy(
+        logits.flatten(0, 1).float(), target_ids.flatten()
+    )
+    loss.backward()
+
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            assert param.grad is not None, f"{name} gradient is None"
+
+    for layer in model.layers:
+        callback._disable(layer)
+
+    model.zero_grad()
+    for name, p in model.named_parameters():
+        assert p.grad is None or p.grad.sum().item() == 0, f"{name} grad not zeroed"
+
+
+def test_gradient_checkpointing_trainer_integration(base_test_env, random_dataset):
+    """Gradient checkpointing runs end-to-end via Trainer."""
+
+    def optimizer_fn(model):
+        return torch.optim.AdamW(model.parameters())
+
+    def scheduler_fn(optim):
+        return SchedulerFactory.create(
+            optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
+        )
+
+    train_config = TrainConfig(
+        model=base_test_env["model"],
+        strategy="seq",
+        dataset=random_dataset,
+        optimizer_fn=optimizer_fn,
+        scheduler_fn=scheduler_fn,
+        ckpt_dir=base_test_env["test_dir"],
+        log_dir=os.path.join(base_test_env["test_dir"], "logs"),
+        n_epoch=1,
+        batch_per_device=2,
+        ckpt_interval=3,
+        grad_accum_steps=1,
+        max_grad_norm=1.0,
+        random_seed=42,
+        device_type=base_test_env["device"],
+        gradient_checkpointing_modules=[DecoderBlock],
+    )
+
+    trainer = Trainer(train_config)
+    trainer.train()
+    # no crash = callback correctly enabled/disabled
+
+
 def test_callback_integration(base_test_env, random_dataset):
     """Test that all callbacks are properly integrated"""
 
@@ -24,10 +146,11 @@ def test_callback_integration(base_test_env, random_dataset):
         optimizer_fn=optimizer_fn,
         scheduler_fn=scheduler_fn,
         ckpt_dir=base_test_env["test_dir"],
+        log_dir=os.path.join(base_test_env["test_dir"], "logs"),
         n_epoch=1,
-        batch_size=2,
+        batch_per_device=2,
         ckpt_interval=3,
-        accumulation_steps=1,
+        grad_accum_steps=1,
         max_grad_norm=1.0,
         random_seed=42,
         device_type=base_test_env["device"],

tests/trainer/test_early_stopping.py

@@ -27,10 +27,11 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
         model=base_test_env["model"],
         dataset=early_stopping_dataset,
         ckpt_dir=base_test_env["test_dir"],
+        log_dir=os.path.join(base_test_env["test_dir"], "logs"),
         n_epoch=2,
-        batch_size=2,
+        batch_per_device=2,
         ckpt_interval=1,
-        accumulation_steps=2,
+        grad_accum_steps=2,
         random_seed=np.random.randint(1e4),
         device_type=base_test_env["device"],
     )

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

tests/trainer/test_trainer.py

@@ -7,45 +7,45 @@ def test_different_batch_sizes(base_test_env, random_dataset, train_config_facto
     """Test training with different batch sizes"""
     batch_sizes = [1, 2, 4, 8]
 
-    for batch_size in batch_sizes:
+    for batch_per_device in batch_sizes:
         train_config = train_config_factory(
             model=base_test_env["model"],
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
-            batch_size=batch_size,
+            batch_per_device=batch_per_device,
         )
 
-        assert train_config.batch_size == batch_size
+        assert train_config.batch_per_device == batch_per_device
 
 
 def test_gradient_accumulation(base_test_env, random_dataset, train_config_factory):
     """Test training with different gradient accumulation steps"""
-    accumulation_steps_list = [1, 2, 4]
+    grad_accum_steps_list = [1, 2, 4]
 
-    for accumulation_steps in accumulation_steps_list:
+    for grad_accum_steps in grad_accum_steps_list:
         train_config = train_config_factory(
             model=base_test_env["model"],
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
-            batch_size=2,
-            accumulation_steps=accumulation_steps,
+            batch_per_device=2,
+            grad_accum_steps=grad_accum_steps,
         )
 
         trainer = Trainer(train_config)
         trainer.train()
 
-        assert train_config.accumulation_steps == accumulation_steps
+        assert train_config.grad_accum_steps == grad_accum_steps
 
 
 def test_memory_efficient_training(base_test_env, random_dataset, train_config_factory):
     """Test training with memory-efficient configurations"""
     # Test with smaller batch sizes and gradient checkpointing
     small_batch_configs = [
-        {"batch_size": 1, "accumulation_steps": 8},
-        {"batch_size": 2, "accumulation_steps": 4},
-        {"batch_size": 4, "accumulation_steps": 2},
+        {"batch_per_device": 1, "grad_accum_steps": 8},
+        {"batch_per_device": 2, "grad_accum_steps": 4},
+        {"batch_per_device": 4, "grad_accum_steps": 2},
     ]
 
     for config in small_batch_configs:
@@ -54,8 +54,9 @@ def test_memory_efficient_training(base_test_env, random_dataset, train_config_f
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
-            batch_size=config["batch_size"],
-            accumulation_steps=config["accumulation_steps"],
+            batch_per_device=config["batch_per_device"],
+            grad_accum_steps=config["grad_accum_steps"],
         )
 
-        assert train_config.accumulation_steps == config["accumulation_steps"]
+        assert train_config.grad_accum_steps == config["grad_accum_steps"]
+        assert train_config.batch_per_device == config["batch_per_device"]