release : v1.3.8

- refactor: 重写 IFD 评估为三层架构，引入 BFD 装箱与自定义 attention mask 批处理打分
- refactor: 重写 HumanEval 评估为函数式流水线，修复测试超时与动态 pass@k
- perf: 替换 paged KV cache 为 ContiguousCache，解码所有 group
- feat: 新增 ROUGE 评估脚本、JSONL 数据集 store、stream_chat 参数
- fix: 修复 IFD token-set 不对称、SFT position_ids 默认值、文档边界保留

.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

@@ -5,8 +5,10 @@
 !*/
 
 # Allow specific file types and root files
-!*.py
-!*.sh
+!astrai/**/*.py
+!scripts/**/*.py
+!scripts/**/*.sh
+!tests/**/*.py
 
 # Allow GitHub files
 !/.github/**

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/ViperEkura/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,21 +18,22 @@ 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 \
     && pip install --no-cache-dir . \
-    --extra-index-url https://download.pytorch.org/whl/cu126
+    --extra-index-url https://download.pytorch.org/whl/cu128
 
 # 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

@@ -9,9 +9,9 @@
 <div align="center">
   <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
   <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?color=76bad9" alt="release">
-  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.stargazers_count&label=stars&suffix=%20stars&color=76bad9" alt="stars">
-  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.forks_count&label=forks&suffix=%20forks&color=76bad9" alt="forks">
+  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
+  <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
 <br>
 
@@ -28,7 +28,8 @@
 ## 📖 Table of Contents
 
 - [Features](#features)
-- [Quick Start](#quick-start)
+- [Getting Started](#getting-started)
+- [Demo](#demo)
 - [Documentation](#documentation)
 - [Contributing](#contributing)
 - [Community](#community)
@@ -49,45 +50,117 @@
 - 🤗 **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
+### Getting Started
 
-#### Installation
+End-to-end walkthrough in 5 steps:
+
+**1. Install**
 
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
 pip install -e .
+# pip install -e ".[dev]"    # optional: dev dependencies (pytest, ruff)
 ```
 
-For development dependencies:
+**2. Download model**
 
 ```bash
-pip install -e ".[dev]"
+python scripts/demo/download.py    # downloads 1B checkpoint to params/
 ```
 
-#### Train a Model
+**3. Preprocess data**
+
+Create `pretrain.json` (preprocessing config for `seq` strategy):
+
+```json
+{
+    "version": 1,
+    "input": {"sections": [{"field": "text", "action": "train"}]},
+    "preprocessing": {"max_seq_len": 2048},
+    "output": {"storage_format": "bin"}
+}
+```
 
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
-    --train_type seq \
-    --data_root_path /path/to/dataset \
-    --param_path /path/to/model \
-    --batch_size 4 \
-    --accumulation_steps 8 \
-    --max_lr 3e-4 \
-    --warmup_steps 1000 \
-    --n_epoch 1
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
 ```
 
-Full reference at [Parameter Guide](assets/docs/params.md).
+**4. Train**
 
-#### Generate Text
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --parallel_mode=ddp \
+    --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 &
+```
+
+**5. Serve & query**
+
+```bash
+# Terminal 1: start server
+python scripts/tools/server.py --param_path ./params --device cuda
+
+# Terminal 2: query
+curl http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+### Demo
+
+Check out the demos in the `scripts/demo/` folder:
+
+```bash
+# Download model weights (required before running demos)
+python scripts/demo/download.py                      # model → params/
+
+# Interactive streaming chat (multi-turn, maintains history)
+python scripts/demo/stream_chat.py
+# Type your message after >>, type !exit to quit
+
+# Batch generation (5 hardcoded prompts, non-streaming)
+python scripts/demo/generate_batch.py
+
+# Single-prompt autoregressive streaming
+python scripts/demo/generate_ar.py
+```
+
+All generation demos use `temperature=0.8`, `top_p=0.95`, `top_k=50`, `max_tokens=2048` by default and require `params/` to contain model weights (run `download.py` first).
+
+Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).
+
+---
+
+See [Documentation](#documentation) for full references beyond the examples above.
+
+#### Text Generation
+
+Batch generation from a JSONL file:
 
 ```bash
 python scripts/tools/generate.py \
-    --param_path /path/to/model \
-    --input_json_file /path/to/input.json \
-    --output_json_file /path/to/output.json
+    --param_path ./params \
+    --input_json_file input.jsonl \
+    --output_json_file output.jsonl
 ```
 
 #### Docker
@@ -101,9 +174,6 @@ docker build -t astrai:latest .
 # Run with GPU support
 docker run --gpus all -it astrai:latest
 
-# Run with specific GPUs
-docker run --gpus '"device=0,1"' -it astrai:latest
-
 # Run inference server
 docker run --gpus all -p 8000:8000 astrai:latest \
   python -m scripts.tools.server --port 8000 --device cuda
@@ -120,87 +190,42 @@ 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
+#### HTTP API Examples
 
-Start the inference server with OpenAI and Anthropic-compatible HTTP API:
+Additional request examples beyond the [Getting Started](#getting-started) flow:
 
 ```bash
-python -m scripts.tools.server --port 8000 --device cuda
-```
-
-Make requests:
-
-```bash
-# OpenAI-compatible
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "Hello"}],
-    "max_tokens": 512
-  }'
-
 # OpenAI-compatible streaming
 curl -X POST http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "Tell a story"}],
-    "stream": true,
-    "max_tokens": 500
-  }'
+  -d '{"messages":[{"role":"user","content":"Tell a story"}],"stream":true,"max_tokens":500}'
 
 # Anthropic-compatible
 curl -X POST http://localhost:8000/v1/messages \
   -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "You are a helpful assistant.",
-    "messages": [{"role": "user", "content": "Hello"}],
-    "max_tokens": 512
-  }'
+  -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"]
-  }'
+  -d '{"model":"astrai","messages":[{"role":"user","content":"Write a story"}],"max_tokens":500,"stream":true,"stop_sequences":["The end"]}'
 
 # Health check
 curl http://localhost:8000/health
 ```
 
-#### Demo
-
-Check out the demos in the `scripts/demo/` folder:
-
-```bash
-# Download pre‑processed data (required before running demos)
-python scripts/demo/download.py
-
-# Interactive streaming chat
-python scripts/demo/stream_chat.py
-
-# Batch generation
-python scripts/demo/generate_batch.py
-
-# Auto‑regressive generation
-python scripts/demo/generate_ar.py
-```
-
-Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd).
+See [Inference Guide](assets/docs/inference.md) for SSE streaming format, error codes, and stats endpoint.
 
 ### 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 |
+| [CLI Reference](./assets/docs/params.md) | Parameters for all CLI tools (train, server, generate, preprocess) |
+| [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 |
+| [Preprocessing](./assets/docs/preprocessing.md) | Declarative JSON-driven data preprocessing |
 
 ### Contributing
 

assets/docs/README-zh-CN.md

@@ -15,9 +15,9 @@
 <div align="center">
   <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
   <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?color=76bad9" alt="release">
-  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.stargazers_count&label=stars&suffix=%20stars&color=76bad9" alt="stars">
-  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.forks_count&label=forks&suffix=%20forks&color=76bad9" alt="forks">
+  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
+  <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
 
 <br>
@@ -34,7 +34,8 @@
 ## 📖 目录
 
 - [特性](#特性)
-- [快速开始](#快速开始)
+- [快速上手](#快速上手)
+- [演示](#演示)
 - [文档](#文档)
 - [贡献](#贡献)
 - [社区](#社区)
@@ -55,45 +56,117 @@
 - 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
 - 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
 
-### 快速开始
+### 快速上手
 
-#### 安装
+端到端演示，只需 5 步：
+
+**1. 安装**
 
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
 pip install -e .
+# pip install -e ".[dev]"    # 可选：开发依赖（pytest, ruff）
 ```
 
-安装开发依赖：
+**2. 下载模型**
 
 ```bash
-pip install -e ".[dev]"
+python scripts/demo/download.py    # 下载 1B 检查点到 params/
 ```
 
-#### 训练模型
+**3. 预处理数据**
+
+创建 `pretrain.json`（`seq` 策略的预处理配置）：
+
+```json
+{
+    "version": 1,
+    "input": {"sections": [{"field": "text", "action": "train"}]},
+    "preprocessing": {"max_seq_len": 2048},
+    "output": {"storage_format": "bin"}
+}
+```
 
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
-    --train_type seq \
-    --data_root_path /path/to/dataset \
-    --param_path /path/to/model \
-    --batch_size 4 \
-    --accumulation_steps 8 \
-    --max_lr 3e-4 \
-    --warmup_steps 1000 \
-    --n_epoch 1
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
 ```
 
-完整参数列表见[参数说明](./params.md)。
+**4. 训练**
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --parallel_mode=ddp \
+    --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 &
+```
+
+**5. 启动服务并调用**
+
+```bash
+# 终端 1：启动服务
+python scripts/tools/server.py --param_path ./params --device cuda
+
+# 终端 2：发起请求
+curl http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"messages":[{"role":"user","content":"你好"}],"max_tokens":512}'
+```
+
+### 演示
+
+查看 `scripts/demo/` 文件夹中的演示：
+
+```bash
+# 下载模型权重（运行演示前必需）
+python scripts/demo/download.py                      # model → params/
+
+# 交互式流式聊天（多轮对话，保持历史记录）
+python scripts/demo/stream_chat.py
+# 在 >> 后输入消息，输入 !exit 退出
+
+# 批量生成（5 条硬编码提示词，非流式）
+python scripts/demo/generate_batch.py
+
+# 单条提示词自回归流式生成
+python scripts/demo/generate_ar.py
+```
+
+所有生成演示默认使用 `temperature=0.8`、`top_p=0.95`、`top_k=50`、`max_tokens=2048`，需要 `params/` 目录包含模型权重（请先运行 `download.py`）。
+
+观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
+
+---
+
+更多选项请参考[文档](#文档)。
 
 #### 文本生成
 
+从 JSONL 文件批量生成：
+
 ```bash
 python scripts/tools/generate.py \
-    --param_path /path/to/model \
-    --input_json_file /path/to/input.json \
-    --output_json_file /path/to/output.json
+    --param_path ./params \
+    --input_json_file input.jsonl \
+    --output_json_file output.jsonl
 ```
 
 #### Docker
@@ -107,9 +180,6 @@ docker build -t astrai:latest .
 # 启用 GPU 运行
 docker run --gpus all -it astrai:latest
 
-# 指定特定 GPU
-docker run --gpus '"device=0,1"' -it astrai:latest
-
 # 运行推理服务
 docker run --gpus all -p 8000:8000 astrai:latest \
   python -m scripts.tools.server --port 8000 --device cuda
@@ -126,87 +196,42 @@ docker compose --profile cpu up -d
 
 > **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。
 
-#### 启动 HTTP 服务
+#### HTTP API 示例
 
-启动推理服务器，支持 OpenAI 和 Anthropic 兼容的 HTTP API：
+除[快速上手](#快速上手)流程外，更多请求示例：
 
 ```bash
-python -m scripts.tools.server --port 8000 --device cuda
-```
-
-发起请求：
-
-```bash
-# OpenAI 兼容
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "你好"}],
-    "max_tokens": 512
-  }'
-
 # OpenAI 兼容流式
 curl -X POST http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "讲个故事"}],
-    "stream": true,
-    "max_tokens": 500
-  }'
+  -d '{"messages":[{"role":"user","content":"讲个故事"}],"stream":true,"max_tokens":500}'
 
 # Anthropic 兼容
 curl -X POST http://localhost:8000/v1/messages \
   -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "你是一个乐于助人的助手。",
-    "messages": [{"role": "user", "content": "你好"}],
-    "max_tokens": 512
-  }'
+  -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": ["结束"]
-  }'
+  -d '{"model":"astrai","messages":[{"role":"user","content":"写个故事"}],"max_tokens":500,"stream":true,"stop_sequences":["结束"]}'
 
 # 健康检查
 curl http://localhost:8000/health
 ```
 
-#### 演示
-
-查看 `scripts/demo/` 文件夹中的演示：
-
-```bash
-# 下载预处理数据（运行演示前必需）
-python scripts/demo/download.py
-
-# 交互式流式聊天
-python scripts/demo/stream_chat.py
-
-# 批量生成
-python scripts/demo/generate_batch.py
-
-# 自回归生成
-python scripts/demo/generate_ar.py
-```
-
-观看 [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd) 上的视频演示。
+SSE 流式格式、错误码和统计端点详见[推理文档](./inference.md)。
 
 ### 文档
 
 | 文档 | 说明 |
 |------|------|
-| [参数说明](./params.md) | 训练与推理参数配置 |
-| [设计文档](./design.md) | 系统架构与模块设计 |
-| [数据流程](./dataflow.md) | 数据处理管道详解 |
-| [模型介绍](./introduction.md) | 模型架构与技术细节 |
+| [CLI 参考](./params.md) | 所有 CLI 工具参数（训练、服务、生成、预处理） |
+| [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
+| [训练文档](./training.md) | 训练循环、策略与公式 |
+| [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
+| [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 |
+| [数据预处理](./preprocessing.md) | 声明式 JSON 驱动数据预处理 |
 
 ### 贡献
 

assets/docs/dataflow.md

@@ -1,237 +1,109 @@
-# 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. For creating preprocessing configs, see [Preprocessing Guide](preprocessing.md).
+
+## Contents
+
+- [Overview](#overview)
+- [Data Preparation](#data-preparation) — tokenization, format detection, backends
+- [Data Keys by Training Type](#data-keys-by-training-type)
+- [Dataset Architecture](#dataset-architecture)
+- [Sampler](#sampler)
+- [DataLoader](#dataloader)
 
 ## 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, callbacks, metric utilities
-- **Inference Module** (`astrai/inference/`): Inference engine with continuous batching, streaming generation
-- **Config Module** (`astrai/config/`): ModelConfig, TrainConfig
-- **Factory Module** (`astrai/factory/`): Registry, BaseFactory for component registration
-- **Parallel Module** (`astrai/parallel/`): Distributed training support
-- **Serialization** (`astrai/serialization.py`): HDF5 data loading, checkpoint management
-
-## Data Flow Diagram
-
-```mermaid
-flowchart LR
-    subgraph A[Data Preparation]
-        direction TB
-        A1[Raw Text] --> A2[AutoTokenizer]
-        A2 --> A3[Tokenized .h5 files]
-        A3 --> A4[BaseDataset]
-        A4 --> A5[ResumableDistributedSampler]
-        A5 --> A6[DataLoader]
-    end
-
-    subgraph B[Training]
-        direction TB
-        B1[DataLoader] --> B2[BaseStrategy]
-        B2 --> B3[Transformer Forward]
-        B3 --> B4[Loss + Backward]
-        B4 --> B5[Gradient Accumulation]
-        B5 -->|every accum_steps| B6[Optimizer Step]
-        B6 --> B7[LR Scheduler]
-        B7 -->|next batch| B2
-        B6 --> B8[CheckpointCallback]
-    end
-
-    subgraph C[Inference]
-        direction TB
-        C1[Checkpoint] --> C2[AutoModel]
-        C1 --> C3[AutoTokenizer]
-        C2 --> C4[InferenceEngine]
-        C3 --> C4
-        C4 --> C5[InferenceScheduler]
-        C5 --> C6[Transformer Forward]
-        C6 --> C7[sample]
-        C7 --> C8{End?}
-        C8 -->|No| C6
-        C8 -->|Yes| C9[Generated Text]
-    end
-
-    A --> B
-    B --> C
+```
+JSONL Lines → Pipeline (mask builder) → Tokenized Tensors
+                                              ↓
+                                      .h5 or .bin storage
+                                              ↓
+                                      Store.load()
+                                              ↓
+                                      Store.fetch(begin, end, keys)
+                                              ↓
+                                      BaseDataset.__getitem__(idx)
+                                              ↓
+                                      Sampler → DataLoader → Training / Inference
 ```
 
-## Detailed Module Descriptions
+## Data Preparation
 
-### 1. Serialization (`astrai/serialization.py`)
+Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
 
-- **`save_h5`**: Saves tensors by groups as HDF5 files (`.h5`), each key maps to a list of tensors
-- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory
-- **`Checkpoint`**: Encapsulates model state dict + epoch + iteration; uses safetensors
+### Tokenization
 
-### 2. Dataset Module
+The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](preprocessing.md)), and produces flat token sequences:
 
-#### 2.1 Dataset (`dataset.py`)
-- **`BaseDataset`**: Abstract base class for windowed sequence sampling
-- **`BaseSegmentFetcher` / `MultiSegmentFetcher`**: Fetch tensor segments by index range
-- **`DatasetFactory`**: Creates dataset instances by `train_type` (`seq`, `sft`, `dpo`, `grpo`)
-- Data keys: `"sequence"` (SEQ), `"loss_mask"` (SFT), `"chosen_mask"/"rejected_mask"` (DPO), `"masks"` (GRPO)
-
-#### 2.2 Sampler (`sampler.py`)
-- **`ResumableDistributedSampler`**: Tracks `epoch` and `iter` for breakpoint resume; supports shuffle and drop_last
-
-### 3. Model Module
-
-#### 3.1 Transformer / AutoModel
-- **`AutoModel`**: Base class with `from_pretrained()` / `save_pretrained()`
-- **`Transformer`**: Decoder-only architecture, registered via `@AutoModel.register('transformer')`
-- Embedding → N×DecoderBlock → RMSNorm → Linear lm_head
-- RoPE position encoding, optional weight tying
-
-#### 3.2 Submodules (`module.py`)
-- **`DecoderBlock`**: GQA attention + residual + MLP + RMSNorm
-- **`GQA`**: Grouped Query Attention (also `MLA` for multi-latent attention)
-- **`MLP`**: `SiLU(gate(x)) * up(x)` → down projection
-- **`RotaryEmbedding`**: RoPE cos/sin cache
-- **`RMSNorm`**: Layer normalization
-
-### 4. Training Module
-
-#### 4.1 Training Context (`train_context.py`)
-- **`TrainContext`**: Dataclass holding model, optimizer, dataloader, strategy, scheduler, checkpoint state
-- **`TrainContextBuilder`**: Builder pattern — takes checkpoint for resume, builds all components
-
-#### 4.2 Trainer (`trainer.py`)
-
-The training loop is nested: **epoch** → **batch** (with step phase interspersed):
-
-```
-on_train_begin
-  on_epoch_begin
-    for each batch:
-      if iteration % accumulation_steps == 0:        ← step phase
-        on_step_begin → optimizer.step() → zero_grad → on_step_end
-                                                      ← batch phase
-      on_batch_begin → strategy(batch) → loss → backward → on_batch_end
-      iteration += 1
-
-    on_epoch_end
-on_train_end
+```python
+# Per JSONL line: messages → chat template → token IDs + loss mask
+tokens = tokenizer.encode(rendered_text)        # List[int]
+loss_mask = [0, 0, 0, 1, 1, 1, 1, 1, 1]        # 0=masked, 1=train
+# Stored as flat tensors, packed with other lines by packing strategy
 ```
 
-Key points:
-- `on_step_*` wraps optimizer step (fires every `accumulation_steps` batches)
-- `on_batch_*` wraps loss computation (fires every batch)
-- `SchedulerCallback` fires on `on_batch_end` — LR scheduler steps every batch
-- `GradientClippingCallback` fires on `on_step_begin`
+The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
 
-#### 4.3 Strategy (`strategy.py`)
-- **`SEQStrategy`**: Next-token prediction, cross-entropy with label smoothing
-- **`SFTStrategy`**: Supervised fine-tuning with loss masking
-- **`DPOStrategy`**: Direct Preference Optimization with reference model
-- **`GRPOStrategy`**: Group Relative Policy Optimization with clipped ratio
+### Format Detection
 
-#### 4.4 Scheduler (`schedule.py`)
-- **`CosineScheduler`**: Cosine decay + linear warmup
-- **`SGDRScheduler`**: Cosine annealing with warm restarts
-- Created by `SchedulerFactory` and bound to optimizer
+`detect_format(load_path)` inspects the path:
 
-#### 4.5 Callbacks
-- **`CheckpointCallback`**: Saves safetensors at `ckpt_interval` iterations
-- **`ProgressBarCallback`**: tqdm progress display
-- **`MetricLoggerCallback`**: Writes JSONL metrics to `{ckpt_dir}/logs/`
-- **`GradientClippingCallback`**: `clip_grad_norm_` on `on_step_begin`
-- **`SchedulerCallback`**: `scheduler.step()` on `on_batch_end`
+- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, unknown suffix raises `ValueError`
+- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, or `*.bin` + `**/meta.json` → `"bin"`
 
-### 5. Inference Module
+### Store Backends
 
-#### 5.1 Inference Engine (`engine.py`)
-- **`InferenceEngine`**: Facade over scheduler; provides `generate()`, `generate_with_request()`, `generate_async()`
-- Accepts `prompt: str | List[str]`, returns generator (stream) or string (non-stream)
-
-#### 5.2 Scheduler 4-Phase Loop (`scheduler.py`)
-
-Background thread runs continuously:
+Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
 
 ```
-1. Cleanup → Remove finished tasks, free KV cache pages
-2. Refill  → Pop from waiting_queue, alloc pages, add to active
-3. Prefill → Group active tasks by prompt_len, run full forward pass
-4. Decode  → Pick largest same-position group, run single-token forward
+StoreFactory.create("h5")  → H5Store
+StoreFactory.create("bin") → MmapStore
 ```
 
-- **`Task`**: Tracks prompt_ids, output_ids, page_table, status (PENDING/RUNNING/FINISHED/ABORTED)
-- **`PagedCache`**: Bitmask-based page allocator with page-table-indirected read/write
-- **`CacheView`**: Batch view bundling cache + page table for attention layers
-- **`sample()`**: Temperature → top-k → top-p → multinomial
+**H5Store**: Reads HDF5 files, supports `share_memory_()` for multi-process DataLoader workers (copies tensors to shared memory).
 
-#### 5.3 Server (`server.py`)
-- FastAPI with OpenAI `/v1/chat/completions` and Anthropic `/v1/messages` endpoints
-- Streaming via SSE, health check at `/health`, stats at `/stats`
+**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
 
-### 6. Tokenizer Module
+Both backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based indexing).
 
-- **`AutoTokenizer`**: Wraps HuggingFace tokenizers (BBPE); `encode`/`decode`/`apply_chat_template`
-- **`ChatTemplate`**: Jinja2-based template rendering for multi-turn chat
+## Data Keys by Training Type
 
-### 7. Factory & Parallel
+| Type | Storage Keys |
+|------|-------------|
+| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) |
+| `sft` | `sequence`, `loss_mask`, `position_ids` |
+| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` |
+| `grpo` | `prompts`, `responses`, `masks`, `rewards` |
 
-- **`Registry` / `BaseFactory`**: Decorator-based component registration
-- **`spawn_parallel_fn`**: Multi-process DDP launcher with NCCL backend
-- **`ParallelModel` / `ColumnParallelLinear` / `RowParallelLinear`**: Tensor model parallelism
+## Dataset Architecture
 
-## Training Data Flow — Detailed Steps
+```
+DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_type=None)
+  → BaseDataset.load(load_path, storage_type=None)
+    → detect_format(load_path)
+    → StoreFactory.create(storage_type)
+    → Store.load(load_path)
+      → _normalize(raw)  # base Store, shared by both backends
+        → Store._data[Dict[str, List[Tensor]]] + _cum[Dict[str, List[int]]]
+          → BaseDataset.__getitem__(idx)
+            → get_index(idx) → [begin, end)
+            → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
+```
 
-1. **Data Preparation**
-   - Raw text → token IDs via `AutoTokenizer.encode()`
-   - Save as `.h5` files (groups of tensor lists per data key)
+`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). `storage_type` defaults to `None` (auto-detect via `detect_format`).
 
-2. **Dataset Loading**
-   - `BaseDataset.load()` calls `load_h5()`, builds `MultiSegmentFetcher`
-   - Sliding window of `window_size` with `stride` determines sample boundaries
+`Store.fetch(begin, end, keys)` accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
 
-3. **Sampling & Batching**
-   - `ResumableDistributedSampler` produces shuffled index sequences
-   - `DataLoader` fetches `[batch_size, window_size]` tensors via `__getitem__`
+## Sampler
 
-4. **Strategy Forward**
-   - Strategy receives batch, calls `Transformer.forward()` for logits
-   - Computes task-specific loss (cross-entropy, DPO, GRPO)
+`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
 
-5. **Backward & Accumulation**
-   - `loss = raw_loss / accumulation_steps`
-   - `loss.backward()` accumulates gradients
-   - Every `accumulation_steps` batches: `optimizer.step()` → `zero_grad()`
-   - Every batch: `scheduler.step()` updates learning rate
+- Tracks `start_epoch` / `start_iter` for resume
+- Shuffle via `torch.Generator(seed + epoch)`
+- Per-replica index slicing for DDP
 
-6. **Checkpoint**
-   - `CheckpointCallback` saves `model.state_dict()` + metadata to safetensors at `ckpt_interval` iterations
-   - Does NOT save optimizer/scheduler state (resume resets those)
+## DataLoader
 
-## Inference Data Flow — Detailed Steps
+Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
 
-1. **Model Loading**
-   - `AutoModel.from_pretrained(path)` loads weights from safetensors
-   - `torch.inference_mode()` wraps generation
-
-2. **Prompt Construction**
-   - Messages → `apply_chat_template(messages, tokenize=False)` → prompt string
-   - `tokenizer.encode(prompt)` → token IDs (truncated to `max_prompt_len`)
-
-3. **Continuous Batching Loop**
-   - **Cleanup**: Finished tasks → `stream_callback(STOP)`, free KV pages
-   - **Refill**: Pop from waiting queue, `PagedCache.alloc_n()` for prompt pages
-   - **Prefill**: Group by prompt length, run full forward with `start_pos=0`
-   - **Decode**: Pick position group with most tasks, single-token forward:
-     - Model forward → `logits` → `sample()` → next token ID
-     - Append to `output_ids`, update `output_tokens`
-     - `_maybe_alloc_page()` grows page table as needed
-     - `stream_callback(token)` for streaming clients
-
-4. **Output**
-   - `tokenizer.decode(output_ids)` → text
-   - Return to caller (streaming: token-by-token; non-streaming: complete string)
-
-## Checkpoint & Serialization
-
-- **Training Checkpoint**: safetensors weights + epoch/iteration metadata. Optimizer/scheduler state is NOT persisted.
-- **Inference Loading**: `AutoModel.from_pretrained()` loads from the same safetensors format.
-- **Dataset Serialization**: HDF5 with shared memory support for large-scale pre-training data.
-
-> Document Update Time: 2026-05-09
+> Document Update Time: 2026-06-19

assets/docs/design.md

@@ -1,719 +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
-            +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 epoch
-            +int iter
-        }
-
-        class DatasetFactory {
-            +Registry _registry
-            +register(name) decorator
-            +create(train_type, window_size, stride) BaseDataset
-            +load(train_type, load_path, window_size, stride) BaseDataset
-        }
-    }
-
-    namespace serialization {
-        class Checkpoint {
-            +dict state_dict
-            +int epoch
-            +int iteration
-            +save(save_dir)
-            +load(save_dir) Checkpoint
-        }
-    }
-
-    namespace model {
-        class AutoModel {
-            +ModelConfig config
-            +Registry _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, paged_cache, 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, paged_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, paged_cache, start_pos) Tensor
-        }
-
-        class MLA {
-            +int n_heads
-            +int n_kv_heads
-            +int head_dim
-            +int kv_lora_rank
-            +int qk_nope_head_dim
-            +int qk_rope_head_dim
-            +Linear q_proj, kv_a_proj, kv_b_proj
-            +Linear o_proj
-            +RMSNorm kv_norm
-            +forward(x, rotary_emb, mask, paged_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[int] 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, system_prompt, **extra_variables) 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
-            +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
-            +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
-            +generate_with_request(request) Union[Generator, str, List[str]]
-            +generate_async(prompt, max_tokens, temperature, top_p, top_k) AsyncGenerator
-            +get_stats() Dict
-            +shutdown()
-        }
-
-        class InferenceScheduler {
-            +nn.Module model
-            +AutoTokenizer tokenizer
-            +PagedCache page_cache
-            +int max_batch_size
-            +int max_seq_len
-            +int max_prompt_len
-            +int page_size
-            +List waiting_queue
-            +List active_tasks
-            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
-            +remove_task(task_id)
-            +start()
-            +stop()
-            +get_stats() Dict
-        }
-
-        class PagedCache {
-            +int page_size
-            +int _free_mask
-            +List[int] _refs
-            +Tensor k_cache
-            +Tensor v_cache
-            +alloc() int
-            +alloc_n(n) List[int]
-            +free(idx)
-            +bind(page_table, total_len) CacheView
-            +write(layer_id, page_table, start_pos, k, v)
-            +gather(layer_id, page_table) Tuple[Tensor, Tensor]
-        }
-
-        class CacheView {
-            +PagedCache _cache
-            +Tensor _page_table
-            +int _total_len
-            +write(layer_id, start_pos, k, v)
-            +gather(layer_id) Tuple[Tensor, Tensor]
-        }
-
-        class Task {
-            +str task_id
-            +List prompt_ids
-            +int max_tokens
-            +float temperature
-            +float top_p
-            +int top_k
-            +TaskStatus status
-            +List output_ids
-            +int input_tokens
-            +int output_tokens
-            +List[int] page_table
-            +int n_pages
-            +float arrival_time
-            +float finish_time
-            +Callable stream_callback
-            +int next_pos
-            +is_finished(stop_ids) bool
-        }
-
-        class TaskStatus {
-            <<enumeration>>
-            PENDING
-            RUNNING
-            FINISHED
-            ABORTED
-        }
-
-        class GenerationRequest {
-            +List[Dict] messages
-            +GenerationParams params
-            +bool stream
-        }
-
-        class GenerationParams {
-            <<value object>>
-            +int top_k
-            +float top_p
-            +float temperature
-            +int max_tokens
-        }
-
-        class BaseSamplingStrategy {
-            <<abstract>>
-            +apply(logits, filter_value) Tensor
-        }
-
-        class TemperatureStrategy {
-            +float temperature
-            +apply(logits, filter_value) Tensor
-        }
-
-        class TopKStrategy {
-            +int top_k
-            +apply(logits, filter_value) Tensor
-        }
-
-        class TopPStrategy {
-            +float top_p
-            +apply(logits, filter_value) Tensor
-        }
-
-        class SamplingPipeline {
-            +List strategies
-            +apply(logits, filter_value) Tensor
-            +sample(logits, filter_value) Tensor
-        }
-
-        class _Result {
-            +List[str] tokens
-            +List[str] results
-            +List[bool] _done
-            +append(token, idx)
-            +get_results() List[str]
-            +pop_all() List[str]
-            +wait(timeout) bool
-        }
-
-        class ChatMessage {
-            +str role
-            +str content
-        }
-
-        class ChatCompletionRequest {
-            +List[ChatMessage] messages
-            +float temperature
-            +float top_p
-            +int top_k
-            +int max_tokens
-            +bool stream
-            +Optional[str] stop
-            +Optional[int] n
-        }
-    }
-
-    namespace parallel {
-        class ParallelFunctions {
-            +spawn_parallel_fn(fn, nprocs)
-            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type)
-        }
-
-        class ParallelModel {
-            +dist.ProcessGroup process_group
-            +int rank
-            +int world_size
-        }
-
-        class ColumnParallelLinear {
-            +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
-    Checkpoint ..> Checkpoint : saves/loads
-    TrainContext --> Checkpoint : manages
-    TrainContext --> BaseStrategy : uses
-    TrainContext --> BaseScheduler : uses
-    SchedulerFactory ..> BaseScheduler : creates
-    BaseScheduler <|-- CosineScheduler
-    BaseScheduler <|-- SGDRScheduler
-    CallbackFactory ..> TrainCallback : creates
-    TrainCallback <|-- GradientClippingCallback
-    TrainCallback <|-- SchedulerCallback
-    TrainCallback <|-- CheckpointCallback
-    TrainCallback <|-- ProgressBarCallback
-    TrainCallback <|-- MetricLoggerCallback
-    InferenceEngine --> InferenceScheduler : uses
-    InferenceEngine --> GenerationRequest : uses
-    GenerationRequest --> GenerationParams : contains
-    InferenceScheduler --> Task : manages
-    Task --> TaskStatus : uses
-    InferenceScheduler --> TaskStatus : uses
-    InferenceScheduler --> PagedCache : uses
-    InferenceScheduler --> Transformer : uses
-    InferenceEngine --> Transformer : uses
-    InferenceEngine --> _Result : uses
-    BaseSamplingStrategy <|-- TemperatureStrategy
-    BaseSamplingStrategy <|-- TopKStrategy
-    BaseSamplingStrategy <|-- TopPStrategy
-    SamplingPipeline --> BaseSamplingStrategy : composes
-    BaseDataset <|-- SEQDataset
-    BaseDataset <|-- SFTDataset
-    BaseDataset <|-- DPODataset
-    BaseDataset <|-- GRPODataset
-    DatasetFactory ..> BaseDataset : creates
-    MultiSegmentFetcher --> BaseSegmentFetcher : uses
-    BaseDataset --> MultiSegmentFetcher : uses
-    AutoModel <|-- Transformer
-    AutoModel --> ModelConfig : contains
-    Transformer --> DecoderBlock : uses
-    Transformer --> RotaryEmbedding : uses
-    Transformer --> Embedding : uses
-    DecoderBlock --> GQA : uses
-    DecoderBlock --> MLP : uses
-    DecoderBlock --> RMSNorm : uses
-    TrainContextBuilder --> ResumableDistributedSampler : creates
-    ResumableDistributedSampler --> BaseDataset : samples
-    ParallelModel <|-- RowParallelLinear
-    ParallelModel <|-- ColumnParallelLinear
-    AutoTokenizer --> ChatTemplate : uses
-    TrainConfig --> DatasetFactory : selects
-    TrainConfig --> SchedulerFactory : selects
-    TrainConfig --> CallbackFactory : selects
-    AutoModel ..> AutoTokenizer : loads with
-    BaseFactory <|-- DatasetFactory
-    BaseFactory <|-- StrategyFactory
-    BaseFactory <|-- SchedulerFactory
-    BaseFactory <|-- CallbackFactory
-```
-
-### Module Overview
-
-| Module | Components | Description |
-|--------|------------|-------------|
-| **astrai.config** | ModelConfig, TrainConfig | Configuration management |
-| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
-| **astrai.serialization** | Checkpoint, save_h5, load_h5 | Model serialization and checkpoint management |
-| **astrai.model** | AutoModel, Transformer, DecoderBlock, GQA, MLA, MLP, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
-| **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
-| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy, StrategyFactory, BaseScheduler, SchedulerFactory, TrainCallback, CallbackFactory | Training workflow management |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, PagedCache, CacheView, Task, TaskStatus, GenerationParams, GenerationRequest, BaseSamplingStrategy, TemperatureStrategy, TopKStrategy, TopPStrategy, SamplingPipeline, ChatMessage, ChatCompletionRequest | Inference service with continuous batching and paged KV cache |
-| **astrai.parallel** | ParallelFunctions, ParallelModel, 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) |
-| **Context** | `TrainContext` | Training process state container with model, optimizer, scheduler and checkpoint |
-| **Registry** | `BaseFactory`, `Registry` | Generic component registration with category and priority support |
-| **Object Pool** | `PagedCache` | Page-based KV cache with O(1) alloc/free via bitmask |
-| **Strategy (Sampling)** | `BaseSamplingStrategy`, `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations with temperature, top-k, top-p |
-| **Producer-Consumer** | `InferenceScheduler`, `Task`, `waiting_queue`, `active_tasks` | Continuous batching with dynamic task queue management |
-| **Event-Driven** | `threading.Event`, `_task_event` | Non-blocking wait mechanism for task scheduling using Python's `threading` module |
-| **AutoModel Registry** | `AutoModel`, `Transformer` | Model type registration and dynamic loading via decorator pattern |
-| **Generator Pattern** | `_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**: `InferenceEngine` → `InferenceScheduler` → `Transformer`, uses `PagedCache` for paged KV cache management and `SamplingPipeline` for efficient continuous batching with streaming/non-streaming
-5. **Distributed Support**: `spawn_parallel_fn` and `setup_parallel` provide multi-process training capability for `Trainer`
-6. **Dataset Loading**: `DatasetFactory` creates datasets (SEQDataset, SFTDataset, DPODataset, GRPODataset), supports HDF5 loading via `BaseSegmentFetcher` and `MultiSegmentFetcher`
-7. **Checkpoint Management**: `Checkpoint` handles model state serialization/deserialization with safetensors
-8. **Scheduler Support**: `SchedulerFactory` creates learning rate schedulers (CosineScheduler, SGDRScheduler)
-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}
-$$
-
-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,241 @@
+# Inference
+
+## Contents
+
+- [KV Cache](#kv-cache)
+- [KVCache System](#kvcache-system)
+- [Continuous Batching](#continuous-batching)
+- [Sampling](#sampling-strategy-pattern)
+- [Protocol Handlers](#protocol-handlers-strategy-pattern)
+- [Engine & GenerateResult](#engine--generateresult)
+- [HTTP API](#http-api) — endpoints, SSE, errors, stats
+- [Engine API](#engine-api)
+
+## 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
+
+Seven classes working together:
+
+```
+KVCache (facade)
+  ├── PagePool         orchestrates page allocation + prefix matching
+  │     ├── Allocator   bitmask-based page allocator + ref-count + LRU eviction (inside PagePool)
+  │     └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) (inside PagePool)
+  ├── 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 (returned by bind())
+```
+
+`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 (ABC)
+  ├── TemperatureStrategy
+  ├── TopKStrategy
+  ├── TopPStrategy
+  └── SamplingPipeline
+```
+
+`SamplingPipeline` composes them: Temperature → Top-K → Top-P → softmax → multinomial.  
+`sample()` is a convenience shortcut for one-shot usage.
+
+## Protocol Handlers (Strategy Pattern)
+
+```python
+class ProtocolHandler:  # concrete orchestrator
+    def __init__(self, request, engine, builder): ...
+    async def handle(self):
+        prompt, ctx, stops = builder.prepare(request, engine)
+        agen = engine.generate_async(prompt, ...)
+        if stream: self._handle_stream(agen, ctx, stops)
+        else:      return await self._handle_non_stream(agen, ctx, stops)
+```
+
+`ResponseBuilder` (ABC): `prepare()`, `format_stream_start()`, `format_chunk()`, `format_stream_end()`, `format_response()`.
+
+`OpenAIResponseBuilder` → `/v1/chat/completions`, `AnthropicResponseBuilder` → `/v1/messages`.
+
+Adding a protocol = one builder file, no handler subclassing needed.
+
+## Engine & GenerateResult
+
+```
+InferenceEngine
+  ├── generate(prompt, stream, ...) → str | List[str] | Generator
+  ├── generate_with_request(req)    → same
+  ├── generate_async(prompt, ...)   → AsyncGenerator
+  ├── get_stats()                   → Dict
+  └── shutdown()
+```
+
+`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",
+  "created": 1717000000,
+  "model": "astrai",
+  "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"}],
+  "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}
+}
+```
+
+Streaming SSE: `object: "chat.completion.chunk"` — starts with role delta, then token chunks, ends with finish chunk + usage stats, then `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) |
+| `top_k` | int | 50 | Top-k count |
+| `top_p` | float | 1.0 | Nucleus threshold |
+| `temperature` | float | 1.0 | Sampling temperature (> 0.0) |
+| `max_tokens` | Optional[int] | None | Max generation length |
+| `stream` | bool | False | Stream output |
+
+### SSE Streaming Format
+
+**OpenAI** (`/v1/chat/completions`, `stream=true`):
+
+```
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai",
+       "choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":0,"model":"astrai",
+       "choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai",
+       "choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
+
+data: {"prompt_tokens":5,"completion_tokens":1,"total_tokens":6}
+
+data: [DONE]
+```
+
+**Anthropic** (`/v1/messages`, `stream=true`):
+
+```
+event: message_start
+data: {"type":"message_start","message":{"id":"msg_...","model":"astrai","role":"assistant",
+       "content":[],"usage":{"input_tokens":0}}}
+
+event: content_block_start
+data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
+
+event: content_block_stop
+data: {"type":"content_block_stop","index":0}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{...}}
+
+event: message_stop
+data: {"type":"message_stop"}
+```
+
+### Error Responses
+
+The server returns standard HTTP status codes. Pydantic validation errors (e.g. missing required fields)
+are handled automatically by FastAPI with 422 status. The only application-level error is engine initialization:
+
+| Status | Meaning |
+|--------|---------|
+| 200 | Success |
+| 422 | Unprocessable entity (Pydantic validation) |
+| 503 | Service unavailable (model not loaded, engine not ready) |
+
+Error response body (503):
+
+```json
+{
+    "detail": "Engine not initialized"
+}
+```
+
+### Stats Endpoint
+
+```
+GET /stats
+```
+
+Response:
+
+```json
+{
+    "total_tasks": 128,
+    "total_tokens": 10240,
+    "active_tasks": 3,
+    "waiting_queue": 2
+}
+```
+
+## 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
+async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[str]
+    print(token)
+```
+
+> Document Update Time: 2026-06-19

assets/docs/introduction.md

@@ -1,334 +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 24 layers of Transformer blocks, with 1.0 billion parameters. Transformer is an autoregressive model that calculates the relationship between all previous tokens to obtain the probability distribution of the next token.
-
-The model now uses the **AutoModel** base class for flexible loading and saving:
-
-```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 (gate)"]
-        U --> W["Linear (up)"]
-        V --> X[SiLU]
-        X --> Y[×]
-        W --> Y
-        Y --> Z["Linear (down)"]
-        Z --> AA[+]
-        T --> AA
-        AA --> BB[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,
-)
-
-# 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 | 1.0 | Sampling temperature (0.0-2.0) |
-| `top_p` | float | 1.0 | Nucleus sampling threshold |
-| `top_k` | int | 50 | Top-k sampling parameter |
-| `max_tokens` | int | 1024 | Maximum tokens to generate |
-| `stream` | bool | false | Enable streaming response |
-
-**Response (non-streaming):**
-```json
-{
-  "id": "chatcmpl-1234567890",
-  "object": "chat.completion",
-  "created": 1234567890,
-  "model": "astrai",
-  "choices": [
-    {
-      "index": 0,
-      "message": {"role": "assistant", "content": "Hello! I'm doing well..."},
-      "finish_reason": "stop"
-    }
-  ],
-  "usage": {
-    "prompt_tokens": 20,
-    "completion_tokens": 15,
-    "total_tokens": 35
-  }
-}
-```
-
-### Streaming Response
-
-Enable streaming for real-time token-by-token output:
-
-```bash
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "Write a story"}],
-    "stream": true,
-    "max_tokens": 500
-  }'
-```
-
-The server uses Server-Sent Events (SSE) with content type `text/event-stream`.
-
-### Anthropic-Compatible Endpoint
-
-The server also provides an Anthropic-compatible endpoint at `/v1/messages`:
-
-```bash
-curl -X POST http://localhost:8000/v1/messages \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "You are a helpful assistant.",
-    "messages": [{"role": "user", "content": "Hello, how are you?"}],
-    "max_tokens": 2048
-  }'
-```
-
-Response:
-```json
-{
-  "id": "msg_abc123...",
-  "type": "message",
-  "role": "assistant",
-  "model": "astrai",
-  "content": [{"type": "text", "text": "Hello! I am doing well..."}],
-  "stop_reason": "end_turn",
-  "stop_sequence": null,
-  "usage": {"input_tokens": 20, "output_tokens": 15}
-}
-```
-
-Streaming:
-```bash
-curl -X POST http://localhost:8000/v1/messages \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "You are a helpful assistant.",
-    "messages": [{"role": "user", "content": "Write a short poem"}],
-    "max_tokens": 500,
-    "stream": true
-  }'
-```
-
-Supports `stop_sequences` for early termination:
-```bash
-curl -X POST http://localhost:8000/v1/messages \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "messages": [{"role": "user", "content": "Write a story"}],
-    "max_tokens": 500,
-    "stop_sequences": ["The end", "THE END"]
-  }'
-```
-
-### Health Check
-
-Monitor server and model status:
-
-```bash
-curl http://localhost:8000/health
-# {"status": "ok", "model_loaded": true}
-
-curl http://localhost:8000/stats
-# {"total_tasks": 10, "total_tokens": 5000, "active_tasks": 1, "waiting_queue": 0}
-```
-
-> Document Update Time: 2026-04-09

assets/docs/params.md

@@ -1,4 +1,11 @@
-# Parameter Documentation
+# CLI Parameter Reference
+
+## Contents
+
+- [Training Parameters](#training-parameters)
+- [Inference Server](#inference-server-serverpy)
+- [Generate](#generate-generatepy)
+- [Preprocess](#preprocess-preprocesspy)
 
 ## Training Parameters
 
@@ -10,14 +17,14 @@
 | `--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 | 1 |
-| `--accumulation_steps` | Gradient accumulation steps between optimizer 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 |
 |-----------|-------------|---------|
-| `--warmup_steps` | Warmup steps | 1000 |
+| `--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 |
 
@@ -46,113 +53,151 @@
 | `--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 |
+| `--start_samples` | Resume from sample count per rank | 0 |
+
+### Validation
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--val_split` | Ratio to split from training dataset for validation (e.g. 0.05) | None |
+| `--val_step` | Number of optimizer steps between validation runs | 1000 |
+
+### Logging
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--log_dir` | Directory for metric logs | checkpoint/logs |
+| `--log_interval` | Number of optimizer steps between metric logs | 1 |
+| `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr", "grad_norm"] |
+
+### Gradient Checkpointing
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--gradient_checkpointing` | Enable activation checkpointing for DecoderBlock modules | False |
 
 ### Distributed Training
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--nprocs` | Number of GPUs / processes | 1 |
+| `--parallel_mode` | Parallel strategy (`none`, `ddp`, or `fsdp`) | none |
 | `--device_type` | Device type | cuda |
+| `--start_method` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) | spawn |
+| `--backend` | Distributed training backend | nccl |
+| `--master_addr` | Master node address | localhost |
+| `--master_port` | Master node port | 29500 |
 
 ### Strategy-specific
 
 | Parameter | Description | Default | Used by |
 |-----------|-------------|---------|---------|
 | `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
-| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.1 | `seq`, `sft` |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `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` |
+| `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` |
+
+### Scheduler
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine |
+| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default) |
+| `--cycle_length` | SGDR first cycle length in steps | None (total_steps - warmup_steps) |
+| `--t_mult` | SGDR cycle length multiplier per restart | 2 |
+| `--stable_steps` | WSD stable plateau steps | None (required for wsd) |
+| `--decay_steps` | WSD decay steps | None (total_steps - warmup_steps - stable_steps) |
 
 ### Usage Example
 
 ```bash
-python scripts/tools/train.py \
-  --train_type seq \
-  --data_root_path /path/to/dataset \
-  --param_path /path/to/model \
-  --n_epoch 3 \
-  --batch_size 4 \
-  --accumulation_steps 8 \
-  --max_lr 3e-4 \
-  --warmup_steps 2000 \
-  --max_grad_norm 1.0 \
-  --ckpt_interval 5000 \
-  --ckpt_dir ./checkpoints \
-  --num_workers 4 \
-  --nprocs 1 \
-  --device_type cuda
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --parallel_mode=ddp \
+    --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 Parameters
+## Inference Server (`server.py`)
 
-### GenerationRequest Parameters
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--host` | str | `0.0.0.0` | Host address |
+| `--port` | int | `8000` | Port number |
+| `--param_path` | path | `project_root/params` | Path to model parameters |
+| `--device` | str | `cuda` | Device to load model on |
+| `--dtype` | str | `bfloat16` | Model weights dtype (`bfloat16`, `float16`, `float32`) |
+| `--max_batch_size` | int | `16` | Maximum batch size for continuous batching |
+| `--reload` | flag | `False` | Enable auto-reload for development |
 
-| 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 |
-
-### Usage Example
-
-```python
-import torch
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
-from astrai.inference import InferenceEngine, GenerationRequest
-
-# Load model using AutoModel
-model = AutoModel.from_pretrained("your_model_dir")
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("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,
-)
+Usage:
+```bash
+python scripts/tools/server.py --param_path ./params --device cuda --dtype bfloat16
 ```
 
-### Generation Modes
+See [Inference Guide](inference.md) for HTTP API documentation.
 
-| Mode | Description |
-|------|-------------|
-| `stream=True` | Streaming output, yields token by token |
-| `stream=False` | Non-streaming output, returns complete result |
+## Generate (`generate.py`)
 
-> Document Update Time: 2026-04-09
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--param_path` | str | required | Path to the model directory |
+| `--input_json_file` | str | required | Path to the input JSONL file |
+| `--output_json_file` | str | required | Path to the output JSONL file |
+| `--question_key` | str | `question` | Key for the question in input JSON |
+| `--response_key` | str | `response` | Key for the response in output JSON |
+| `--temperature` | float | `0.60` | Sampling temperature |
+| `--top_k` | int | `30` | Top-k filtering |
+| `--top_p` | float | `0.95` | Nucleus sampling threshold |
+| `--batch_size` | int | `1` | Batch size for generation |
+| `--max_tokens` | int | model config `max_len` | Maximum tokens to generate |
+
+Usage:
+```bash
+python scripts/tools/generate.py \
+    --param_path ./params \
+    --input_json_file input.jsonl \
+    --output_json_file output.jsonl
+```
+
+## Preprocess (`preprocess.py`)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) |
+| `--output_dir`, `-o` | path | required | Output directory for processed data |
+| `--config`, `-c` | path | required | Preprocessing pipeline config (JSON) |
+| `--tokenizer_path` | str | `params` | Path to tokenizer directory |
+
+Usage:
+```bash
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c sft.json
+```
+
+See [Preprocessing Guide](preprocessing.md) for config file format and examples.
+
+---
+
+> Document Update Time: 2026-06-19

assets/docs/preprocessing.md

@@ -0,0 +1,364 @@
+# Preprocessing Pipeline
+
+Declarative JSON-driven data preprocessing. One `SectionedMaskBuilder` handles all formats via `input.sections` (single-output) or `input.sources` (multi-output).
+
+## Contents
+
+- [Philosophy](#philosophy)
+- [Config Structure](#config-structure)
+- [Quick Start](#quick-start) — SFT Chat, SFT Instruction, Pretrain, DPO, GRPO examples
+- [Configuration Reference](#configuration-reference) — all fields
+- [Mask Algorithm](#mask-algorithm)
+- [Output Layout](#output-layout)
+- [CLI](#cli)
+- [Python API](#python-api)
+
+## Philosophy
+
+| Component | Responsibility |
+|-----------|---------------|
+| `tokenizer_config.json` (`chat_template`) | Formatting -- how roles become tokens |
+| `pipeline.json` (`mask`) | Masking -- which roles participate in training |
+
+A single config file captures the entire pipeline, reusable and version-controllable.
+
+## Config Structure
+
+```json
+{
+  "version":       1,
+  "input":         {},   // sections (single) or sources (multi)
+  "mask":          {},   // role -> "train" | "mask"
+  "mask_default":  "mask",
+  "preprocessing": {},
+  "output":        {}
+}
+```
+
+### Section Fields
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `field` | str | -- | JSONL key to read |
+| `action` | str | -- | `"train"` / `"mask"` / `"$role"` |
+| `template` | bool | `false` | Apply `chat_template` per message |
+| `add_special_tokens` | bool | `true` for first non-template section | Add special tokens during encode |
+
+### Source Fields (multi-output mode)
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `sections` | list[dict] | -- | Same as single-output section list |
+| `list_field` | bool | `false` | JSONL field holds a list; tokenise each element |
+| `mask_key` | str | `"{key}_mask"` | Explicit output key for loss mask |
+
+---
+
+## Quick Start
+
+### SFT Chat
+
+Input JSONL:
+
+```json
+{"messages": [{"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sections": [
+      {"field": "messages", "action": "$role", "template": true}
+    ]
+  },
+  "mask": {
+    "system": "mask",
+    "user": "mask",
+    "assistant": "train"
+  },
+  "mask_default": "mask",
+  "preprocessing": {
+    "max_seq_len": 2048
+  },
+  "output": {
+    "storage_format": "bin",
+    "dtype": {"loss_mask": "bool"}
+  }
+}
+```
+
+Output keys: `sequence` (int32), `loss_mask` (bool)
+
+### SFT Instruction
+
+Input JSONL:
+
+```json
+{"prompt": "Translate to French: Hello", "response": "Bonjour"}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sections": [
+      {"field": "prompt",   "action": "mask", "add_special_tokens": true},
+      {"field": "response", "action": "train"}
+    ]
+  },
+  "mask_default": "mask",
+  "preprocessing": {
+    "max_seq_len": 2048
+  }
+}
+```
+
+Output keys: `sequence`, `loss_mask`
+
+### Pretrain
+
+Input JSONL:
+
+```json
+{"text": "Artificial Intelligence is a field of computer science..."}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sections": [
+      {"field": "text", "action": "train"}
+    ]
+  },
+  "preprocessing": {
+    "max_seq_len": 8192,
+    "min_chars": 100
+  }
+}
+```
+
+Output keys: `sequence` (no `loss_mask` — all tokens trained)
+
+### DPO
+
+Input JSONL:
+
+```json
+{"chosen": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"}], "rejected": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "5"}]}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sources": {
+      "chosen": {
+        "sections": [
+          {"field": "chosen", "action": "$role", "template": true}
+        ]
+      },
+      "rejected": {
+        "sections": [
+          {"field": "rejected", "action": "$role", "template": true}
+        ]
+      }
+    }
+  },
+  "mask": {
+    "user": "mask",
+    "assistant": "train"
+  },
+  "mask_default": "mask"
+}
+```
+
+Output keys: `chosen`, `chosen_mask`, `rejected`, `rejected_mask`
+
+### GRPO
+
+Input JSONL:
+
+```json
+{"prompt": [{"role": "user", "content": "What is 2+2?"}], "responses": ["4", "Five", "Four"], "rewards": [1.0, 0.3, 0.8]}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sources": {
+      "prompts": {
+        "sections": [
+          {"field": "prompt", "action": "mask", "template": true}
+        ]
+      },
+      "responses": {
+        "sections": [
+          {"field": "responses", "action": "train"}
+        ],
+        "list_field": true,
+        "mask_key": "masks"
+      },
+      "rewards": {
+        "sections": [
+          {"field": "rewards", "action": "value"}
+        ]
+      }
+    }
+  },
+  "mask": {
+    "user": "mask",
+    "assistant": "train"
+  },
+  "mask_default": "mask"
+}
+```
+
+Output keys: `prompts`, `prompts_mask`, `responses`, `masks`, `rewards` (float32)
+
+- `action: "value"` — extract raw values from JSONL without tokenisation
+- `list_field: true` — tokenise each list element independently, then concatenate
+- `mask_key: "masks"` — rename the auto-generated mask key (default: `responses_mask`)
+- `prompts_mask` is auto-generated (all masked) and unused by GRPOStrategy
+
+---
+
+## Configuration Reference
+
+### `input`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `sections` | list[dict] or null | `null` | Section specs for single-output mode |
+| `sources` | dict[str, dict] or null | `null` | Source specs for multi-output mode (DPO/GRPO) |
+
+When `sources` is set, `sections` is ignored.
+
+### `mask`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `mask` | dict | `{}` | `{role: "train" \| "mask"}` |
+| `mask_default` | str | `"mask"` | Default action for unlisted roles |
+
+### `preprocessing`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `max_seq_len` | int | `2048` | Truncate sequences to this length |
+| `min_chars` | int | `50` | Skip text-mode items shorter than this |
+| `max_chars` | int | `2000000` | Skip text-mode items longer than this |
+| `max_items` | int or null | `null` | Stop after N documents |
+| `packing_strategy` | str | `"simple"` | Packing strategy: `"simple"`, `"bfd"`, `"bfd_split"` |
+| `max_packed_len` | int | `8192` | Maximum length of a packed bin |
+| `truncation_mode` | str | `"keep_start"` | How to truncate sequences: `"keep_start"` or `"keep_end"` |
+
+### `output`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `domain_key` | str or null | `null` | JSONL key for domain grouping |
+| `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` |
+| `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens |
+| `dtype` | dict[str, str] | `{}` | Per-key tensor dtype override (e.g. `{"loss_mask": "bool"}`) |
+| `position_ids_mode` | str | `"doc_reset"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` |
+
+---
+
+## Mask Algorithm
+
+### Template mode (`template: true`)
+
+1. Prepend BOS token (masked)
+2. For each message in the field's array:
+   1. Render through `chat_template` for that single message
+   2. Encode rendered text
+   3. Apply mask rule for the message's role
+
+### Non-template mode
+
+Encode the field value as text. Mask value is 1 (train) or 0 (mask) per the section's `action`.
+
+### Text config detection
+
+When no section uses `template` and all sections have `action: "train"`, the builder omits `loss_mask` from the output — all tokens are trained.
+
+---
+
+## Output Layout
+
+### Single-Shard (`bin`)
+
+```
+output/
+  __default__/
+    shard_0000/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+  wiki/
+    shard_0000/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+```
+
+### Multi-Shard (`bin`)
+
+When `max_tokens_per_shard` is exceeded:
+
+```
+output/
+  __default__/
+    shard_0000/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+    shard_0001/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+```
+
+For `bin` format, `MmapStore` discovers all shards under the domain directory via `rglob("meta.json")`. For `h5` format, `H5Store` discovers `.h5`/`.hdf5` files via recursive glob.
+
+---
+
+## CLI
+
+```bash
+# SFT
+python scripts/tools/preprocess.py data/sft/*.jsonl -o output/sft/ -c configs/sft_chat.json
+
+# DPO
+python scripts/tools/preprocess.py data/dpo/*.jsonl -o output/dpo/ -c configs/dpo.json --tokenizer_path params
+
+# GRPO
+python scripts/tools/preprocess.py data/grpo/*.jsonl -o output/grpo/ -c configs/grpo.json
+```
+
+---
+
+## Python API
+
+```python
+from astrai.preprocessing.pipeline import Pipeline
+from astrai.config.preprocess_config import PipelineConfig
+
+config = PipelineConfig.from_file("sft.json")
+Pipeline(
+    config,
+    ["data_part1.jsonl", "data_part2.jsonl"],
+    output_dir="output/",
+    tokenizer_path="params",
+).run()
+```
+
+> Document Update Time: 2026-06-03

assets/docs/training.md

@@ -0,0 +1,217 @@
+# Training
+
+## Contents
+
+- [Autoregression](#autoregression)
+- [Causal Mask](#causal-mask)
+- [Rotary Position Embedding (RoPE)](#rotary-position-embedding-rope)
+- [Training Loop](#training-loop)
+- [Strategies](#strategies) — SEQ, SFT, DPO, GRPO
+- [LR Schedulers](#lr-schedulers)
+- [Gradient Checkpointing](#gradient-checkpointing)
+- [Checkpoint](#checkpoint)
+- [TrainContextBuilder](#traincontextbuilder-builder-pattern)
+- [Training CLI](#training-cli)
+
+### 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
+  model.train()
+  on_epoch_begin
+    for batch in dataloader:
+      on_batch_begin
+      with executor.accumulate(model):
+        loss = strategy.compute_loss(batch)
+        context.loss = loss.item()
+        stand_loss = loss / executor.grad_accum_steps
+        executor.backward(stand_loss)
+        context.consumed_samples += (
+            context.config.batch_per_device * context.world_size
+        )
+        on_batch_end
+
+        if executor.sync_gradients:
+          on_optimizer_step
+          optimizer.step()
+          optimizer.zero_grad()
+          if scheduler:
+            scheduler.step()
+    on_epoch_end
+on_train_end
+```
+
+### Callback Lifecycle
+
+| Hook | Fires | Default callback |
+|------|-------|-----------------|
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
+| `on_batch_begin` | Every batch | — |
+| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricLoggerCallback`, `ValidationCallback` |
+| `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
+| `on_epoch_end` | End of each epoch | `ProgressBarCallback` |
+| `on_error` | On exception during training | `CheckpointCallback`, `MetricLoggerCallback` |
+| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricLoggerCallback`, `GradientCheckpointingCallback` |
+
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `validation` (periodic validation on val_dataset), `metric_logger` (JSONL, rank-0), `progress_bar` (tqdm), `gradient_clipping`.
+
+## 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`. Optional: `label_smoothing`.
+
+### 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`. Optional: `label_smoothing`.
+
+### 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`, `reduction="mean"`. 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`, `reduction="mean"`.
+
+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`) |
+| WSD | `WSDScheduler` | Warmup-Stable-Decay with sqrt cooldown |
+
+Created by `SchedulerFactory.create(schedule_type, optimizer, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`, `"wsd"`. Omit to use no scheduler.
+
+## 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, consumed_samples, extra, meta, config)
+  ├── save(save_dir)    rank-0 only: meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
+  └── load(save_dir, broadcast=False)    loads from local disk; set broadcast=True to broadcast metadata from rank-0
+```
+
+Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
+Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.
+
+## TrainContextBuilder (Builder Pattern)
+
+```python
+context = (
+    TrainContextBuilder(config)
+        .with_resume_dir(resume_dir)
+        .build()
+)
+# Returns TrainContext with model, strategy, optimizer, scheduler, dataloader, checkpoint
+```
+
+- Loads checkpoint weights if provided
+- Creates executor via `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.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, model, device, **kwargs)`
+
+## Training CLI
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --parallel_mode=ddp \
+    --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-30

astrai/__init__.py

@@ -1,32 +1,98 @@
-__version__ = "1.3.4"
+__version__ = "1.3.8"
 __author__ = "ViperEkura"
 
 from astrai.config import (
-    ModelConfig,
+    AutoRegressiveLMConfig,
+    BaseModelConfig,
+    ConfigFactory,
+    EncoderConfig,
+    PipelineConfig,
     TrainConfig,
 )
-from astrai.dataset import DatasetFactory
+from astrai.dataset import (
+    BaseDataset,
+    DatasetFactory,
+    ResumableDistributedSampler,
+    Store,
+    StoreFactory,
+)
 from astrai.factory import BaseFactory
 from astrai.inference import (
     GenerationRequest,
     InferenceEngine,
+    ProtocolHandler,
+    SamplingPipeline,
+    get_app,
+    run_server,
+    sample,
+)
+from astrai.model import (
+    AutoModel,
+    AutoRegressiveLM,
+    EmbeddingEncoder,
+    LoRAConfig,
+    inject_lora,
+)
+from astrai.parallel import (
+    ExecutorFactory,
+    get_rank,
+    get_world_size,
+    only_on_rank,
+    spawn_parallel_fn,
+)
+from astrai.preprocessing import Pipeline, filter_by_length
+from astrai.serialization import Checkpoint
+from astrai.tokenize import AutoTokenizer, ChatTemplate
+from astrai.trainer import (
+    BaseScheduler,
+    BaseStrategy,
+    CallbackFactory,
+    SchedulerFactory,
+    StrategyFactory,
+    TrainCallback,
+    Trainer,
 )
-from astrai.model import AutoModel, Transformer
-from astrai.tokenize import AutoTokenizer
-from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer
 
 __all__ = [
-    "Transformer",
-    "ModelConfig",
-    "TrainConfig",
-    "DatasetFactory",
+    "AutoRegressiveLM",
+    "AutoRegressiveLMConfig",
+    "AutoModel",
     "AutoTokenizer",
+    "BaseDataset",
+    "BaseFactory",
+    "BaseModelConfig",
+    "BaseScheduler",
+    "BaseStrategy",
+    "CallbackFactory",
+    "ChatTemplate",
+    "Checkpoint",
+    "ConfigFactory",
+    "DatasetFactory",
+    "EmbeddingEncoder",
+    "EncoderConfig",
+    "ExecutorFactory",
     "GenerationRequest",
     "InferenceEngine",
-    "Trainer",
-    "CallbackFactory",
-    "StrategyFactory",
+    "LoRAConfig",
+    "Pipeline",
+    "PipelineConfig",
+    "ProtocolHandler",
+    "ResumableDistributedSampler",
+    "SamplingPipeline",
     "SchedulerFactory",
-    "BaseFactory",
-    "AutoModel",
+    "Store",
+    "StoreFactory",
+    "StrategyFactory",
+    "TrainCallback",
+    "TrainConfig",
+    "Trainer",
+    "filter_by_length",
+    "get_app",
+    "get_rank",
+    "get_world_size",
+    "inject_lora",
+    "only_on_rank",
+    "run_server",
+    "sample",
+    "spawn_parallel_fn",
 ]

astrai/config/__init__.py

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

astrai/config/base.py

@@ -0,0 +1,98 @@
+import json
+from dataclasses import MISSING, dataclass, fields
+from pathlib import Path
+from typing import Any, Dict, Optional, Self, Union, 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, tuple)):
+                try:
+                    val = list(v) if isinstance(v, tuple) else v
+                    json.dumps(val)
+                    d[fld.name] = val
+                except (TypeError, ValueError):
+                    pass
+            elif isinstance(v, BaseConfig):
+                d[fld.name] = v.to_dict()
+            elif hasattr(v, "__dataclass_fields__"):
+                sub = {}
+                for f in fields(v):
+                    a = getattr(v, f.name)
+                    sub[f.name] = list(a) if isinstance(a, tuple) else a
+                d[fld.name] = sub
+        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
+        if isinstance(value, dict) and issubclass(target_type, BaseConfig):
+            return target_type.from_dict(value)
+        raise TypeError
+
+    @classmethod
+    def from_file(cls, path: Union[str, Path]) -> Self:
+        with open(path, "r", encoding="utf-8") as f:
+            return cls.from_dict(json.load(f))
+
+    def to_file(self, path: Union[str, Path]):
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)

astrai/config/model_config.py

@@ -1,42 +1,82 @@
-import json
-from dataclasses import asdict, dataclass
-from typing import Optional, Self
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+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
+    neftune_alpha: float = 0.0
+
+
+@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
+    rope_scaling: Optional[dict] = 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
+    rope_scaling: Optional[dict] = None
+
+    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
+
+    ffn_type: str = "mlp"
+    pooling_type: Optional[str] = None
+    normalize_embeddings: Optional[bool] = None

astrai/config/preprocess_config.py

@@ -0,0 +1,109 @@
+"""Pipeline configuration for JSONL preprocessing.
+
+Supports single-sequence (SFT/pretrain) and multi-output (DPO/GRPO)
+modes, both driven declaratively through ``input.sections`` or
+``input.sources``.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+from astrai.config.base import BaseConfig
+
+
+@dataclass
+class InputConfig(BaseConfig):
+    """Declarative input mapping.
+
+    Single-output mode (backward-compatible)::
+
+        {"input": {"sections": [{"field": "messages", ...}]}}
+
+    Multi-output mode (DPO / GRPO)::
+
+        {"input": {"sources": {
+            "chosen": {"sections": [{"field": "chosen", ...}]},
+            "rejected": {"sections": [{"field": "rejected", ...}]},
+        }}}
+    """
+
+    sections: Optional[List[Dict]] = None
+    sources: Optional[Dict[str, Dict]] = None
+
+
+@dataclass
+class ProcessingConfig(BaseConfig):
+    """Processing configuration.
+
+    Parameters
+    ----------
+    max_seq_len : int
+        Maximum sequence length (default: 2048).
+    min_chars : int
+        Minimum number of characters to keep (default: 50).
+    max_chars : int
+        Maximum number of characters to keep (default: 2_000_000).
+    max_items : Optional[int]
+        Maximum number of items to process (default: None, unlimited).
+    packing_strategy : str
+        How to pack sequences into a contiguous stream.
+
+        - ``"simple"``: sequential concatenation (default, backward compatible).
+        - ``"bfd"``: best-fit decreasing bin packing, minimises wasted tokens.
+        - ``"bfd_split"``: BFD with over-length sequences split into chunks.
+    max_packed_len : int
+        Maximum length of a packed bin. Sequences longer than this are
+        truncated or split depending on ``packing_strategy`` (default: 8192).
+    truncation_mode : str
+        How to truncate sequences longer than ``max_packed_len``.
+
+        - ``"keep_start"``: keep the first ``max_packed_len`` tokens (default).
+        - ``"keep_end"``: keep the last ``max_packed_len`` tokens.
+    """
+
+    max_seq_len: int = 2048
+    min_chars: int = 50
+    max_chars: int = 2_000_000
+    max_items: Optional[int] = None
+    packing_strategy: str = "simple"
+    max_packed_len: int = 8192
+    truncation_mode: str = "keep_start"
+
+
+@dataclass
+class OutputConfig(BaseConfig):
+    """Output configuration.
+
+    Parameters
+    ----------
+    domain_key : Optional[str]
+        Domain key for the output store (default: None).
+    storage_format : str
+        Storage format, one of ``"bin"``, ``"jsonl"`` (default: ``"bin"``).
+    max_tokens_per_shard : int
+        Maximum tokens per shard before splitting (default: 100_000_000).
+    dtype : Dict[str, str]
+        Per-key dtype overrides, e.g. ``{"input_ids": "int32"}`` (default: {}).
+    position_ids_mode : Optional[str]
+        How to compute position_ids in packed sequences.
+
+         - ``"none"``: do not generate (default).
+        - ``"doc_reset"``: reset to 0 at each document boundary.
+        - ``"continuous"``: sequential 0, 1, 2, ... (pretrain, single doc).
+    """
+
+    domain_key: Optional[str] = None
+    storage_format: str = "bin"
+    max_tokens_per_shard: int = 100_000_000
+    dtype: Dict[str, str] = field(default_factory=dict)
+    position_ids_mode: str = "doc_reset"
+
+
+@dataclass
+class PipelineConfig(BaseConfig):
+    version: int = 1
+    input: InputConfig = field(default_factory=InputConfig)
+    mask: Dict[str, str] = field(default_factory=dict)
+    mask_default: str = "mask"
+    preprocessing: ProcessingConfig = field(default_factory=ProcessingConfig)
+    output: OutputConfig = field(default_factory=OutputConfig)

astrai/config/train_config.py

@@ -1,43 +1,79 @@
-from dataclasses import dataclass, field
-from typing import Callable, Optional
+from dataclasses import dataclass, field, fields
+from typing import Any, Callable, Dict, List, Optional
 
 import torch.nn as nn
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import Dataset
 
+from astrai.config.base import BaseConfig
+from astrai.model.components.lora import LoRAConfig
+
+
+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_fn: Callable[[], nn.Module] = field(
+        default=None, metadata=required(help="Model factory 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[str] = 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."})
-    start_batch: int = field(
-        default=0, metadata={"help": "Start batch iteration for training."}
+    start_samples: int = field(
+        default=0,
+        metadata={
+            "help": "Start samples count (per rank). Superseded by checkpoint consumed_samples."
+        },
     )
     ckpt_dir: str = field(
         default="./checkpoint", metadata={"help": "Checkpoint directory."}
     )
     ckpt_interval: int = field(
-        default=5000, metadata={"help": "Number of iterations between checkpoints."}
+        default=5000,
+        metadata={"help": "Number of optimizer steps between checkpoints."},
+    )
+
+    # lora setting
+    lora: Optional[LoRAConfig] = field(
+        default=None,
+        metadata={"help": "LoRA config. None means full fine-tuning."},
+    )
+
+    # metric setting
+    log_dir: str = field(
+        default="./checkpoint/logs", metadata={"help": "Directory for metric logs."}
+    )
+    metrics: List[str] = field(
+        default_factory=lambda: ["loss", "lr", "grad_norm"],
+        metadata={"help": "Metrics to record during training."},
     )
 
     # dataloader setting
@@ -66,18 +102,42 @@ 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_type: str = field(
         default="cuda", metadata={"help": "Device type for distributed training."}
     )
-    extra_kwargs: dict = field(
+    val_dataset: Optional[Dataset] = field(
+        default=None, metadata={"help": "Dataset for validation."}
+    )
+    val_split: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "Ratio to split from training dataset for validation (e.g. 0.05). Ignored if val_dataset is set."
+        },
+    )
+    val_step: int = field(
+        default=1000,
+        metadata={"help": "Number of optimizer steps between validation runs."},
+    )
+    neftune_alpha: float = field(
+        default=0.0,
+        metadata={"help": "NEFTune noise alpha (0=disabled, typical: 5.0)."},
+    )
+
+    executor_kwargs: Dict[str, Any] = field(
+        default_factory=dict,
+        metadata={"help": "Extra kwargs passed to ExecutorFactory.create()."},
+    )
+    extra_kwargs: Dict[str, Any] = field(
         default_factory=dict, metadata={"help": "Other arguments."}
     )
 
@@ -85,14 +145,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 (
+    H5Store,
+    JsonlStore,
+    MmapStore,
+    Store,
+    StoreFactory,
+    detect_format,
+)
+from astrai.serialization import (
+    load_bin,
+    load_h5,
+    save_bin,
+    save_h5,
+)
 
 __all__ = [
-    # Base classes
     "BaseDataset",
-    # Factory
     "DatasetFactory",
-    # Fetchers
-    "BaseSegmentFetcher",
-    "MultiSegmentFetcher",
-    # Sampler
+    "Store",
+    "StoreFactory",
+    "H5Store",
+    "MmapStore",
+    "JsonlStore",
+    "detect_format",
+    "save_h5",
+    "load_h5",
+    "save_bin",
+    "load_bin",
     "ResumableDistributedSampler",
 ]

astrai/dataset/dataset.py

@@ -1,140 +1,88 @@
 """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 (
+    Store,
+    StoreFactory,
+    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[Store] = 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, **kwargs):
+        """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", "bin", "jsonl"),
+                          or None for auto-detection
+            **kwargs: Extra arguments forwarded to the store constructor and
+                      to ``store.load()``.
+
+        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 = StoreFactory.create(storage_type, **kwargs)
+        self._load_path = load_path
+        self.storage.load(load_path, **kwargs)
+        self._validate_keys()
+
+    @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 +93,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 +115,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"]):
@@ -182,26 +138,6 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         dataset = DatasetFactory.create("custom", window_size, stride)
     """
 
-    @classmethod
-    def _validate_component(cls, dataset_cls: type) -> None:
-        """Validate that the dataset class inherits from BaseDataset."""
-        if not issubclass(dataset_cls, BaseDataset):
-            raise TypeError(f"{dataset_cls.__name__} must inherit from BaseDataset")
-
-    @classmethod
-    def create(cls, train_type: str, window_size: int, stride: int) -> "BaseDataset":
-        """Create a dataset instance.
-
-        Args:
-            train_type: Type of training ("seq", "sft", "dpo", "grpo")
-            window_size: Window size for data sampling
-            stride: Stride between consecutive samples
-
-        Returns:
-            Dataset instance
-        """
-        return super().create(train_type, window_size, stride)
-
     @classmethod
     def load(
         cls,
@@ -209,6 +145,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         load_path: str,
         window_size: int,
         stride: Optional[int] = None,
+        storage_type: Optional[str] = None,
+        **kwargs,
     ) -> "BaseDataset":
         """Create and load a dataset in one step.
 
@@ -217,6 +155,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", "bin", "jsonl") or None for auto-detection
+            **kwargs: Extra arguments forwarded to ``dataset.load()``.
 
         Returns:
             Loaded dataset instance
@@ -225,29 +165,21 @@ 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, **kwargs)
 
         return dataset
 
-    @classmethod
-    def available_types(cls) -> list:
-        """Return list of registered dataset type names."""
-        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."""
 
-    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)
@@ -262,35 +194,39 @@ class SEQDataset(BaseDataset):
 class SFTDataset(BaseDataset):
     """Dataset for supervised fine-tuning with loss masking."""
 
-    def __init__(self, window_size: int, stride: int):
-        super().__init__(window_size, stride)
+    @property
+    def required_keys(self) -> List[str]:
+        return ["sequence", "loss_mask", "position_ids"]
 
     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)
 
-        x = self._fetch_data(begin_idx, end_idx, "sequence").to(dtype=torch.long)
-        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(
-            dtype=torch.long
-        )
-        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask").to(
-            dtype=torch.bool
-        )
+        x = self._fetch_data(begin_idx, end_idx, "sequence")
+        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence")
+        position_ids = self._fetch_data(begin_idx, end_idx, "position_ids")
+        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask")
 
-        return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask}
+        return {
+            "input_ids": x.to(dtype=torch.long),
+            "target_ids": y.to(dtype=torch.long),
+            "position_ids": position_ids.to(dtype=torch.long),
+            "loss_mask": loss_mask.to(dtype=torch.bool),
+        }
 
 
 @DatasetFactory.register("dpo")
 class DPODataset(BaseDataset):
     """Dataset for Direct Preference Optimization training."""
 
-    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)
@@ -316,18 +252,21 @@ class DPODataset(BaseDataset):
 class GRPODataset(BaseDataset):
     """Dataset for Group Relative Policy Optimization training."""
 
-    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)
 
-        prompts = self._fetch_data(begin_idx, end_idx, "prompts")
-        responses = self._fetch_data(begin_idx, end_idx, "responses")
-        masks = self._fetch_data(begin_idx, end_idx, "masks")
+        prompts = self._fetch_data(begin_idx, end_idx, "prompts").to(dtype=torch.long)
+        responses = self._fetch_data(begin_idx, end_idx, "responses").to(
+            dtype=torch.long
+        )
+        masks = self._fetch_data(begin_idx, end_idx, "masks").to(dtype=torch.bool)
         rewards = self._fetch_data(begin_idx, end_idx, "rewards")
 
         return {

astrai/dataset/sampler.py

@@ -43,6 +43,7 @@ class ResumableDistributedSampler(Sampler[int]):
         offset = 0 if drop_last else self.num_replicas - 1
         self.num_samples_per_replica = (self.num_samples + offset) // self.num_replicas
         self.total_size = self.num_samples_per_replica * self.num_replicas
+        self.iter = self.iter % self.num_samples_per_replica
 
         self._indices = None
 
@@ -73,6 +74,12 @@ class ResumableDistributedSampler(Sampler[int]):
 
         self.epoch += 1
         self._indices = None
+        self.iter = self.iter % self.num_samples_per_replica
+
+    @property
+    def _remaining(self):
+        remaining = self.num_samples_per_replica - self.iter
+        return max(remaining, 0)
 
     def __len__(self):
-        return self.num_samples_per_replica
+        return self._remaining

astrai/dataset/storage.py

@@ -0,0 +1,312 @@
+"""Storage backends for different data formats.
+
+Layers:
+  - I/O layer:       save_* / load_* functions, read/write raw files (HDF5/bin)
+                      return Dict[str, List[Tensor]] — format-specific, no state
+  - Store (ABC):     central abstraction, normalizes multi-segment into
+                      Dict[str, List[Tensor]] per key via _normalize(),
+                      fetch() uses bisect across segments — no forced concat
+  - Dataset layer:   BaseDataset owns a Store, only calls store.fetch(begin, end, key)
+
+Key properties:
+  - Multi-segment:   segments kept as-is, no forced concatenation — safe for
+                      datasets larger than RAM
+  - Explicit length: _length = min(total elements across keys), set at load,
+                      __len__ returns O(1)
+  - Zero-copy mmap:  MmapStore wraps np.memmap(mode="r"), all DataLoader
+                       workers share OS page-cache pages
+"""
+
+import bisect
+import glob
+import json
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Dict, List, Union
+
+import torch
+from torch import Tensor
+
+from astrai.config.preprocess_config import PipelineConfig
+from astrai.factory import BaseFactory
+from astrai.preprocessing.builder import MaskBuilderFactory
+from astrai.preprocessing.position_id import PositionIdStrategyFactory
+from astrai.serialization import (
+    load_bin,
+    load_h5,
+)
+from astrai.tokenize import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+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", "bin", or "jsonl")
+
+    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 == ".jsonl":
+            return "jsonl"
+        raise ValueError(f"Unsupported file format: {suffix}")
+
+    h5_files = [
+        Path(p)
+        for pattern in ("*.h5", "*.hdf5")
+        for p in glob.glob(str(root / "**" / pattern), recursive=True)
+    ]
+    if h5_files:
+        return "h5"
+    bin_files = [Path(p) for p in glob.glob(str(root / "**" / "*.bin"), recursive=True)]
+    if bin_files:
+        has_meta = (root / "meta.json").exists() or len(
+            [Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)]
+        ) > 0
+        if has_meta:
+            return "bin"
+    jsonl_files = [
+        Path(p) for p in glob.glob(str(root / "**" / "*.jsonl"), recursive=True)
+    ]
+    if jsonl_files:
+        return "jsonl"
+    raise FileNotFoundError(f"No supported data files found at {load_path}")
+
+
+class Store(ABC):
+    """String keys -> segmented tensors with ``fetch(begin, end, keys)``.
+
+    Each key maps to one or more tensor segments (no forced concatenation).
+    ``len(store)`` returns ``self._length`` (explicit, O(1)), the minimum
+    total element count across all keys.
+
+    Subclasses fill ``self._data`` and ``self._cum`` during ``load()``
+    via ``_normalize()``.
+    """
+
+    def __init__(self):
+        self._data: Dict[str, List[Tensor]] = {}
+        self._cum: Dict[str, List[int]] = {}
+        self._length: int = 0
+
+    @abstractmethod
+    def load(self, path: str) -> None:
+        raise NotImplementedError
+
+    @property
+    def keys(self) -> List[str]:
+        return list(self._data.keys())
+
+    def __len__(self) -> int:
+        return self._length
+
+    def fetch(
+        self,
+        begin: int,
+        end: int,
+        keys: Union[str, List[str]],
+    ):
+        if not self._data:
+            raise RuntimeError("Store not loaded")
+        if not (0 <= begin < self._length and 0 <= end <= self._length):
+            raise ValueError(
+                f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
+            )
+        if isinstance(keys, str):
+            return self._fetch_key(keys, begin, end)
+        return {k: self._fetch_key(k, begin, end) for k in keys}
+
+    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
+        """Fetch slice [begin, end) across potentially multiple segments."""
+        segments = self._data[key]
+        cum = self._cum[key]
+        seg_start = bisect.bisect_right(cum, begin)
+        seg_end = bisect.bisect_left(cum, end)
+
+        results = []
+        for i in range(seg_start, seg_end + 1):
+            prev = cum[i - 1] if i > 0 else 0
+            s = max(begin - prev, 0)
+            e = min(end - prev, segments[i].shape[0])
+            results.append(segments[i][s:e])
+
+        return results[0] if len(results) == 1 else torch.cat(results, dim=0)
+
+    def _normalize(self, raw: Dict[str, List[Tensor]]):
+        """Register segments and pre-compute cumulative lengths.
+
+        Does NOT concatenate — segments are kept as-is to avoid OOM on
+        large datasets.  Sets ``self._length`` to the minimum total
+        element count across all keys.
+        """
+        for key, tensors in raw.items():
+            self._data[key] = tensors
+            cum = []
+            total = 0
+            for t in tensors:
+                total += t.shape[0]
+                cum.append(total)
+            self._cum[key] = cum
+        self._length = (
+            min((cum[-1] if cum else 0) for cum in self._cum.values())
+            if self._cum
+            else 0
+        )
+
+
+class StoreFactory(BaseFactory["Store"]):
+    """Factory for creating Store instances by type name.
+
+    Example::
+
+        @StoreFactory.register("custom")
+        class CustomStore(Store):
+            ...
+    """
+
+
+@StoreFactory.register("h5")
+class H5Store(Store):
+    """HDF5-based storage backend (pre-tokenized data)."""
+
+    def load(self, path: str):
+        self._normalize(load_h5(path))
+
+
+@StoreFactory.register("bin")
+class MmapStore(Store):
+    """Memory-mapped binary storage backend.
+
+    Each key is a single .bin file backed by ``np.memmap(mode="r")``.
+    No per-process memory duplication — all DataLoader workers share the
+    same OS page-cache pages.
+
+    Format on disk::
+
+        data_root/
+          meta.json          # {key: {shape, dtype}, ...}
+          <key>.bin          # raw numpy array, one per key
+    """
+
+    def load(self, path: str):
+        self._mmap_refs = []
+        root = Path(path)
+        all_raw: Dict[str, List[Tensor]] = {}
+        meta_paths = [
+            Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)
+        ]
+        for meta_path in meta_paths:
+            raw = load_bin(str(meta_path.parent))
+            for key, tensors in raw.items():
+                if key not in all_raw:
+                    all_raw[key] = []
+                all_raw[key].extend(tensors)
+        if not meta_paths:
+            raise FileNotFoundError(f"No meta.json found under {path}")
+        self._normalize(all_raw)
+        for tensors in self._data.values():
+            self._mmap_refs.extend(tensors)
+
+
+@StoreFactory.register("jsonl")
+class JsonlStore(Store):
+    """On-the-fly tokenization store for raw JSONL files.
+
+    A JSONL dataset directory contains ``*.jsonl`` files plus a
+    ``dataset_config.json`` file that follows the same schema as
+    :class:`PipelineConfig` with an additional ``tokenizer_path`` field.
+    Records are tokenized when the store is loaded and concatenated into
+    segmented tensors matching the key layout expected by the dataset
+    classes (``sequence``, ``loss_mask``, ``position_ids``, ...).
+    """
+
+    CONFIG_NAME = "dataset_config.json"
+
+    def load(self, path: str):
+        root = Path(path)
+        config_path = root / self.CONFIG_NAME
+        if not config_path.exists():
+            raise FileNotFoundError(
+                f"JSONL dataset config not found: {config_path}. "
+                f"Expected {self.CONFIG_NAME} alongside *.jsonl files."
+            )
+
+        with open(config_path, "r", encoding="utf-8") as f:
+            raw_config = json.load(f)
+
+        tokenizer_path = raw_config.pop("tokenizer_path", None)
+        if tokenizer_path is None:
+            raise ValueError(
+                f"JSONL dataset config must specify 'tokenizer_path': {config_path}"
+            )
+
+        self.config = PipelineConfig.from_dict(raw_config)
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+        mask_builder = MaskBuilderFactory.create("sectioned")
+        position_strategy = PositionIdStrategyFactory.create(
+            self.config.output.position_ids_mode
+        )
+
+        raw: Dict[str, List[Tensor]] = {}
+        doc_sequences: List[List[int]] = []
+
+        for jsonl_path in sorted(root.glob("*.jsonl")):
+            with open(jsonl_path, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        item = json.loads(line)
+                    except json.JSONDecodeError:
+                        logger.warning(
+                            "Failed to parse JSON line in %s, skipping", jsonl_path
+                        )
+                        continue
+
+                    result = mask_builder.build(item, self.config, tokenizer)
+                    if result is None:
+                        continue
+
+                    result.pop("domain", None)
+                    primary_ids = self._primary_ids(result)
+                    if not primary_ids:
+                        continue
+
+                    doc_sequences.append(primary_ids)
+                    for key, ids in result.items():
+                        if key not in raw:
+                            raw[key] = []
+                        raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
+
+        pos_ids = position_strategy.generate(doc_sequences)
+        if pos_ids:
+            raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
+
+        self._normalize(raw)
+
+    @staticmethod
+    def _primary_ids(result: dict) -> List[int]:
+        """Return the first integer list in *result* as the primary id sequence."""
+        for val in result.values():
+            if isinstance(val, list) and val and isinstance(val[0], int):
+                return val
+        return []
+
+    @staticmethod
+    def _infer_dtype(ids: List) -> torch.dtype:
+        """Infer tensor dtype from the first element of a token/value list."""
+        if ids and isinstance(ids[0], float):
+            return torch.float32
+        return torch.int32

astrai/factory.py

@@ -1,190 +1,144 @@
-"""Base factory class for extensible component registration."""
+"""Base factory with decorator-based registration and kwarg-filtered instantiation."""
 
+import inspect
+import sys
 from abc import ABC
-from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+from typing import (
+    Callable,
+    Dict,
+    ForwardRef,
+    Generic,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
+from typing import get_args as _get_args
+from typing import get_origin as _get_origin
 
 T = TypeVar("T")
 
 
-class Registry:
-    """Flexible registry for component classes with category and priority support.
+def _resolve_type(
+    arg: Union[Type, str, ForwardRef], factory_cls: type
+) -> Optional[Type]:
+    """Resolve a generic type-arg (str forward-ref, ForwardRef, or class)."""
+    if not isinstance(arg, (str, ForwardRef)):
+        return arg
 
-    This registry stores component classes with optional metadata (category, priority).
-    It provides methods for registration, retrieval, and listing with filtering.
-    """
+    name = arg if isinstance(arg, str) else arg.__forward_arg__
+    if name == factory_cls.__name__:
+        return factory_cls
 
-    def __init__(self):
-        self._entries = {}  # name -> (component_cls, category, priority)
+    mod = sys.modules.get(factory_cls.__module__)
+    if mod is None:
+        return None
+    ns = vars(mod)
 
-    def register(
-        self,
-        name: str,
-        component_cls: Type,
-        category: Optional[str] = None,
-        priority: int = 0,
-    ) -> None:
-        """Register a component class with optional category and priority."""
-        if name in self._entries:
-            raise ValueError(f"Component '{name}' is already registered")
-        self._entries[name] = (component_cls, category, priority)
+    if isinstance(arg, ForwardRef):
+        return arg._evaluate(ns, None, recursive_guard=frozenset())
 
-    def get(self, name: str) -> Type:
-        """Get component class by name."""
-        if name not in self._entries:
-            raise KeyError(f"Component '{name}' not found in registry")
-        return self._entries[name][0]
-
-    def get_with_metadata(self, name: str) -> Tuple[Type, Optional[str], int]:
-        """Get component class with its metadata."""
-        entry = self._entries.get(name)
-        if entry is None:
-            raise KeyError(f"Component '{name}' not found in registry")
-        return entry
-
-    def contains(self, name: str) -> bool:
-        """Check if a name is registered."""
-        return name in self._entries
-
-    def list_names(self) -> List[str]:
-        """Return list of registered component names."""
-        return sorted(self._entries.keys())
-
-    def list_by_category(self, category: str) -> List[str]:
-        """Return names of components belonging to a specific category."""
-        return sorted(
-            name for name, (_, cat, _) in self._entries.items() if cat == category
-        )
-
-    def list_by_priority(self, reverse: bool = False) -> List[str]:
-        """Return names sorted by priority (default ascending)."""
-        return sorted(
-            self._entries.keys(),
-            key=lambda name: self._entries[name][2],
-            reverse=reverse,
-        )
-
-    def entries(self) -> Dict[str, Tuple[Type, Optional[str], int]]:
-        """Return raw entries dictionary."""
-        return self._entries.copy()
+    return ns.get(name)
 
 
 class BaseFactory(ABC, Generic[T]):
-    """Generic factory class for component registration and creation.
+    """Generic factory with decorator-based component registration.
 
-    This base class provides a decorator-based registration pattern
-    for creating extensible component factories.
-
-    Example usage:
-        class MyFactory(BaseFactory[MyBaseClass]):
+        class MyFactory(BaseFactory[MyBase]):
             pass
 
         @MyFactory.register("custom")
-        class CustomComponent(MyBaseClass):
+        class CustomComponent(MyBase):
             ...
 
-        component = MyFactory.create("custom", *args, **kwargs)
+        obj = MyFactory.create("custom", *args, **kwargs)
+
+    ``create()`` filters kwargs to match the component's ``__init__``
+    signature so components don't need ``**kwargs`` just to absorb
+    unrelated parameters.
     """
 
-    _registry: Registry
+    _entries: Dict[str, Type[T]]
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
-        cls._registry = Registry()
+        for orig_base in getattr(cls, "__orig_bases__", ()):
+            if _get_origin(orig_base) is BaseFactory:
+                (arg,) = _get_args(orig_base)
+                cls._entries = {}
+                cls._component_base = _resolve_type(arg, cls)
+                return
 
     @classmethod
-    def register(
-        cls, name: str, category: Optional[str] = None, priority: int = 0
-    ) -> Callable[[Type[T]], Type[T]]:
-        """Decorator to register a component class with optional category and priority.
+    def register(cls, name: str) -> Callable[[Type[T]], Type[T]]:
+        """Decorator to register a component class.
 
-        Args:
-            name: Registration name for the component
-            category: Optional category for grouping components
-            priority: Priority for ordering (default 0)
-
-        Returns:
-            Decorator function that registers the component class
-
-        Raises:
-            TypeError: If the decorated class doesn't inherit from the base type
+        Validates that the decorated class inherits from the generic
+        type parameter ``T`` declared on the factory.
         """
 
         def decorator(component_cls: Type[T]) -> Type[T]:
             cls._validate_component(component_cls)
-            cls._registry.register(
-                name, component_cls, category=category, priority=priority
-            )
+            if name in cls._entries:
+                raise ValueError(f"Component '{name}' is already registered")
+            cls._entries[name] = component_cls
             return component_cls
 
         return decorator
 
     @classmethod
     def create(cls, name: str, *args, **kwargs) -> T:
-        """Create a component instance by name.
-
-        Args:
-            name: Registered name of the component
-            *args: Positional arguments passed to component constructor
-            **kwargs: Keyword arguments passed to component constructor
-
-        Returns:
-            Component instance
-
-        Raises:
-            ValueError: If the component name is not registered
+        """Create a component instance by name, filtering kwargs to match
+        the component's ``__init__`` signature.
         """
-        if not cls._registry.contains(name):
+        entry = cls._entries.get(name)
+        if entry is None:
             raise ValueError(
-                f"Unknown component: '{name}'. "
-                f"Supported types: {sorted(cls._registry.list_names())}"
+                f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
             )
-        component_cls = cls._registry.get(name)
+        component_cls = entry
+        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
-    def _validate_component(cls, component_cls: Type[T]) -> None:
-        """Validate that the component class is valid for this factory.
+    def _validate_component(cls, component_cls: Type[T]):
+        """Validate the decorated class inherits from the factory's base type.
 
-        Override this method in subclasses to add custom validation.
-
-        Args:
-            component_cls: Component class to validate
-
-        Raises:
-            TypeError: If the component class is invalid
+        Override for custom validation beyond ``issubclass``.
         """
-        pass
+        base = cls._component_base
+        if base is not None and not issubclass(component_cls, base):
+            raise TypeError(
+                f"{component_cls.__name__} must inherit from {base.__name__}"
+            )
 
     @classmethod
-    def list_registered(cls) -> list:
-        """List all registered component names.
+    def get_component_class(cls, name: str) -> Type[T]:
+        """Get the registered component class without instantiating it."""
+        entry = cls._entries.get(name)
+        if entry is None:
+            raise ValueError(
+                f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
+            )
+        return entry
 
-        Returns:
-            List of registered component names
-        """
-        return cls._registry.list_names()
+    @classmethod
+    def list_registered(cls) -> List[str]:
+        """List all registered component names."""
+        return sorted(cls._entries)
 
     @classmethod
     def is_registered(cls, name: str) -> bool:
-        """Check if a component name is registered.
-
-        Args:
-            name: Component name to check
-
-        Returns:
-            True if registered, False otherwise
-        """
-        return cls._registry.contains(name)
-
-    @classmethod
-    def list_by_category(cls, category: str) -> List[str]:
-        """List registered component names in a category."""
-        return cls._registry.list_by_category(category)
-
-    @classmethod
-    def list_by_priority(cls, reverse: bool = False) -> List[str]:
-        """List registered component names sorted by priority."""
-        return cls._registry.list_by_priority(reverse)
-
-
-__all__ = ["Registry", "BaseFactory"]
+        """Check if a component name is registered."""
+        return name in cls._entries

astrai/inference/__init__.py

@@ -1,19 +1,54 @@
 """Inference module for continuous batching.
 
 Layers:
-  - engine.py:    Facade (InferenceEngine), Value Object (GenerationParams, GenerationRequest)
-  - scheduler.py: Continuous-batching loop, Task state machine, TaskStatus enum
-  - cache.py:     PagedCache (page-table-indirected KV cache with alloc/free)
-  - sampling.py:  Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
-  - server.py:    FastAPI HTTP server (OpenAI-compatible endpoints)
+  - core/:        Core inference loop (cache, executor, scheduler, task)
+  - api/:         HTTP orchestration (ProtocolHandler, server)
+  - protocols/:   Response builders (OpenAI, Anthropic)
+  - transport/:   SSE transport utilities
+  - engine.py:    Facade (InferenceEngine), Value Object (GenerationRequest)
+  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
 """
 
-from astrai.inference.engine import (
-    GenerationParams,
-    GenerationRequest,
-    InferenceEngine,
+from astrai.inference.api import (
+    AnthropicMessage,
+    BaseToolParser,
+    ChatCompletionRequest,
+    ChatMessage,
+    FunctionDef,
+    GenContext,
+    MessagesRequest,
+    ProtocolHandler,
+    SimpleJsonToolParser,
+    StopChecker,
+    ToolDef,
+    ToolParserFactory,
+    get_app,
+    run_server,
 )
-from astrai.inference.sampling import (
+from astrai.inference.api.anthropic import AnthropicResponseBuilder
+from astrai.inference.api.openai import OpenAIResponseBuilder
+from astrai.inference.core import (
+    STOP,
+    Allocator,
+    CacheView,
+    ContiguousCache,
+    ContiguousCacheView,
+    Executor,
+    InferenceScheduler,
+    KVCache,
+    PageCache,
+    PageCacheView,
+    PagePool,
+    PrefixCache,
+    Storage,
+    Task,
+    TaskManager,
+    TaskStatus,
+    TaskTable,
+    page_hash,
+)
+from astrai.inference.engine import GenerationRequest, InferenceEngine
+from astrai.inference.sample import (
     BaseSamplingStrategy,
     SamplingPipeline,
     TemperatureStrategy,
@@ -21,26 +56,48 @@ from astrai.inference.sampling import (
     TopPStrategy,
     sample,
 )
-from astrai.inference.scheduler import (
-    InferenceScheduler,
-    Task,
-    TaskStatus,
-)
 
 __all__ = [
-    # Engine / Requests
     "InferenceEngine",
     "GenerationRequest",
-    "GenerationParams",
-    # Scheduler
     "InferenceScheduler",
+    "Executor",
+    "STOP",
     "Task",
+    "TaskManager",
     "TaskStatus",
-    # Sampling (Strategy pattern)
+    "Allocator",
+    "CacheView",
+    "KVCache",
+    "ContiguousCache",
+    "ContiguousCacheView",
+    "PageCache",
+    "PageCacheView",
+    "PagePool",
+    "PrefixCache",
+    "Storage",
+    "TaskTable",
+    "page_hash",
     "sample",
     "BaseSamplingStrategy",
     "TemperatureStrategy",
     "TopKStrategy",
     "TopPStrategy",
     "SamplingPipeline",
+    "ProtocolHandler",
+    "StopChecker",
+    "GenContext",
+    "BaseToolParser",
+    "SimpleJsonToolParser",
+    "ToolParserFactory",
+    "OpenAIResponseBuilder",
+    "AnthropicResponseBuilder",
+    "ChatMessage",
+    "ChatCompletionRequest",
+    "FunctionDef",
+    "ToolDef",
+    "AnthropicMessage",
+    "MessagesRequest",
+    "get_app",
+    "run_server",
 ]

astrai/inference/api/__init__.py

@@ -0,0 +1,39 @@
+"""Inference API: protocol handler, stop checker, tool parsers, and FastAPI server.
+
+``app`` is no longer a module-level global. Use :func:`get_app` to access the
+lazy singleton FastAPI instance.
+"""
+
+from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
+from astrai.inference.api.server import (
+    AnthropicMessage,
+    ChatCompletionRequest,
+    ChatMessage,
+    FunctionDef,
+    MessagesRequest,
+    ToolDef,
+    get_app,
+    run_server,
+)
+from astrai.inference.api.tool_parser import (
+    BaseToolParser,
+    SimpleJsonToolParser,
+    ToolParserFactory,
+)
+
+__all__ = [
+    "ProtocolHandler",
+    "StopChecker",
+    "GenContext",
+    "BaseToolParser",
+    "SimpleJsonToolParser",
+    "ToolParserFactory",
+    "AnthropicMessage",
+    "ChatCompletionRequest",
+    "ChatMessage",
+    "FunctionDef",
+    "ToolDef",
+    "MessagesRequest",
+    "get_app",
+    "run_server",
+]

astrai/inference/api/anthropic.py

@@ -0,0 +1,142 @@
+"""Anthropic message completion response builder."""
+
+import time
+import uuid
+from typing import Any, Dict, List, Tuple, Union
+
+from pydantic import BaseModel
+
+from astrai.inference.api.protocol import (
+    GenContext,
+    ResponseBuilder,
+    StopInfo,
+    sse_event,
+)
+from astrai.inference.engine import InferenceEngine
+
+
+def _extract_text(content: Union[str, List[Dict[str, Any]]]) -> str:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                return block.get("text", "")
+    return ""
+
+
+class AnthropicResponseBuilder(ResponseBuilder):
+    def prepare(
+        self, request: BaseModel, engine: InferenceEngine
+    ) -> Tuple[str, GenContext, List[str]]:
+        messages: List[Dict[str, str]] = []
+        system = getattr(request, "system", None)
+        if system:
+            messages.append({"role": "system", "content": system})
+        for m in request.messages:
+            text = _extract_text(m.content)
+            if text:
+                messages.append({"role": m.role, "content": text})
+        prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
+        ctx = GenContext(
+            resp_id=f"msg_{uuid.uuid4().hex[:24]}",
+            created=int(time.time()),
+            model=request.model,
+        )
+        stop_sequences = getattr(request, "stop_sequences", None) or []
+        return prompt, ctx, stop_sequences
+
+    def format_stream_start(self, ctx: GenContext) -> 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_chunk(self, token: str, **kwargs) -> List[str]:
+        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: GenContext, stop: StopInfo) -> List[str]:
+        events: List[str] = []
+        if stop.matched:
+            trimmed = stop.body[: stop.body.rfind(stop.matched)]
+            unyielded = trimmed[len(stop.yielded) :]
+            if unyielded:
+                events.append(
+                    sse_event(
+                        {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": unyielded},
+                        },
+                        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 stop.matched else "end_turn",
+                        "stop_sequence": stop.matched,
+                    },
+                    "usage": {"output_tokens": ctx.completion_tokens},
+                },
+                event="message_delta",
+            )
+        )
+        events.append(sse_event({"type": "message_stop"}, event="message_stop"))
+        return events
+
+    def format_response(
+        self, ctx: GenContext, content: str, stop: StopInfo
+    ) -> Dict[str, Any]:
+        if stop.matched:
+            content = content[: content.rfind(stop.matched)]
+        return {
+            "id": ctx.resp_id,
+            "type": "message",
+            "role": "assistant",
+            "model": ctx.model,
+            "content": [{"type": "text", "text": content}],
+            "stop_reason": "stop_sequence" if stop.matched else "end_turn",
+            "stop_sequence": stop.matched,
+            "usage": {
+                "input_tokens": ctx.prompt_tokens,
+                "output_tokens": ctx.completion_tokens,
+            },
+        }

astrai/inference/api/openai.py

@@ -0,0 +1,278 @@
+"""OpenAI chat completion response builder."""
+
+import logging
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from pydantic import BaseModel
+
+from astrai.inference.api.protocol import (
+    GenContext,
+    ResponseBuilder,
+    StopInfo,
+    sse_event,
+)
+from astrai.inference.api.tool_parser import BaseToolParser, ToolParserFactory
+from astrai.inference.engine import InferenceEngine
+
+logger = logging.getLogger(__name__)
+
+_UNSUPPORTED_PARAMS = (
+    "n",
+    "presence_penalty",
+    "frequency_penalty",
+    "logit_bias",
+    "user",
+)
+
+
+def _resolve_tool_choice(
+    request: BaseModel,
+) -> Union[str, Dict[str, Any]]:
+    tc = getattr(request, "tool_choice", None)
+    if tc is None:
+        return "auto"
+    if isinstance(tc, str):
+        return tc
+    if isinstance(tc, dict):
+        return tc
+    return "auto"
+
+
+def _resolve_tools(request: BaseModel) -> Optional[List[Dict[str, Any]]]:
+    raw = getattr(request, "tools", None)
+    if not raw:
+        return None
+    if isinstance(raw, list):
+        return [t.model_dump() if hasattr(t, "model_dump") else t for t in raw]
+    return None
+
+
+class OpenAIResponseBuilder(ResponseBuilder):
+    def prepare(
+        self, request: BaseModel, engine: InferenceEngine
+    ) -> Tuple[str, GenContext, List[str]]:
+        messages = [{"role": m.role, "content": m.content} for m in request.messages]
+        tools = _resolve_tools(request)
+        prompt = engine.tokenizer.apply_chat_template(
+            messages, tokenize=False, tools=tools or []
+        )
+
+        self._resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+        self._model = request.model
+
+        for param in _UNSUPPORTED_PARAMS:
+            value = getattr(request, param, None)
+            fields = getattr(type(request), "model_fields", {})
+            default = fields[param].default if param in fields else None
+            if value is not None and value != default:
+                logger.warning(
+                    "ChatCompletionRequest param '%s'=%r is not supported"
+                    " and will be ignored",
+                    param,
+                    value,
+                )
+
+        self._parser: Optional[BaseToolParser] = None
+        if tools:
+            tool_choice = _resolve_tool_choice(request)
+            self._parser = ToolParserFactory.create(
+                "simple_json", tools=tools, tool_choice=tool_choice
+            )
+        self._content_started = False
+
+        ctx = GenContext(
+            resp_id=self._resp_id,
+            created=int(time.time()),
+            model=self._model,
+        )
+        stop = request.stop
+        stop_sequences = (
+            [] if stop is None else [stop] if isinstance(stop, str) else stop
+        )
+        return prompt, ctx, stop_sequences
+
+    def format_stream_start(self, ctx: GenContext) -> List[str]:
+        return [
+            sse_event(
+                {
+                    "id": self._resp_id,
+                    "object": "chat.completion.chunk",
+                    "created": ctx.created,
+                    "model": self._model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"role": "assistant"},
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+            )
+        ]
+
+    def format_chunk(self, token: str, **kwargs) -> List[str]:
+        body = kwargs.get("body", "")
+        if self._parser is not None:
+            return self._format_tool_chunk(body, **kwargs)
+
+        return [
+            sse_event(
+                {
+                    "id": self._resp_id,
+                    "object": "chat.completion.chunk",
+                    "created": 0,
+                    "model": self._model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": token},
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+            )
+        ]
+
+    def _format_tool_chunk(self, body: str, **kwargs) -> List[str]:
+        deltas = self._parser.feed(
+            body,
+            current_token_ids=kwargs.get("current_token_ids"),
+            delta_token_ids=kwargs.get("delta_token_ids"),
+        )
+        events: List[str] = []
+        for d in deltas:
+            if "content" in d:
+                if not self._content_started:
+                    events.append(self._role_chunk())
+                    self._content_started = True
+                events.append(
+                    sse_event(
+                        {
+                            "id": self._resp_id,
+                            "object": "chat.completion.chunk",
+                            "created": 0,
+                            "model": self._model,
+                            "choices": [
+                                {
+                                    "index": 0,
+                                    "delta": {"content": d["content"]},
+                                    "finish_reason": None,
+                                }
+                            ],
+                        }
+                    )
+                )
+            elif "tool_calls" in d:
+                if not self._content_started:
+                    events.append(self._role_chunk())
+                    self._content_started = True
+                events.append(
+                    sse_event(
+                        {
+                            "id": self._resp_id,
+                            "object": "chat.completion.chunk",
+                            "created": 0,
+                            "model": self._model,
+                            "choices": [
+                                {
+                                    "index": 0,
+                                    "delta": {"tool_calls": d["tool_calls"]},
+                                    "finish_reason": None,
+                                }
+                            ],
+                        }
+                    )
+                )
+        return events
+
+    def _role_chunk(self) -> str:
+        return sse_event(
+            {
+                "id": self._resp_id,
+                "object": "chat.completion.chunk",
+                "created": 0,
+                "model": self._model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {"role": "assistant"},
+                        "finish_reason": None,
+                    }
+                ],
+            }
+        )
+
+    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
+        finish_reason = "stop"
+        if self._parser is not None and self._parser.has_tool_calls:
+            finish_reason = "tool_calls"
+        return [
+            sse_event(
+                {
+                    "id": self._resp_id,
+                    "object": "chat.completion.chunk",
+                    "created": ctx.created,
+                    "model": self._model,
+                    "choices": [
+                        {"index": 0, "delta": {}, "finish_reason": finish_reason}
+                    ],
+                }
+            ),
+            sse_event(
+                {
+                    "prompt_tokens": ctx.prompt_tokens,
+                    "completion_tokens": ctx.completion_tokens,
+                    "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
+                }
+            ),
+        ]
+
+    def format_response(
+        self, ctx: GenContext, content: str, stop: StopInfo
+    ) -> Dict[str, Any]:
+        if self._parser is not None:
+            parsed = self._parser.parse_complete(content)
+            if parsed and parsed.get("tool_calls"):
+                return {
+                    "id": self._resp_id,
+                    "object": "chat.completion",
+                    "created": ctx.created,
+                    "model": self._model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "message": {
+                                "role": "assistant",
+                                "content": parsed.get("content"),
+                                "tool_calls": parsed["tool_calls"],
+                            },
+                            "finish_reason": "tool_calls",
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": ctx.prompt_tokens,
+                        "completion_tokens": ctx.completion_tokens,
+                        "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
+                    },
+                }
+
+        return {
+            "id": self._resp_id,
+            "object": "chat.completion",
+            "created": ctx.created,
+            "model": self._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,
+            },
+        }

astrai/inference/api/protocol.py

@@ -0,0 +1,199 @@
+"""Orchestration layer: ProtocolHandler, StopChecker, GenContext, StopInfo, ResponseBuilder, SSE utils.
+
+ProtocolHandler orchestrates the async generation loop and delegates
+protocol-specific formatting to a ResponseBuilder.
+"""
+
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, 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 GenContext:
+    """Per-generation metadata passed to builder format methods."""
+
+    resp_id: str
+    created: int
+    model: str
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+
+@dataclass
+class StopInfo:
+    """Stop-check result passed to format_stream_end / format_response."""
+
+    matched: Optional[str] = None
+    body: str = ""
+    yielded: 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
+
+
+class ResponseBuilder(ABC):
+    """Interface for protocol-specific response formatting.
+
+    A new protocol requires one concrete builder implementing 5 methods.
+    """
+
+    @abstractmethod
+    def prepare(
+        self, request: BaseModel, engine: InferenceEngine
+    ) -> Tuple[str, GenContext, List[str]]:
+        """Return (prompt, ctx, stop_sequences) for a generation request."""
+
+    @abstractmethod
+    def format_stream_start(self, ctx: GenContext) -> List[str]:
+        """SSE events that open the stream."""
+
+    @abstractmethod
+    def format_chunk(self, token: str, **kwargs) -> List[str]:
+        """SSE events for a single generated token.
+
+        ``body`` (the full accumulated text so far) is always provided
+        as a keyword argument. Additional keyword arguments such as
+        ``current_token_ids`` and ``delta_token_ids`` may be included
+        for tool parsers that need token-level information.
+        Returns a list of SSE event strings (may be empty).
+        """
+
+    @abstractmethod
+    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
+        """SSE events that close the stream."""
+
+    @abstractmethod
+    def format_response(
+        self, ctx: GenContext, content: str, stop: StopInfo
+    ) -> Dict[str, Any]:
+        """JSON response body for non-streaming mode."""
+
+
+class ProtocolHandler:
+    """Orchestrates the generation loop, delegates formatting to a builder.
+
+    Usage::
+
+        handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
+        response = await handler.handle()
+    """
+
+    def __init__(
+        self, request: BaseModel, engine: InferenceEngine, builder: ResponseBuilder
+    ):
+        self.request = request
+        self.engine = engine
+        self.builder = builder
+
+    async def handle(self) -> Union[StreamingResponse, Dict[str, Any]]:
+        prompt, ctx, stop_sequences = self.builder.prepare(self.request, self.engine)
+        ctx.prompt_tokens = len(self.engine.tokenizer.encode(prompt))
+
+        agen = self.engine.generate_async(
+            prompt=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, stop_sequences)
+        else:
+            return await self._handle_non_stream(agen, ctx, stop_sequences)
+
+    def _handle_stream(
+        self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
+    ) -> StreamingResponse:
+        checker = StopChecker(stop_sequences)
+
+        async def event_stream():
+            for event in self.builder.format_stream_start(ctx):
+                yield event
+
+            body = ""
+            yielded = ""
+            matched = None
+            token_ids: List[int] = []
+            async for token in agen:
+                body += token
+
+                new_ids = self.engine.tokenizer.encode(token)
+                token_ids.extend(new_ids)
+
+                matched = checker.check(body)
+                if matched:
+                    break
+
+                ctx.completion_tokens += 1
+                for event in self.builder.format_chunk(
+                    token,
+                    body=body,
+                    current_token_ids=token_ids,
+                    delta_token_ids=new_ids,
+                ):
+                    yield event
+                yielded += token
+
+            stop = StopInfo(matched=matched, body=body, yielded=yielded)
+            for event in self.builder.format_stream_end(ctx, stop):
+                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: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
+    ) -> Dict[str, Any]:
+        checker = StopChecker(stop_sequences)
+        chunks: List[str] = []
+        body = ""
+        matched = None
+
+        async for token in agen:
+            chunks.append(token)
+            body += token
+
+            matched = checker.check(body)
+            if matched:
+                break
+
+            ctx.completion_tokens += 1
+
+        content = "".join(chunks)
+        stop = StopInfo(matched=matched, body=body)
+        return self.builder.format_response(ctx, content, stop)

astrai/inference/api/server.py

@@ -0,0 +1,202 @@
+"""
+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.
+
+``app`` is lazily constructed — importing this module does NOT create a FastAPI instance.
+Use :func:`get_app` to access the singleton.
+"""
+
+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 APIRouter, FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+from astrai.inference.api.anthropic import AnthropicResponseBuilder
+from astrai.inference.api.openai import OpenAIResponseBuilder
+from astrai.inference.api.protocol import ProtocolHandler
+from astrai.inference.engine import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+_app_instance: Optional[FastAPI] = None
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    tool_call_id: Optional[str] = None
+
+
+class FunctionDef(BaseModel):
+    name: str
+    description: Optional[str] = None
+    parameters: Optional[Dict[str, Any]] = None
+
+
+class ToolDef(BaseModel):
+    type: str = "function"
+    function: FunctionDef
+
+
+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
+    tools: Optional[List[ToolDef]] = None
+    tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
+
+
+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")
+
+
+router = APIRouter()
+
+
+def _create_engine(
+    param_path: Path,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    max_batch_size: int = 16,
+) -> InferenceEngine:
+    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_app() -> FastAPI:
+    """Return the singleton FastAPI instance (lazily created on first call)."""
+    global _app_instance
+    if _app_instance is None:
+        _app_instance = FastAPI(
+            title="AstrAI Inference Server",
+            version="0.2.0",
+            lifespan=lifespan,
+        )
+        _app_instance.include_router(router)
+        _app_instance.state.server_config = {}
+        _app_instance.state.engine = None
+    return _app_instance
+
+
+def _get_engine() -> InferenceEngine:
+    engine = get_app().state.engine
+    if engine is None:
+        raise HTTPException(status_code=503, detail="Engine not initialized")
+    return engine
+
+
+@router.get("/health")
+async def health():
+    app = get_app()
+    return {
+        "status": "ok",
+        "model_loaded": app.state.engine is not None,
+    }
+
+
+@router.get("/stats")
+async def get_stats():
+    return _get_engine().get_stats()
+
+
+@router.post("/v1/chat/completions")
+async def chat_completion(request: ChatCompletionRequest):
+    engine = _get_engine()
+    handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
+    return await handler.handle()
+
+
+@router.post("/v1/messages")
+async def create_message(request: MessagesRequest):
+    engine = _get_engine()
+    handler = ProtocolHandler(request, engine, AnthropicResponseBuilder())
+    return await handler.handle()
+
+
+def run_server(
+    param_path: Path,
+    host: str = "0.0.0.0",
+    port: int = 8000,
+    reload: bool = False,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    max_batch_size: int = 16,
+):
+    app = get_app()
+    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/api/tool_parser.py

@@ -0,0 +1,325 @@
+"""Tool call parsers for extracting structured tool calls from model output.
+
+Patterned after vLLM's ToolParser abstraction. Each parser knows how to
+detect and incrementally extract tool calls from raw generated text.
+
+Subclasses may optionally consume ``token_ids`` for token-level parsing
+(e.g. Harmony / VLM-style parsers).
+"""
+
+import re
+import uuid
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional
+
+from astrai.factory import BaseFactory
+
+
+class BaseToolParser(ABC):
+    """Abstract tool call parser — one instance per request.
+
+    Maintains streaming state internally so that each call to :meth:`feed`
+    can diff against previously emitted content.
+
+    Parameters
+    ----------
+    tools : list of dict, optional
+        Tool definitions from the request.
+    tool_choice : str
+        ``"auto"`` / ``"required"`` / ``"none"`` or a named tool choice
+        dict.
+    """
+
+    def __init__(self, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
+        self.tools = tools or []
+        self.tool_choice = tool_choice
+
+    @abstractmethod
+    def feed(
+        self,
+        body: str,
+        current_token_ids: Optional[List[int]] = None,
+        delta_token_ids: Optional[List[int]] = None,
+    ) -> List[Dict]:
+        """Feed the *full* accumulated text each step.
+
+        Returns a list of delta dicts to emit. Each delta is one of:
+
+        - ``{"content": "text"}``      — plain text delta
+        - ``{"tool_calls": [...]}``    — tool-call delta (OpenAI format)
+
+        Returns an empty list when nothing new should be emitted.
+
+        Parameters
+        ----------
+        body : str
+            The complete accumulated generated text so far.
+        current_token_ids : list of int, optional
+            All token IDs decoded into *body* (cumulative).
+        delta_token_ids : list of int, optional
+            Only the token IDs for this chunk.
+        """
+
+    @abstractmethod
+    def parse_complete(self, body: str) -> Optional[Dict]:
+        """Parse the *complete* generated text after generation ends.
+
+        Returns ``None`` when no tool calls were found, otherwise a dict
+        with ``content`` (str or None) and ``tool_calls`` (list of dicts).
+        """
+
+    @property
+    @abstractmethod
+    def has_tool_calls(self) -> bool:
+        """True if the parser detected at least one tool call in the stream."""
+
+
+class ToolParserFactory(BaseFactory["BaseToolParser"]):
+    pass
+
+
+_TOOL_CALL_HEAD_RE = re.compile(r'\{\s*"name"\s*:')
+
+
+def _scan_json(text: str, start: int = 0):
+    """Scan for a complete JSON object starting at *start*.
+
+    Returns ``(end, complete)`` where *end* is one-past the closing
+    brace (or ``len(text)`` if unclosed), and *complete* is a bool.
+    """
+    depth = 0
+    in_string = False
+    escape = False
+    for i in range(start, len(text)):
+        c = text[i]
+        if escape:
+            escape = False
+            continue
+        if c == "\\":
+            escape = True
+            continue
+        if c == '"':
+            in_string = not in_string
+            continue
+        if in_string:
+            continue
+        if c == "{":
+            depth += 1
+        elif c == "}":
+            depth -= 1
+            if depth == 0:
+                return i + 1, True
+    return len(text), False
+
+
+def _parse_tool_call_json(json_str: str, complete: bool):
+    """Extract *name* and *arguments* from a tool-call JSON string.
+
+    Returns ``(name, args, valid)``.
+    """
+    name_match = re.search(r'"name"\s*:\s*"([^"]*)"', json_str)
+    if not name_match:
+        return None, "", False
+    name = name_match.group(1)
+
+    args_match = re.search(r'"arguments"\s*:\s*(.*)', json_str, re.DOTALL)
+    if not args_match:
+        return name, "", True
+
+    raw = args_match.group(1).rstrip()
+    if complete and raw.endswith("}"):
+        raw = raw[:-1].rstrip()
+    if raw.startswith("{"):
+        inner = raw[1:].rstrip()
+        if inner.endswith("}"):
+            inner = inner[:-1].rstrip()
+        raw = inner
+    return name, raw, True
+
+
+def _find_tool_calls(text: str, start_pos: int = 0):
+    """Find all complete ``{...}`` tool-call objects in *text*.
+
+    Returns a list of dicts with keys *start*, *end*, *name*, *args*,
+    *complete*.
+    """
+    results = []
+    pos = start_pos
+
+    while True:
+        brace = text.find("{", pos)
+        if brace == -1:
+            break
+
+        end, complete = _scan_json(text, brace)
+        if not complete:
+            break
+
+        json_str = text[brace:end]
+        if not _TOOL_CALL_HEAD_RE.search(json_str):
+            pos = end
+            continue
+
+        name, args, valid = _parse_tool_call_json(json_str, complete=True)
+        if not valid or name is None:
+            pos = end
+            continue
+
+        results.append(
+            {
+                "start": brace,
+                "end": end,
+                "name": name,
+                "args": args,
+                "complete": True,
+            }
+        )
+        pos = end
+
+    return results
+
+
+def _find_partial_tool_call(text: str, start_pos: int = 0):
+    """Find one incomplete (still-generating) tool-call JSON object."""
+    brace = text.find("{", start_pos)
+    if brace == -1:
+        return None
+
+    json_str = text[brace:]
+    if not _TOOL_CALL_HEAD_RE.search(json_str):
+        return None
+
+    name, args, valid = _parse_tool_call_json(json_str, complete=False)
+    if not valid or name is None:
+        return None
+
+    return {
+        "start": brace,
+        "name": name,
+        "args": args,
+        "complete": False,
+    }
+
+
+@ToolParserFactory.register("simple_json")
+class SimpleJsonToolParser(BaseToolParser):
+    """Parser for models that output tool calls as plain JSON objects.
+
+    Detects ``{"name": "<func>", "arguments": {...}}`` anywhere in the
+    generated text.  Handles single and (non-overlapping) multiple tool
+    calls.  Text preceding the first tool call is emitted as plain
+    ``content`` deltas.
+    """
+
+    def __init__(self, tools=None, tool_choice="auto"):
+        super().__init__(tools, tool_choice)
+        self._emitted_content_len = 0
+        self._tc_state: List[Dict] = []
+        self._has_tool_calls = False
+
+    # -------------------------------------------------------------- feed
+
+    def feed(
+        self,
+        body: str,
+        current_token_ids: Optional[List[int]] = None,
+        delta_token_ids: Optional[List[int]] = None,
+    ) -> List[Dict]:
+        deltas: List[Dict] = []
+
+        completed = _find_tool_calls(body)
+
+        if not completed:
+            partial = _find_partial_tool_call(body)
+            if not partial:
+                return self._emit_plain_content(body, deltas)
+            all_tcs = [partial]
+        else:
+            all_tcs = completed
+            partial = _find_partial_tool_call(body, completed[-1]["end"])
+            if partial:
+                all_tcs = completed + [partial]
+
+        first_start = all_tcs[0]["start"]
+        if first_start > self._emitted_content_len:
+            content = body[self._emitted_content_len : first_start]
+            self._emitted_content_len = first_start
+            if content:
+                deltas.append({"content": content})
+
+        for i, tc in enumerate(all_tcs):
+            if i >= len(self._tc_state):
+                self._tc_state.append(
+                    {
+                        "id": f"call_{uuid.uuid4().hex[:12]}",
+                        "name_emitted": False,
+                        "args_emitted_len": 0,
+                    }
+                )
+                self._has_tool_calls = True
+            st = self._tc_state[i]
+
+            if not st["name_emitted"]:
+                st["name_emitted"] = True
+                deltas.append(
+                    {
+                        "tool_calls": [
+                            {
+                                "index": i,
+                                "id": st["id"],
+                                "type": "function",
+                                "function": {"name": tc["name"], "arguments": ""},
+                            }
+                        ]
+                    }
+                )
+
+            new_args = tc["args"]
+            if len(new_args) > st["args_emitted_len"]:
+                diff = new_args[st["args_emitted_len"] :]
+                st["args_emitted_len"] = len(new_args)
+                deltas.append(
+                    {
+                        "tool_calls": [
+                            {
+                                "index": i,
+                                "function": {"arguments": diff},
+                            }
+                        ]
+                    }
+                )
+
+        return deltas
+
+    def _emit_plain_content(self, body: str, deltas: List[Dict]) -> List[Dict]:
+        new_content = body[self._emitted_content_len :]
+        if new_content:
+            self._emitted_content_len = len(body)
+            deltas.append({"content": new_content})
+        return deltas
+
+    # -------------------------------------------------------- complete
+
+    def parse_complete(self, body: str) -> Optional[Dict]:
+        completed = _find_tool_calls(body)
+        if not completed:
+            return None
+
+        content = body[: completed[0]["start"]].strip() or None
+        tool_calls = []
+        for i, tc in enumerate(completed):
+            tool_calls.append(
+                {
+                    "id": f"call_{uuid.uuid4().hex[:12]}",
+                    "type": "function",
+                    "function": {
+                        "name": tc["name"],
+                        "arguments": tc["args"],
+                    },
+                }
+            )
+        return {"content": content, "tool_calls": tool_calls}
+
+    @property
+    def has_tool_calls(self) -> bool:
+        return self._has_tool_calls

astrai/inference/cache.py

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

astrai/inference/core/__init__.py

@@ -0,0 +1,40 @@
+"""Inference core: cache, executor, scheduler, task management."""
+
+from astrai.inference.core.cache import (
+    Allocator,
+    CacheView,
+    ContiguousCache,
+    ContiguousCacheView,
+    KVCache,
+    PageCache,
+    PageCacheView,
+    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",
+    "CacheView",
+    "KVCache",
+    "ContiguousCache",
+    "ContiguousCacheView",
+    "PageCache",
+    "PageCacheView",
+    "PagePool",
+    "PrefixCache",
+    "Storage",
+    "TaskTable",
+    "page_hash",
+    "Executor",
+    "InferenceScheduler",
+    "STOP",
+    "Task",
+    "TaskManager",
+    "TaskStatus",
+]

astrai/inference/core/cache.py

@@ -0,0 +1,499 @@
+import threading
+from abc import ABC, abstractmethod
+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):
+        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):
+        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):
+        with self._lock:
+            if idx in self._lru:
+                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):
+        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):
+        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):
+        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):
+        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):
+        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):
+        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,
+    ):
+        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 CacheView(ABC):
+    """Abstract view passed to attention layers for KV-cache I/O."""
+
+    @abstractmethod
+    def write(self, layer_id: int, k: Tensor, v: Tensor): ...
+
+    @abstractmethod
+    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: ...
+
+
+class KVCache(ABC):
+    """Abstract KV-cache facade for scheduler/executor."""
+
+    @abstractmethod
+    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: ...
+
+    @abstractmethod
+    def task_free(self, task_id: str): ...
+
+    @abstractmethod
+    def task_extend(self, task_id: str, pos: int) -> bool: ...
+
+    @abstractmethod
+    def bind_tasks(
+        self, task_ids: List[str], total_len: int, device: torch.device
+    ) -> CacheView: ...
+
+    def task_cached(self, task_id: str) -> int:
+        return 0
+
+    def task_record_hashes(
+        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
+    ): ...
+
+
+class PageCacheView(CacheView):
+    """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):
+        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 PageCache(KVCache):
+    """Paged KV-cache with prefix sharing."""
+
+    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):
+        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
+    ):
+        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 bind_tasks(
+        self, task_ids: List[str], total_len: int, device: torch.device
+    ) -> PageCacheView:
+        page_table = self._table.table_tensor(task_ids, device)
+        return PageCacheView(self._storage, page_table, total_len)
+
+
+class ContiguousCacheView(CacheView):
+    """Contiguous KV-cache view for attention layers."""
+
+    def __init__(
+        self, cache: "ContiguousCache", batch_indices: Tensor, total_len: int = 0
+    ):
+        self._cache = cache
+        self._batch_indices = batch_indices
+        self._total_len = total_len
+
+    def write(self, layer_id: int, k: Tensor, v: Tensor):
+        seq_len = k.size(1)
+        start_pos = self._total_len - seq_len
+        indices = self._batch_indices
+        self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k
+        self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v
+        new_len = start_pos + seq_len
+        for s in indices.tolist():
+            cur = self._cache._slot_len.get(s, 0)
+            if new_len > cur:
+                self._cache._slot_len[s] = new_len
+
+    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
+        max_len = max(
+            self._cache._slot_len.get(int(s), 0) for s in self._batch_indices.tolist()
+        )
+        indices = self._batch_indices
+        k = self._cache.k[layer_id, indices, :max_len]
+        v = self._cache.v[layer_id, indices, :max_len]
+        return k, v
+
+
+class ContiguousCache(KVCache):
+    """Contiguous per-slot KV cache (default implementation)."""
+
+    def __init__(
+        self,
+        n_layers: int,
+        max_batch_size: int,
+        max_seq_len: int,
+        n_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        self.max_seq_len = max_seq_len
+        self.k = torch.zeros(
+            n_layers,
+            max_batch_size,
+            max_seq_len,
+            n_kv_heads,
+            head_dim,
+            device=device,
+            dtype=dtype,
+        )
+        self.v = torch.zeros(
+            n_layers,
+            max_batch_size,
+            max_seq_len,
+            n_kv_heads,
+            head_dim,
+            device=device,
+            dtype=dtype,
+        )
+        self._slot_len: Dict[int, int] = {}
+        self._task_slot: Dict[str, int] = {}
+        self._free_slots = list(range(max_batch_size))
+        self._device = device
+
+    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
+        if not self._free_slots:
+            return False
+        slot = self._free_slots.pop(0)
+        self._task_slot[task_id] = slot
+        self._slot_len[slot] = 0
+        return True
+
+    def task_free(self, task_id: str):
+        slot = self._task_slot.pop(task_id, None)
+        if slot is not None:
+            self._slot_len.pop(slot, None)
+            self._free_slots.append(slot)
+
+    def task_extend(self, task_id: str, pos: int) -> bool:
+        return pos < self.max_seq_len
+
+    def bind_tasks(
+        self, task_ids: List[str], total_len: int, device: torch.device
+    ) -> ContiguousCacheView:
+        slots = [self._task_slot[tid] for tid in task_ids]
+        batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
+        return ContiguousCacheView(self, batch_indices, total_len)

astrai/inference/core/executor.py

@@ -0,0 +1,92 @@
+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,
+        kv_cache: KVCache,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.kv_cache = kv_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):
+        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]
+
+        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.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
+            )
+
+    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]
+
+        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.kv_cache.bind_tasks(task_ids, total_len, self.device),
+                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 ContiguousCache, 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:
+    """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
+
+    def __init__(
+        self,
+        model: AutoModel,
+        tokenizer: AutoTokenizer,
+        max_batch_size: int = 16,
+        max_seq_len: Optional[int] = None,
+        max_prompt_len: int = 2048,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+        cache: Optional[KVCache] = 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
+
+        head_dim = config.dim // config.n_heads
+
+        if cache is not None:
+            self._cache = cache
+        else:
+            self._cache = ContiguousCache(
+                config.n_layers,
+                max_batch_size,
+                self.max_seq_len,
+                config.n_kv_heads,
+                head_dim,
+                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,
+            kv_cache=self._cache,
+            device=self.device,
+            dtype=self.dtype,
+        )
+
+        self._stop_event = threading.Event()
+        self._loop_thread: Optional[threading.Thread] = None
+
+    def add_task(self, prompt: str, **kwargs) -> str:
+        return self._task_mgr.add_task(prompt, **kwargs)
+
+    def remove_task(self, task_id: str):
+        for task in self._task_mgr.remove_task(task_id):
+            self._cache.task_free(task.task_id)
+
+    def get_stats(self) -> Dict[str, Any]:
+        return self._task_mgr.get_stats()
+
+    def _run_generation_loop(self):
+        stop_ids = self._task_mgr.tokenizer.stop_ids
+        cache = self._cache
+        try:
+            while not self._stop_event.is_set():
+                finished = self._task_mgr.remove_finished_tasks(stop_ids)
+                for task in finished:
+                    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 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
+                    and cache.task_cached(t.task_id) < len(t.prompt_ids)
+                ]
+                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),
+                            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 // getattr(
+                            cache, "page_size", 64
+                        )
+                        for t in group:
+                            cache.task_record_hashes(
+                                t.task_id, t.prompt_ids, 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)
+
+                for next_pos in sorted(pos_groups.keys()):
+                    group = sorted(pos_groups[next_pos], key=lambda t: t.task_id)
+
+                    valid: List[Task] = []
+                    for t in group:
+                        if cache.task_extend(t.task_id, t.next_pos):
+                            valid.append(t)
+                        else:
+                            t.status = TaskStatus.ABORTED
+                            self._task_mgr.invoke_callback(t.task_id, 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
+                            self._task_mgr.invoke_callback(
+                                t.task_id,
+                                self._task_mgr.tokenizer.decode([ntok]),
+                            )
+
+                        for t in valid:
+                            if t.is_finished(stop_ids):
+                                self._task_mgr.invoke_callback(t.task_id, STOP)
+
+        except Exception as e:
+            self._stop_event.set()
+            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
+            for task in self._task_mgr.get_active_tasks():
+                self._task_mgr.invoke_callback(task.task_id, STOP)
+                cache.task_free(task.task_id)
+            for task in self._task_mgr.get_waiting_tasks():
+                self._task_mgr.invoke_callback(task.task_id, STOP)
+            self._task_mgr.clear_queues()
+
+    def start(self):
+        if self._loop_thread is not None and self._loop_thread.is_alive():
+            return
+        self._stop_event.clear()
+        t = threading.Thread(target=self._run_generation_loop, daemon=True)
+        t.start()
+        self._loop_thread = t
+
+    def stop(self):
+        self._stop_event.set()
+        self._task_mgr.wake()
+        if self._loop_thread is not None:
+            self._loop_thread.join(timeout=2.0)
+        self._loop_thread = None
+        for task in self._task_mgr.get_active_tasks():
+            self._task_mgr.invoke_callback(task.task_id, STOP)
+            self._cache.task_free(task.task_id)
+        for task in self._task_mgr.get_waiting_tasks():
+            self._task_mgr.invoke_callback(task.task_id, STOP)
+        self._task_mgr.clear_queues()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()

astrai/inference/core/task.py

@@ -0,0 +1,216 @@
+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,
+    ):
+        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
+
+    @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._callbacks: Dict[str, Callable[[str], None]] = {}
+
+        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,
+        )
+
+        with self._lock:
+            self.waiting_queue.append(task)
+            self._total_tasks += 1
+            if stream_callback:
+                self._callbacks[task_id] = stream_callback
+
+        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]
+            self._callbacks.pop(task_id, None)
+        return removed_active
+
+    def invoke_callback(self, task_id: str, token: str):
+        cb = self._callbacks.get(task_id)
+        if cb:
+            cb(token)
+
+    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):
+        task.status = TaskStatus.RUNNING
+        with self._lock:
+            self.active_tasks.append(task)
+
+    def return_to_waiting(self, tasks: List[Task]):
+        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):
+        with self._lock:
+            if self.waiting_queue or self.active_tasks:
+                return
+            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):
+        with self._lock:
+            self.waiting_queue.clear()
+            self.active_tasks.clear()
+            self._callbacks.clear()
+
+    def wake(self):
+        self._task_event.set()

astrai/inference/engine.py

@@ -1,132 +1,34 @@
-"""Unified inference engine for continuous batching.
-
-Layers:
-  - GenerationParams:    Immutable value object for sampling parameters.
-  - GenerationRequest:   User-facing request DTO with validation.
-  - _Result:             Thread-safe token accumulator (Observer pattern).
-  - InferenceEngine:     Facade over InferenceScheduler + async wrapper.
-"""
+"""Unified inference engine for continuous batching."""
 
 import asyncio
 import gc
 import threading
-from dataclasses import dataclass
-from typing import Any, AsyncGenerator, 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.cache import STOP
-from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.core.cache import KVCache
+from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.core.task import STOP
 from astrai.tokenize import AutoTokenizer
 
 
-@dataclass(frozen=True)
-class GenerationParams:
-    """Immutable value object for sampling hyperparameters."""
-
-    top_k: int = 50
-    top_p: float = 1.0
-    temperature: float = 1.0
-    max_tokens: int = 1024
-
-
-class GenerationRequest:
-    """Request parameters for text generation.
-
-    Encapsulates messages, sampling parameters (via GenerationParams),
-    and streaming preference for a single generation request.
-    """
-
-    def __init__(
-        self,
-        messages: List[Dict[str, str]],
-        top_k: int = 50,
-        top_p: float = 1.0,
-        temperature: float = 1.0,
-        max_len: int = 1024,
-        stream: bool = False,
-    ):
-        """Initializes a generation request.
-
-        Args:
-            messages: Conversation history as list of {"role": ..., "content": ...}.
-            top_k: Top-k sampling count (0 disables).
-            top_p: Nucleus sampling probability threshold.
-            temperature: Sampling temperature.
-            max_len: Maximum tokens to generate.
-            stream: Whether to return output as a token stream.
-        """
-        self.messages = messages
-        self.params = GenerationParams(
-            top_k=top_k,
-            top_p=top_p,
-            temperature=temperature,
-            max_tokens=max_len,
-        )
-        self.stream = stream
-        self._validate()
-
-    @property
-    def top_k(self) -> int:
-        return self.params.top_k
-
-    @property
-    def top_p(self) -> float:
-        return self.params.top_p
-
-    @property
-    def temperature(self) -> float:
-        return self.params.temperature
-
-    @property
-    def max_len(self) -> int:
-        return self.params.max_tokens
-
-    def _validate(self):
-        """Validates sampling parameter ranges."""
-        if not (isinstance(self.top_k, int) and self.top_k >= 0):
-            raise ValueError("top_k must be a non-negative integer")
-        if not (0.0 <= self.top_p <= 1.0):
-            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:
-    """Thread-safe token accumulator for streaming and non-streaming modes.
-
-    Supports multiple concurrent generation tasks with per-index result tracking.
-    Uses a threading.Condition for efficient completion notification
-    and a threading.Event for streaming wakeup.
-    """
+class GenerateResult:
+    """Thread-safe token accumulator for streaming and non-streaming modes."""
 
     def __init__(self, count: int = 1):
-        """Initializes the accumulator.
-
-        Args:
-            count: Number of concurrent generation tasks to track.
-        """
         self._cond = threading.Condition()
         self._event = threading.Event()
-        self.tokens: List[str] = []
+        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):
-        """Appends a token to the result buffer.
-
-        In non-streaming mode, tokens are concatenated into results[idx].
-        The sentinel STOP marks a task as complete.
-
-        Args:
-            token: The decoded token string, or STOP sentinel.
-            idx: Index of the generation task this token belongs to.
-        """
         with self._cond:
-            self.tokens.append(token)
+            self.tokens.append((idx, token))
             if token is not STOP:
                 self.results[idx] += token
             else:
@@ -136,12 +38,7 @@ class _Result:
                     self._cond.notify_all()
             self._event.set()
 
-    def pop_all(self) -> List[str]:
-        """Returns and clears all accumulated tokens.
-
-        Returns:
-            List of token strings since the last call.
-        """
+    def pop_all(self) -> List[Tuple[int, str]]:
         with self._cond:
             out = self.tokens.copy()
             self.tokens.clear()
@@ -150,45 +47,52 @@ class _Result:
             return out
 
     def wait(self, timeout: Optional[float] = None) -> bool:
-        """Blocks until new tokens arrive or the timeout expires.
-
-        Args:
-            timeout: Maximum wait time in seconds (None = infinite).
-
-        Returns:
-            True if the event was set (new data available), False on timeout.
-        """
         return self._event.wait(timeout=timeout)
 
-    def wait_completion(self) -> None:
-        """Blocks until all tasks complete (non-streaming).
-
-        Uses a Condition to sleep efficiently instead of busy-waiting.
-        The calling thread is parked until a STOP signal arrives.
-        """
+    def wait_completion(self, timeout: float = 300.0):
         with self._cond:
-            self._cond.wait_for(lambda: self._completed >= self._total)
+            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]:
-        """Returns all accumulated results for non-streaming mode.
-
-        Returns:
-            List of complete generated strings, one per task index.
-        """
         with self._cond:
             return self.results.copy()
 
 
+class GenerationRequest:
+    """Request parameters for text generation."""
+
+    def __init__(
+        self,
+        messages: List[Dict[str, str]],
+        top_k: int = 50,
+        top_p: float = 1.0,
+        temperature: float = 1.0,
+        max_tokens: Optional[int] = None,
+        stream: bool = False,
+    ):
+        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 positive number")
+
+        self.messages = messages
+        self.top_k = top_k
+        self.top_p = top_p
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.stream = stream
+
+
 class InferenceEngine:
-    """Unified inference engine backed by continuous-batching scheduler.
-
-    Usage:
-        with InferenceEngine(model, tokenizer) as engine:
-            for token in engine.generate("hello", stream=True):
-                print(token, end="")
-
-            text = engine.generate("hello")
-    """
+    """Unified inference engine backed by continuous-batching scheduler."""
 
     def __init__(
         self,
@@ -198,18 +102,8 @@ class InferenceEngine:
         max_seq_len: Optional[int] = None,
         max_prompt_len: int = 2048,
         page_size: int = 128,
+        cache: Optional[KVCache] = None,
     ):
-        """Initializes the inference engine.
-
-        Args:
-            model: The model instance.
-            tokenizer: The tokenizer instance.
-            max_batch_size: Maximum number of concurrent tasks.
-            max_seq_len: Maximum sequence length.
-            max_prompt_len: Maximum prompt tokens.
-            compile: Whether to compile the model with torch.compile.
-            page_size: Number of tokens per KV cache page.
-        """
         self.model = model
         self.tokenizer = tokenizer
         self.scheduler = InferenceScheduler(
@@ -218,7 +112,7 @@ class InferenceEngine:
             max_batch_size=max_batch_size,
             max_seq_len=max_seq_len,
             max_prompt_len=max_prompt_len,
-            page_size=page_size,
+            cache=cache,
         )
 
         self.scheduler.start()
@@ -234,25 +128,11 @@ 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,
-    ) -> Union[Generator[str, None, None], str, List[str]]:
-        """Generates text from a prompt.
-
-        Args:
-            prompt: Single string or list of strings for batch generation.
-            stream: If True, returns a generator yielding tokens one by one.
-            max_tokens: Maximum number of tokens to generate.
-            temperature: Sampling temperature.
-            top_p: Nucleus sampling probability threshold.
-            top_k: Top-k sampling count (0 disables).
-
-        Returns:
-            Generator (stream=True), single string (non-stream, single prompt),
-            or list of strings (non-stream, batch prompts).
-        """
+    ) -> Union[Generator, str, List[str]]:
         is_batch = isinstance(prompt, list)
         prompts = prompt if is_batch else [prompt]
 
@@ -268,26 +148,11 @@ class InferenceEngine:
     def generate_async(
         self,
         prompt: str,
-        max_tokens: int = 1024,
+        max_tokens: Optional[int] = None,
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = 50,
     ) -> AsyncGenerator[str, None]:
-        """Async streaming generator that does not block the event loop.
-
-        Runs the synchronous generator in a background thread pool executor,
-        yielding tokens to the async consumer as they arrive.
-
-        Args:
-            prompt: Input text to generate from.
-            max_tokens: Maximum tokens to generate.
-            temperature: Sampling temperature.
-            top_p: Nucleus sampling threshold.
-            top_k: Top-k sampling count.
-
-        Yields:
-            Decoded token strings as they are generated.
-        """
         sync_gen = self._generate_streaming(
             [prompt], False, max_tokens, temperature, top_p, top_k
         )
@@ -304,14 +169,6 @@ class InferenceEngine:
 
     @staticmethod
     def _next_token(gen: Generator) -> Optional[str]:
-        """Retrieves the next token from a synchronous generator.
-
-        Args:
-            gen: A synchronous generator yielding token strings.
-
-        Returns:
-            The next token, or None if the generator is exhausted.
-        """
         try:
             return next(gen)
         except StopIteration:
@@ -320,77 +177,80 @@ class InferenceEngine:
     def generate_with_request(
         self, request: GenerationRequest
     ) -> Union[Generator[str, None, None], str, List[str]]:
-        """Generates text from a structured GenerationRequest.
-
-        Applies the chat template to the request's messages before generation.
-
-        Args:
-            request: A GenerationRequest with messages and parameters.
-
-        Returns:
-            Generator, string, or list of strings (see generate()).
-        """
         prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
         return self.generate(
             prompt=prompt,
             stream=request.stream,
-            max_tokens=request.params.max_tokens,
-            temperature=request.params.temperature,
-            top_p=request.params.top_p,
-            top_k=request.params.top_k,
+            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,
-    ) -> Generator[str, None, None]:
-        """Internal streaming generator.
-
-        Polls the _Result accumulator in a loop, yielding tokens as they arrive.
-        Cleans up the scheduler task on GeneratorExit.
-
-        Args:
-            prompts: List of prompts (only first is used; batch not yet supported).
-            is_batch: If True, raises NotImplementedError.
-            max_tokens: Maximum tokens to generate.
-            temperature: Sampling temperature.
-            top_p: Nucleus sampling threshold.
-            top_k: Top-k sampling count.
-
-        Yields:
-            Decoded token strings.
-        """
-        if is_batch:
-            raise NotImplementedError("Batch streaming not yet supported")
-
-        result = _Result()
-
-        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=lambda tok: result.append(tok, 0),
+    ) -> 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:
+                while remaining > 0:
+                    items = result.pop_all()
+                    for idx, token in items:
                         if token is STOP:
-                            return
-                        yield token
-                    if not result.wait(timeout=0.05):
-                        pass
+                            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:
-                self.scheduler.remove_task(task_id)
+                for tid in task_ids:
+                    self.scheduler.remove_task(tid)
 
         return gen()
 
@@ -398,62 +258,32 @@ class InferenceEngine:
         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]]:
-        """Internal non-streaming generator.
+        result, task_ids = self._submit_tasks(
+            prompts, max_tokens, temperature, top_p, top_k
+        )
 
-        Submits all prompts to the scheduler and waits for all to complete.
+        try:
+            result.wait_completion()
+        except TimeoutError:
+            for tid in task_ids:
+                self.scheduler.remove_task(tid)
+            raise
 
-        Args:
-            prompts: List of prompt strings.
-            is_batch: Whether multiple prompts were provided.
-            max_tokens: Maximum tokens to generate.
-            temperature: Sampling temperature.
-            top_p: Nucleus sampling threshold.
-            top_k: Top-k sampling count.
-
-        Returns:
-            Single string for one prompt, list of strings for batch.
-        """
-        result = _Result(count=len(prompts))
-        task_ids = []
-
-        for i, p in enumerate(prompts):
-
-            def make_cb(idx):
-                return lambda tok: result.append(tok, idx)
-
-            task_id = self.scheduler.add_task(
-                prompt=p,
-                max_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                stream_callback=make_cb(i),
-            )
-            task_ids.append(task_id)
-
-        result.wait_completion()
-
-        for task_id in task_ids:
-            self.scheduler.remove_task(task_id)
+        for tid in task_ids:
+            self.scheduler.remove_task(tid)
 
         res = result.get_results()
         return res if is_batch else res[0]
 
     def get_stats(self) -> Dict[str, Any]:
-        """Returns current engine statistics.
-
-        Returns:
-            Dict with total_tasks, total_tokens, active_tasks, waiting_queue.
-        """
         return self.scheduler.get_stats()
 
-    def shutdown(self) -> None:
-        """Shuts down the engine, stops the scheduler, and frees GPU memory."""
+    def shutdown(self):
         self.scheduler.stop()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

astrai/inference/sample.py

@@ -29,6 +29,7 @@ class BaseSamplingStrategy(ABC):
         Returns:
             Transformed logits tensor.
         """
+        raise NotImplementedError
 
 
 class TemperatureStrategy(BaseSamplingStrategy):
@@ -41,13 +42,15 @@ class TemperatureStrategy(BaseSamplingStrategy):
     def __init__(self, temperature: Union[float, Tensor] = 1.0):
         self.temperature = temperature
 
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
         t = self.temperature
         if isinstance(t, Tensor):
+            t = t.to(logits.device, non_blocking=True).view(-1, 1)
+            t = torch.clamp(t, min=1e-8)
             if (t != 1.0).any():
-                logits = logits / t.to(logits.device, non_blocking=True).view(-1, 1)
+                logits = logits / t
         elif t != 1.0:
-            logits = logits / t
+            logits = logits / max(t, 1e-8)
         return logits
 
 
@@ -61,19 +64,29 @@ class TopKStrategy(BaseSamplingStrategy):
     def __init__(self, top_k: Union[int, Tensor] = 0):
         self.top_k = top_k
 
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
         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
-            k = min(max_k, logits.size(-1))
-        elif tk > 0:
-            k = min(tk, logits.size(-1))
-        else:
+            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
-        thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
-        logits[logits < thresholds] = filter_value
+        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
 
 
@@ -88,7 +101,9 @@ class TopPStrategy(BaseSamplingStrategy):
     def __init__(self, top_p: Union[float, Tensor] = 1.0):
         self.top_p = top_p
 
-    def _apply(self, logits, top_p, filter_value):
+    def _apply(
+        self, logits: Tensor, top_p: Union[float, Tensor], filter_value: float
+    ) -> Tensor:
         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
@@ -99,7 +114,7 @@ class TopPStrategy(BaseSamplingStrategy):
         logits[mask] = filter_value
         return logits
 
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
         tp = self.top_p
         if isinstance(tp, Tensor):
             tp = tp.to(logits.device, non_blocking=True)
@@ -130,7 +145,7 @@ class SamplingPipeline(BaseSamplingStrategy):
     def __init__(self, strategies: List[BaseSamplingStrategy]):
         self.strategies = strategies
 
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
         for strategy in self.strategies:
             logits = strategy.apply(logits, filter_value)
         return logits

astrai/inference/scheduler.py

@@ -1,411 +0,0 @@
-"""Inference scheduler for single-GPU continuous batching with paged KV cache."""
-
-import logging
-import threading
-import time
-import uuid
-from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple
-
-import torch
-from torch import Tensor
-
-from astrai.inference.cache import STOP, PagedCache
-from astrai.inference.sampling import sample
-from astrai.model.automodel import AutoModel
-from astrai.tokenize.tokenizer import AutoTokenizer
-
-logger = logging.getLogger(__name__)
-
-
-class TaskStatus(Enum):
-    """Task states in the continuous batching lifecycle."""
-
-    PENDING = "pending"
-    RUNNING = "running"
-    FINISHED = "finished"
-    ABORTED = "aborted"
-
-
-class Task:
-    """Represents a single generation request with paged KV cache tracking."""
-
-    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.page_table: List[int] = []
-        self.n_pages: int = 0
-        self._prefix_cached_tokens: int = 0
-        self.arrival_time = time.time()
-        self.finish_time: Optional[float] = None
-        self.stream_callback = stream_callback
-        self._pages_freed: bool = False
-
-    @property
-    def next_pos(self) -> int:
-        return self.input_tokens + len(self.output_ids)
-
-    def is_finished(self, stop_ids: List[int]) -> bool:
-        if 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 InferenceScheduler:
-    """Continuous batching scheduler with paged KV cache.
-
-    Runs a background generation loop with four phases per iteration:
-      1. Cleanup finished tasks and release resources.
-      2. Refill active batch from the waiting queue.
-      3. Prefill newly activated tasks.
-      4. Decode the largest same-position group of active tasks.
-    """
-
-    def __init__(
-        self,
-        model: AutoModel,
-        tokenizer: AutoTokenizer,
-        max_batch_size: int = 16,
-        max_seq_len: Optional[int] = None,
-        max_prompt_len: int = 512,
-        page_size: int = 64,
-        device: Optional[str] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        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_prompt_len = max_prompt_len
-        self.page_size = page_size
-        self.device = device or next(model.parameters()).device
-        self.dtype = dtype or next(model.parameters()).dtype
-
-        n_kv_heads = config.n_kv_heads
-        head_dim = config.dim // config.n_heads
-        n_layers = config.n_layers
-        n_pages = (
-            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
-        ) // page_size
-
-        self.page_cache = PagedCache(
-            n_layers,
-            n_pages,
-            page_size,
-            n_kv_heads,
-            head_dim,
-            self.device,
-            self.dtype,
-        )
-
-        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 _n_pages_for(self, n_tokens: int) -> int:
-        return (n_tokens + self.page_size - 1) // self.page_size
-
-    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:
-        task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
-        prompt_ids = self.tokenizer.encode(prompt)
-        if len(prompt_ids) > self.max_prompt_len:
-            prompt_ids = prompt_ids[-self.max_prompt_len :]
-
-        task = Task(
-            task_id=task_id,
-            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) -> None:
-        with self._lock:
-            removed_active = [t for t in self.active_tasks if t.task_id == task_id]
-            self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
-            self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
-
-        for task in removed_active:
-            if not task._pages_freed:
-                self._free_pages(task.page_table)
-                task.page_table.clear()
-                task.n_pages = 0
-                task._pages_freed = True
-
-    def _free_pages(self, indices: List[int]) -> None:
-        for idx in indices:
-            self.page_cache.free(idx)
-
-    def _record_page_hashes(self, task: Task, start_logical_page: int = 0) -> None:
-        full_pages = len(task.prompt_ids) // self.page_size
-        for i in range(start_logical_page, full_pages):
-            self.page_cache.record_page(task.page_table[i], task.prompt_ids, i)
-
-    def _remove_finished_tasks(self) -> None:
-        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:
-            if not task._pages_freed:
-                self._free_pages(task.page_table)
-                task.page_table.clear()
-                task.n_pages = 0
-                task._pages_freed = True
-
-        self.active_tasks = [
-            t for t in self.active_tasks if t.status != TaskStatus.FINISHED
-        ]
-
-    def _refill_active_batch(self) -> None:
-        available = self.max_batch_size - len(self.active_tasks)
-        if available <= 0:
-            return
-
-        to_add: List[Task] = []
-        with self._lock:
-            n = min(available, len(self.waiting_queue))
-            for _ in range(n):
-                to_add.append(self.waiting_queue.pop(0))
-
-        failed: List[Task] = []
-        for task in to_add:
-            prompt_len = len(task.prompt_ids)
-
-            hit_pages = self.page_cache.lookup_prefix(task.prompt_ids)
-            cached_tokens = len(hit_pages) * self.page_size
-            for p in hit_pages:
-                self.page_cache.inc_ref(p)
-
-            remaining = prompt_len - cached_tokens
-            n_new = self._n_pages_for(remaining) if remaining > 0 else 0
-            new_pages = self.page_cache.alloc_n(n_new) if n_new > 0 else []
-
-            if remaining > 0 and not new_pages:
-                for p in hit_pages:
-                    self.page_cache.free(p)
-                failed.append(task)
-                continue
-
-            task.page_table = hit_pages + new_pages
-            task.n_pages = len(task.page_table)
-            task._prefix_cached_tokens = cached_tokens
-            task.status = TaskStatus.RUNNING
-            self.active_tasks.append(task)
-
-        if failed:
-            with self._lock:
-                self.waiting_queue[:0] = failed
-
-    def _execute_prefill(
-        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
-    ) -> None:
-        tasks = sorted(tasks, key=lambda t: t.task_id)
-        batch_sz = len(tasks)
-
-        seq_len = prompt_len - start_pos
-        input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
-        input_mask = torch.ones(batch_sz, seq_len, dtype=torch.bool, device=self.device)
-
-        for i, t in enumerate(tasks):
-            input_ids[i] = torch.tensor(
-                t.prompt_ids[start_pos:prompt_len], device=self.device
-            )
-
-        page_tables = self._make_page_table_tensor(tasks)
-
-        with torch.inference_mode():
-            self.model(
-                input_ids,
-                input_mask=input_mask,
-                start_pos=start_pos,
-                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
-            )
-
-        start_logical_page = start_pos // self.page_size
-        for t in tasks:
-            self._record_page_hashes(t, start_logical_page=start_logical_page)
-
-    def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
-        if not tasks:
-            return
-
-        tasks = sorted(tasks, key=lambda t: t.task_id)
-        batch_sz = len(tasks)
-
-        for t in tasks:
-            self._maybe_alloc_page(t, start_pos)
-
-        input_ids = torch.tensor(
-            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
-            dtype=torch.long,
-            device=self.device,
-        )
-
-        active_mask = torch.ones((batch_sz, 1), dtype=torch.bool, device=self.device)
-
-        page_tables = self._make_page_table_tensor(tasks)
-        total_len = start_pos + 1
-
-        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
-        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
-        top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
-
-        with torch.inference_mode():
-            outputs = self.model(
-                input_ids.unsqueeze(1),
-                input_mask=active_mask,
-                paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
-                start_pos=start_pos,
-            )
-            logits = outputs["logits"][:, -1, :]
-
-        next_tokens = sample(
-            logits,
-            temperature=temperatures,
-            top_k=top_ks,
-            top_p=top_ps,
-        ).tolist()
-
-        for t, ntok in zip(tasks, next_tokens):
-            t.output_ids.append(ntok)
-            t.output_tokens += 1
-            pos = t.input_tokens + t.output_tokens
-            self._maybe_alloc_page(t, pos)
-            if t.stream_callback:
-                t.stream_callback(self.tokenizer.decode([ntok]))
-
-        for t in tasks:
-            if t.is_finished(self.tokenizer.stop_ids):
-                if t.stream_callback:
-                    t.stream_callback(STOP)
-
-    def _make_page_table_tensor(self, tasks: List[Task]) -> Tensor:
-        max_pages = max(t.n_pages for t in tasks)
-        rows = [t.page_table + [-1] * (max_pages - t.n_pages) for t in tasks]
-        return torch.tensor(rows, dtype=torch.long, device=self.device)
-
-    def _maybe_alloc_page(self, task: Task, pos: int) -> None:
-        needed = self._n_pages_for(pos + 1)
-        while task.n_pages < needed:
-            p = self.page_cache.alloc()
-            if p < 0:
-                break
-            task.page_table.append(p)
-            task.n_pages += 1
-
-    def _run_generation_loop(self) -> None:
-        try:
-            while self._running:
-                self._remove_finished_tasks()
-                self._refill_active_batch()
-
-                if not self.active_tasks and not self.waiting_queue:
-                    self._task_event.clear()
-                    self._task_event.wait(timeout=1.0)
-                    continue
-
-                to_prefill = [t for t in self.active_tasks if t.output_tokens == 0]
-                if to_prefill:
-                    for t in to_prefill:
-                        t.input_tokens = len(t.prompt_ids)
-
-                    groups: Dict[Tuple[int, int], List[Task]] = {}
-                    for t in to_prefill:
-                        key = (len(t.prompt_ids), t._prefix_cached_tokens)
-                        groups.setdefault(key, []).append(t)
-
-                    for (prompt_len, start_pos), group in groups.items():
-                        if start_pos < prompt_len:
-                            self._execute_prefill(group, prompt_len, start_pos)
-
-                pos_groups: Dict[int, List[Task]] = {}
-                for t in self.active_tasks:
-                    pos_groups.setdefault(t.next_pos, []).append(t)
-
-                if pos_groups:
-                    best_pos = max(pos_groups, key=lambda p: len(pos_groups[p]))
-                    self._execute_decode(pos_groups[best_pos], best_pos)
-        except Exception as e:
-            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
-            for task in self.active_tasks:
-                if task.stream_callback:
-                    task.stream_callback(STOP)
-            for task in self.waiting_queue:
-                if task.stream_callback:
-                    task.stream_callback(STOP)
-            raise
-
-    def start(self) -> None:
-        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_event.set()
-        if hasattr(self, "_loop_thread"):
-            self._loop_thread.join(timeout=2.0)
-        self.waiting_queue.clear()
-        self.active_tasks.clear()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    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),
-        }

astrai/inference/server.py

@@ -1,486 +0,0 @@
-"""
-OpenAI / Anthropic-compatible chat completion server backed by continuous-batching inference.
-"""
-
-import json
-import logging
-import time
-import uuid
-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 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__)
-
-_project_root = Path(__file__).parent.parent.parent
-
-
-class ServerState:
-    def __init__(self):
-        self.engine: Optional[InferenceEngine] = None
-        self.config: Dict[str, Any] = {
-            "device": "cuda",
-            "dtype": torch.bfloat16,
-            "param_path": None,
-            "max_batch_size": 16,
-        }
-
-
-_state = ServerState()
-
-
-class ChatMessage(BaseModel):
-    role: str
-    content: str
-
-
-class ChatCompletionRequest(BaseModel):
-    """OpenAI Chat Completion API request body."""
-
-    model: str = "astrai"
-    messages: List[ChatMessage]
-    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
-    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
-    top_k: Optional[int] = Field(default=50, ge=1)
-    stream: Optional[bool] = False
-    stop: Optional[Union[str, List[str]]] = None
-    max_tokens: Optional[int] = Field(default=2048, ge=1)
-    n: Optional[int] = Field(default=1, ge=1)
-    presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
-    frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
-    logit_bias: Optional[Dict[int, float]] = None
-    user: Optional[str] = None
-
-
-class AnthropicMessage(BaseModel):
-    role: str
-    content: Union[str, List[Dict[str, Any]]]
-
-
-class MessagesRequest(BaseModel):
-    """Anthropic Messages API request body."""
-
-    model: str = "astrai"
-    max_tokens: int = Field(default=1024, ge=1)
-    messages: List[AnthropicMessage]
-    system: Optional[str] = None
-    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
-    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
-    top_k: Optional[int] = Field(default=50, ge=1)
-    stream: Optional[bool] = False
-    stop_sequences: Optional[List[str]] = None
-
-
-def configure_server(
-    device: str = "cuda",
-    dtype: torch.dtype = torch.bfloat16,
-    param_path: Optional[Path] = None,
-    max_batch_size: int = 16,
-):
-    _state.config.update(
-        device=device,
-        dtype=dtype,
-        param_path=param_path,
-        max_batch_size=max_batch_size,
-    )
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    try:
-        load_model(
-            param_path=_state.config["param_path"],
-            device=_state.config["device"],
-            dtype=_state.config["dtype"],
-            max_batch_size=_state.config["max_batch_size"],
-        )
-    except Exception as e:
-        logger.error(f"Failed to load model: {e}")
-        raise
-    yield
-    if _state.engine:
-        _state.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,
-):
-    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}")
-
-    _state.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}")
-
-
-def _get_engine() -> InferenceEngine:
-    if _state.engine is None:
-        raise HTTPException(status_code=503, detail="Engine not initialized")
-    return _state.engine
-
-
-def _make_chunk(
-    delta: Dict[str, str],
-    finish_reason: Optional[str] = None,
-    *,
-    resp_id: str,
-    created: int,
-    model: str,
-    index: int = 0,
-) -> str:
-    """Build a single SSE ``data:`` chunk matching OpenAI streaming format."""
-    data = {
-        "id": resp_id,
-        "object": "chat.completion.chunk",
-        "created": created,
-        "model": model,
-        "choices": [
-            {
-                "index": index,
-                "delta": delta,
-                "finish_reason": finish_reason,
-            }
-        ],
-    }
-    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
-
-
-@app.get("/health")
-async def health():
-    return {
-        "status": "ok",
-        "model_loaded": _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):
-    """OpenAI-compatible chat completion endpoint (streaming + non-streaming)."""
-    engine = _get_engine()
-    resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
-    created = int(time.time())
-    model = request.model
-
-    prompt = engine.tokenizer.apply_chat_template(
-        [{"role": m.role, "content": m.content} for m in request.messages],
-        tokenize=False,
-    )
-    prompt_tokens = len(engine.tokenizer.encode(prompt))
-
-    if request.stream:
-        agen = engine.generate_async(
-            prompt=prompt,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            top_k=request.top_k,
-        )
-
-        async def event_stream():
-            yield _make_chunk(
-                {"role": "assistant"},
-                finish_reason=None,
-                resp_id=resp_id,
-                created=created,
-                model=model,
-            )
-
-            completion_tokens = 0
-            async for token in agen:
-                yield _make_chunk(
-                    {"content": token},
-                    finish_reason=None,
-                    resp_id=resp_id,
-                    created=created,
-                    model=model,
-                )
-                completion_tokens += 1
-
-            yield _make_chunk(
-                {},
-                finish_reason="stop",
-                resp_id=resp_id,
-                created=created,
-                model=model,
-            )
-
-            usage = {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
-                "total_tokens": prompt_tokens + completion_tokens,
-            }
-            yield f"data: {json.dumps(usage, ensure_ascii=False)}\n\n"
-            yield "data: [DONE]\n\n"
-
-        return StreamingResponse(
-            event_stream(),
-            media_type="text/event-stream",
-            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
-        )
-
-    completion_tokens = 0
-    chunks: List[str] = []
-    agen = engine.generate_async(
-        prompt=prompt,
-        max_tokens=request.max_tokens,
-        temperature=request.temperature,
-        top_p=request.top_p,
-        top_k=request.top_k,
-    )
-    async for token in agen:
-        chunks.append(token)
-        completion_tokens += 1
-    content = "".join(chunks)
-
-    return {
-        "id": resp_id,
-        "object": "chat.completion",
-        "created": created,
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "message": {"role": "assistant", "content": content},
-                "finish_reason": "stop",
-            }
-        ],
-        "usage": {
-            "prompt_tokens": prompt_tokens,
-            "completion_tokens": completion_tokens,
-            "total_tokens": prompt_tokens + completion_tokens,
-        },
-    }
-
-
-def _make_anthropic_sse(event: str, data: Dict[str, Any]) -> str:
-    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
-
-
-def _check_stop_sequence(text: str, stop_sequences: List[str]) -> Optional[str]:
-    for seq in stop_sequences:
-        if seq and seq in text:
-            return seq
-    return None
-
-
-def _extract_text_content(content: Union[str, List[Dict[str, Any]]]) -> str:
-    if isinstance(content, str):
-        return content
-    if isinstance(content, list):
-        for block in content:
-            if isinstance(block, dict) and block.get("type") == "text":
-                return block.get("text", "")
-    return ""
-
-
-def _build_anthropic_messages(
-    messages: List[AnthropicMessage], system: Optional[str]
-) -> List[Dict[str, str]]:
-    result: List[Dict[str, str]] = []
-    if system:
-        result.append({"role": "system", "content": system})
-    for m in messages:
-        content = _extract_text_content(m.content)
-        if content:
-            result.append({"role": m.role, "content": content})
-    return result
-
-
-@app.post("/v1/messages")
-async def create_message(request: MessagesRequest):
-    """Anthropic-compatible Messages API endpoint (streaming + non-streaming)."""
-    engine = _get_engine()
-    resp_id = f"msg_{uuid.uuid4().hex[:24]}"
-    model = request.model
-
-    chat_messages = _build_anthropic_messages(request.messages, request.system)
-    prompt = engine.tokenizer.apply_chat_template(chat_messages, tokenize=False)
-    prompt_tokens = len(engine.tokenizer.encode(prompt))
-
-    stop_sequences = request.stop_sequences or []
-
-    if request.stream:
-        agen = engine.generate_async(
-            prompt=prompt,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            top_k=request.top_k,
-        )
-
-        async def event_stream():
-            yield _make_anthropic_sse(
-                "message_start",
-                {
-                    "type": "message_start",
-                    "message": {
-                        "id": resp_id,
-                        "type": "message",
-                        "role": "assistant",
-                        "model": model,
-                        "content": [],
-                        "usage": {"input_tokens": prompt_tokens},
-                    },
-                },
-            )
-
-            yield _make_anthropic_sse(
-                "content_block_start",
-                {
-                    "type": "content_block_start",
-                    "index": 0,
-                    "content_block": {"type": "text", "text": ""},
-                },
-            )
-
-            completion_tokens = 0
-            accumulated = ""
-            stopped_seq: Optional[str] = None
-            async for token in agen:
-                accumulated += token
-                completion_tokens += 1
-
-                matched = _check_stop_sequence(accumulated, stop_sequences)
-                if matched:
-                    text = accumulated[: accumulated.rfind(matched)]
-                    stopped_seq = matched
-                    if text:
-                        yield _make_anthropic_sse(
-                            "content_block_delta",
-                            {
-                                "type": "content_block_delta",
-                                "index": 0,
-                                "delta": {"type": "text_delta", "text": text},
-                            },
-                        )
-                    break
-
-                yield _make_anthropic_sse(
-                    "content_block_delta",
-                    {
-                        "type": "content_block_delta",
-                        "index": 0,
-                        "delta": {"type": "text_delta", "text": token},
-                    },
-                )
-
-            yield _make_anthropic_sse(
-                "content_block_stop",
-                {"type": "content_block_stop", "index": 0},
-            )
-
-            stop_reason = "stop_sequence" if stopped_seq else "end_turn"
-            yield _make_anthropic_sse(
-                "message_delta",
-                {
-                    "type": "message_delta",
-                    "delta": {"stop_reason": stop_reason, "stop_sequence": stopped_seq},
-                    "usage": {"output_tokens": completion_tokens},
-                },
-            )
-
-            yield _make_anthropic_sse(
-                "message_stop",
-                {"type": "message_stop"},
-            )
-
-        return StreamingResponse(
-            event_stream(),
-            media_type="text/event-stream",
-            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
-        )
-
-    completion_tokens = 0
-    chunks: List[str] = []
-    agen = engine.generate_async(
-        prompt=prompt,
-        max_tokens=request.max_tokens,
-        temperature=request.temperature,
-        top_p=request.top_p,
-        top_k=request.top_k,
-    )
-    stopped_seq: Optional[str] = None
-    accumulated = ""
-    async for token in agen:
-        chunks.append(token)
-        completion_tokens += 1
-        accumulated += token
-        matched = _check_stop_sequence(accumulated, stop_sequences)
-        if matched:
-            stopped_seq = matched
-            break
-
-    content = "".join(chunks)
-    if stopped_seq:
-        idx = content.rfind(stopped_seq)
-        if idx != -1:
-            content = content[:idx]
-
-    return {
-        "id": resp_id,
-        "type": "message",
-        "role": "assistant",
-        "model": model,
-        "content": [{"type": "text", "text": content}],
-        "stop_reason": "stop_sequence" if stopped_seq else "end_turn",
-        "stop_sequence": stopped_seq,
-        "usage": {
-            "input_tokens": prompt_tokens,
-            "output_tokens": completion_tokens,
-        },
-    }
-
-
-def run_server(
-    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,
-):
-    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

@@ -4,18 +4,22 @@ AutoModel base class for model loading and saving.
 
 from contextlib import contextmanager
 from pathlib import Path
-from typing import 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.factory import Registry
+from astrai.config.model_config import BaseModelConfig, ConfigFactory
+from astrai.factory import BaseFactory
+from astrai.serialization import load_model_config, load_model_weights, save_model
 
 
 @contextmanager
 def _disable_random_init(enable: bool = True):
-    init_functions = [
+    if not enable:
+        yield
+        return
+
+    names = (
         "xavier_normal_",
         "xavier_uniform_",
         "kaiming_normal_",
@@ -25,60 +29,27 @@ def _disable_random_init(enable: bool = True):
         "constant_",
         "normal_",
         "uniform_",
-    ]
-    original_funcs = {}
-    for name in init_functions:
-        if enable and hasattr(nn.init, name):
-            original_funcs[name] = getattr(nn.init, name)
-            setattr(nn.init, name, lambda *args, **kwargs: None)
+    )
+    orig = {n: getattr(nn.init, n) for n in names if hasattr(nn.init, n)}
+    for n in orig:
+        setattr(nn.init, n, lambda *a, **kw: None)
     try:
         yield
     finally:
-        if enable:
-            for name, orig_func in original_funcs.items():
-                setattr(nn.init, name, orig_func)
+        for n, fn in orig.items():
+            setattr(nn.init, n, fn)
 
 
-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.
     """
 
-    _registry = Registry()
-
-    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.register(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 not cls._registry.contains(model_type):
-            available = cls._registry.list_names()
-            raise ValueError(
-                f"Unknown model_type: {model_type}. Available: {available}"
-            )
-        return cls._registry.get(model_type)
-
     @classmethod
     def from_pretrained(
         cls,
@@ -89,24 +60,22 @@ class AutoModel(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))
-        else:
+        if not config_path.exists():
             raise FileNotFoundError(f"Config file not found: {config_path}")
 
-        model_type = config.model_type or "transformer"
-        actual_cls = cls.get_model_class(model_type)
+        raw = load_model_config(str(model_path))
+        config = ConfigFactory.load(raw)
+        model_type = config.model_type or "autoregressive_lm"
+
+        actual_cls = AutoModel.get_component_class(model_type)
 
         with _disable_random_init(enable=disable_random_init):
             model = actual_cls(config)
 
-        # Load weights
         weights_path = model_path / "model.safetensors"
         if weights_path.exists():
-            state_dict = st.load_file(str(weights_path))
+            state_dict = load_model_weights(str(model_path))
             model.load_state_dict(state_dict, strict=strict)
 
         return model
@@ -114,15 +83,12 @@ class AutoModel(nn.Module):
     def save_pretrained(
         self,
         save_directory: Union[str, Path],
-    ) -> None:
-        save_path = Path(save_directory)
-        save_path.mkdir(parents=True, exist_ok=True)
-
-        # Save config
-        self.config.save(str(save_path / "config.json"))
-
-        # Save weights
-        st.save_file(self.state_dict(), str(save_path / "model.safetensors"))
+    ):
+        save_model(
+            config=self.config.to_dict(),
+            state_dict=self.state_dict(),
+            save_directory=str(save_directory),
+        )
 
     def to(self, *args, **kwargs) -> Self:
         """Move model to device/dtype."""

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,214 @@
+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 CacheView
+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]):
+    pass
+
+
+@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,
+        n_layers: int = 1,
+    ):
+        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, init_std=0.02 / (2 * n_layers) ** 0.5)
+
+        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[CacheView] = 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,
+        n_layers: int = 1,
+    ):
+        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, init_std=0.02 / (2 * n_layers) ** 0.5
+        )
+
+        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[CacheView] = 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,39 @@
+from dataclasses import asdict
+from typing import Optional
+
+import torch.nn as nn
+from torch import Tensor
+
+from astrai.inference.core.cache import CacheView
+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, config, layer_id: int):
+        super().__init__()
+        cfg = asdict(config)
+        cfg["down_init_std"] = 0.02 / (2 * config.n_layers) ** 0.5
+        self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
+        self.input_norm = RMSNorm(config.dim, config.norm_eps)
+        self.post_attention_norm = RMSNorm(config.dim, config.norm_eps)
+        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
+
+    def forward(
+        self,
+        x: Tensor,
+        rotary_emb: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        paged_cache: Optional[CacheView] = 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,26 @@
+import math
+
+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, neftune_alpha: float = 0.0):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
+        self.neftune_noise_alpha = neftune_alpha
+
+    def set_neftune_alpha(self, alpha: float):
+        self.neftune_noise_alpha = alpha
+
+    def reset_parameters(self):
+        nn.init.normal_(self.weight, mean=0.0, std=0.02)
+
+    def forward(self, x: Tensor) -> Tensor:
+        out = F.embedding(x, self.weight)
+        if self.training and self.neftune_noise_alpha > 0.0:
+            eps = self.neftune_noise_alpha / math.sqrt(out.size(1))
+            out = out + eps * torch.randn_like(out)
+        return out

astrai/model/components/linear.py

@@ -0,0 +1,24 @@
+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, init_std: float = 0.02
+    ):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
+        self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
+        self.init_std = init_std
+
+    def reset_parameters(self):
+        nn.init.normal_(self.weight, mean=0.0, std=self.init_std)
+        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,194 @@
+import logging
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Optional, Set
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from astrai.model.components.linear import Linear
+from astrai.serialization import (
+    load_json,
+    load_safetensors,
+    save_json,
+    save_safetensors,
+)
+
+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)
+    save_safetensors(lora_sd, path / "adapter_model.safetensors")
+    save_json(asdict(config), path / "adapter_config.json")
+    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)
+    raw = load_json(path / "adapter_config.json")
+    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 = load_safetensors(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,100 @@
+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]):
+    pass
+
+
+@FFNFactory.register("mlp")
+class MLP(nn.Module):
+    def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02):
+        super().__init__()
+        self.up = Linear(dim, dim_ffn)
+        self.gate = Linear(dim, dim_ffn)
+        self.down = Linear(dim_ffn, dim, init_std=down_init_std)
+
+    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",
+        n_layers: int = 1,
+    ):
+        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)
+        moe_scale = 1 / max(n_shared_experts, 1) + 1 / n_activated_experts
+        down_init_std = 0.02 / (2 * n_layers * moe_scale) ** 0.5
+
+        self.shared_experts = nn.ModuleList(
+            [
+                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                for _ in range(n_shared_experts)
+            ]
+        )
+        self.routed_experts = nn.ModuleList(
+            [
+                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                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,71 @@
+from typing import Dict, 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 ntk_base(base: float, dim: int, factor: float) -> float:
+    return base * (factor ** (dim / (dim - 2)))
+
+
+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,
+        rope_scaling: Optional[Dict] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_len = max_len
+        self.base = base
+        self.rope_scaling = rope_scaling
+
+        if rope_scaling is not None:
+            scaling_type = rope_scaling.get("type", "ntk")
+            factor = rope_scaling.get("factor", 1.0)
+            if scaling_type == "ntk":
+                self.base = ntk_base(base, dim, factor)
+
+        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,89 @@
+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, rope_scaling=config.rope_scaling
+        )
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+        )
+
+        self.layers = nn.ModuleList(
+            [DecoderBlock(config, 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)
+
+        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,337 +0,0 @@
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-
-from astrai.inference.cache import CacheView
-
-
-def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
-    """Repeat KV heads n_rep times for GQA."""
-    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]:
-    """Precompute cos/sin for RoPE."""
-    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 via cos/sin (shape-preserving)."""
-    dtype = x.dtype
-    cos, sin = rotary_emb
-    cos = cos.unsqueeze(0).unsqueeze(2)
-    sin = sin.unsqueeze(0).unsqueeze(2)
-    x_real = x[..., 0::2]
-    x_imag = x[..., 1::2]
-    x_real_rot = x_real * cos - x_imag * sin
-    x_imag_rot = x_real * sin + x_imag * cos
-    x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1)
-    x_out = x_out.view(*x_out.shape[:-2], -1)
-    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,
-        paged_cache: Optional[CacheView] = None,
-        start_pos: int = 0,
-    ) -> Tensor:
-        bsz, seq_len, _ = x.size()
-        is_causal = mask is None
-
-        # (bsz, seq_len, dim) -> (bsz, seq_len, n_heads, head_dim)
-        q = self._split_heads(self.q_proj(x), self.n_heads)
-        k = self._split_heads(self.k_proj(x), self.n_kv_heads)
-        v = self._split_heads(self.v_proj(x), self.n_kv_heads)
-        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, start_pos, k, v)
-            k, v = paged_cache.gather(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)
-        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)
-
-        # fused KV: (k_nope, k_rope, v)
-        self.kv_b_proj = Linear(
-            kv_lora_rank,
-            n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
-        )
-
-        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,
-        paged_cache: Optional[CacheView] = 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 paged_cache is not None:
-            paged_cache.write(self.layer_id, start_pos, 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, 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,
-        paged_cache: Optional[CacheView] = None,
-        start_pos: int = 0,
-    ) -> Tensor:
-        attn_output = self.attention(
-            self.input_norm(x),
-            rotary_emb,
-            attention_mask,
-            paged_cache,
-            start_pos,
-        )
-        x = attn_output + x
-
-        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,98 +1,83 @@
-from typing import Any, Mapping, Optional
+from typing import Any, Dict, Mapping, Optional
 
 import torch
 import torch.nn as nn
 from torch import Tensor
 
-from astrai.config.model_config import ModelConfig
-from astrai.inference.cache import CacheView
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.inference.core.cache import CacheView
 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:
-    """Build 4D attention mask from 2D seq_mask, with optional causal masking."""
+) -> 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 = input_tensor.size(0)
+    T = position_ids.max().item() + 1
 
-    if seq_mask is None:
-        if start_pos != 0:
-            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
-
-    if seq_mask.dim() > 2:
-        return seq_mask
-
-    batch_size = seq_mask.size(0)
-    seq_mask = seq_mask[:, : start_pos + seq_len].to(device=device, dtype=torch.bool)
-    expanded_mask = seq_mask.unsqueeze(1).expand(
-        batch_size, seq_len, start_pos + seq_len
-    )
+        attend = torch.ones(B, 1, T, dtype=torch.bool, device=device)
+    else:
+        attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1)
 
     if is_causal:
-        expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)
+        causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
+        attend = attend & causal
 
-    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)
-
-    return attention_mask
+    return attend.unsqueeze(1)
 
 
-@AutoModel.register("transformer")
-class Transformer(AutoModel):
-    """Transformer language model with paged KV cache."""
+@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
+        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(
-            config.dim // config.n_heads, config.max_len
+            rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
+        )
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
         )
-        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)
-            ]
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
         )
 
         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"
@@ -100,7 +85,7 @@ class Transformer(AutoModel):
 
         state_dict = dict(state_dict)
 
-        if self.config.tie_weight:
+        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]
@@ -116,7 +101,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]
@@ -128,17 +113,16 @@ class Transformer(AutoModel):
         input_ids: Tensor,
         input_mask: Optional[Tensor] = None,
         paged_cache: Optional[CacheView] = None,
-        start_pos: int = 0,
-    ) -> Tensor:
+        position_ids: Optional[Tensor] = None,
+    ) -> Dict[str, 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, paged_cache, 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,286 @@
+"""Unified training executor — parallel strategy + gradient accumulation."""
+
+import contextlib
+import logging
+import os
+from contextlib import contextmanager
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch.distributed.fsdp import FullStateDictConfig, StateDictType
+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):
+        return model.state_dict()
+
+    @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
+
+    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
+        total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
+        if isinstance(total_norm, torch.Tensor):
+            return total_norm.item()
+        return total_norm
+
+
+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 = int(os.environ.get("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):
+        if isinstance(model, DDP):
+            return model.module.state_dict()
+        return model.state_dict()
+
+
+@ExecutorFactory.register("fsdp")
+class FSDPExecutor(BaseExecutor):
+    def __init__(
+        self,
+        grad_accum_steps: int = 1,
+        process_group=None,
+        sharding_strategy=None,
+        cpu_offload=None,
+        auto_wrap_policy=None,
+        backward_prefetch=None,
+        mixed_precision=None,
+        ignored_modules=None,
+        param_init_fn=None,
+        sync_module_states: bool = False,
+        forward_prefetch: bool = False,
+        limit_all_gathers: bool = True,
+        ignored_states=None,
+        device_mesh=None,
+    ):
+        super().__init__(grad_accum_steps=grad_accum_steps)
+        self._fsdp_kwargs = {
+            k: v
+            for k, v in dict(
+                process_group=process_group,
+                sharding_strategy=sharding_strategy,
+                cpu_offload=cpu_offload,
+                auto_wrap_policy=auto_wrap_policy,
+                backward_prefetch=backward_prefetch,
+                mixed_precision=mixed_precision,
+                ignored_modules=ignored_modules,
+                param_init_fn=param_init_fn,
+                sync_module_states=sync_module_states,
+                forward_prefetch=forward_prefetch,
+                limit_all_gathers=limit_all_gathers,
+                use_orig_params=True,
+                ignored_states=ignored_states,
+                device_mesh=device_mesh,
+            ).items()
+            if v is not None
+        }
+        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 clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
+        if isinstance(model, FSDP) and self.use_distributed:
+            total_norm = model.clip_grad_norm_(max_norm)
+            if isinstance(total_norm, torch.Tensor):
+                return total_norm.item()
+            return total_norm
+        return super().clip_grad_norm(model, max_norm)
+
+    def unwrap_model(self, model: nn.Module):
+        if isinstance(model, FSDP) and self.use_distributed:
+            with FSDP.state_dict_type(
+                model,
+                StateDictType.FULL_STATE_DICT,
+                FullStateDictConfig(offload_to_cpu=True, rank0_only=False),
+            ):
+                return model.state_dict()
+
+        return model.state_dict()

astrai/parallel/setup.py

@@ -1,4 +1,5 @@
 import os
+from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from functools import wraps
 from typing import Callable
@@ -30,6 +31,7 @@ def get_rank() -> int:
 def setup_parallel(
     rank: int,
     world_size: int,
+    local_rank: int,
     backend: str = "nccl",
     master_addr: str = "localhost",
     master_port: str = "29500",
@@ -41,20 +43,26 @@ def setup_parallel(
         return
 
     if world_size <= 1:
+        device_id = torch.device(device_type, local_rank)
+        os.environ["LOCAL_RANK"] = str(local_rank)
+        os.environ["WORLD_SIZE"] = "1"
+        os.environ["LOCAL_DEVICE"] = str(device_id)
         yield None
         return
 
-    device_id = torch.device(device_type, rank)
+    device_id = torch.device(device_type, local_rank)
 
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port
-    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(local_rank)
     os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["LOCAL_DEVICE"] = str(device_id)
 
-    dist.init_process_group(
-        rank=rank, world_size=world_size, backend=backend, device_id=device_id
-    )
+    pg_kwargs = dict(rank=rank, world_size=world_size, backend=backend)
+    if backend in ("nccl", "ccl"):
+        pg_kwargs["device_id"] = device_id
+
+    dist.init_process_group(**pg_kwargs)
 
     try:
         if backend == "nccl" and torch.cuda.is_available():
@@ -90,7 +98,7 @@ def only_on_rank(rank, sync=False):
     return decorator
 
 
-def wrapper_spawn_func(
+def _run_single_rank(
     rank: int,
     world_size: int,
     backend: str,
@@ -100,20 +108,108 @@ def wrapper_spawn_func(
     func: Callable,
     kwargs: dict,
 ):
-    try:
+    with setup_parallel(
+        rank=rank,
+        world_size=world_size,
+        local_rank=rank,
+        backend=backend,
+        master_addr=master_addr,
+        master_port=master_port,
+        device_type=device_type,
+    ):
+        func(**kwargs)
+
+
+class LaunchStrategy(ABC):
+    """Strategy for launching a function in a distributed context."""
+
+    def __init__(
+        self,
+        world_size: int,
+        backend: str,
+        master_addr: str,
+        master_port: str,
+        device_type: str,
+        start_method: str,
+    ):
+        self.world_size = world_size
+        self.backend = backend
+        self.master_addr = master_addr
+        self.master_port = master_port
+        self.device_type = device_type
+        self.start_method = start_method
+
+    @abstractmethod
+    def launch(self, func: Callable, **kwargs):
+        raise NotImplementedError
+
+
+class TorchrunStrategy(LaunchStrategy):
+    """External orchestrator (torchrun, SLURM, K8s) — env vars pre-set."""
+
+    def launch(self, func: Callable, **kwargs):
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+        local_rank = int(os.environ.get("LOCAL_RANK", rank))
         with setup_parallel(
             rank=rank,
             world_size=world_size,
-            backend=backend,
-            master_addr=master_addr,
-            master_port=master_port,
-            device_type=device_type,
+            local_rank=local_rank,
+            backend=self.backend,
+            master_addr=os.environ.get("MASTER_ADDR", self.master_addr),
+            master_port=os.environ.get("MASTER_PORT", self.master_port),
+            device_type=self.device_type,
         ):
             func(**kwargs)
 
-    except Exception as e:
-        print(f"Error in rank {rank}: {e}")
-        raise
+
+class LocalStrategy(LaunchStrategy):
+    """Local launcher — single-process or mp.start_processes."""
+
+    def launch(self, func: Callable, **kwargs):
+        args = (
+            self.world_size,
+            self.backend,
+            self.master_addr,
+            self.master_port,
+            self.device_type,
+            func,
+            kwargs,
+        )
+
+        if self.world_size == 1:
+            _run_single_rank(0, *args)
+            return
+
+        ctx = mp.start_processes(
+            _run_single_rank,
+            args=args,
+            nprocs=self.world_size,
+            start_method=self.start_method,
+            join=False,
+        )
+        try:
+            while not ctx.join():
+                pass
+        except BaseException:
+            for p in ctx.processes:
+                p.terminate()
+            ctx.join()
+            raise
+
+
+def _detect_launcher() -> str:
+    """Detect the distributed launcher from environment.
+
+    Returns one of: "torchelastic", "torchrun", "external", "local".
+    """
+    if dist.is_torchelastic_launched():
+        return "torchelastic"
+    if "LOCAL_WORLD_SIZE" in os.environ:
+        return "torchrun"
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        return "external"
+    return "local"
 
 
 def spawn_parallel_fn(
@@ -123,39 +219,16 @@ def spawn_parallel_fn(
     master_addr: str = "localhost",
     master_port: str = "29500",
     device_type: str = "cuda",
+    start_method: str = "spawn",
     **kwargs,
 ):
-    # clear environment variables
-    for key in [
-        "MASTER_ADDR",
-        "MASTER_PORT",
-        "RANK",
-        "WORLD_SIZE",
-        "LOCAL_RANK",
-        "LOCAL_DEVICE",
-    ]:
-        if key in os.environ:
-            del os.environ[key]
-
-    if world_size == 1:
-        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)
-        return
-
-    wrapper_spawn_func_args = (
-        world_size,
-        backend,
-        master_addr,
-        master_port,
-        device_type,
-        func,
-        kwargs,
-    )
-
-    mp.spawn(
-        wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
-    )
+    launcher = _detect_launcher()
+    if launcher in ("torchelastic", "torchrun", "external"):
+        strategy = TorchrunStrategy(
+            world_size, backend, master_addr, master_port, device_type, start_method
+        )
+    else:
+        strategy = LocalStrategy(
+            world_size, backend, master_addr, master_port, device_type, start_method
+        )
+    strategy.launch(func, **kwargs)

astrai/preprocessing/__init__.py

@@ -0,0 +1,32 @@
+from astrai.preprocessing.builder import (
+    BaseMaskBuilder,
+    MaskBuilderFactory,
+    SectionedMaskBuilder,
+)
+from astrai.preprocessing.packing import (
+    PackingStrategy,
+    PackingStrategyFactory,
+)
+from astrai.preprocessing.pipeline import Pipeline, filter_by_length
+from astrai.preprocessing.position_id import (
+    PositionIdStrategy,
+    PositionIdStrategyFactory,
+)
+from astrai.preprocessing.writer import (
+    StoreWriter,
+    StoreWriterFactory,
+)
+
+__all__ = [
+    "BaseMaskBuilder",
+    "MaskBuilderFactory",
+    "PackingStrategy",
+    "PackingStrategyFactory",
+    "Pipeline",
+    "PositionIdStrategy",
+    "PositionIdStrategyFactory",
+    "SectionedMaskBuilder",
+    "StoreWriter",
+    "StoreWriterFactory",
+    "filter_by_length",
+]

astrai/preprocessing/builder.py

@@ -0,0 +1,315 @@
+"""Mask building for preprocessing pipeline.
+
+:class:`SectionRenderer` converts section specs into token ids and loss
+masks (template / text / value extraction).  :class:`SectionedMaskBuilder`
+orchestrates single-output / multi-output (DPO / GRPO) assembly.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from astrai.factory import BaseFactory
+
+
+def _extract_domain(item: dict, domain_key: Optional[str]) -> str:
+    if not domain_key:
+        return "__default__"
+    val = item.get(domain_key, "__default__")
+    return val if isinstance(val, str) else "__default__"
+
+
+def _resolve_action(action: str, role: str, config) -> str:
+    if action == "$role":
+        return config.mask.get(role, config.mask_default)
+    return action
+
+
+class SectionRenderer:
+    """Render section specs into ``(ids, loss_mask)`` tuples."""
+
+    def process_sections(
+        self,
+        item: dict,
+        sections: list,
+        config,
+        tokenizer,
+        *,
+        is_top_level: bool = False,
+    ):
+        all_ids: list[int] = []
+        loss_mask: list[int] = []
+
+        has_template = any(s.get("template") for s in sections)
+        is_text_config = not has_template and all(
+            s["action"] == "train" for s in sections
+        )
+
+        if is_top_level and has_template and tokenizer.bos_token_id is not None:
+            all_ids.append(tokenizer.bos_token_id)
+            loss_mask.append(0)
+
+        first_section = True
+        for sec in sections:
+            field = sec["field"]
+            action = sec["action"]
+            use_template = sec.get("template", False)
+            add_special = sec.get(
+                "add_special_tokens", not use_template and first_section
+            )
+
+            if use_template:
+                success = self._append_template(
+                    item, field, action, tokenizer, config, all_ids, loss_mask
+                )
+                if not success:
+                    continue
+            else:
+                success = self._append_text(
+                    item,
+                    field,
+                    action,
+                    tokenizer,
+                    add_special,
+                    is_text_config,
+                    config,
+                    all_ids,
+                    loss_mask,
+                )
+                if not success:
+                    continue
+
+            first_section = False
+
+        max_len = config.preprocessing.max_seq_len
+        all_ids = all_ids[:max_len]
+        loss_mask = loss_mask[: len(all_ids)]
+
+        if not all_ids:
+            return None, None
+
+        if is_top_level and has_template and len(all_ids) <= 1:
+            return None, None
+
+        return all_ids, loss_mask
+
+    def process_list_field(self, item: dict, sections: list, config, tokenizer):
+        all_ids: list[int] = []
+        loss_mask: list[int] = []
+
+        for sec in sections:
+            field = sec["field"]
+            action = sec["action"]
+            use_template = sec.get("template", False)
+
+            values = item.get(field)
+            if not isinstance(values, list):
+                continue
+
+            for val in values:
+                if use_template:
+                    if isinstance(val, list):
+                        wrapper = {field: val}
+                        self._append_template(
+                            wrapper,
+                            field,
+                            action,
+                            tokenizer,
+                            config,
+                            all_ids,
+                            loss_mask,
+                        )
+                else:
+                    wrapper = {field: str(val)}
+                    self._append_text(
+                        wrapper,
+                        field,
+                        action,
+                        tokenizer,
+                        False,
+                        False,
+                        config,
+                        all_ids,
+                        loss_mask,
+                    )
+
+        max_len = config.preprocessing.max_seq_len
+        all_ids = all_ids[:max_len]
+        loss_mask = loss_mask[: len(all_ids)]
+
+        if not all_ids:
+            return None, None
+        return all_ids, loss_mask
+
+    @staticmethod
+    def is_value_section(sections: list) -> bool:
+        return len(sections) == 1 and sections[0].get("action") == "value"
+
+    @staticmethod
+    def extract_raw_value(item: dict, sections: list):
+        sec = sections[0]
+        field = sec["field"]
+        raw = item.get(field)
+        if raw is None:
+            return None
+        if isinstance(raw, list):
+            return [float(v) for v in raw]
+        return [float(raw)]
+
+    def _append_template(
+        self, item, field, action, tokenizer, config, all_ids, loss_mask
+    ):
+        messages = item.get(field)
+        if not isinstance(messages, list) or not messages:
+            return False
+        for msg in messages:
+            role = msg.get("role", "")
+            act = _resolve_action(action, role, config)
+            rendered = tokenizer.apply_chat_template(
+                [msg], tokenize=False, add_generation_prompt=False
+            )
+            ids = tokenizer.encode(rendered, add_special_tokens=False)
+            all_ids.extend(ids)
+            val = 1 if act == "train" else 0
+            loss_mask.extend([val] * len(ids))
+        return True
+
+    def _append_text(
+        self,
+        item,
+        field,
+        action,
+        tokenizer,
+        add_special,
+        is_text_config,
+        config,
+        all_ids,
+        loss_mask,
+    ):
+        text = str(item.get(field, ""))
+        if not text.strip():
+            return False
+        if is_text_config:
+            pp = config.preprocessing
+            if pp.min_chars > 0 and len(text) < pp.min_chars:
+                return False
+            if len(text) > pp.max_chars:
+                return False
+        ids = tokenizer.encode(text, add_special_tokens=add_special)
+        all_ids.extend(ids)
+        val = 1 if action == "train" else 0
+        loss_mask.extend([val] * len(ids))
+        return True
+
+
+class BaseMaskBuilder(ABC):
+    """Convert a JSONL item into token ids and optional loss_mask."""
+
+    @abstractmethod
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]: ...
+
+
+class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
+    pass
+
+
+@MaskBuilderFactory.register("sectioned")
+class SectionedMaskBuilder(BaseMaskBuilder):
+    """Config-driven builder supporting single and multi-output modes.
+
+    Single-output::
+
+        {"input": {"sections": [
+            {"field": "messages", "action": "$role", "template": true}
+        ]}}
+        → {"sequence": [...], "loss_mask": [...], "domain": "..."}
+
+    Multi-output (DPO / GRPO)::
+
+        {"input": {"sources": {
+            "chosen": {"sections": [{"field": "chosen", "action": "$role", "template": true}]},
+            "rejected": {"sections": [{"field": "rejected", "action": "$role", "template": true}]},
+        }}}
+        → {"chosen": [...], "chosen_mask": [...], "rejected": [...], "rejected_mask": [...], "domain": "..."}
+
+    Output spec fields::
+
+        sections      – list of section specs (same format as single-output)
+        list_field    – True when JSONL field holds a list (GRPO responses)
+        mask_key      – explicit loss-mask output key (default: ``"{output_key}_mask"``)
+    """
+
+    def __init__(self):
+        self.renderer = SectionRenderer()
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        sources_spec = getattr(config.input, "sources", None)
+        if sources_spec:
+            return self._build_multi(item, sources_spec, config, tokenizer)
+        return self._build_single(item, config, tokenizer)
+
+    def _build_single(self, item: dict, config, tokenizer) -> Optional[dict]:
+        sections = config.input.sections
+        if not sections:
+            return None
+
+        ids, mask = self.renderer.process_sections(
+            item, sections, config, tokenizer, is_top_level=True
+        )
+        if ids is None:
+            return None
+
+        result: dict = {
+            "sequence": ids,
+            "domain": _extract_domain(item, config.output.domain_key),
+        }
+        if not all(m == 1 for m in mask):
+            result["loss_mask"] = mask
+        return result
+
+    def _build_multi(
+        self, item: dict, sources_spec: dict, config, tokenizer
+    ) -> Optional[dict]:
+        result: dict = {}
+        any_output = False
+
+        for output_key, spec in sources_spec.items():
+            sections = spec.get("sections", [])
+            if not sections:
+                continue
+
+            if self.renderer.is_value_section(sections):
+                ids = self.renderer.extract_raw_value(item, sections)
+                if ids is None:
+                    continue
+                result[output_key] = ids
+                any_output = True
+                continue
+
+            list_field = spec.get("list_field", False)
+            mask_key = spec.get("mask_key", f"{output_key}_mask")
+
+            if list_field:
+                ids, mask = self.renderer.process_list_field(
+                    item, sections, config, tokenizer
+                )
+            else:
+                ids, mask = self.renderer.process_sections(
+                    item, sections, config, tokenizer, is_top_level=True
+                )
+
+            if ids is None:
+                continue
+
+            result[output_key] = ids
+            if not all(m == 1 for m in mask):
+                result[mask_key] = mask
+            elif "mask_key" in spec:
+                result[mask_key] = mask
+
+            any_output = True
+
+        if not any_output:
+            return None
+
+        result["domain"] = _extract_domain(item, config.output.domain_key)
+        return result

astrai/preprocessing/packing.py

@@ -0,0 +1,121 @@
+"""Sequence packing strategies for shard-level reordering and truncation.
+
+Each strategy receives the accumulated ``{key: [list of token lists]}``
+dict for a shard and returns a reordered / truncated version.  The
+pipeline later flattens the result into contiguous tensors.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+from astrai.factory import BaseFactory
+
+
+def _truncate(seq: List[int], max_len: int, mode: str) -> List[int]:
+    if len(seq) <= max_len:
+        return seq
+    if mode == "keep_end":
+        return seq[-max_len:]
+    return seq[:max_len]
+
+
+class PackingStrategy(ABC):
+    """Reorder and truncate sequences within a shard."""
+
+    @abstractmethod
+    def apply(
+        self,
+        keys: Dict[str, List[List[int]]],
+        max_packed_len: int,
+        truncation_mode: str,
+    ) -> Dict[str, List[List[int]]]:
+        raise NotImplementedError
+
+
+class PackingStrategyFactory(BaseFactory["PackingStrategy"]):
+    pass
+
+
+@PackingStrategyFactory.register("simple")
+class SimplePacking(PackingStrategy):
+    def apply(
+        self,
+        keys: Dict[str, List[List[int]]],
+        max_packed_len: int,
+        truncation_mode: str,
+    ) -> Dict[str, List[List[int]]]:
+        return {
+            k: [_truncate(v, max_packed_len, truncation_mode) for v in vals]
+            for k, vals in keys.items()
+        }
+
+
+@PackingStrategyFactory.register("bfd")
+class BFDPacking(PackingStrategy):
+    """Best-Fit Decreasing bin packing.
+
+    Assigns sequences to bins using a best-fit heuristic (sorted by
+    decreasing length) and concatenates sequences within each bin into
+    a single packed sequence.  Packed sequences are truncated to
+    *max_packed_len* so that each packed bin fits within one context
+    window during training.
+    """
+
+    def apply(
+        self,
+        keys: Dict[str, List[List[int]]],
+        max_packed_len: int,
+        truncation_mode: str,
+    ) -> Dict[str, List[List[int]]]:
+        sequences = keys.get("sequence", [])
+        if not sequences:
+            return keys
+        bins = self._plan(sequences, max_packed_len, truncation_mode)
+
+        packed: Dict[str, List[List[int]]] = {}
+        for k, vals in keys.items():
+            packed[k] = [
+                _truncate(
+                    self._concat_bin(vals, bin_indices),
+                    max_packed_len,
+                    truncation_mode,
+                )
+                for bin_indices in bins
+            ]
+        return packed
+
+    @staticmethod
+    def _concat_bin(vals: List[List[int]], indices: List[int]) -> List[int]:
+        result: List[int] = []
+        for i in indices:
+            result.extend(vals[i])
+        return result
+
+    @staticmethod
+    def _plan(
+        sequences: List[List[int]], max_packed_len: int, truncation_mode: str
+    ) -> List[List[int]]:
+        n = len(sequences)
+        order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
+        bins: List[List[int]] = []
+        bin_lengths: List[int] = []
+
+        for orig_idx in order:
+            seq_len = len(
+                _truncate(sequences[orig_idx], max_packed_len, truncation_mode)
+            )
+            best_bin = None
+            best_remain = max_packed_len + 1
+            for i, bl in enumerate(bin_lengths):
+                remain = max_packed_len - bl
+                if seq_len <= remain < best_remain:
+                    best_remain = remain
+                    best_bin = i
+            if best_bin is not None:
+                bins[best_bin].append(orig_idx)
+                bin_lengths[best_bin] += seq_len
+            else:
+                bins.append([orig_idx])
+                bin_lengths.append(seq_len)
+
+        return bins

astrai/preprocessing/pipeline.py

@@ -0,0 +1,192 @@
+"""Config-driven JSONL preprocessing pipeline.
+
+Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
+sharding and flush to ``.h5`` / ``.bin`` storage.  Packing, position-id
+generation and storage writing are each delegated to pluggable strategies,
+dispatched by configuration keys.
+"""
+
+import json
+import logging
+import os
+from collections import defaultdict
+from itertools import chain
+from typing import Dict, List, Optional
+
+import torch
+import tqdm
+
+from astrai.config.preprocess_config import PipelineConfig
+from astrai.preprocessing.builder import MaskBuilderFactory
+from astrai.preprocessing.packing import PackingStrategyFactory
+from astrai.preprocessing.position_id import PositionIdStrategyFactory
+from astrai.preprocessing.writer import StoreWriterFactory
+from astrai.tokenize import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+_STR_TO_DTYPE: dict[str, torch.dtype] = {
+    "bool": torch.bool,
+    "uint8": torch.uint8,
+    "int8": torch.int8,
+    "int16": torch.int16,
+    "int32": torch.int32,
+    "int64": torch.int64,
+    "float16": torch.float16,
+    "float32": torch.float32,
+    "float64": torch.float64,
+}
+
+
+def filter_by_length(text: str, min_len: int = 50, max_len: int = 2_000_000) -> bool:
+    return min_len <= len(text) <= max_len
+
+
+class Pipeline:
+    """Tokenization pipeline driven by a declarative :class:`PipelineConfig`.
+
+    Usage::
+
+        config = PipelineConfig.from_file("sft_pipeline.json")
+        Pipeline(config, ["data.jsonl"], output_dir="out", tokenizer_path="params").run()
+    """
+
+    def __init__(
+        self,
+        config: PipelineConfig,
+        input_paths: list[str],
+        output_dir: str,
+        tokenizer_path: str,
+    ):
+        os.makedirs(output_dir, exist_ok=True)
+        self.config = config
+        self.paths = input_paths
+        self.output_dir = output_dir
+        self.tokenizer_path = tokenizer_path
+
+        self.mask_builder = MaskBuilderFactory.create("sectioned")
+        self._packer = PackingStrategyFactory.create(
+            config.preprocessing.packing_strategy
+        )
+        self._position_id = PositionIdStrategyFactory.create(
+            config.output.position_ids_mode
+        )
+        self._writer = StoreWriterFactory.create(config.output.storage_format)
+
+    def transform(self, item: dict) -> Optional[dict]:
+        return self.mask_builder.build(item, self.config, self._tokenizer)
+
+    def run(self):
+        self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path)
+        domains: dict = defaultdict(lambda: defaultdict(list))
+        total_tokens = 0
+        shard_idx: dict[str, int] = defaultdict(int)
+        count = 0
+
+        pp = self.config.preprocessing
+
+        for item in tqdm.tqdm(
+            self._iter_items(), desc="Tokenizing", unit="docs", mininterval=0.5
+        ):
+            if pp.max_items and count >= pp.max_items:
+                break
+
+            try:
+                result = self.transform(item)
+            except Exception:
+                logger.warning(
+                    "Failed to process item #%d, skipping", count + 1, exc_info=True
+                )
+                continue
+            if result is None:
+                continue
+
+            domain = result.pop("domain", "__default__")
+
+            is_multi = bool(getattr(self.config.input, "sources", None))
+            if is_multi:
+                ids = self._primary_ids(result)
+            else:
+                ids = result.pop("sequence")
+                result["sequence"] = ids
+
+            if not ids:
+                continue
+
+            bucket = domains[domain]
+            self._align_bucket(bucket, result, ids)
+            for key, val in result.items():
+                bucket[key].append(val)
+
+            count += 1
+            total_tokens += len(ids)
+
+            if total_tokens >= self.config.output.max_tokens_per_shard:
+                self._flush(domains, shard_idx)
+                domains.clear()
+                total_tokens = 0
+
+        if total_tokens > 0:
+            self._flush(domains, shard_idx)
+
+    @staticmethod
+    def _primary_ids(result: dict) -> list:
+        """Return the first list-valued entry in *result* as the primary id
+        sequence for token counting."""
+        for val in result.values():
+            if isinstance(val, list) and val and isinstance(val[0], int):
+                return val
+        return []
+
+    @staticmethod
+    def _align_bucket(bucket: dict, result: dict, ids: list):
+        """Pad previously-accumulated keys that are missing from *result*."""
+        for key in list(bucket.keys()):
+            if key in result:
+                continue
+            bucket[key].append([0] * len(ids))
+
+    def _iter_items(self):
+        for path in self.paths:
+            with open(path, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    yield json.loads(line)
+
+    def _flush(self, domains, shard_idx):
+        for domain, keys in domains.items():
+            idx = shard_idx[domain]
+
+            pp = self.config.preprocessing
+            original_sequences = keys.get("sequence", [])
+            mode = self.config.output.position_ids_mode
+
+            if mode == "doc_reset" and original_sequences:
+                keys["position_ids"] = [list(range(len(s))) for s in original_sequences]
+
+            keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode)
+
+            tensors: Dict[str, List[torch.Tensor]] = {}
+            for key, ids_list in keys.items():
+                dt = _STR_TO_DTYPE.get(
+                    self.config.output.dtype.get(key, "int32"), torch.int32
+                )
+                tensors[key] = [
+                    torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
+                ]
+
+            if mode == "continuous" and original_sequences:
+                pos_ids = self._position_id.generate(keys.get("sequence", []))
+                if pos_ids:
+                    tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
+
+            self._writer.save(self.output_dir, domain, idx, tensors)
+            shard_idx[domain] = idx + 1
+
+            first_key = "sequence" if "sequence" in tensors else next(iter(tensors))
+            tqdm.tqdm.write(
+                f"  saved {domain}/shard_{idx:04d}  "
+                f"({tensors[first_key][0].numel():,} tokens)"
+            )

astrai/preprocessing/position_id.py

@@ -0,0 +1,46 @@
+"""Position-id generation strategies for packed sequences.
+
+Each strategy takes the list of per-document token sequences after packing
+and returns a flat list of position ids (same total length as all
+sequences combined).  The pipeline wraps the result into a tensor and
+attaches it as ``position_ids``.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from astrai.factory import BaseFactory
+
+
+class PositionIdStrategy(ABC):
+    """Generate ``position_ids`` for packed sequences."""
+
+    @abstractmethod
+    def generate(self, sequences: List[List[int]]) -> List[int]:
+        raise NotImplementedError
+
+
+class PositionIdStrategyFactory(BaseFactory["PositionIdStrategy"]):
+    pass
+
+
+@PositionIdStrategyFactory.register("none")
+class NoPositionId(PositionIdStrategy):
+    def generate(self, sequences: List[List[int]]) -> List[int]:
+        return []
+
+
+@PositionIdStrategyFactory.register("doc_reset")
+class DocResetPositionId(PositionIdStrategy):
+    def generate(self, sequences: List[List[int]]) -> List[int]:
+        pos_ids = []
+        for seq in sequences:
+            pos_ids.extend(range(len(seq)))
+        return pos_ids
+
+
+@PositionIdStrategyFactory.register("continuous")
+class ContinuousPositionId(PositionIdStrategy):
+    def generate(self, sequences: List[List[int]]) -> List[int]:
+        total = sum(len(seq) for seq in sequences)
+        return list(range(total))

astrai/preprocessing/writer.py

@@ -0,0 +1,75 @@
+"""Storage writer strategies for pipeline output.
+
+The :class:`StoreWriter` abstraction decouples the pipeline from the
+concrete storage format (bin / h5).  The pipeline builds a ``{key:
+List[Tensor]}`` dict and delegates the write to the writer selected
+by ``output.storage_format``.
+"""
+
+import logging
+import os
+import shutil
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+import torch
+
+from astrai.factory import BaseFactory
+from astrai.serialization import save_bin, save_h5
+
+logger = logging.getLogger(__name__)
+
+
+class StoreWriter(ABC):
+    """Write pre-tokenized tensors to disk in a format-specific way."""
+
+    @abstractmethod
+    def save(
+        self,
+        output_dir: str,
+        domain: str,
+        shard_idx: int,
+        tensors: Dict[str, List[torch.Tensor]],
+    ) -> None: ...
+
+
+class StoreWriterFactory(BaseFactory["StoreWriter"]):
+    pass
+
+
+@StoreWriterFactory.register("bin")
+class BinWriter(StoreWriter):
+    def save(self, output_dir, domain, shard_idx, tensors):
+        shard_path = os.path.join(output_dir, domain, f"shard_{shard_idx:04d}")
+        try:
+            save_bin(shard_path, tensors)
+        except Exception:
+            if os.path.exists(shard_path):
+                shutil.rmtree(shard_path, ignore_errors=True)
+            logger.error(
+                "Failed to write shard %s/%s_%04d, cleaned up partial output",
+                domain,
+                "shard",
+                shard_idx,
+                exc_info=True,
+            )
+            raise
+
+
+@StoreWriterFactory.register("h5")
+class H5Writer(StoreWriter):
+    def save(self, output_dir, domain, shard_idx, tensors):
+        chunk_dir = os.path.join(output_dir, domain)
+        file_path = os.path.join(chunk_dir, f"data_{shard_idx:04d}.h5")
+        try:
+            save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
+        except Exception:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            logger.error(
+                "Failed to write shard %s/data_%04d.h5, cleaned up partial output",
+                domain,
+                shard_idx,
+                exc_info=True,
+            )
+            raise

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,116 +0,0 @@
-import json
-import os
-from pathlib import Path
-from typing import Any, Dict, List, 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,
-    ):
-        self.state_dict = state_dict
-        self.epoch = epoch
-        self.iteration = iteration
-        self.extra = extra or {}
-
-    def save(
-        self,
-        save_dir: str,
-    ) -> None:
-
-        save_path = Path(save_dir)
-        save_path.mkdir(parents=True, exist_ok=True)
-
-        rank = get_rank()
-        if rank == 0:
-            meta = {
-                "epoch": self.epoch,
-                "iteration": self.iteration,
-            }
-            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:
-                torch.save(self.extra, save_path / "extra.pt")
-
-    @classmethod
-    def load(
-        cls,
-        save_dir: str,
-    ) -> "Checkpoint":
-
-        rank = get_rank()
-        save_path = Path(save_dir)
-
-        meta = {}
-        if rank == 0:
-            with open(Path(save_dir) / "meta.json", "r") as f:
-                meta = json.load(f)
-
-        if dist.is_initialized():
-            meta_list = [meta]
-            dist.broadcast_object_list(meta_list, src=0)
-            meta = meta_list[0]
-
-        state_dict = st.load_file(save_path / "state_dict.safetensors")
-
-        extra = None
-        extra_path = save_path / "extra.pt"
-        if extra_path.exists():
-            extra = torch.load(extra_path, map_location="cpu", weights_only=False)
-
-        return cls(
-            state_dict=state_dict,
-            epoch=meta["epoch"],
-            iteration=meta["iteration"],
-            extra=extra,
-        )

astrai/serialization/__init__.py

@@ -0,0 +1,43 @@
+"""Serialization utilities for models and datasets.
+
+This package re-exports checkpoint helpers and dataset storage helpers so
+that existing imports from ``astrai.serialization`` continue to work.
+"""
+
+from astrai.serialization.checkpoint import (
+    Checkpoint,
+    load_json,
+    load_model_config,
+    load_model_weights,
+    load_safetensors,
+    load_state_dict,
+    load_torch,
+    save_json,
+    save_model,
+    save_safetensors,
+    save_torch,
+)
+from astrai.serialization.dataset import (
+    load_bin,
+    load_h5,
+    save_bin,
+    save_h5,
+)
+
+__all__ = [
+    "Checkpoint",
+    "load_json",
+    "load_model_config",
+    "load_model_weights",
+    "load_safetensors",
+    "load_state_dict",
+    "load_torch",
+    "save_json",
+    "save_model",
+    "save_safetensors",
+    "save_torch",
+    "load_bin",
+    "load_h5",
+    "save_bin",
+    "save_h5",
+]

astrai/serialization/checkpoint.py

@@ -0,0 +1,204 @@
+"""Model checkpoint serialization helpers."""
+
+import io
+import json
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import safetensors.torch as st
+import torch
+import torch.distributed as dist
+
+from astrai.parallel.setup import get_rank
+
+_META_FILE = "meta.json"
+_CONFIG_FILE = "config.json"
+_WEIGHTS_FILE = "model.safetensors"
+
+
+def save_safetensors(state_dict: dict, path: Union[str, Path]):
+    st.save_file(state_dict, str(path))
+
+
+def load_safetensors(path: Union[str, Path], broadcast: bool = False) -> dict:
+    if not broadcast or not dist.is_initialized():
+        return st.load_file(str(path))
+
+    rank = get_rank()
+    if rank == 0:
+        state_dict = st.load_file(str(path))
+    else:
+        state_dict = {}
+    tmp = [state_dict]
+    dist.broadcast_object_list(tmp, src=0)
+    return tmp[0]
+
+
+def save_json(data: dict, path: Union[str, Path]):
+    with open(str(path), "w") as f:
+        json.dump(data, f, indent=2)
+
+
+def load_json(path: Union[str, Path], broadcast: bool = False) -> dict:
+    if not broadcast or not dist.is_initialized():
+        with open(str(path), "r") as f:
+            return json.load(f)
+
+    rank = get_rank()
+    if rank == 0:
+        with open(str(path), "r") as f:
+            data = json.load(f)
+    else:
+        data = {}
+    tmp = [data]
+    dist.broadcast_object_list(tmp, src=0)
+    return tmp[0]
+
+
+def save_torch(obj: Any, path: Union[str, Path]):
+    torch.save(obj, str(path))
+
+
+def load_torch(path: Union[str, Path], broadcast: bool = False) -> Any:
+    if not broadcast or not dist.is_initialized():
+        return torch.load(str(path), map_location="cpu", weights_only=False)
+
+    path = Path(path)
+    rank = get_rank()
+
+    if rank == 0:
+        with open(path, "rb") as f:
+            raw = f.read()
+        data_tensor = torch.frombuffer(bytearray(raw), dtype=torch.uint8)
+        num_bytes = torch.tensor([len(raw)], dtype=torch.long)
+    else:
+        num_bytes = torch.tensor([0], dtype=torch.long)
+
+    dist.broadcast(num_bytes, src=0)
+
+    if rank != 0:
+        data_tensor = torch.empty(num_bytes.item(), dtype=torch.uint8)
+
+    dist.broadcast(data_tensor, src=0)
+
+    buf = io.BytesIO(data_tensor.numpy().tobytes())
+    return torch.load(buf, map_location="cpu", weights_only=False)
+
+
+def save_model(config: dict, state_dict: dict, save_directory: str):
+    save_path = Path(save_directory)
+    save_path.mkdir(parents=True, exist_ok=True)
+    save_json(config, save_path / _CONFIG_FILE)
+    save_safetensors(state_dict, save_path / _WEIGHTS_FILE)
+
+
+def load_model_config(save_directory: str) -> dict:
+    return load_json(Path(save_directory) / _CONFIG_FILE)
+
+
+def load_model_weights(save_directory: str) -> dict:
+    return load_state_dict(Path(save_directory) / _WEIGHTS_FILE)
+
+
+def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict:
+    path = Path(path)
+    if not broadcast or not dist.is_initialized():
+        return load_safetensors(path)
+
+    rank = get_rank()
+    if rank == 0:
+        state_dict = load_safetensors(path)
+        specs = [
+            (k, list(state_dict[k].shape), str(state_dict[k].dtype).split(".")[-1])
+            for k in sorted(state_dict)
+        ]
+    else:
+        state_dict = {}
+        specs = []
+
+    specs_list = [specs]
+    dist.broadcast_object_list(specs_list, src=0)
+    specs = specs_list[0]
+
+    for key, shape, dtype_name in specs:
+        dtype = getattr(torch, dtype_name)
+        if rank != 0:
+            tensor = torch.empty(shape, dtype=dtype, device="cpu")
+        else:
+            tensor = state_dict[key].contiguous().cpu()
+        dist.broadcast(tensor, src=0)
+        if rank != 0:
+            state_dict[key] = tensor
+    return state_dict
+
+
+@dataclass
+class Checkpoint:
+    state_dict: Dict[str, Any] = field(default_factory=dict)
+    epoch: int = 0
+    consumed_samples: int = 0
+    extra: Dict[str, Any] = field(default_factory=dict)
+    meta: Dict[str, Any] = field(default_factory=dict)
+    config: Dict[str, Any] = field(default_factory=dict)
+
+    def save(self, save_dir: str):
+        save_path = Path(save_dir)
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        if get_rank() != 0:
+            return
+
+        meta = {
+            "epoch": self.epoch,
+            "consumed_samples": self.consumed_samples,
+            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+            **self.meta,
+        }
+        save_json(meta, save_path / _META_FILE)
+        save_json(self.config, save_path / _CONFIG_FILE)
+        save_safetensors(self.state_dict, save_path / _WEIGHTS_FILE)
+        for key, value in self.extra.items():
+            save_torch(value, save_path / f"{key}.pt")
+
+    @classmethod
+    def load(cls, save_dir: str, broadcast: bool = False) -> "Checkpoint":
+        save_path = Path(save_dir)
+
+        meta = load_json(save_path / _META_FILE, broadcast)
+        config = load_json(save_path / _CONFIG_FILE, broadcast)
+        state_dict = load_state_dict(save_path / _WEIGHTS_FILE, broadcast=broadcast)
+
+        extra = {}
+        for f in sorted(save_path.iterdir()):
+            if f.suffix == ".pt":
+                extra[f.stem] = load_torch(f, broadcast=broadcast)
+
+        return cls(
+            state_dict=state_dict,
+            epoch=meta.get("epoch", 0),
+            consumed_samples=meta.get("consumed_samples", 0),
+            extra=extra,
+            config=config,
+        )
+
+    @classmethod
+    def load_any(cls, save_dir: str, broadcast: bool = False) -> Optional["Checkpoint"]:
+        save_path = Path(save_dir)
+        meta_path = save_path / _META_FILE
+        weights_path = save_path / _WEIGHTS_FILE
+
+        if meta_path.exists():
+            return cls.load(save_dir, broadcast=broadcast)
+
+        if weights_path.exists():
+            state_dict = load_state_dict(weights_path, broadcast=broadcast)
+            config = {}
+            config_path = save_path / _CONFIG_FILE
+            if config_path.exists():
+                config = load_json(config_path, broadcast)
+            return cls(state_dict=state_dict, config=config)
+
+        return None

astrai/serialization/dataset.py

@@ -0,0 +1,73 @@
+"""Dataset storage serialization helpers (HDF5 / memory-mapped binary)."""
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, List
+
+import h5py
+import numpy as np
+import torch
+from torch import Tensor
+
+
+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_bin(file_path: str, tensor_group: Dict[str, List[Tensor]]):
+    os.makedirs(file_path, exist_ok=True)
+    meta = {}
+    for key, tensors in tensor_group.items():
+        cat = torch.cat(tensors, dim=0)
+        meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]}
+        np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
+    with open(os.path.join(file_path, "meta.json"), "w") as f:
+        json.dump(meta, f)
+
+
+def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
+    with open(os.path.join(file_path, "meta.json"), "r") as f:
+        meta = json.load(f)
+    segments: Dict[str, List[Tensor]] = {}
+    for key, info in meta.items():
+        arr = np.memmap(
+            os.path.join(file_path, f"{key}.bin"),
+            dtype=info["dtype"],
+            mode="r+",
+            shape=tuple(info["shape"]),
+        )
+        segments[key] = [torch.from_numpy(arr)]
+    return segments

astrai/tokenize/chat_template.py

@@ -1,13 +1,10 @@
-from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
 from jinja2 import Template
 
-# Message type for chat messages
 type MessageType = Dict[str, Any]
 
 
-@dataclass
 class ChatTemplate:
     """A chat template with Jinja2 rendering support.
 
@@ -15,23 +12,24 @@ class ChatTemplate:
         name: Unique identifier for the template.
         template_str: Jinja2 template string.
         description: Optional description.
-        default_variables: Optional dictionary of default variable values
-            that will be passed to the template if not overridden during rendering.
+        default_variables: Optional dictionary of default variable values.
         special_tokens: Optional dictionary mapping token names to their string values.
-            These tokens are automatically added to the template variables.
     """
 
-    name: str
-    template_str: str
-    description: str = ""
-    default_variables: Dict[str, Any] = None
-    special_tokens: Dict[str, str] = None
-
-    def __post_init__(self):
-        if self.default_variables is None:
-            self.default_variables = {}
-        if self.special_tokens is None:
-            self.special_tokens = {}
+    def __init__(
+        self,
+        name: str = "",
+        template_str: str = "",
+        description: str = "",
+        default_variables: Optional[Dict[str, Any]] = None,
+        special_tokens: Optional[Dict[str, str]] = None,
+    ):
+        self.name = name
+        self.template_str = template_str
+        self.description = description
+        self.default_variables = default_variables or {}
+        self.special_tokens = special_tokens or {}
+        self._compiled: Template = Template(template_str)
 
     @classmethod
     def from_string(
@@ -43,7 +41,7 @@ class ChatTemplate:
     ) -> "ChatTemplate":
         """Create a ChatTemplate instance directly from a template string."""
         return cls(
-            name="",  # empty name for ad‑hoc templates
+            name="",
             template_str=template_str,
             description=description,
             default_variables=default_variables,
@@ -73,5 +71,4 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        jinja_template = Template(self.template_str)
-        return jinja_template.render(**variables)
+        return self._compiled.render(**variables)

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):

astrai/trainer/metric_util.py

@@ -1,75 +1,25 @@
 from typing import 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 = {}
-    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
+def grad_norm(model: nn.Module, per_param: bool = False) -> float | Dict[str, float]:
+    grads = [p.grad.detach() for p in model.parameters() if p.grad is not None]
+    if not grads:
+        return 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
-
-
-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
-
-
-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
-
-
-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
-
-
-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
+    total_sq = torch.stack([g.pow(2).sum() for g in grads]).sum()
+    if per_param:
+        norms = {}
+        for name, param in model.named_parameters():
+            if param.grad is not None:
+                norms[name] = param.grad.norm(2).item()
+            else:
+                norms[name] = 0.0
+        norms["total"] = total_sq.sqrt().item()
+        return norms
+    return total_sq.sqrt().item()
 
 
 def ctx_get_loss(ctx):
@@ -80,25 +30,9 @@ 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)
-
-
-def ctx_get_grad_std(ctx):
-    return grad_std(ctx.model)
-
-
-def ctx_get_grad_max(ctx):
-    return grad_max(ctx.model)
-
-
-def ctx_get_grad_min(ctx):
-    return grad_min(ctx.model)
-
-
-def ctx_get_grad_mean(ctx):
-    return grad_mean(ctx.model)
-
-
-def ctx_get_grad_nan_num(ctx):
-    return grad_nan_num(ctx.model)
+    return ctx.grad_norm

astrai/trainer/schedule.py

@@ -2,7 +2,7 @@
 
 import math
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Type
+from typing import Any, Dict, List
 
 from torch.optim.lr_scheduler import LRScheduler
 
@@ -31,7 +31,6 @@ class SchedulerFactory(BaseFactory["BaseScheduler"]):
     """Factory class for creating learning rate schedulers.
 
     Supports decorator-based registration for extensible scheduler types.
-    Also supports creation from ScheduleConfig objects.
 
     Example usage:
         @SchedulerFactory.register("custom")
@@ -41,33 +40,6 @@ class SchedulerFactory(BaseFactory["BaseScheduler"]):
         scheduler = SchedulerFactory.create("custom", optimizer, **kwargs)
     """
 
-    @classmethod
-    def _validate_component(cls, scheduler_cls: Type[BaseScheduler]) -> None:
-        """Validate that the scheduler class inherits from BaseScheduler."""
-        if not issubclass(scheduler_cls, BaseScheduler):
-            raise TypeError(f"{scheduler_cls.__name__} must inherit from BaseScheduler")
-
-    @classmethod
-    def create(
-        cls, optimizer, schedule_type: str = "none", **kwargs
-    ) -> "BaseScheduler":
-        """Create a scheduler instance by type name.
-
-        Args:
-            optimizer: PyTorch optimizer
-            schedule_type: Type of scheduler ("cosine", "sgdr")
-            **kwargs: Arguments passed to the scheduler constructor
-
-        Returns:
-            Scheduler instance
-        """
-        return super().create(schedule_type, optimizer, **kwargs)
-
-    @classmethod
-    def available_types(cls) -> list:
-        """Return list of registered scheduler type names."""
-        return cls.list_registered()
-
 
 # ----------- Scheduler implementations -----------
 
@@ -81,7 +53,7 @@ class CosineScheduler(BaseScheduler):
         optimizer,
         warmup_steps: int,
         lr_decay_steps: int,
-        min_rate: float = 0.05,
+        min_rate: float = 0.01,
         last_epoch: int = -1,
     ):
         self.warmup_steps = warmup_steps
@@ -93,11 +65,15 @@ class CosineScheduler(BaseScheduler):
     def get_lr(self) -> List[float]:
         # warmup
         if self.last_epoch < self.warmup_steps:
-            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
+            warmup_factor = max(
+                self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
+            )
             return [base_lr * warmup_factor for base_lr in self.base_lrs]
 
         # cosine decay
-        decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
+        decay_progress = (self.last_epoch - self.warmup_steps) / max(
+            self.lr_decay_steps, 1
+        )
         decay_progress = min(decay_progress, 1.0)
         cosine_decay = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
         decay_factor = max(self.min_rate, cosine_decay)
@@ -132,7 +108,7 @@ class SGDRScheduler(BaseScheduler):
         optimizer,
         warmup_steps: int,
         cycle_length: int,
-        min_rate: float = 0.05,
+        min_rate: float = 0.01,
         t_mult: int = 2,
         last_epoch: int = -1,
     ):
@@ -146,7 +122,9 @@ class SGDRScheduler(BaseScheduler):
     def get_lr(self):
         # warmup
         if self.last_epoch < self.warmup_steps:
-            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
+            warmup_factor = max(
+                self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
+            )
             return [base_lr * warmup_factor for base_lr in self.base_lrs]
 
         # SGDR
@@ -192,3 +170,66 @@ class SGDRScheduler(BaseScheduler):
         self.min_rate = state_dict.pop("min_rate")
         self.t_mult = state_dict.pop("t_mult")
         super().load_state_dict(state_dict)
+
+
+@SchedulerFactory.register("wsd")
+class WSDScheduler(BaseScheduler):
+    """WSD (Warmup-Stable-Decay) scheduler with sqrt cooldown.
+
+    warmup_steps:  linear warmup from min_rate to 1.0
+    stable_steps:   constant at base_lr
+    decay_steps:    sqrt decay from base_lr to min_rate
+    min_rate:       minimum lr as fraction of base_lr (default 0.0)
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_steps: int,
+        stable_steps: int,
+        decay_steps: int,
+        min_rate: float = 0.01,
+        last_epoch: int = -1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.stable_steps = stable_steps
+        self.decay_steps = decay_steps
+        self.min_rate = min_rate
+        self.total_steps = warmup_steps + stable_steps + decay_steps
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> List[float]:
+        if self.last_epoch < self.warmup_steps:
+            factor = max(self.min_rate, self.last_epoch / max(self.warmup_steps, 1))
+            return [base_lr * factor for base_lr in self.base_lrs]
+
+        offset = self.last_epoch - self.warmup_steps
+
+        if offset < self.stable_steps:
+            return list(self.base_lrs)
+
+        decay_ratio = (offset - self.stable_steps) / max(self.decay_steps, 1)
+        decay_ratio = min(decay_ratio, 1.0)
+        factor = (1.0 - self.min_rate) * (1.0 - decay_ratio) ** 2 + self.min_rate
+        return [base_lr * factor for base_lr in self.base_lrs]
+
+    def state_dict(self):
+        state = super().state_dict()
+        state.update(
+            {
+                "warmup_steps": self.warmup_steps,
+                "stable_steps": self.stable_steps,
+                "decay_steps": self.decay_steps,
+                "min_rate": self.min_rate,
+                "total_steps": self.total_steps,
+            }
+        )
+        return state
+
+    def load_state_dict(self, state_dict):
+        self.warmup_steps = state_dict.pop("warmup_steps")
+        self.stable_steps = state_dict.pop("stable_steps")
+        self.decay_steps = state_dict.pop("decay_steps")
+        self.min_rate = state_dict.pop("min_rate")
+        self.total_steps = state_dict.pop("total_steps")
+        super().load_state_dict(state_dict)

astrai/trainer/strategy.py

@@ -1,39 +1,28 @@
 """Training strategy implementations with factory pattern."""
 
-import copy
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Union
+from typing import Callable, Dict, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
-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
-    return model
-
-
-def create_ref_model(model: nn.Module) -> nn.Module:
-    """Create a reference model for DPO/GRPO training.
-
-    Handles DDP-wrapped models safely by unwrapping first,
-    then creating a deep copy with frozen gradients.
-    """
-    original_model = unwrap_model(model)
-    ref_model = copy.deepcopy(original_model)
+def create_ref_model(
+    model_fn: Callable[[], nn.Module], state_dict: Dict[str, Tensor]
+) -> nn.Module:
+    """Create a frozen reference model from model_fn + full state dict."""
+    ref_model = model_fn()
+    ref_model.load_state_dict(state_dict)
     ref_model.requires_grad_(False)
     ref_model.eval()
     return ref_model
 
 
-def move_to_device(batch: Dict[str, Tensor], device: str) -> Any:
+def move_to_device(batch: Dict[str, Tensor], device: str) -> Dict[str, Tensor]:
     """Move batch tensors to specified device with non-blocking transfer."""
     return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
 
@@ -43,7 +32,7 @@ def get_logprobs(
     input_ids: Tensor,
     mask: Tensor,
     reduction: str,
-):
+) -> Tensor:
     """Compute token-wise log probabilities from model outputs.
 
     Args:
@@ -81,14 +70,35 @@ def get_logprobs(
         return token_logprobs * shifted_mask
 
 
+def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
+    S = position_ids.size(1)
+    device = position_ids.device
+    boundaries = position_ids[:, 1:] <= position_ids[:, :-1]
+    doc_ids = torch.cat(
+        [
+            torch.zeros(position_ids.size(0), 1, dtype=torch.long, device=device),
+            boundaries.long().cumsum(dim=1),
+        ],
+        dim=1,
+    )
+    same_doc = doc_ids.unsqueeze(-1) == doc_ids.unsqueeze(-2)
+    causal = torch.tril(torch.ones(S, S, dtype=torch.bool, device=device))
+    return (same_doc & causal).unsqueeze(1)
+
+
 class BaseStrategy(ABC):
     """Abstract base class for training strategies."""
 
     def __init__(
-        self, model: Union[Callable[..., Dict[str, Tensor]]], device: str, **kwargs
+        self,
+        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
+        device: str,
+        **kwargs,
     ):
         self.model = model
         self.device = device
+        self.executor = kwargs.pop("executor", None)
+        self.model_fn = kwargs.pop("model_fn", None)
         self.extra_kwargs = kwargs
 
     @abstractmethod
@@ -122,32 +132,6 @@ class StrategyFactory(BaseFactory["BaseStrategy"]):
         strategy = StrategyFactory.create("custom", model, device)
     """
 
-    @classmethod
-    def _validate_component(cls, strategy_cls: type) -> None:
-        """Validate that the strategy class inherits from BaseStrategy."""
-        if not issubclass(strategy_cls, BaseStrategy):
-            raise TypeError(f"{strategy_cls.__name__} must inherit from BaseStrategy")
-
-    @classmethod
-    def create(cls, train_type: str, model, device: str, **kwargs) -> "BaseStrategy":
-        """Create a strategy instance based on training type.
-
-        Args:
-            train_type: Type of training ("seq", "sft", "dpo", "grpo")
-            model: Model instance for the strategy
-            device: Device to run the strategy on
-            **kwargs: Additional arguments passed to strategy constructor
-
-        Returns:
-            Strategy instance
-        """
-        return super().create(train_type, model, device, **kwargs)
-
-    @classmethod
-    def available_strategies(cls) -> list:
-        """Return list of registered strategy names."""
-        return cls.list_registered()
-
 
 # ============== Strategy Classes ==============
 # All strategies are registered at class definition time using the decorator
@@ -160,7 +144,13 @@ class SEQStrategy(BaseStrategy):
     Computes cross-entropy loss for next token prediction.
     """
 
-    def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
+    def __init__(
+        self,
+        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
+        device: str,
+        label_smoothing: float = 0.0,
+        **kwargs,
+    ):
         super().__init__(model, device, **kwargs)
         self.label_smoothing = label_smoothing
 
@@ -185,21 +175,31 @@ class SFTStrategy(BaseStrategy):
     Applies cross-entropy loss only to tokens where loss_mask is True.
     """
 
-    def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
+    def __init__(
+        self,
+        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
+        device: str,
+        label_smoothing: float = 0.0,
+        **kwargs,
+    ):
         super().__init__(model, device, **kwargs)
         self.label_smoothing = label_smoothing
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
         batch = move_to_device(batch, self.device)
-        input_ids, target_ids, loss_mask = (
+        input_ids, target_ids, position_ids, loss_mask = (
             batch["input_ids"],
             batch["target_ids"],
+            batch["position_ids"],
             batch["loss_mask"],
         )
 
         ignore_index = -100
-        logits = self.model(input_ids=input_ids)["logits"]
-        target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
+        input_mask = make_doc_boundary_mask(position_ids)
+        target_ids = target_ids.masked_fill(~loss_mask, ignore_index)
+        logits = self.model(
+            input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
+        )["logits"]
 
         loss = F.cross_entropy(
             input=logits.flatten(0, 1).float(),
@@ -228,7 +228,9 @@ class DPOStrategy(BaseStrategy):
         **kwargs,
     ):
         super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(model)
+        self.ref_model = create_ref_model(
+            self.model_fn, self.executor.unwrap_model(model)
+        ).to(device=self.device)
         self.beta = beta
         self.reduction = reduction
 
@@ -282,7 +284,9 @@ class GRPOStrategy(BaseStrategy):
         **kwargs,
     ):
         super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(model)
+        self.ref_model = create_ref_model(
+            self.model_fn, self.executor.unwrap_model(model)
+        ).to(device=self.device)
         self.clip_eps = clip_eps
         self.kl_coef = kl_coef
         self.group_size = group_size
@@ -292,8 +296,7 @@ class GRPOStrategy(BaseStrategy):
 
     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)
+        self.ref_model.load_state_dict(self.executor.unwrap_model(self.model))
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
         self._step += 1

astrai/trainer/train_callback.py

@@ -1,28 +1,31 @@
 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, get_rank
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
-    ctx_get_grad_max,
-    ctx_get_grad_mean,
-    ctx_get_grad_min,
-    ctx_get_grad_nan_num,
     ctx_get_grad_norm,
-    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 +45,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 +69,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 +79,45 @@ 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
-        clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
+    def on_optimizer_step(self, context: TrainContext):
+        context.grad_norm = context.executor.clip_grad_norm(
+            context.model, 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,54 +126,62 @@ 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,
     ):
         self.save_dir = save_dir
         self.interval = interval
         self.weight_only = weight_only
-        self.state_dict_fn = state_dict_fn
-        self.save_extra_fn = save_extra_fn
-        self.last_ckpt_iter = 0
+        self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
+        self.last_ckpt_step = 0
 
-    @only_on_rank(0)
     def _save_checkpoint(self, context: TrainContext):
-        save_path = os.path.join(
-            self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
-        )
-        state_dict = (
-            self.state_dict_fn(context.model)
-            if self.state_dict_fn
-            else context.model.state_dict()
-        )
+        state_dict = context.executor.unwrap_model(context.model)
+        self.last_ckpt_step = context.optimizer_step
 
-        extra = self.save_extra_fn(context) if self.save_extra_fn else None
-        context.checkpoint = Checkpoint(
-            state_dict=state_dict,
-            epoch=context.epoch,
-            iteration=context.iteration,
-            extra=extra,
-        )
-
-        context.checkpoint.save(save_path)
-        self.last_ckpt_iter = context.iteration
+        if get_rank() == 0:
+            save_path = os.path.join(
+                self.save_dir,
+                f"epoch_{context.epoch}_step_{context.optimizer_step}",
+            )
+            extra = self.save_extra_fn(context)
+            meta = context.config.to_dict()
+            context.checkpoint = Checkpoint(
+                state_dict=state_dict,
+                epoch=context.epoch,
+                consumed_samples=context.consumed_samples,
+                config=context.model_config,
+                extra=extra,
+                meta=meta,
+            )
+            context.checkpoint.save(save_path)
 
     def on_batch_end(self, context: TrainContext):
-        if context.iteration - self.last_ckpt_iter >= self.interval:
+        if context.optimizer_step - self.last_ckpt_step >= self.interval:
             self._save_checkpoint(context)
 
     def on_train_end(self, context: TrainContext):
-        if context.iteration != self.last_ckpt_iter:
+        if context.optimizer_step != self.last_ckpt_step:
             self._save_checkpoint(context)
 
     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
+
 
 @CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):
@@ -170,27 +189,37 @@ 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: Optional[IO[str]] = None
+    ):
         self.num_epoch = num_epoch
+        self.log_interval = log_interval
+        self.file = file
         self.progress_bar: tqdm = None
 
     @only_on_rank(0)
     def on_epoch_begin(self, context: TrainContext):
+        total_steps = len(context.dataloader) // context.executor.grad_accum_steps
         self.progress_bar = tqdm(
-            context.dataloader,
+            total=total_steps,
             desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
             dynamic_ncols=True,
+            file=self.file or sys.stdout,
         )
 
     @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}",
-            }
-        )
+    def on_optimizer_step(self, context: TrainContext):
         self.progress_bar.update(1)
+        postfix = {
+            "step": context.optimizer_step,
+            "loss": f"{context.loss:.4f}",
+            "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
+        }
+        if context.grad_norm is not None:
+            postfix["grad_norm"] = f"{context.grad_norm:.2f}"
+        if context.val_loss is not None:
+            postfix["val_loss"] = f"{context.val_loss:.4f}"
+        self.progress_bar.set_postfix(postfix)
 
     @only_on_rank(0)
     def on_epoch_end(self, context: TrainContext):
@@ -199,19 +228,20 @@ class ProgressBarCallback(TrainCallback):
             self.progress_bar.close()
 
 
-@CallbackFactory.register("metric_logger")
-class MetricLoggerCallback(TrainCallback):
+@CallbackFactory.register("metric")
+class MetricCallback(TrainCallback):
     def __init__(
         self,
         log_dir: str,
         save_interval: int,
-        log_interval: int = 10,
         metrics: List[str] = None,
+        val_step: int = 0,
     ):
-        self.last_log_iter = 0
+        self.last_log_flush_step = 0
         self.save_interval = save_interval
-        self.log_interval = log_interval
         self.metrics = metrics or ["loss", "lr"]
+        self.val_step = val_step
+        self._next_val_step = 0
 
         self.log_dir = Path(log_dir) if log_dir else Path.cwd() / "logs"
         self.log_dir.mkdir(parents=True, exist_ok=True)
@@ -221,46 +251,84 @@ 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,
-            "grad_min": ctx_get_grad_min,
-            "grad_mean": ctx_get_grad_mean,
-            "grad_nan_num": ctx_get_grad_nan_num,
         }
 
-    def _get_log_data(self, context: TrainContext):
+    def _metrics(self, context: TrainContext, names):
         return {
-            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
-            "epoch": context.epoch,
-            "iter": context.iteration,
-            **{m: self._metric_funcs[m](context) for m in self.metrics},
+            m: self._metric_funcs[m](context)
+            for m in names
+            if self._metric_funcs[m](context) is not None
         }
 
     @only_on_rank(0)
-    def _add_log(self, log_data):
-        self.log_cache.append(log_data)
+    def _append(self, event_type: str, context: TrainContext, **extra):
+        entry = {
+            "type": event_type,
+            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+            "epoch": context.epoch,
+            "step": context.optimizer_step,
+            "consumed_samples": context.consumed_samples,
+            **extra,
+        }
+        self.log_cache.append(entry)
+
+    def _run_validation(self, context: TrainContext) -> float:
+        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
+
+        if context.world_size > 1 and dist.is_initialized():
+            stats = torch.tensor(
+                [total_loss, float(num_batches)], device=get_current_device()
+            )
+            dist.all_reduce(stats, op=dist.ReduceOp.SUM)
+            avg_loss = (stats[0] / stats[1]).item()
+        else:
+            avg_loss = total_loss / max(num_batches, 1)
+
+        context.model.train()
+        return avg_loss
 
     @only_on_rank(0)
-    def _save_log(self, epoch, iter):
-        log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
-
+    def _flush(self, epoch, step):
+        log_file = self.log_dir / f"epoch_{epoch}_step_{step}_metric.jsonl"
+        log_file.parent.mkdir(parents=True, exist_ok=True)
         with open(log_file, "w") as f:
             for log in self.log_cache:
                 f.write(json.dumps(log) + "\n")
 
-    def on_batch_end(self, context):
-        if context.iteration % self.log_interval == 0:
-            log_data = self._get_log_data(context)
-            self._add_log(log_data)
+    def on_optimizer_step(self, context):
+        if (
+            context.val_dataloader is not None
+            and self.val_step > 0
+            and context.optimizer_step >= self._next_val_step
+        ):
+            context.val_loss = self._run_validation(context)
+            self._next_val_step = context.optimizer_step + self.val_step
+            self._append("validation", context, val_loss=context.val_loss)
 
-        if context.iteration - self.last_log_iter >= self.save_interval:
-            self._save_log(context.epoch, context.iteration)
-            self.last_log_iter = context.iteration
+        step_metrics = [m for m in self.metrics if m != "val_loss"]
+        self._append("step", context, **self._metrics(context, step_metrics))
+
+        if context.optimizer_step - self.last_log_flush_step >= self.save_interval:
+            self._flush(context.epoch, context.optimizer_step)
+            self.last_log_flush_step = context.optimizer_step
+
+    def on_epoch_end(self, context):
+        self._append("epoch", context)
 
     def on_train_end(self, context):
-        if context.iteration != self.last_log_iter:
-            self._save_log(context.epoch, context.iteration)
+        if context.optimizer_step != self.last_log_flush_step:
+            self._flush(context.epoch, context.optimizer_step)
 
     def on_error(self, context):
-        self._save_log(context.epoch, context.iteration)
+        self._flush(context.epoch, context.optimizer_step)

astrai/trainer/train_context.py

@@ -1,15 +1,18 @@
 from dataclasses import dataclass, field
-from typing import Callable, Optional, Self
+from pathlib import Path
+from typing import Any, Dict, Optional, Self
 
+import torch
 import torch.nn as nn
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LRScheduler
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, random_split
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import ResumableDistributedSampler
+from astrai.model.components.lora import inject_lora
+from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
-from astrai.serialization import Checkpoint
+from astrai.protocols import OptimizerProtocol, SchedulerProtocol
+from astrai.serialization import Checkpoint, load_json
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
 
@@ -18,84 +21,169 @@ 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)
+    model_config: dict = field(default_factory=dict)
+    executor: BaseExecutor = field(default=None)
 
     epoch: int = field(default=0)
-    iteration: int = field(default=0)
+    consumed_samples: int = field(default=0)
     loss: float = field(default=0.0)
+    grad_norm: Optional[float] = field(default=None)
+    val_dataloader: Optional[DataLoader] = field(default=None)
+    val_loss: Optional[float] = field(default=None)
 
     world_size: int = field(default=1)
     rank: int = field(default=0)
-    kwargs: dict = field(default_factory=dict)
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def optimizer_step(self) -> int:
+        return self.consumed_samples // (
+            self.config.batch_per_device
+            * self.world_size
+            * self.config.grad_accum_steps
+        )
 
 
 class TrainContextBuilder:
     def __init__(
         self,
         config: TrainConfig,
-        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
     ):
         self.config = config
-        self._checkpoint: Optional[Checkpoint] = None
-        self._load_extra_fn = load_extra_fn
+        self._resume_dir: Optional[str] = None
 
-    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
-        self._checkpoint = checkpoint
+    def with_resume_dir(self, resume_dir: Optional[str]) -> Self:
+        self._resume_dir = resume_dir
         return self
 
     def build(self) -> TrainContext:
-        context = TrainContext(
-            model=self.config.model,
-            world_size=get_world_size(),
-            rank=get_rank(),
+        cfg = self.config
+        device = get_current_device()
+
+        executor = ExecutorFactory.create(
+            cfg.parallel_mode,
+            grad_accum_steps=cfg.grad_accum_steps,
+            **cfg.executor_kwargs,
         )
 
-        device = get_current_device()
-        context.model = context.model.to(device=device)
+        model = cfg.model_fn()
+        model = model.to(device=device)
 
-        if self.config.nprocs > 1 and self.config.parallel_wrapper:
-            context.model = self.config.parallel_wrapper(context.model)
+        model_config = {}
+        if self._resume_dir:
+            config_path = Path(self._resume_dir) / "config.json"
+            if config_path.exists():
+                model_config = load_json(config_path)
 
-        if self._checkpoint is not None:
-            context.epoch = max(self._checkpoint.epoch, self.config.start_epoch)
-            context.iteration = max(self._checkpoint.iteration, self.config.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(),
+        if not model_config and hasattr(model, "config"):
+            model_config = model.config.to_dict()
+
+        context = TrainContext(
+            model=model,
+            world_size=get_world_size(),
+            rank=get_rank(),
+            config=cfg,
+            model_config=model_config,
+            executor=executor,
+        )
+
+        if self._resume_dir:
+            checkpoint = Checkpoint.load_any(self._resume_dir)
+            if checkpoint is not None:
+                model.load_state_dict(checkpoint.state_dict, strict=False)
+                if checkpoint.config:
+                    context.model_config = checkpoint.config
+                context.epoch = checkpoint.epoch or cfg.start_epoch
+                if checkpoint.consumed_samples > 0:
+                    context.consumed_samples = checkpoint.consumed_samples
+                else:
+                    context.consumed_samples = cfg.start_samples * context.world_size
+                context.checkpoint = checkpoint
+
+        if cfg.lora is not None:
+            inject_lora(
+                model,
+                r=cfg.lora.r,
+                alpha=cfg.lora.alpha,
+                target_modules=set(cfg.lora.target_modules),
             )
 
-        context.optimizer = self.config.optimizer_fn(context.model)
-        context.scheduler = self.config.scheduler_fn(context.optimizer)
+        context.optimizer = cfg.optimizer_fn(model)
+        context.scheduler = cfg.scheduler_fn(context.optimizer)
 
-        if self._checkpoint and self._checkpoint.extra and self._load_extra_fn:
-            self._load_extra_fn(self._checkpoint.extra, context)
+        train_dataset = cfg.dataset
+        val_dataset = cfg.val_dataset
 
-        cfg = self.config
-        sampler_offset = context.iteration * cfg.batch_size
+        if val_dataset is None and cfg.val_split is not None:
+            n_total = len(cfg.dataset)
+            n_val = max(1, int(n_total * cfg.val_split))
+            n_train = n_total - n_val
+            generator = torch.Generator().manual_seed(cfg.random_seed)
+            train_dataset, val_dataset = random_split(
+                cfg.dataset, [n_train, n_val], generator=generator
+            )
+
+        sampler_offset = context.consumed_samples // context.world_size
         sampler = ResumableDistributedSampler(
-            data_source=cfg.dataset,
+            data_source=train_dataset,
             start_epoch=context.epoch,
             start_iter=sampler_offset,
             seed=cfg.random_seed,
         )
         context.dataloader = DataLoader(
-            cfg.dataset,
-            batch_size=cfg.batch_size,
+            train_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 val_dataset is not None:
+            val_sampler = ResumableDistributedSampler(
+                data_source=val_dataset,
+                start_epoch=0,
+                start_iter=0,
+                seed=cfg.random_seed,
+                shuffle=False,
+            )
+            context.val_dataloader = DataLoader(
+                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(
+                model,
+                context.optimizer,
+                context.dataloader,
+                context.scheduler,
+            )
+        )
+
+        if context.checkpoint and context.checkpoint.extra:
+            extra = context.checkpoint.extra
+            for name in ("optimizer", "scheduler"):
+                if name in extra:
+                    obj = getattr(context, name, None)
+                    if obj is not None:
+                        obj.load_state_dict(extra[name])
+
         context.strategy = StrategyFactory.create(
+            cfg.strategy,
             model=context.model,
-            train_type=self.config.strategy,
             device=device,
-            **self.config.extra_kwargs,
+            executor=executor,
+            model_fn=cfg.model_fn,
+            **cfg.extra_kwargs,
         )
 
         return context

astrai/trainer/trainer.py

@@ -3,7 +3,6 @@ from typing import List, Optional
 
 from astrai.config import TrainConfig
 from astrai.parallel.setup import spawn_parallel_fn
-from astrai.serialization import Checkpoint
 from astrai.trainer.train_callback import (
     CallbackFactory,
     TrainCallback,
@@ -25,18 +24,27 @@ 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",
+                log_dir=cfg.log_dir,
+                save_interval=cfg.ckpt_interval,
+                metrics=cfg.metrics,
+                val_step=cfg.val_step,
+            ),
             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"),
         ]
-
-    def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
-        return (
-            TrainContextBuilder(self.train_config).with_checkpoint(checkpoint).build()
-        )
+        return callbacks
 
     def _call_callbacks(self, method_name: str, context: TrainContext):
         for callback in self.callbacks:
@@ -44,55 +52,58 @@ 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,
-            checkpoint=checkpoint,
+    def _trainer_loop(self, resume_dir: Optional[str] = None):
+        context = (
+            TrainContextBuilder(self.train_config).with_resume_dir(resume_dir).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)
 
-                accumulation_steps = max(self.train_config.accumulation_steps, 1)
                 for batch in context.dataloader:
-                    if context.iteration % 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)
+                    with executor.accumulate(context.model):
+                        self._call_callbacks("on_batch_begin", context)
+                        loss = context.strategy(batch)
+                        context.loss = loss.item()
+                        stand_loss = loss / executor.grad_accum_steps
+                        executor.backward(stand_loss)
+                        context.consumed_samples += (
+                            context.config.batch_per_device * context.world_size
+                        )
+                        self._call_callbacks("on_batch_end", context)
 
-                    # 3. batch
-                    self._call_callbacks("on_batch_begin", context)
-                    loss = context.strategy(batch)
-                    context.loss = loss.item()
-                    context.iteration += 1
+                        if executor.sync_gradients:
+                            self._call_callbacks("on_optimizer_step", context)
+                            context.optimizer.step()
+                            context.optimizer.zero_grad()
 
-                    # to make the loss normalized by accumulation steps
-                    stand_loss = loss / accumulation_steps
-                    stand_loss.backward()
-
-                    self._call_callbacks("on_batch_end", context)
+                            if context.scheduler:
+                                context.scheduler.step()
 
                 self._call_callbacks("on_epoch_end", context)
 
         except Exception as e:
-            logger.error(f"Training failed: {str(e)}", exc_info=True)
+            logger.error("Training failed: %s", str(e), exc_info=True)
             self._call_callbacks("on_error", context)
             raise
         finally:
             self._call_callbacks("on_train_end", context)
+
+    def train(self, resume_dir: Optional[str] = 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,
+            resume_dir=resume_dir,
+        )

docker-compose.yml

@@ -1,12 +1,13 @@
 services:
   server:
-    build: .
-    image: astrai:latest
+    build:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
       - ./params:/app/params:ro
-      - ./checkpoints:/app/checkpoints
     command: python -m scripts.tools.server --port 8000 --device cuda
     deploy:
       resources:
@@ -25,13 +26,14 @@ services:
 
   server-cpu:
     profiles: [cpu]
-    build: .
-    image: astrai:latest
+    build:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
       - ./params:/app/params:ro
-      - ./checkpoints:/app/checkpoints
     command: python -m scripts.tools.server --port 8000 --device cpu
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8000/health"]

pyproject.toml

@@ -9,8 +9,8 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "h5py==3.15.1",
-    "numpy==2.3.2",
-    "torch==2.7.1",
+    "numpy==2.4.4",
+    "torch==2.11.0",
     "tokenizers==0.21.4",
     "tqdm==4.67.1",
     "safetensors==0.5.3",
@@ -37,7 +37,7 @@ dev = ["pytest==9.0.2", "ruff"]
 where = ["."]
 
 [tool.pip]
-extra-index-url = "https://download.pytorch.org/whl/cu126"
+extra-index-url = "https://download.pytorch.org/whl/cu128"
 
 [tool.setuptools.dynamic]
 version = { attr = "astrai.__version__" }

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

@@ -1,3 +1,4 @@
+from argparse import ArgumentParser
 from pathlib import Path
 
 import torch
@@ -7,42 +8,82 @@ from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
-PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+
+
+def parse_args():
+    parser = ArgumentParser(description="Interactive streaming chat")
+    parser.add_argument(
+        "--model_path",
+        type=Path,
+        default=PROJECT_ROOT / "params",
+        help="Path to model weights (params/ or checkpoint/epoch_N_step_M/)",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.8,
+        help="Sampling temperature (default: 0.8)",
+    )
+    parser.add_argument(
+        "--top_p",
+        type=float,
+        default=0.95,
+        help="Top-p sampling threshold",
+    )
+    parser.add_argument(
+        "--top_k",
+        type=int,
+        default=50,
+        help="Top-k sampling threshold",
+    )
+    parser.add_argument(
+        "--max_tokens",
+        type=int,
+        default=2048,
+        help="Maximum tokens to generate",
+    )
+    parser.add_argument(
+        "--system_prompt",
+        type=str,
+        default="You are a helpful assistant.",
+        help="Optional system prompt",
+    )
+    return parser.parse_args()
 
 
 def chat():
-    model = AutoModel.from_pretrained(PARAMETER_ROOT)
-    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
-    model.to(device="cuda", dtype=torch.bfloat16)
+    args = parse_args()
+    model_path = args.model_path
 
-    messages = [{"role": "system", "content": "You are a helpful assistant."}]
+    model = AutoModel.from_pretrained(model_path)
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model.to(device="cuda", dtype=torch.bfloat16)
     engine = InferenceEngine(model=model, tokenizer=tokenizer)
 
+    messages = [{"role": "system", "content": args.system_prompt}]
+
     while True:
         query = input(">> ")
         if query == "!exit":
             break
 
-        # Add user message
         messages.append({"role": "user", "content": query})
 
-        # Generate response
         full_response = ""
         prompt = tokenizer.apply_chat_template(messages, tokenize=False)
 
         for token in engine.generate(
             prompt=prompt,
             stream=True,
-            max_tokens=2048,
-            temperature=0.8,
-            top_p=0.95,
-            top_k=50,
+            max_tokens=args.max_tokens,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            top_k=args.top_k,
         ):
             print(token, end="", flush=True)
             full_response += token
 
         print()
-        # Add assistant response to messages
         messages.append({"role": "assistant", "content": full_response.strip()})
 
 

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/eval/analyze_weights.py

@@ -0,0 +1,307 @@
+"""SVD effective rank & weight statistics analysis for model checkpoints."""
+
+import argparse
+import json
+from pathlib import Path
+
+import safetensors.torch
+import torch
+
+
+def effective_rank_metrics(w: torch.Tensor) -> dict:
+    if w.ndim == 1:
+        return {"shape": tuple(w.shape), "is_1d": True}
+
+    w = w.float()
+    s = torch.linalg.svdvals(w)
+    s_sq = s**2
+    total = s_sq.sum()
+    cumsum = torch.cumsum(s_sq, dim=0) / total
+
+    min_dim = min(w.shape[0], w.shape[1])
+    er_90 = (cumsum < 0.90).sum().item() + 1
+    er_95 = (cumsum < 0.95).sum().item() + 1
+    er_99 = (cumsum < 0.99).sum().item() + 1
+
+    p = s_sq / total
+    p = p[p > 1e-30]
+    entropy = -(p * torch.log(p)).sum()
+    entropic_rank = torch.exp(entropy).item()
+
+    return {
+        "shape": tuple(w.shape),
+        "min_dim": min_dim,
+        "er_90": er_90,
+        "er_95": er_95,
+        "er_99": er_99,
+        "er_99_norm": er_99 / min_dim,
+        "er_95_norm": er_95 / min_dim,
+        "entropic_rank": entropic_rank,
+        "entropic_rank_norm": entropic_rank / min_dim,
+        "top1_ratio": s[0].item() / s.sum().item(),
+        "top5_ratio": s[:5].sum().item() / s.sum().item(),
+        "decay_ratio": s[-1].item() / s[0].item(),
+        "condition_number": s[0].item() / s[-1].item(),
+        "mean": w.mean().item(),
+        "std": w.std().item(),
+        "min": w.min().item(),
+        "max": w.max().item(),
+    }
+
+
+def format_header(headers: list[str], widths: list[int]) -> str:
+    return "".join(h.ljust(w) for h, w in zip(headers, widths))
+
+
+def format_row(values: list[str], widths: list[int]) -> str:
+    return "".join(v.ljust(w) for v, w in zip(values, widths))
+
+
+def group_by_component(results: dict[str, dict]) -> dict[str, list[dict]]:
+    groups: dict[str, list[dict]] = {}
+    for key, r in results.items():
+        parts = key.split(".")
+        if parts[0] == "layers" and len(parts) >= 3:
+            sub = parts[2:]
+            if sub[0] == "attention":
+                comp = f"attn.{sub[1]}"
+            elif sub[0] == "mlp":
+                comp = f"mlp.{sub[1]}"
+            elif sub[0] == "input_norm":
+                comp = "input_norm"
+            elif sub[0] == "post_attention_norm":
+                comp = "post_attn_norm"
+            else:
+                comp = ".".join(sub)
+        else:
+            comp = key
+        groups.setdefault(comp, []).append(r)
+    return groups
+
+
+def print_component_summary(results: dict[str, dict], title: str):
+    groups = group_by_component(results)
+    matrix_groups = {
+        k: [v for v in vs if not v.get("is_1d")]
+        for k, vs in groups.items()
+        if any(not v.get("is_1d") for v in vs)
+    }
+
+    widths = [20, 12, 12, 12, 12, 12]
+    print(f"\n{title}")
+    print(
+        format_header(
+            ["Component", "N", "ER@99%", "EntRank%", "Top1 σ(%)", "Cond. Num"], widths
+        )
+    )
+    print("-" * sum(widths))
+
+    for name in sorted(matrix_groups.keys()):
+        items = matrix_groups[name]
+        n = len(items)
+        print(
+            format_row(
+                [
+                    name,
+                    str(n),
+                    f"{sum(r['er_99_norm'] for r in items) / n:.4f}",
+                    f"{sum(r['entropic_rank_norm'] for r in items) / n:.4f}",
+                    f"{sum(r['top1_ratio'] for r in items) / n:.4f}",
+                    f"{sum(r['condition_number'] for r in items) / n:.1f}",
+                ],
+                widths,
+            )
+        )
+
+    all_er = [
+        r["er_99_norm"]
+        for vs in matrix_groups.values()
+        for r in vs
+        if "_norm" not in r or not r.get("is_1d")
+    ]
+    if all_er:
+        m = sum(all_er) / len(all_er)
+        print(f"\n  Overall Mean ER@99: {m:.4f} ({m * 100:.1f}% of dimension)")
+        if m > 0.85:
+            print("  → HIGH utilization: model near capacity → need more params")
+        elif m > 0.5:
+            print("  → MODERATE utilization: some headroom left")
+        else:
+            print("  → LOW utilization: significant unused capacity")
+
+
+def print_layer_grid(results: dict[str, dict]):
+    comps = [
+        "attn.q_proj",
+        "attn.k_proj",
+        "attn.v_proj",
+        "attn.o_proj",
+        "mlp.up",
+        "mlp.gate",
+        "mlp.down",
+    ]
+    widths = [6] + [10] * len(comps)
+    metric = "er_99_norm"
+
+    print(f"\n--- Per-Layer Effective Rank (99% energy) ---")
+    print(format_header(["Layer"] + comps, widths))
+    print("-" * sum(widths))
+
+    layer_data: dict[int, dict[str, dict]] = {}
+    for key, r in results.items():
+        parts = key.split(".")
+        if parts[0] != "layers":
+            continue
+        li = int(parts[1])
+        sub = parts[2:]
+        if sub[0] == "attention":
+            cname = f"attn.{sub[1]}"
+        elif sub[0] == "mlp":
+            cname = f"mlp.{sub[1]}"
+        else:
+            continue
+        layer_data.setdefault(li, {})[cname] = r
+
+    for li in sorted(layer_data):
+        values = [str(li)]
+        for c in comps:
+            v = layer_data[li].get(c, {}).get(metric, 0)
+            values.append(f"{v:.4f}")
+        print(format_row(values, widths))
+
+
+def print_weight_stats(results: dict[str, dict]):
+    groups = group_by_component(results)
+    widths = [20, 12, 12, 12, 12]
+    print(f"\n--- Weight Value Statistics ---")
+    print(format_header(["Component", "Mean", "Std", "Min", "Max"], widths))
+    print("-" * sum(widths))
+
+    for name in sorted(groups.keys()):
+        items = groups[name]
+        means = [r.get("mean", 0) for r in items]
+        stds = [r.get("std", 0) for r in items]
+        mins = [r.get("min", 0) for r in items]
+        maxs = [r.get("max", 0) for r in items]
+        g_mean = sum(means) / len(means)
+        g_std = sum(stds) / len(stds)
+        g_min = min(mins)
+        g_max = max(maxs)
+        print(
+            format_row(
+                [
+                    name,
+                    f"{g_mean:.6f}",
+                    f"{g_std:.6f}",
+                    f"{g_min:.6f}",
+                    f"{g_max:.6f}",
+                ],
+                widths,
+            )
+        )
+
+
+def print_params_summary(results: dict[str, dict]):
+    total_2d = sum(
+        r["shape"][0] * r["shape"][1] for r in results.values() if not r.get("is_1d")
+    )
+    total_1d = sum(r["shape"][0] for r in results.values() if r.get("is_1d"))
+    print(f"\n  Total 2D params: {total_2d:,}")
+    print(f"  Total 1D params: {total_1d:,}")
+    print(f"  Total params:     {total_2d + total_1d:,}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="SVD effective rank & weight statistics of a model checkpoint."
+    )
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        required=True,
+        help="Path to checkpoint directory (containing model.safetensors + config.json).",
+    )
+    parser.add_argument(
+        "--compare",
+        type=str,
+        nargs="*",
+        help="Additional checkpoint directories to compare against.",
+    )
+    parser.add_argument(
+        "--no_svd",
+        action="store_true",
+        help="Skip SVD analysis, only show weight statistics (mean/std/min/max).",
+    )
+    args = parser.parse_args()
+
+    def analyze_one(ckpt_dir: str, label: str):
+        ckpt_dir = Path(ckpt_dir)
+        weights_path = ckpt_dir / "model.safetensors"
+        if not weights_path.exists():
+            print(f"ERROR: {weights_path} not found")
+            return {}
+
+        meta = {}
+        meta_path = ckpt_dir / "meta.json"
+        if meta_path.exists():
+            with open(meta_path) as f:
+                meta = json.load(f)
+
+        print(f"\n{'=' * 70}")
+        print(f"  {label}: {ckpt_dir}")
+        if meta:
+            print(
+                f"  Iteration: {meta.get('iteration', '?')}, "
+                f"Strategy: {meta.get('strategy', '?')}, "
+                f"nprocs={meta.get('nprocs', '?')}"
+            )
+        print(f"{'=' * 70}")
+
+        print(f"Loading weights...")
+        sd = safetensors.torch.load_file(str(weights_path))
+        print(f"  {len(sd)} keys loaded")
+
+        weight_keys = [
+            k
+            for k in sd
+            if ".weight" in k and "rotary_embedding" not in k and "freqs_cis" not in k
+        ]
+
+        results = {}
+        if not args.no_svd:
+            print(f"Computing SVD on {len(weight_keys)} tensors...")
+            for i, k in enumerate(sorted(weight_keys)):
+                print(f"  [{i + 1}/{len(weight_keys)}] {k:<60s}", end="\r")
+                results[k] = effective_rank_metrics(sd[k])
+            print()
+        else:
+            print(f"Computing stats on {len(weight_keys)} tensors (no SVD)...")
+            for i, k in enumerate(sorted(weight_keys)):
+                t = sd[k]
+                results[k] = {
+                    "shape": tuple(t.shape),
+                    "is_1d": t.ndim == 1,
+                    "mean": t.float().mean().item(),
+                    "std": t.float().std().item(),
+                    "min": t.float().min().item(),
+                    "max": t.float().max().item(),
+                }
+
+        print_params_summary(results)
+        if not args.no_svd:
+            print_component_summary(
+                results, "\n=== SVD Effective Rank by Component ==="
+            )
+            print_layer_grid(results)
+        print_weight_stats(results)
+        return results
+
+    analyze_one(args.ckpt_dir, "Primary")
+
+    if args.compare:
+        for cdir in args.compare:
+            analyze_one(cdir, "Compare")
+
+
+if __name__ == "__main__":
+    main()

scripts/eval/evaluate_humaneval.py

@@ -0,0 +1,393 @@
+"""HumanEval benchmark — functional pipeline design.
+
+Pipeline:
+    load -> generate -> extract -> test -> score -> report
+
+Each stage is a pure function (except GPU/CPU-bound I/O stages).
+Config is a single dataclass; side effects are isolated at pipeline boundaries.
+"""
+
+import argparse
+import itertools
+import json
+import os
+import re
+import subprocess
+import sys
+from dataclasses import dataclass
+from math import prod
+from typing import Dict, Iterator, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import tqdm
+
+from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+HUMANEVAL_URL = (
+    "https://github.com/openai/human-eval/raw/master/data/HumanEval.jsonl.gz"
+)
+
+STOP_SEQUENCES = [
+    "\nclass ",
+    "\ndef ",
+    "\n# ",
+    "\nif __name__",
+    "\nprint(",
+    "\n\n\n",
+]
+
+
+@dataclass
+class EvalConfig:
+    param_path: str = "./params"
+    data_path: str = "./humaneval/HumanEval.jsonl"
+    output: Optional[str] = None
+
+    test_only: Optional[str] = None
+    generate_only: bool = False
+
+    num_samples: int = 200
+    max_tokens: int = 512
+    temperature: float = 0.8
+    top_p: float = 0.95
+    top_k: int = 50
+    batch_size: int = 32
+    test_timeout: float = 3.0
+    test_workers: int = 8
+    k_values: Tuple[int, ...] = (1, 10, 100)
+    problem_indices: Optional[List[int]] = None
+
+
+def download(url: str, path: str):
+    if os.path.exists(path):
+        return
+    import gzip
+    import urllib.request
+
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    print(f"Downloading {url} ...")
+    tmp = path + ".tmp"
+    urllib.request.urlretrieve(url, tmp)
+    with gzip.open(tmp, "rb") as f_in:
+        with open(path, "wb") as f_out:
+            f_out.write(f_in.read())
+    os.remove(tmp)
+    print(f"  saved to {path}")
+
+
+def load_jsonl(path: str) -> List[dict]:
+    rows = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                rows.append(json.loads(line))
+    return rows
+
+
+def save_json(path: str, data):
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+
+
+def create_engine(param_path: str, batch_size: int) -> InferenceEngine:
+    model = AutoModel.from_pretrained(param_path)
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
+    model.to(device="cuda", dtype=torch.bfloat16)
+    return InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=batch_size,
+    )
+
+
+def trim_stop(text: str) -> str:
+    for stop in STOP_SEQUENCES:
+        idx = text.find(stop)
+        if idx != -1:
+            text = text[:idx]
+    return text
+
+
+def extract_body(code: str, entry_point: str) -> Optional[str]:
+    pattern = rf"def\s+{re.escape(entry_point)}\b[^:]*:"
+    match = re.search(pattern, code)
+    if not match:
+        return code
+
+    lines = code[match.end() :].split("\n")
+    body_lines = []
+    started = False
+
+    for line in lines:
+        stripped = line.rstrip()
+        if not stripped and not started:
+            continue
+        if not stripped and started:
+            body_lines.append("")
+            continue
+        if not started:
+            started = True
+        if stripped.lstrip() == stripped and started:
+            break
+        body_lines.append(stripped)
+
+    body = "\n".join(body_lines)
+    return body if body.strip() else None
+
+
+def deduplicate(seq: Sequence[str]) -> List[str]:
+    seen = set()
+    return [x for x in seq if not (x in seen or seen.add(x))]
+
+
+def generate_batch(
+    engine: InferenceEngine,
+    prompt: str,
+    n: int,
+    batch_size: int,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+) -> List[str]:
+    completions = []
+    remaining = n
+    while remaining > 0:
+        current = min(batch_size, remaining)
+        outputs = engine.generate(
+            prompt=[prompt] * current,
+            stream=False,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+        )
+        completions.extend(outputs if isinstance(outputs, list) else [outputs])
+        remaining -= current
+    return deduplicate(completions)
+
+
+def extract_completions(
+    raw: Sequence[str],
+    entry_point: str,
+) -> List[str]:
+    bodies = []
+    for r in raw:
+        t = trim_stop(r)
+        body = extract_body(t, entry_point)
+        if body:
+            bodies.append(body)
+    return bodies
+
+
+def generate_all(
+    engine: InferenceEngine,
+    problems: Sequence[dict],
+    cfg: EvalConfig,
+) -> List[dict]:
+    results = []
+    for problem in tqdm.tqdm(problems, desc="Generating", unit="problem"):
+        raw = generate_batch(
+            engine,
+            problem["prompt"],
+            cfg.num_samples,
+            cfg.batch_size,
+            cfg.max_tokens,
+            cfg.temperature,
+            cfg.top_p,
+            cfg.top_k,
+        )
+        bodies = extract_completions(raw, problem["entry_point"])
+        results.append(
+            dict(
+                task_id=problem["task_id"],
+                entry_point=problem["entry_point"],
+                prompt=problem["prompt"],
+                test=problem["test"],
+                completions=bodies,
+            )
+        )
+    return results
+
+
+def execute_one(args: tuple) -> bool:
+    full_code, entry_point, timeout = args
+    try:
+        r = subprocess.run(
+            [sys.executable, "-c", full_code],
+            capture_output=True,
+            timeout=timeout,
+        )
+        return r.returncode == 0
+    except subprocess.TimeoutExpired:
+        return False
+    except Exception:
+        return False
+
+
+def test_one(item: dict, cfg: EvalConfig) -> Tuple[str, int, int]:
+    from concurrent.futures import ProcessPoolExecutor
+
+    task_id = item["task_id"]
+    completions = item["completions"]
+    codes = [
+        (
+            item["prompt"] + c + "\n" + item["test"],
+            item["entry_point"],
+            cfg.test_timeout,
+        )
+        for c in completions
+    ]
+    n = len(codes)
+    passed = 0
+    with ProcessPoolExecutor(max_workers=cfg.test_workers) as pool:
+        for ok in pool.map(execute_one, codes):
+            if ok:
+                passed += 1
+    return task_id, n, passed
+
+
+def test_all(
+    items: Sequence[dict],
+    cfg: EvalConfig,
+) -> Iterator[Tuple[str, int, int]]:
+    for item in tqdm.tqdm(items, desc="Testing", unit="problem"):
+        yield test_one(item, cfg)
+
+
+def pass_at_k(n: int, c: int, k: int) -> float:
+    if n - c < k:
+        return 1.0
+    return 1.0 - float(prod(1.0 - k / np.arange(n - c + 1, n + 1)))
+
+
+def score_results(
+    results: Iterator[Tuple[str, int, int]],
+    k_values: Tuple[int, ...],
+) -> Dict:
+    # filter to k <= n (peek first result to get n)
+    first = next(results)
+    results = itertools.chain([first], results)
+    n = first[1]
+    k_values = tuple(k for k in k_values if k <= n)
+
+    scores = {k: [] for k in k_values}
+    output = {}
+    for task_id, n, passed in results:
+        entry = {"task_id": task_id, "n": n, "passed": passed}
+        for k in k_values:
+            pk = round(pass_at_k(n, passed, k), 4)
+            entry[f"pass@{k}"] = pk
+            scores[k].append(pk)
+        output[task_id] = entry
+
+    summary = {}
+    for k in k_values:
+        vals = scores[k]
+        summary[f"pass@{k}"] = round(float(np.mean(vals)), 4)
+    output["_summary"] = summary
+    return output
+
+
+def run_pipeline(cfg: EvalConfig) -> Dict:
+    if cfg.test_only:
+        with open(cfg.test_only, encoding="utf-8") as f:
+            generated = json.load(f)
+    else:
+        download(HUMANEVAL_URL, cfg.data_path)
+
+        problems = load_jsonl(cfg.data_path)
+        if cfg.problem_indices:
+            problems = [problems[i] for i in cfg.problem_indices if i < len(problems)]
+
+        engine = create_engine(cfg.param_path, cfg.batch_size)
+
+        try:
+            generated = generate_all(engine, problems, cfg)
+        finally:
+            engine.shutdown()
+
+        if cfg.output:
+            mid = cfg.output.replace(".json", "_completions.json")
+            save_json(mid, generated)
+            print(f"Completions saved to {mid}")
+
+        if cfg.generate_only:
+            return {}
+
+    results = test_all(generated, cfg)
+    scored = score_results(results, cfg.k_values)
+    return scored
+
+
+def parse_args(argv: Optional[List[str]] = None) -> EvalConfig:
+    p = argparse.ArgumentParser(description="HumanEval benchmark")
+    p.add_argument("--param_path", type=str, default="./params")
+    p.add_argument("--data_path", type=str, default="./humaneval/HumanEval.jsonl")
+    p.add_argument("--output", type=str, default=None)
+    p.add_argument(
+        "--test_only",
+        type=str,
+        default=None,
+        help="Skip generation, test existing completions JSON",
+    )
+    p.add_argument(
+        "--generate_only", action="store_true", help="Only generate, skip testing"
+    )
+    p.add_argument("--num_samples", type=int, default=200)
+    p.add_argument("--max_tokens", type=int, default=512)
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top_p", type=float, default=0.95)
+    p.add_argument("--top_k", type=int, default=50)
+    p.add_argument("--batch_size", type=int, default=32)
+    p.add_argument("--test_workers", type=int, default=8)
+    p.add_argument("--test_timeout", type=float, default=3.0)
+    p.add_argument("--problems", type=int, nargs="+", default=None)
+    args = p.parse_args(argv)
+
+    return EvalConfig(
+        param_path=args.param_path,
+        data_path=args.data_path,
+        output=args.output,
+        test_only=args.test_only,
+        generate_only=args.generate_only,
+        num_samples=args.num_samples,
+        max_tokens=args.max_tokens,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        top_k=args.top_k,
+        batch_size=args.batch_size,
+        test_workers=args.test_workers,
+        test_timeout=args.test_timeout,
+        problem_indices=args.problems,
+    )
+
+
+def report(scored: Dict):
+    summary = scored.pop("_summary", {})
+    print(f"\n{'=' * 60}")
+    for k, v in summary.items():
+        print(f"  {k}: {v:.2%}")
+    print(f"{'=' * 60}")
+    scored["_summary"] = summary
+
+
+def main():
+    cfg = parse_args()
+    scored = run_pipeline(cfg)
+    report(scored)
+    if cfg.output:
+        save_json(cfg.output, scored)
+        print(f"Results saved to {cfg.output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/eval/evaluate_ifd.py

@@ -0,0 +1,474 @@
+"""IFD (Instruction Following Difficulty) data quality scoring.
+
+IFD = conditional_NLL / unconditional_NLL
+
+- Messages format: plain text concatenation (no chat template)
+- Plain format: raw instr_key + resp_key fields
+
+v2 changelog:
+  - Same token set: unconditional pass prefixes resp with a plain-text sentinel
+    (default ``\\n``; use ``--sentinel_text ""`` for bos/pad fallback).
+    Both branches predict the identical N resp tokens.
+    Single-token answers (rl=1) are now supported.
+  - ctx_len tracked in output
+  - skip_reason for None samples (no more silent None)
+  - --per_token for per-token IFD breakdown
+"""
+
+import argparse
+import json
+import statistics
+
+import torch
+import torch.nn.functional as F
+import tqdm
+
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+
+def _pack_bins(pairs, max_len):
+    """BFD bin packing: pack (c+r) into bins of max total length."""
+    indexed = sorted(enumerate(pairs), key=lambda x: -(len(x[1][0]) + len(x[1][1])))
+    bins = []
+    lengths = []
+    for orig_idx, (c, r) in indexed:
+        size = len(c) + len(r)
+        best_bin = -1
+        for bi, rem in enumerate(lengths):
+            if rem >= size:
+                if best_bin < 0 or rem < lengths[best_bin]:
+                    best_bin = bi
+        if best_bin >= 0:
+            bins[best_bin].append((orig_idx, c, r))
+            lengths[best_bin] -= size
+        else:
+            bins.append([(orig_idx, c, r)])
+            lengths.append(max_len - size)
+    return bins
+
+
+def _resolve_sentinel_ids(tokenizer, sentinel_text):
+    """Tokenize the sentinel text for the unconditional pass prefix.
+
+    Falls back to bos/pad_token_id when sentinel_text is empty or
+    cannot be encoded.
+    """
+    if sentinel_text:
+        ids = tokenizer.encode(sentinel_text, add_special_tokens=False)
+        if ids:
+            return ids
+    for attr in ("bos_token_id", "pad_token_id", "eos_token_id"):
+        tid = getattr(tokenizer, attr, None)
+        if tid is not None:
+            return [tid]
+    return [0]
+
+
+@torch.inference_mode()
+def _score_batch(
+    pairs, model, device, max_len=2048, sentinel_ids=None, per_token=False
+):
+    """BFD-packed IFD with text-sentinel-anchored unconditional pass.
+
+    Conditional:   (ctx + resp[0..i-1]) → resp[i],  i = 0..N-1
+    Unconditional: (<sentinel> + resp[0..i-1]) → resp[i],  i = 0..N-1
+
+    Both branches predict the identical N response tokens.  A short
+    plain-text sentinel gives the unconditional pass a prefix so that
+    every response token can be predicted.  Single-token answers (rl=1)
+    are supported.
+    """
+    if not pairs:
+        return []
+
+    if sentinel_ids is None:
+        sentinel_ids = [0]
+
+    bins = _pack_bins(pairs, max_len)
+    result = [None] * len(pairs)
+
+    # ---- conditional pass (packed, per-document position IDs) ----
+    for bin_items in bins:
+        seq_ids = []
+        global_pos = []
+        doc_ids = []
+        doc_offsets = []
+
+        for di, (orig_idx, c, r) in enumerate(bin_items):
+            ctx_len = len(c)
+            start = len(seq_ids)
+            item_len = len(c) + len(r)
+            seq_ids.extend(c)
+            seq_ids.extend(r)
+            end = len(seq_ids)
+            global_pos.extend(range(item_len))
+            doc_ids.extend([di] * item_len)
+            doc_offsets.append((start, end, orig_idx, ctx_len))
+
+        full_ids = torch.tensor([seq_ids], device=device, dtype=torch.long)
+        pos_ids = torch.tensor([global_pos], device=device, dtype=torch.long)
+        seq_len = len(seq_ids)
+        causal = torch.tril(
+            torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
+        )
+        doc_t = torch.tensor([doc_ids], device=device)
+        doc_mask = doc_t.unsqueeze(-1) == doc_t.unsqueeze(-2)
+        attn_mask = (causal & doc_mask[0]).unsqueeze(0).unsqueeze(0)
+        logits_full = model(full_ids, position_ids=pos_ids, input_mask=attn_mask)[
+            "logits"
+        ][0]
+
+        for start, end, orig_idx, ctx_len in doc_offsets:
+            rl = end - start - ctx_len
+            resp_start = start + ctx_len - 1
+            resp_logits = logits_full[resp_start : end - 1]
+            resp_targets = torch.tensor(
+                seq_ids[start + ctx_len : end], device=device, dtype=torch.long
+            )
+            cond_losses = F.cross_entropy(
+                resp_logits, resp_targets, reduction="none"
+            ).cpu()
+            result[orig_idx] = {
+                "_cond_losses": cond_losses,
+                "_rl": rl,
+                "_ctx_len": ctx_len,
+            }
+
+    # ---- unconditional pass (sentinel-prefixed, batched 2D) ----
+    valid_items = [
+        (
+            i,
+            result[i]["_rl"],
+            result[i]["_ctx_len"],
+            result[i]["_cond_losses"],
+            pairs[i][1],
+        )
+        for i in range(len(pairs))
+        if result[i] is not None and "_cond_losses" in result[i]
+    ]
+    if not valid_items:
+        return result
+
+    valid_items.sort(key=lambda x: -x[1])
+    prefix_len = len(sentinel_ids)
+    max_rl = prefix_len + max(rl for _, rl, _, _, _ in valid_items)
+    bsz = len(valid_items)
+
+    u_batch = torch.zeros(bsz, max_rl, dtype=torch.long, device=device)
+    for ri, (_, rl, _, _, r_ids) in enumerate(valid_items):
+        u_batch[ri, :prefix_len] = torch.tensor(sentinel_ids, dtype=torch.long)
+        u_batch[ri, prefix_len : prefix_len + rl] = torch.tensor(
+            r_ids, dtype=torch.long
+        )
+
+    logits_resp = model(u_batch)["logits"]
+
+    for ri, (orig_idx, rl, ctx_len, cond_losses, _) in enumerate(valid_items):
+        unp_logits = logits_resp[ri, prefix_len - 1 : prefix_len - 1 + rl]
+        unp_targets = u_batch[ri, prefix_len : prefix_len + rl]
+        uncond_losses = F.cross_entropy(unp_logits, unp_targets, reduction="none").cpu()
+
+        L_cond = cond_losses.mean().item()
+        L_uncond = uncond_losses.mean().item()
+        ifd = L_cond / L_uncond if L_uncond > 0 else None
+
+        out = {
+            "L_cond": round(L_cond, 6),
+            "L_uncond": round(L_uncond, 6),
+            "ifd": round(ifd, 6) if ifd is not None else None,
+            "ctx_len": ctx_len,
+            "resp_len": rl,
+        }
+        if per_token:
+            per = [
+                (round(c.item() / u.item(), 6) if u.item() > 0 else None)
+                for c, u in zip(cond_losses, uncond_losses)
+            ]
+            out["ifd_per_token"] = per
+        result[orig_idx] = out
+
+    return result
+
+
+def _trim(context_ids, resp_ids, max_len):
+    """Truncate to fit max_len, keeping response intact if possible."""
+    if len(resp_ids) > max_len // 2:
+        resp_ids = resp_ids[: max_len // 2]
+    full_ids = context_ids + resp_ids
+    if len(full_ids) <= max_len:
+        return context_ids, resp_ids
+    overflow = len(full_ids) - max_len
+    if overflow >= len(context_ids):
+        return [], resp_ids[:max_len]
+    return context_ids[overflow:], resp_ids
+
+
+def score_plain(
+    model,
+    tokenizer,
+    instruction,
+    response,
+    device,
+    max_len=2048,
+    sentinel_ids=None,
+    per_token=False,
+):
+    """Compute IFD for a single instruction-response pair (plain format)."""
+    ctx_ids = tokenizer.encode(instruction, add_special_tokens=False)
+    resp_ids = tokenizer.encode(response, add_special_tokens=False)
+    ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
+    if not ctx_ids or not resp_ids:
+        return {
+            "L_cond": None,
+            "L_uncond": None,
+            "ifd": None,
+            "skip_reason": "empty ctx or resp",
+        }
+    return _score_batch(
+        [(ctx_ids, resp_ids)],
+        model,
+        device,
+        max_len,
+        sentinel_ids=sentinel_ids,
+        per_token=per_token,
+    )[0]
+
+
+def score_messages(
+    model, tokenizer, messages, device, max_len=2048, sentinel_ids=None, per_token=False
+):
+    """Compute IFD for each assistant turn in a messages array."""
+    turns = []
+    for i, msg in enumerate(messages):
+        if msg.get("role") != "assistant":
+            continue
+        ctx_text = "\n\n".join(m["content"] for m in messages[:i])
+        ctx_ids = tokenizer.encode(ctx_text)
+        resp_ids = tokenizer.encode(msg["content"], add_special_tokens=False)
+        ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
+        if ctx_ids and resp_ids:
+            turns.append((ctx_ids, resp_ids))
+    if not turns:
+        return None
+    raw_scores = _score_batch(
+        turns, model, device, max_len, sentinel_ids=sentinel_ids, per_token=per_token
+    )
+    valid = [s for s in raw_scores if s is not None and s.get("ifd") is not None]
+    if not valid:
+        return {"ifd": None, "ifd_turns": raw_scores}
+    avg = sum(s["ifd"] for s in valid) / len(valid)
+    return {
+        "ifd": avg,
+        "ifd_detail": valid[0] if len(valid) == 1 else None,
+        "ifd_turns": raw_scores,
+    }
+
+
+def process_file(
+    param_path,
+    input_file,
+    output_file,
+    instr_key,
+    resp_key,
+    max_len=2048,
+    data_format="plain",
+    batch_size=1,
+    device=None,
+    sentinel_text="\n",
+    per_token=False,
+):
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16 if "cuda" in device else torch.float32
+
+    model = AutoModel.from_pretrained(param_path)
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
+    model.to(device=device, dtype=dtype)
+    model.eval()
+
+    sentinel_ids = _resolve_sentinel_ids(tokenizer, sentinel_text)
+
+    with open(input_file, encoding="utf-8") as f:
+        data = [json.loads(line) for line in f if line.strip()]
+
+    results = []
+    all_ifds = []
+    buffer = []
+
+    for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
+        if data_format == "messages":
+            turns = []
+            for i, msg in enumerate(item.get("messages", [])):
+                if msg.get("role") != "assistant":
+                    continue
+                ctx_text = "\n\n".join(m["content"] for m in item["messages"][:i])
+                ctx_ids = tokenizer.encode(ctx_text)
+                resp_ids = tokenizer.encode(msg["content"], add_special_tokens=False)
+                ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
+                if ctx_ids and resp_ids:
+                    turns.append((ctx_ids, resp_ids))
+            if not turns:
+                results.append(
+                    {
+                        **item,
+                        "ifd": None,
+                        "skip_reason": "no valid assistant turns",
+                        "ifd_turns": [],
+                    }
+                )
+                continue
+            buffer.append((item, turns, "messages"))
+        else:
+            ctx_ids = tokenizer.encode(item[instr_key], add_special_tokens=False)
+            resp_ids = tokenizer.encode(item[resp_key], add_special_tokens=False)
+            ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
+            if not ctx_ids or not resp_ids:
+                results.append(
+                    {
+                        **item,
+                        "ifd": None,
+                        "ifd_detail": {"skip_reason": "empty ctx or resp"},
+                    }
+                )
+                continue
+            buffer.append((item, [(ctx_ids, resp_ids)], "plain"))
+
+        if len(buffer) >= batch_size:
+            _flush_buffer(
+                buffer,
+                results,
+                all_ifds,
+                model,
+                device,
+                max_len,
+                sentinel_ids,
+                per_token,
+            )
+
+    if buffer:
+        _flush_buffer(
+            buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
+        )
+
+    with open(output_file, "w", encoding="utf-8") as f:
+        for item in results:
+            f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+    valid_ifd = [v for v in all_ifds if v is not None]
+    if valid_ifd:
+        print(f"\n{'=' * 50}")
+        print(f"  Samples:       {len(data)}")
+        print(f"  Valid IFD:     {len(valid_ifd)}")
+        print(f"  Skipped:       {len(data) - len(valid_ifd)}")
+        print(f"  Mean IFD:      {statistics.mean(valid_ifd):.4f}")
+        print(f"  Median IFD:    {statistics.median(valid_ifd):.4f}")
+        if len(valid_ifd) > 1:
+            print(f"  Stdev IFD:     {statistics.stdev(valid_ifd):.4f}")
+        print(f"  Min IFD:       {min(valid_ifd):.4f}")
+        print(f"  Max IFD:       {max(valid_ifd):.4f}")
+        print(f"{'=' * 50}")
+    print(f"Results saved to {output_file}")
+
+
+def _flush_buffer(
+    buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
+):
+    all_pairs = []
+    indices = []
+    for item, turns, fmt in buffer:
+        start = len(all_pairs)
+        all_pairs.extend(turns)
+        indices.append((item, turns, fmt, start, len(all_pairs)))
+
+    raw = _score_batch(
+        all_pairs,
+        model,
+        device,
+        max_len,
+        sentinel_ids=sentinel_ids,
+        per_token=per_token,
+    )
+
+    for item, turns, fmt, start, end in indices:
+        turn_scores = raw[start:end]
+        if fmt == "messages":
+            valid = [
+                s for s in turn_scores if s is not None and s.get("ifd") is not None
+            ]
+            if not valid:
+                results.append({**item, "ifd": None, "ifd_turns": turn_scores})
+            else:
+                avg = sum(s["ifd"] for s in valid) / len(valid)
+                all_ifds.append(avg)
+                results.append(
+                    {
+                        **item,
+                        "ifd": avg,
+                        "ifd_detail": valid[0] if len(valid) == 1 else None,
+                        "ifd_turns": turn_scores,
+                    }
+                )
+        else:
+            score = turn_scores[0]
+            all_ifds.append(score.get("ifd"))
+            results.append({**item, "ifd": score.get("ifd"), "ifd_detail": score})
+
+    buffer.clear()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compute IFD scores for instruction-response data"
+    )
+    parser.add_argument("--param_path", type=str, required=True, help="Model directory")
+    parser.add_argument("--input", type=str, required=True, help="Input JSONL file")
+    parser.add_argument("--output", type=str, required=True, help="Output JSONL file")
+    parser.add_argument("--max_len", type=int, default=2048, help="Max token length")
+    parser.add_argument(
+        "--format",
+        type=str,
+        default="plain",
+        choices=["plain", "messages"],
+        help="Input format",
+    )
+    parser.add_argument(
+        "--instr_key", type=str, default="instruction", help="Key for instruction field"
+    )
+    parser.add_argument(
+        "--resp_key", type=str, default="response", help="Key for response field"
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=8, help="Batch size for model forward passes"
+    )
+    parser.add_argument("--device", type=str, default=None, help="Device (e.g. cuda:0)")
+    parser.add_argument(
+        "--sentinel_text",
+        type=str,
+        default="\n",
+        help='Plain-text prefix for unconditional pass (default: "\\n"). Use "" for bos/pad fallback.',
+    )
+    parser.add_argument(
+        "--per_token",
+        action="store_true",
+        help="Include per-token IFD breakdown in output",
+    )
+    args = parser.parse_args()
+
+    process_file(
+        args.param_path,
+        args.input,
+        args.output,
+        args.instr_key,
+        args.resp_key,
+        args.max_len,
+        data_format=args.format,
+        batch_size=args.batch_size,
+        device=args.device,
+        sentinel_text=args.sentinel_text,
+        per_token=args.per_token,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/eval/evaluate_ifeval.py

@@ -0,0 +1,609 @@
+"""IFEval instruction-following evaluation benchmark.
+
+Evaluates model responses against regex-based constraint verifiers.
+Supports all IFEval constraint types except language detection.
+
+Usage::
+
+    python scripts/tools/evaluate_ifeval.py --param_path ./params \
+        --data_path ifeval.jsonl --output results.json \
+        --temperature 0.1 --max_tokens 512
+"""
+
+import argparse
+import json
+import os
+import re
+import urllib.request
+from typing import Callable, Dict, List, Optional
+
+import torch
+import tqdm
+
+from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+IFEVAL_URL = (
+    "https://raw.githubusercontent.com/google-research/"
+    "google-research/master/instruction_following_eval/data/input_data.jsonl"
+)
+
+CONSTRAINT_VERIFIERS: Dict[str, Callable[[str, dict], bool]] = {}
+
+
+def register(instruction_id: str):
+    def decorator(fn):
+        CONSTRAINT_VERIFIERS[instruction_id] = fn
+        return fn
+
+    return decorator
+
+
+@register("keywords:existence")
+def check_keyword_existence(response: str, kwargs: dict) -> bool:
+    for kw in kwargs["keywords"]:
+        if not re.search(re.escape(kw), response, re.IGNORECASE):
+            return False
+    return True
+
+
+@register("keywords:frequency")
+def check_keyword_frequency(response: str, kwargs: dict) -> bool:
+    keyword = kwargs["keyword"]
+    frequency = kwargs.get("frequency", 1)
+    relation = kwargs.get("relation", "at least")
+    count = len(re.findall(re.escape(keyword), response, re.IGNORECASE))
+    if relation == "less than":
+        return count < frequency
+    return count >= frequency
+
+
+@register("keywords:forbidden_words")
+def check_forbidden_words(response: str, kwargs: dict) -> bool:
+    for word in kwargs["forbidden_words"]:
+        if re.search(r"\b" + re.escape(word) + r"\b", response, re.IGNORECASE):
+            return False
+    return True
+
+
+@register("keywords:letter_frequency")
+def check_letter_frequency(response: str, kwargs: dict) -> bool:
+    letter = kwargs["letter"].lower()
+    frequency = kwargs.get("let_frequency", 1)
+    relation = kwargs.get("let_relation", "at least")
+    count = response.lower().count(letter)
+    if relation == "less than":
+        return count < frequency
+    return count >= frequency
+
+
+@register("detectable_content:number_placeholders")
+def check_placeholders(response: str, kwargs: dict) -> bool:
+    num = kwargs.get("num_placeholders", 1)
+    placeholders = re.findall(r"\[.*?\]", response)
+    return len(placeholders) >= num
+
+
+@register("detectable_content:postscript")
+def check_postscript(response: str, kwargs: dict) -> bool:
+    marker = kwargs.get("postscript_marker", "P.S.")
+    response_lower = response.lower()
+    if marker == "P.P.S":
+        return bool(re.search(r"p\.\s?p\.\s?s", response_lower))
+    elif marker == "P.S.":
+        return bool(re.search(r"p\.\s?s\.", response_lower))
+    else:
+        return bool(re.search(re.escape(marker.lower()), response_lower))
+
+
+@register("detectable_format:number_bullet_lists")
+def check_bullet_lists(response: str, kwargs: dict) -> bool:
+    num = kwargs.get("num_bullets", 1)
+    bullets = re.findall(r"^\s*\*[^\*].*$", response, re.MULTILINE)
+    dashes = re.findall(r"^\s*-.*$", response, re.MULTILINE)
+    return len(bullets) + len(dashes) == num
+
+
+@register("detectable_format:number_highlighted_sections")
+def check_highlighted_sections(response: str, kwargs: dict) -> bool:
+    num = kwargs.get("num_highlights", 1)
+    highlights = re.findall(r"\*[^\n\*]+\*", response)
+    count = 0
+    for h in highlights:
+        if h.strip("*").strip():
+            count += 1
+    return count >= num
+
+
+@register("detectable_format:multiple_sections")
+def check_multiple_sections(response: str, kwargs: dict) -> bool:
+    splitter = kwargs.get("section_spliter", "Section")
+    num = kwargs.get("num_sections", 1)
+    pattern = r"\s?" + re.escape(splitter) + r"\s?\d+\s?"
+    sections = re.split(pattern, response)
+    return len(sections) - 1 >= num
+
+
+@register("detectable_format:title")
+def check_title(response: str, kwargs: dict) -> bool:
+    titles = re.findall(r"<<[^>\n]+>>", response)
+    for title in titles:
+        if title.strip("<>").strip():
+            return True
+    return False
+
+
+@register("detectable_format:json_format")
+def check_json_format(response: str, kwargs: dict) -> bool:
+    value = response.strip()
+    for prefix in ("```json", "```Json", "```JSON", "```"):
+        if value.lower().startswith(prefix.lower()):
+            value = value[len(prefix) :].strip()
+    if value.endswith("```"):
+        value = value[:-3].strip()
+    try:
+        json.loads(value)
+        return True
+    except (ValueError, json.JSONDecodeError):
+        return False
+
+
+@register("detectable_format:general_punctuation")
+def check_general_punctuation(response: str, kwargs: dict) -> bool:
+    punctuation_blacklist = kwargs.get("punctuation_blacklist", [])
+    for punct in punctuation_blacklist:
+        if punct in response:
+            return False
+    return True
+
+
+@register("detectable_format:number_highlighted_words")
+def check_highlighted_words(response: str, kwargs: dict) -> bool:
+    num = kwargs.get("num_highlights", 1)
+    highlights = re.findall(r"\*[^\s\*][^\*]*[^\s\*]\*", response)
+    return len(highlights) >= num
+
+
+@register("startend:end_checker")
+def check_end_checker(response: str, kwargs: dict) -> bool:
+    end_phrase = kwargs["end_phrase"]
+    return (
+        response.strip()
+        .rstrip('"')
+        .rstrip()
+        .lower()
+        .endswith(end_phrase.strip().lower())
+    )
+
+
+@register("startend:quotation")
+def check_quotation(response: str, kwargs: dict) -> bool:
+    value = response.strip()
+    return value.startswith('"') and value.endswith('"')
+
+
+@register("startend:start_checker")
+def check_start_checker(response: str, kwargs: dict) -> bool:
+    starter = kwargs["starter"]
+    return bool(re.search(r"^\s*" + re.escape(starter), response, re.MULTILINE))
+
+
+@register("change_case:english_capital")
+def check_english_capital(response: str, kwargs: dict) -> bool:
+    return response.isupper()
+
+
+@register("change_case:english_lowercase")
+def check_english_lowercase(response: str, kwargs: dict) -> bool:
+    return response.islower()
+
+
+@register("change_case:capital_word_frequency")
+def check_capital_word_frequency(response: str, kwargs: dict) -> bool:
+    frequency = kwargs.get("capital_frequency", 1)
+    relation = kwargs.get("capital_relation", "at least")
+    capital_words = re.findall(r"\b[A-Z]{2,}\b", response)
+    count = len(capital_words)
+    if relation == "less than":
+        return count < frequency
+    return count >= frequency
+
+
+@register("punctuation:no_comma")
+def check_no_comma(response: str, kwargs: dict) -> bool:
+    return "," not in response
+
+
+def count_words(text: str) -> int:
+    return len(re.findall(r"\b\w+\b", text))
+
+
+def count_sentences(text: str) -> int:
+    text = text.strip()
+    if not text:
+        return 0
+    sentences = re.split(r"(?<=[.!?])\s+", text)
+    return len([s for s in sentences if s.strip()])
+
+
+@register("length_constraints:number_words")
+def check_number_words(response: str, kwargs: dict) -> bool:
+    num = kwargs.get("num_words", 100)
+    relation = kwargs.get("relation", "at least")
+    cnt = count_words(response)
+    if relation == "less than":
+        return cnt < num
+    return cnt >= num
+
+
+@register("length_constraints:number_sentences")
+def check_number_sentences(response: str, kwargs: dict) -> bool:
+    num = kwargs.get("num_sentences", 5)
+    relation = kwargs.get("relation", "at least")
+    cnt = count_sentences(response)
+    if relation == "less than":
+        return cnt < num
+    return cnt >= num
+
+
+@register("length_constraints:number_paragraphs")
+def check_number_paragraphs(response: str, kwargs: dict) -> bool:
+    num = kwargs.get("num_paragraphs", 1)
+    if "***" in response:
+        paragraphs = re.split(r"\s?\*\*\*\s?", response)
+    else:
+        paragraphs = re.split(r"\n\n+", response)
+    actual = len([p for p in paragraphs if p.strip()])
+    return actual == num
+
+
+@register("length_constraints:nth_paragraph_first_word")
+def check_nth_paragraph_first_word(response: str, kwargs: dict) -> bool:
+    num_paragraphs = kwargs.get("num_paragraphs", 1)
+    nth = kwargs.get("nth_paragraph", 1)
+    first_word = kwargs.get("first_word", "").lower()
+
+    paragraphs = re.split(r"\n\n+", response)
+    paragraphs = [p.strip() for p in paragraphs if p.strip()]
+
+    if len(paragraphs) != num_paragraphs:
+        return False
+    if nth > len(paragraphs):
+        return False
+
+    target = paragraphs[nth - 1]
+    words = target.split()
+    if not words:
+        return False
+
+    word = words[0].strip().lstrip("'\"").rstrip(".,!?:;\"'")
+    return word.lower() == first_word
+
+
+@register("length_constraints:nth_word_checker")
+def check_nth_word(response: str, kwargs: dict) -> bool:
+    nth = kwargs.get("nth_word", 1)
+    target = kwargs.get("target_word", "").lower()
+    words = re.findall(r"\b\w+\b", response)
+    if nth > len(words):
+        return False
+    return words[nth - 1].lower() == target
+
+
+@register("combination:repeat_prompt")
+def check_repeat_prompt(response: str, kwargs: dict) -> bool:
+    prompt = kwargs["prompt_to_repeat"]
+    return response.strip().lower().startswith(prompt.strip().lower())
+
+
+@register("combination:two_responses")
+def check_two_responses(response: str, kwargs: dict) -> bool:
+    parts = response.split("******")
+    valid = [p for p in parts if p.strip()]
+    if len(valid) != 2:
+        return False
+    return valid[0].strip() != valid[1].strip()
+
+
+def download_ifeval(data_path: str):
+    if os.path.exists(data_path):
+        return
+    os.makedirs(os.path.dirname(data_path) or ".", exist_ok=True)
+    print(f"Downloading IFEval from {IFEVAL_URL} ...")
+    tmp = data_path + ".tmp"
+    urllib.request.urlretrieve(IFEVAL_URL, tmp)
+    with open(tmp, "rb") as f_in:
+        content = f_in.read()
+    with open(data_path, "wb") as f_out:
+        f_out.write(content)
+    os.remove(tmp)
+    print(f"  saved to {data_path}")
+
+
+def load_problems(data_path: str) -> List[dict]:
+    problems = []
+    with open(data_path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                problems.append(json.loads(line))
+    return problems
+
+
+def verify_response(response: str, instruction_id: str, kwargs: dict) -> Optional[bool]:
+    verifier = CONSTRAINT_VERIFIERS.get(instruction_id)
+    if verifier is None:
+        return None
+    try:
+        return verifier(response, kwargs)
+    except Exception:
+        return False
+
+
+def generate_one(
+    engine: InferenceEngine,
+    tokenizer: AutoTokenizer,
+    prompt: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+) -> str:
+    formatted = tokenizer.apply_chat_template(
+        [{"role": "user", "content": prompt}],
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+    output = engine.generate(
+        prompt=formatted,
+        stream=False,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+    )
+    if isinstance(output, list):
+        return output[0]
+    return output
+
+
+def evaluate(
+    engine: InferenceEngine,
+    tokenizer: AutoTokenizer,
+    problems: List[dict],
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    num_samples: int = 1,
+) -> Dict:
+    results = {}
+    constraint_stats: Dict[str, Dict[str, int]] = {}
+    total_constraints = 0
+    total_passed = 0
+
+    for problem in tqdm.tqdm(problems, desc="IFEval", unit="problem"):
+        key = problem["key"]
+        prompt = problem["prompt"]
+        instruction_ids = problem["instruction_id_list"]
+        kwargs_list = problem["kwargs"]
+
+        samples = []
+        for _ in range(num_samples):
+            response = generate_one(
+                engine, tokenizer, prompt, max_tokens, temperature, top_p, top_k
+            )
+            samples.append(response)
+
+        constraint_results = []
+        passed = 0
+        verified = 0
+
+        for idx, instruction_id in enumerate(instruction_ids):
+            kwargs = kwargs_list[idx] if idx < len(kwargs_list) else {}
+            best_pass = False
+            for response in samples:
+                result = verify_response(response, instruction_id, kwargs)
+                if result is None:
+                    continue
+                if result:
+                    best_pass = True
+                    break
+
+            verifier_exists = instruction_id in CONSTRAINT_VERIFIERS
+            if verifier_exists:
+                verified += 1
+                if best_pass:
+                    passed += 1
+
+            constraint_results.append(
+                {
+                    "instruction_id": instruction_id,
+                    "passed": best_pass,
+                    "supported": verifier_exists,
+                    "kwargs": kwargs,
+                }
+            )
+
+            if verifier_exists:
+                if instruction_id not in constraint_stats:
+                    constraint_stats[instruction_id] = {
+                        "total": 0,
+                        "passed": 0,
+                    }
+                constraint_stats[instruction_id]["total"] += 1
+                if best_pass:
+                    constraint_stats[instruction_id]["passed"] += 1
+
+        total_constraints += verified
+        total_passed += passed
+
+        accuracy = passed / verified if verified > 0 else None
+        results[str(key)] = {
+            "key": key,
+            "prompt": prompt,
+            "response": samples[0],
+            "num_samples": num_samples,
+            "num_constraints": len(instruction_ids),
+            "num_verified": verified,
+            "num_passed": passed,
+            "accuracy": round(accuracy, 4) if accuracy is not None else None,
+            "constraints": constraint_results,
+        }
+
+    overall_accuracy = (
+        round(total_passed / total_constraints, 4) if total_constraints > 0 else 0.0
+    )
+
+    type_summary = {}
+    for inst_id, stats in sorted(constraint_stats.items()):
+        type_summary[inst_id] = {
+            "total": stats["total"],
+            "passed": stats["passed"],
+            "accuracy": round(stats["passed"] / stats["total"], 4)
+            if stats["total"] > 0
+            else 0.0,
+        }
+
+    unsupported_count = sum(
+        1
+        for p in problems
+        for iid in p["instruction_id_list"]
+        if iid not in CONSTRAINT_VERIFIERS
+    )
+
+    results["_summary"] = {
+        "total_problems": len(problems),
+        "total_constraints": total_constraints,
+        "total_passed": total_passed,
+        "overall_accuracy": overall_accuracy,
+        "unsupported_constraints": unsupported_count,
+        "supported_types": sorted(CONSTRAINT_VERIFIERS.keys()),
+        "per_type_accuracy": type_summary,
+    }
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="IFEval benchmark")
+    parser.add_argument(
+        "--param_path", type=str, default="./params", help="Model directory"
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        default="./ifeval/input_data.jsonl",
+        help="IFEval JSONL file (auto-download if missing)",
+    )
+    parser.add_argument("--output", type=str, default=None, help="Output JSON path")
+    parser.add_argument(
+        "--max_tokens", type=int, default=512, help="Max generation tokens"
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.1,
+        help="Sampling temperature",
+    )
+    parser.add_argument("--top_p", type=float, default=0.95, help="Top-p sampling")
+    parser.add_argument("--top_k", type=int, default=50, help="Top-k sampling")
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=1,
+        help="Number of samples per problem (best-of-n scoring)",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=1, help="Inference batch size"
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Limit to first N problems (for quick testing)",
+    )
+    parser.add_argument(
+        "--dump_responses",
+        type=str,
+        default=None,
+        help="Path to dump raw model responses (JSONL)",
+    )
+    args = parser.parse_args()
+
+    download_ifeval(args.data_path)
+    problems = load_problems(args.data_path)
+    if args.limit:
+        problems = problems[: args.limit]
+
+    print(f"Loaded {len(problems)} problems")
+    print(f"Supported constraint types: {len(CONSTRAINT_VERIFIERS)}")
+
+    model = AutoModel.from_pretrained(args.param_path)
+    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
+    model.to(device="cuda", dtype=torch.bfloat16)
+    model.eval()
+
+    engine = InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=args.batch_size,
+    )
+
+    results = evaluate(
+        engine=engine,
+        tokenizer=tokenizer,
+        problems=problems,
+        max_tokens=args.max_tokens,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        top_k=args.top_k,
+        num_samples=args.num_samples,
+    )
+
+    summary = results.pop("_summary")
+    print(f"\n{'=' * 60}")
+    print(f"  Problems:       {summary['total_problems']}")
+    print(f"  Constraints:    {summary['total_constraints']}")
+    print(f"  Passed:         {summary['total_passed']}")
+    print(f"  Accuracy:       {summary['overall_accuracy']:.2%}")
+    print(f"  Unsupported:    {summary['unsupported_constraints']}")
+    print(f"{'=' * 60}")
+
+    print("\nPer-type accuracy:")
+    for inst_id, stats in sorted(summary["per_type_accuracy"].items()):
+        print(
+            f"  {inst_id:50s}  {stats['accuracy']:.2%}  "
+            f"({stats['passed']}/{stats['total']})"
+        )
+
+    if args.output:
+        results["_summary"] = summary
+        with open(args.output, "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2, ensure_ascii=False)
+        print(f"\nResults saved to {args.output}")
+
+    if args.dump_responses:
+        with open(args.dump_responses, "w", encoding="utf-8") as f:
+            for k, v in results.items():
+                if k.startswith("_"):
+                    continue
+                f.write(
+                    json.dumps(
+                        {
+                            "key": v["key"],
+                            "prompt": v["prompt"],
+                            "response": v["response"],
+                        },
+                        ensure_ascii=False,
+                    )
+                    + "\n"
+                )
+        print(f"Responses dumped to {args.dump_responses}")
+
+    engine.shutdown()
+
+
+if __name__ == "__main__":
+    main()

scripts/eval/evaluate_mmlu.py

@@ -0,0 +1,319 @@
+"""MMLU evaluation via log-likelihood ranking."""
+
+import argparse
+import csv
+import json
+import os
+import shutil
+import tarfile
+
+import requests
+import torch
+import torch.nn.functional as F
+import tqdm
+
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+MMLU_URL = "https://people.eecs.berkeley.edu/~hendrycks/data.tar"
+MMLU_SUBJECTS = [
+    "abstract_algebra",
+    "anatomy",
+    "astronomy",
+    "business_ethics",
+    "clinical_knowledge",
+    "college_biology",
+    "college_chemistry",
+    "college_computer_science",
+    "college_mathematics",
+    "college_medicine",
+    "college_physics",
+    "computer_security",
+    "conceptual_physics",
+    "econometrics",
+    "electrical_engineering",
+    "elementary_mathematics",
+    "formal_logic",
+    "global_facts",
+    "high_school_biology",
+    "high_school_chemistry",
+    "high_school_computer_science",
+    "high_school_european_history",
+    "high_school_geography",
+    "high_school_government_and_politics",
+    "high_school_macroeconomics",
+    "high_school_mathematics",
+    "high_school_microeconomics",
+    "high_school_physics",
+    "high_school_psychology",
+    "high_school_statistics",
+    "high_school_us_history",
+    "high_school_world_history",
+    "human_aging",
+    "human_sexuality",
+    "international_law",
+    "jurisprudence",
+    "logical_fallacies",
+    "machine_learning",
+    "management",
+    "marketing",
+    "medical_genetics",
+    "miscellaneous",
+    "moral_disputes",
+    "moral_scenarios",
+    "nutrition",
+    "philosophy",
+    "prehistory",
+    "professional_accounting",
+    "professional_law",
+    "professional_medicine",
+    "professional_psychology",
+    "public_relations",
+    "security_studies",
+    "sociology",
+    "us_foreign_policy",
+    "virology",
+    "world_religions",
+]
+
+
+def _download_and_extract(url: str, data_dir: str):
+    tar_path = os.path.join(data_dir, "data.tar")
+    os.makedirs(data_dir, exist_ok=True)
+    print(f"Downloading MMLU data from {url}...")
+    resp = requests.get(url, stream=True, timeout=300)
+    resp.raise_for_status()
+    total = int(resp.headers.get("content-length", 0))
+    with tqdm.tqdm(total=total, unit="B", unit_scale=True, desc="  Download") as bar:
+        with open(tar_path, "wb") as f:
+            for chunk in resp.iter_content(chunk_size=8192):
+                f.write(chunk)
+                bar.update(len(chunk))
+    print("Extracting...")
+    with tarfile.open(tar_path, "r") as tf:
+        tf.extractall(data_dir)
+    os.remove(tar_path)
+
+
+def download_mmlu(data_dir: str):
+    _download_and_extract(MMLU_URL, data_dir)
+    src = os.path.join(data_dir, "data")
+    if os.path.exists(src):
+        for item in os.listdir(src):
+            src_item = os.path.join(src, item)
+            dst_item = os.path.join(data_dir, item)
+            if os.path.exists(dst_item):
+                if os.path.isdir(dst_item):
+                    shutil.rmtree(dst_item)
+                else:
+                    os.remove(dst_item)
+            os.rename(src_item, dst_item)
+        os.rmdir(src)
+    print(f"MMLU data saved to {data_dir}")
+
+
+def _strip_prefix(text: str, prefix: str) -> str:
+    if text.startswith(prefix):
+        return text[len(prefix) :].strip()
+    return text
+
+
+def load_csv(path: str) -> list[dict]:
+    data = []
+    with open(path, "r", encoding="utf-8") as f:
+        for row in csv.reader(f):
+            if len(row) < 6:
+                continue
+            if row[0].strip().lower() == "question":
+                continue
+            data.append(
+                {
+                    "question": row[0].strip(),
+                    "A": _strip_prefix(row[1].strip(), "A)"),
+                    "B": _strip_prefix(row[2].strip(), "B)"),
+                    "C": _strip_prefix(row[3].strip(), "C)"),
+                    "D": _strip_prefix(row[4].strip(), "D)"),
+                    "answer": row[5].strip(),
+                }
+            )
+    return data
+
+
+def build_prompt(
+    question: str, choices: dict, subject: str, n_shot: int, dev_data: list[dict]
+) -> str:
+    prompt = ""
+    if n_shot > 0 and dev_data:
+        prompt = f"The following are multiple choice questions (with answers) about {subject}.\n\n"
+        for item in dev_data[:n_shot]:
+            prompt += f"Question: {item['question']}\n"
+            for k in ("A", "B", "C", "D"):
+                prompt += f"{k}. {item[k]}\n"
+            prompt += f"Answer: {item['answer']}\n\n"
+    prompt += f"Question: {question}\n"
+    for k in ("A", "B", "C", "D"):
+        prompt += f"{k}. {choices[k]}\n"
+    prompt += "Answer:"
+    return prompt
+
+
+def apply_chat(
+    tokenizer, raw_prompt: str, n_shot: int, dev_data: list[dict] | None
+) -> str:
+    """Wrap raw MMLU prompt in the model's chat template format.
+
+    For few-shot, prepend example Q&A pairs as a second user/assistant exchange.
+    """
+    messages = []
+    if n_shot > 0 and dev_data:
+        for item in dev_data[:n_shot]:
+            q = f"Question: {item['question']}\n"
+            for k in ("A", "B", "C", "D"):
+                q += f"{k}. {item[k]}\n"
+            q += "Answer:"
+            messages.append({"role": "user", "content": q})
+            messages.append({"role": "assistant", "content": item["answer"]})
+    messages.append({"role": "user", "content": raw_prompt})
+    return tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+
+def choice_logprob(
+    model, tokenizer, context_ids: list[int], choice_letter: str, device: str
+) -> float:
+    choice_text = choice_letter
+    choice_ids = tokenizer.encode(choice_text, add_special_tokens=False)
+    input_ids = context_ids + choice_ids
+    max_len = model.config.max_len
+    if len(input_ids) > max_len:
+        overflow = len(input_ids) - max_len
+        input_ids = input_ids[overflow:]
+        ctx_len = len(input_ids) - len(choice_ids)
+    else:
+        ctx_len = len(context_ids)
+
+    input_tensor = torch.tensor([input_ids], device=device, dtype=torch.long)
+    with torch.inference_mode():
+        logits = model(input_tensor)["logits"][0]
+
+    score = 0.0
+    for i, tid in enumerate(choice_ids):
+        pos = ctx_len - 1 + i
+        if pos >= len(logits):
+            break
+        score += F.log_softmax(logits[pos], dim=-1)[tid].item()
+    return score
+
+
+def evaluate_subject(
+    model,
+    tokenizer,
+    subject: str,
+    test_data: list[dict],
+    dev_data: list[dict] | None,
+    device: str,
+    n_shot: int,
+) -> tuple[float, int, int]:
+    correct = 0
+    total = 0
+    for item in tqdm.tqdm(test_data, desc=f"{subject:40s}", leave=False):
+        raw_prompt = build_prompt(
+            item["question"], item, subject, n_shot, dev_data or []
+        )
+        context = apply_chat(tokenizer, raw_prompt, n_shot, dev_data or [])
+        context_ids = tokenizer.encode(context)
+        scores = {
+            c: choice_logprob(model, tokenizer, context_ids, c, device)
+            for c in ("A", "B", "C", "D")
+        }
+        if max(scores, key=scores.get) == item["answer"]:
+            correct += 1
+        total += 1
+    return correct / total, correct, total
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MMLU evaluation")
+    parser.add_argument(
+        "--param_path", type=str, default="./params", help="Model directory"
+    )
+    parser.add_argument(
+        "--data_dir", type=str, default="./mmlu_data", help="MMLU data directory"
+    )
+    parser.add_argument("--download", action="store_true", help="Download MMLU data")
+    parser.add_argument(
+        "--n_shot", type=int, default=5, help="Few-shot examples (0 for zero-shot)"
+    )
+    parser.add_argument(
+        "--subjects", type=str, nargs="+", help="Specific subjects (default: all)"
+    )
+    parser.add_argument("--output", type=str, help="Output JSON path")
+    parser.add_argument("--split", type=str, default="test", choices=["test", "val"])
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device",
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16" if torch.cuda.is_available() else "float32",
+        help="Torch dtype",
+    )
+    args = parser.parse_args()
+
+    if args.download or not os.path.exists(args.data_dir):
+        download_mmlu(args.data_dir)
+
+    model = AutoModel.from_pretrained(args.param_path)
+    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
+    device = args.device
+    dtype = getattr(torch, args.dtype)
+    model.to(device=device, dtype=dtype)
+    model.eval()
+
+    subjects = args.subjects or MMLU_SUBJECTS
+    results = {}
+    total_correct = 0
+    total_questions = 0
+
+    for subject in subjects:
+        dev_path = os.path.join(args.data_dir, "dev", f"{subject}_dev.csv")
+        test_path = os.path.join(
+            args.data_dir, args.split, f"{subject}_{args.split}.csv"
+        )
+
+        if not os.path.exists(test_path):
+            print(f"  Skipping {subject}: test file not found")
+            continue
+
+        dev_data = load_csv(dev_path) if os.path.exists(dev_path) else None
+        test_data = load_csv(test_path)
+
+        acc, corr, tot = evaluate_subject(
+            model, tokenizer, subject, test_data, dev_data, device, args.n_shot
+        )
+        results[subject] = {"accuracy": round(acc, 4), "correct": corr, "total": tot}
+        total_correct += corr
+        total_questions += tot
+        print(f"  {subject:40s}  {acc:.2%}  ({corr}/{tot})")
+
+    overall = total_correct / total_questions if total_questions else 0
+    print(f"\n{'=' * 70}")
+    print(f"  Overall: {overall:.2%}  ({total_correct}/{total_questions})")
+    results["_overall"] = {
+        "accuracy": round(overall, 4),
+        "correct": total_correct,
+        "total": total_questions,
+    }
+
+    if args.output:
+        with open(args.output, "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2)
+        print(f"Results saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/eval/evaluate_ppl.py

@@ -10,11 +10,11 @@ from astrai.tokenize import AutoTokenizer
 
 
 def process_file(
-    model_dir: str, input_file: str, output_file: str, batch_size: int, text_key: str
+    param_path: str, input_file: str, output_file: str, batch_size: int, text_key: str
 ):
     # 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)
 
     with open(input_file, "r", encoding="utf-8") as f:
@@ -44,8 +44,8 @@ def process_file(
 
         for seq in batch_encoded:
             pad_len = max_len - len(seq)
-            padded_seq = [tokenizer.pad_id] * pad_len + seq
-            mask = [False] * pad_len + [True] * len(seq)
+            padded_seq = seq + [tokenizer.pad_id] * pad_len
+            mask = [True] * len(seq) + [False] * pad_len
             padded_ids.append(padded_seq)
             masks.append(mask)
 
@@ -86,9 +86,9 @@ def process_file(
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run perplexity with a Khaosz model.")
+    parser = argparse.ArgumentParser(description="Perplexity evaluation on JSONL text.")
     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_file", type=str, required=True, help="Path to the input file."

scripts/eval/evaluate_rouge.py

@@ -0,0 +1,153 @@
+"""ROUGE evaluation (manual implementation, no external deps).
+
+Computes ROUGE-1, ROUGE-2, ROUGE-L precision, recall, and F1.
+
+Usage::
+
+    # Batch evaluation from JSONL (each line: {"reference": ..., "candidate": ...})
+    python scripts/eval/evaluate_rouge.py --data_path preds.jsonl --output results.json
+
+    # As a library
+    from scripts.eval.evaluate_rouge import compute_rouge
+    scores = compute_rouge("the cat sat on the mat", "the cat sat")
+"""
+
+import argparse
+import json
+from collections import Counter
+from typing import Dict, List, Tuple
+
+
+def _tokenize(text: str) -> List[str]:
+    return text.split()
+
+
+def _ngrams(tokens: List[str], n: int) -> Counter:
+    return Counter(zip(*[tokens[i:] for i in range(n)]))
+
+
+def _lcs(x: List[str], y: List[str]) -> int:
+    m, n = len(x), len(y)
+    dp = [[0] * (n + 1) for _ in range(m + 1)]
+    for i in range(1, m + 1):
+        xi = x[i - 1]
+        dpi = dp[i]
+        dpi_1 = dp[i - 1]
+        for j in range(1, n + 1):
+            if xi == y[j - 1]:
+                dpi[j] = dpi_1[j - 1] + 1
+            else:
+                dpi[j] = dpi_1[j] if dpi_1[j] > dpi[j - 1] else dpi[j - 1]
+    return dp[m][n]
+
+
+def _f1(precision: float, recall: float) -> float:
+    if precision + recall == 0:
+        return 0.0
+    return 2 * precision * recall / (precision + recall)
+
+
+def _rouge_n(ref_tokens: List[str], cand_tokens: List[str], n: int) -> Dict[str, float]:
+    ref_ngrams = _ngrams(ref_tokens, n)
+    cand_ngrams = _ngrams(cand_tokens, n)
+
+    overlap = sum((cand_ngrams & ref_ngrams).values())
+    cand_total = sum(cand_ngrams.values())
+    ref_total = sum(ref_ngrams.values())
+
+    precision = overlap / cand_total if cand_total > 0 else 0.0
+    recall = overlap / ref_total if ref_total > 0 else 0.0
+    f1 = _f1(precision, recall)
+
+    return {"precision": precision, "recall": recall, "f1": f1}
+
+
+def _rouge_l(ref_tokens: List[str], cand_tokens: List[str]) -> Dict[str, float]:
+    lcs_len = _lcs(ref_tokens, cand_tokens)
+    ref_len = len(ref_tokens)
+    cand_len = len(cand_tokens)
+
+    recall = lcs_len / ref_len if ref_len > 0 else 0.0
+    precision = lcs_len / cand_len if cand_len > 0 else 0.0
+    f1 = _f1(precision, recall)
+
+    return {"precision": precision, "recall": recall, "f1": f1}
+
+
+def compute_rouge(
+    reference: str, candidate: str, n: int = 2
+) -> Dict[str, Dict[str, float]]:
+    """Compute ROUGE-N (1..n) and ROUGE-L scores.
+
+    Returns::
+
+        {
+            "rouge-1": {"precision": ..., "recall": ..., "f1": ...},
+            "rouge-2": {"precision": ..., "recall": ..., "f1": ...},
+            "rouge-l": {"precision": ..., "recall": ..., "f1": ...},
+        }
+    """
+    ref_tokens = _tokenize(reference)
+    cand_tokens = _tokenize(candidate)
+
+    results = {}
+    for i in range(1, n + 1):
+        results[f"rouge-{i}"] = _rouge_n(ref_tokens, cand_tokens, i)
+    results["rouge-l"] = _rouge_l(ref_tokens, cand_tokens)
+    return results
+
+
+def evaluate_file(data_path: str) -> Dict:
+    with open(data_path, "r", encoding="utf-8") as f:
+        pairs = [json.loads(line) for line in f if line.strip()]
+
+    agg = {
+        k: {"precision": 0.0, "recall": 0.0, "f1": 0.0}
+        for k in ("rouge-1", "rouge-2", "rouge-l")
+    }
+    per_item = []
+
+    for item in pairs:
+        ref = item["reference"]
+        cand = item["candidate"]
+        scores = compute_rouge(ref, cand)
+        per_item.append({**item, "scores": scores})
+        for k, v in scores.items():
+            agg[k]["precision"] += v["precision"]
+            agg[k]["recall"] += v["recall"]
+            agg[k]["f1"] += v["f1"]
+
+    n = len(pairs)
+    for k in agg:
+        agg[k] = {m: v / n for m, v in agg[k].items()}
+
+    return {"num_samples": n, "aggregate": agg, "per_item": per_item}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="ROUGE evaluation")
+    parser.add_argument(
+        "--data_path", required=True, help="JSONL with reference/candidate per line"
+    )
+    parser.add_argument("--output", type=str, default=None, help="Output JSON path")
+    args = parser.parse_args()
+
+    results = evaluate_file(args.data_path)
+    agg = results["aggregate"]
+
+    print(f"Samples: {results['num_samples']}")
+    print()
+    for metric in ("rouge-1", "rouge-2", "rouge-l"):
+        s = agg[metric]
+        print(
+            f"  {metric:8s}  P={s['precision']:.4f}  R={s['recall']:.4f}  F1={s['f1']:.4f}"
+        )
+
+    if args.output:
+        with open(args.output, "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2, ensure_ascii=False)
+        print(f"\nSaved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/tools/benchmark.py

@@ -1,14 +1,13 @@
-"""Benchmark Transformer with PagedCache (replaces old persistent_key_values)."""
+"""Benchmark AutoRegressiveLM with KVCache"""
 
 from dataclasses import dataclass
 from typing import Any, Dict
 
 import torch
-from torch import Tensor
 
-from astrai.config import ModelConfig
-from astrai.inference.cache import PagedCache
-from astrai.model.transformer import Transformer
+from astrai.config import AutoRegressiveLMConfig
+from astrai.inference import KVCache
+from astrai.model.transformer import AutoRegressiveLM
 
 
 @dataclass
@@ -22,7 +21,7 @@ class BenchmarkResult:
 class GenerationBenchmark:
     def __init__(
         self,
-        config: ModelConfig,
+        config: AutoRegressiveLMConfig,
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
         page_size: int = 128,
@@ -30,11 +29,11 @@ class GenerationBenchmark:
         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()
         head_dim = config.dim // config.n_heads
         n_pages = (config.max_len * 4 + page_size - 1) // page_size
-        self._page_cache = PagedCache(
+        self._page_cache = KVCache(
             config.n_layers,
             n_pages,
             page_size,
@@ -61,9 +60,6 @@ class GenerationBenchmark:
         )
         return prompt_ids, gen_ids
 
-    def _make_mask(self, batch_size: int, seq_len: int) -> Tensor:
-        return torch.ones(batch_size, seq_len, dtype=torch.bool, device=self.device)
-
     @torch.inference_mode()
     def run_prefill_benchmark(
         self,
@@ -134,7 +130,12 @@ class GenerationBenchmark:
             )
 
             n_pages = (prompt_length + gen_length + page_size - 1) // page_size
-            pages = self._page_cache.alloc_n(n_pages * batch_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,
@@ -145,8 +146,11 @@ class GenerationBenchmark:
             _ = self.model(
                 prompt_ids,
                 paged_cache=cv,
-                start_pos=0,
-                input_mask=self._make_mask(batch_size, prompt_length),
+                position_ids=torch.arange(
+                    prompt_length, dtype=torch.long, device=self.device
+                )
+                .unsqueeze(0)
+                .expand(batch_size, -1),
             )
 
             torch.cuda.synchronize()
@@ -162,8 +166,12 @@ class GenerationBenchmark:
                 _ = self.model(
                     input_token,
                     paged_cache=cv,
-                    start_pos=current_pos,
-                    input_mask=self._make_mask(batch_size, 1),
+                    position_ids=torch.full(
+                        (batch_size, 1),
+                        current_pos,
+                        dtype=torch.long,
+                        device=self.device,
+                    ),
                 )
                 current_pos += 1
             end.record()
@@ -173,7 +181,7 @@ class GenerationBenchmark:
             total_time += trial_time
 
             for idx in pages:
-                self._page_cache.free(idx)
+                self._page_cache._pool.free(idx)
 
             print(
                 f"  Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
@@ -208,7 +216,7 @@ def print_benchmark_result(result: BenchmarkResult):
 
 
 if __name__ == "__main__":
-    config = ModelConfig(
+    config = AutoRegressiveLMConfig(
         vocab_size=10000,
         dim=1536,
         n_heads=24,
@@ -222,7 +230,7 @@ if __name__ == "__main__":
     benchmark = GenerationBenchmark(config)
 
     print("=" * 80)
-    print("Running Transformer Generation Benchmark (PagedCache)")
+    print("Running AutoRegressiveLM Generation Benchmark (KVCache)")
     print("=" * 80)
 
     prefill_result = benchmark.run_prefill_benchmark(

scripts/tools/generate.py

@@ -1,5 +1,6 @@
 import argparse
 import json
+from typing import Optional
 
 import torch
 
@@ -17,7 +18,8 @@ def processor(
     top_p: float,
     question_key: str,
     response_key: str,
-    max_tokens: int,
+    max_tokens: Optional[int],
+    batch_size: int,
 ):
     # Load model and tokenizer
     model = AutoModel.from_pretrained(param_path)
@@ -25,7 +27,9 @@ def processor(
     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]
@@ -69,7 +73,7 @@ def processor(
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run generate with a Khaosz model.")
+    parser = argparse.ArgumentParser(description="Batch generation from JSONL file.")
 
     parser.add_argument(
         "--param_path", type=str, required=True, help="Path to the model directory."
@@ -90,36 +94,42 @@ if __name__ == "__main__":
         "--question_key",
         type=str,
         default="question",
-        help="Key for the question in the input JSON.",
+        help="Key for the question in the input JSON (default: question).",
     )
     parser.add_argument(
         "--response_key",
         type=str,
         default="response",
-        help="Key for the response in the output JSON.",
+        help="Key for the response in the output JSON (default: response).",
     )
     parser.add_argument(
         "--temperature",
         type=float,
         default=0.60,
-        help="Temperature for generating responses.",
+        help="Temperature for generating responses (default: 0.60).",
     )
     parser.add_argument(
-        "--top_k", type=int, default=30, help="Top-k value for generating responses."
+        "--top_k",
+        type=int,
+        default=30,
+        help="Top-k value for generating responses (default: 30).",
     )
     parser.add_argument(
         "--top_p",
         type=float,
         default=0.95,
-        help="Top-p value for generating responses.",
+        help="Top-p value for generating responses (default: 0.95).",
     )
     parser.add_argument(
-        "--batch_size", type=int, default=1, help="Batch size for generating responses."
+        "--batch_size",
+        type=int,
+        default=1,
+        help="Batch size for generating responses (default: 1).",
     )
     parser.add_argument(
         "--max_tokens",
         type=int,
-        default=2048,
+        default=None,
         help="Maximum tokens to generate (default: model config max_len).",
     )
 

scripts/tools/preprocess.py

@@ -0,0 +1,38 @@
+"""CLI: JSONL → tokenized .h5/.bin via config-driven Pipeline."""
+
+import argparse
+
+from astrai.config.preprocess_config import PipelineConfig
+from astrai.preprocessing.pipeline import Pipeline
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Raw JSONL → tokenized .h5/.bin via config-driven Pipeline"
+    )
+    parser.add_argument(
+        "inputs", nargs="+", metavar="JSONL", help="One or more JSONL files"
+    )
+    parser.add_argument("--output_dir", "-o", required=True, help="Output directory")
+    parser.add_argument(
+        "--config", "-c", required=True, help="Path to pipeline config JSON"
+    )
+    parser.add_argument(
+        "--tokenizer_path",
+        default="params",
+        help="Path to tokenizer directory (default: params)",
+    )
+    args = parser.parse_args()
+
+    config = PipelineConfig.from_file(args.config)
+
+    Pipeline(
+        config=config,
+        input_paths=args.inputs,
+        output_dir=args.output_dir,
+        tokenizer_path=args.tokenizer_path,
+    ).run()
+
+
+if __name__ == "__main__":
+    main()

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():
@@ -18,7 +18,7 @@ def main():
         "--reload", action="store_true", help="Enable auto-reload for development"
     )
     parser.add_argument(
-        "--param-path",
+        "--param_path",
         type=Path,
         default=None,
         help="Path to model parameters (default: project_root/params)",

scripts/tools/train.py

@@ -2,22 +2,19 @@ import argparse
 import os
 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.model.components.decoder_block import DecoderBlock
 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",
@@ -42,18 +39,20 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--n_epoch", type=int, default=1, help="Number of epochs to train."
     )
-    parser.add_argument("--batch_size", type=int, default=1, help="Batch size per GPU.")
     parser.add_argument(
-        "--accumulation_steps",
+        "--batch_per_device", type=int, default=1, help="Batch size per GPU."
+    )
+    parser.add_argument(
+        "--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 warmup steps for LR scheduler.",
+        "--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."
@@ -68,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",
@@ -114,9 +113,15 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--label_smoothing",
         type=float,
-        default=0.1,
+        default=0.0,
         help="cross_entropy function label smoothing parameter",
     )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Enable activation checkpointing for DecoderBlock modules.",
+    )
 
     parser.add_argument(
         "--ckpt_interval",
@@ -130,6 +135,30 @@ def parse_args() -> argparse.Namespace:
         default="checkpoint",
         help="Directory to save checkpoints.",
     )
+    parser.add_argument(
+        "--val_split",
+        type=float,
+        default=None,
+        help="Ratio to split from training dataset for validation (e.g. 0.05).",
+    )
+    parser.add_argument(
+        "--val_step",
+        type=int,
+        default=1000,
+        help="Number of optimizer steps between validation runs.",
+    )
+    parser.add_argument(
+        "--metrics",
+        nargs="*",
+        default=["loss", "lr", "grad_norm"],
+        help="Metrics to log (e.g. --metrics loss lr val_loss). Default: loss lr grad_norm.",
+    )
+    parser.add_argument(
+        "--log_dir",
+        type=str,
+        default="checkpoint/logs",
+        help="Directory for metric logs.",
+    )
     parser.add_argument(
         "--grpo_sync_interval",
         type=int,
@@ -140,43 +169,141 @@ def parse_args() -> argparse.Namespace:
         "--start_epoch", type=int, default=0, help="Start epoch for training."
     )
     parser.add_argument(
-        "--start_batch", type=int, default=0, help="Start batch for training."
+        "--start_samples",
+        type=int,
+        default=0,
+        help="Start samples (per rank) for training.",
     )
 
+    parser.add_argument(
+        "--master_addr",
+        type=str,
+        default="localhost",
+        help="Master node address for distributed training.",
+    )
+    parser.add_argument(
+        "--master_port",
+        type=str,
+        default="29500",
+        help="Master node port for distributed training.",
+    )
+    parser.add_argument(
+        "--backend",
+        type=str,
+        default="nccl",
+        help="Distributed training backend.",
+    )
     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", "fsdp"],
+        help="Parallel training strategy (none, ddp, fsdp).",
+    )
     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.",
+    )
+    parser.add_argument(
+        "--neftune_alpha",
+        type=float,
+        default=0.0,
+        help="NEFTune noise alpha (0=disabled, typical: 5.0).",
+    )
+
+    parser.add_argument(
+        "--schedule_type",
+        type=str,
+        default="cosine",
+        choices=["cosine", "sgdr", "wsd"],
+        help="Learning rate scheduler type.",
+    )
+    parser.add_argument(
+        "--min_rate",
+        type=float,
+        default=None,
+        help="Minimum LR as fraction of base LR. Uses scheduler default if not set (cosine/sgdr: 0.05, wsd: 0.0).",
+    )
+    parser.add_argument(
+        "--cycle_length",
+        type=int,
+        default=None,
+        help="SGDR first cycle length in steps. Defaults to total_steps - warmup_steps.",
+    )
+    parser.add_argument(
+        "--t_mult",
+        type=int,
+        default=2,
+        help="SGDR cycle length multiplier per restart.",
+    )
+    parser.add_argument(
+        "--stable_steps",
+        type=int,
+        default=None,
+        help="WSD stable plateau steps. Required when --schedule_type wsd.",
+    )
+    parser.add_argument(
+        "--decay_steps",
+        type=int,
+        default=None,
+        help="WSD decay steps. Defaults to total_steps - warmup_steps - stable_steps.",
+    )
 
     args = parser.parse_args()
 
     return args
 
 
-def ddp_wrap(model: nn.Module):
-    local_rank = get_rank()
-    model = model.to(dtype=torch.bfloat16)
-    ddp_model = DDP(
-        model,
-        device_ids=[local_rank],
-        output_device=local_rank,
-        find_unused_parameters=False,
-    )
-    return ddp_model
+def create_model(config):
+    return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
 
 
-def create_optimizer(model: nn.Module, **kwargs) -> optim.Optimizer:
-    return optim.AdamW(model.parameters(), **kwargs)
+def create_optimizer(model, **kwargs) -> optim.Optimizer:
+    decay_params = []
+    no_decay_params = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if param.dim() < 2 or "norm" in name or "bias" in name:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+    param_groups = [
+        {"params": decay_params, "weight_decay": kwargs.pop("weight_decay", 0.01)},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+    return optim.AdamW(param_groups, fused=True, **kwargs)
 
 
 def create_scheduler(
     optimizer: optim.Optimizer, **kwargs
 ) -> optim.lr_scheduler.LRScheduler:
-    return SchedulerFactory.create(optimizer, **kwargs)
+    schedule_type = kwargs.pop("schedule_type")
+    return SchedulerFactory.create(schedule_type, 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(
@@ -185,13 +312,17 @@ 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,
+    start_samples: int,
+    grad_accum_steps: int,
+    warmup_ratio: float,
     ckpt_interval: int,
     ckpt_dir: str,
+    val_split: float,
+    val_step: int,
+    metrics: list[str],
+    log_dir: str,
     dpo_beta: float,
     grpo_clip_eps: float,
     grpo_kl_coef: float,
@@ -205,34 +336,39 @@ def train(
     random_seed: int,
     num_workers: int,
     pin_memory: bool,
+    gradient_checkpointing: bool,
     window_size: int,
     stride: int,
     nprocs: int,
+    parallel_mode: str,
     device_type: str,
+    backend: str,
+    master_addr: str,
+    master_port: str,
+    start_method: str,
+    neftune_alpha: float,
+    schedule_type: str,
+    min_rate: float,
+    cycle_length: int,
+    t_mult: int,
+    stable_steps: int,
+    decay_steps: int,
 ):
     assert train_type in ["seq", "sft", "dpo", "grpo"]
     assert os.path.exists(param_path)
+    if nprocs > 1 and parallel_mode == "none":
+        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")
 
     # 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)
+    config.neftune_alpha = neftune_alpha
 
     if window_size is None:
         window_size = config.max_len
 
-    # Create bare Transformer (for training, no tokenizer needed)
-    model = Transformer(config)
-
-    # Load weights if available
-    weights_path = os.path.join(param_path, "model.safetensors")
-    if os.path.exists(weights_path):
-        state_dict = st.load_file(weights_path)
-        model.load_state_dict(state_dict, strict=False)
-
     strategy_kwargs = {
-        "dpo_beta": dpo_beta,
+        "beta": dpo_beta,
         "label_smoothing": label_smoothing,
         "clip_eps": grpo_clip_eps,
         "kl_coef": grpo_kl_coef,
@@ -240,6 +376,12 @@ def train(
         "sync_interval": grpo_sync_interval,
     }
 
+    executor_kwargs = {
+        "gradient_as_bucket_view": True,
+        "broadcast_buffers": False,
+    }
+
+    model_fn = partial(create_model, config)
     dataset = DatasetFactory.load(
         train_type=train_type,
         load_path=data_root_path,
@@ -256,42 +398,74 @@ 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)
+    warmup_steps = min(warmup_steps, total_steps)
+
+    scheduler_kwargs = {"warmup_steps": warmup_steps}
+
+    if schedule_type == "cosine":
+        scheduler_kwargs["lr_decay_steps"] = total_steps - warmup_steps
+    elif schedule_type == "sgdr":
+        scheduler_kwargs["cycle_length"] = cycle_length or (total_steps - warmup_steps)
+        scheduler_kwargs["t_mult"] = t_mult
+    elif schedule_type == "wsd":
+        remaining = total_steps - warmup_steps
+        stable_steps_ = stable_steps or max(1, int(remaining * 0.8))
+        scheduler_kwargs["stable_steps"] = stable_steps_
+        scheduler_kwargs["decay_steps"] = max(
+            1, decay_steps or (remaining - stable_steps_)
+        )
+
+    if min_rate is not None:
+        scheduler_kwargs["min_rate"] = min_rate
+
     scheduler_fn = partial(
         create_scheduler,
-        **{
-            "schedule_type": "cosine",
-            "warmup_steps": warmup_steps,
-            "lr_decay_steps": total_steps - warmup_steps,
-        },
+        schedule_type=schedule_type,
+        **scheduler_kwargs,
     )
 
+    grad_ckpt_modules = [DecoderBlock] if gradient_checkpointing else []
+
     train_config = TrainConfig(
-        model=model,
+        model_fn=model_fn,
         strategy=train_type,
         dataset=dataset,
         optimizer_fn=optimizer_fn,
         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,
+        start_samples=start_samples,
         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,
+        backend=backend,
+        master_addr=master_addr,
+        master_port=master_port,
+        parallel_mode=parallel_mode,
         device_type=device_type,
+        start_method=start_method,
+        val_split=val_split,
+        val_step=val_step,
+        metrics=metrics,
+        log_dir=log_dir,
+        gradient_checkpointing_modules=grad_ckpt_modules,
+        executor_kwargs=executor_kwargs,
         extra_kwargs=strategy_kwargs,
+        neftune_alpha=neftune_alpha,
     )
 
     trainer = Trainer(train_config)
-    trainer.train()
+    trainer.train(resume_dir=param_path)
 
 
 if __name__ == "__main__":

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
 
@@ -72,7 +75,7 @@ class MultiTurnDataset(Dataset):
 
 
 class EarlyStoppingDataset(Dataset):
-    """Dataset that triggers early stopping after a specified number of iterations."""
+    """Dataset that triggers early stopping after consuming a specified number of samples."""
 
     def __init__(self, length=10, stop_after=5):
         self.length = length
@@ -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/conftest.py

@@ -0,0 +1,235 @@
+import json
+import os
+import tempfile
+
+import pytest
+from tokenizers import Tokenizer, models, pre_tokenizers, trainers
+
+from astrai.config.preprocess_config import (
+    InputConfig,
+    PipelineConfig,
+    ProcessingConfig,
+)
+from astrai.preprocessing.builder import SectionedMaskBuilder
+from astrai.tokenize import AutoTokenizer
+
+_SPECIAL_TOKENS_CONFIG = {
+    "bos_token": "<|begin_of_sentence|>",
+    "eos_token": "<|end_of_sentence|>",
+    "pad_token": "<|_pad_|>",
+    "unk_token": "<|_unk_|>",
+    "im_start": "<|im_start|>",
+    "im_end": "<|im_end|>",
+}
+
+_SPECIAL_TOKENS = list(_SPECIAL_TOKENS_CONFIG.values())
+
+_CHAT_TEMPLATE = (
+    "{% for message in messages %}"
+    "{% if message['role'] == 'system' %}"
+    "<|im_start|>system\n{{ message['content'] }}<|im_end|>\n"
+    "{% elif message['role'] == 'user' %}"
+    "<|im_start|>user\n{{ message['content'] }}<|im_end|>\n"
+    "{% elif message['role'] == 'assistant' %}"
+    "<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n"
+    "{% endif %}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+)
+
+_CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
+
+_INSTRUCTION_SECTIONS = [
+    {"field": "prompt", "action": "mask", "add_special_tokens": True},
+    {"field": "response", "action": "train"},
+]
+
+_TEXT_SECTIONS = [{"field": "text", "action": "train"}]
+
+_GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}]
+
+
+def _build_chat_tokenizer():
+    tok = Tokenizer(models.BPE())
+    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tr = trainers.BpeTrainer(
+        vocab_size=512,
+        min_frequency=1,
+        special_tokens=_SPECIAL_TOKENS,
+    )
+    train_data = [
+        "hello world",
+        "Hi there!",
+        "You are helpful.",
+        "What is 2+2?",
+        "Tell me a story about dragons and knights.",
+        "Sure, here is a tale.",
+        "Translate to French: Hello",
+        "Bonjour",
+        "Artificial Intelligence is a field of computer science.",
+        "system",
+        "user",
+        "assistant",
+        "<|im_start|>",
+        "<|im_end|>",
+        *[chr(i) for i in range(32, 127)],
+    ]
+    tok.train_from_iterator(train_data, tr)
+
+    auto_tok = AutoTokenizer()
+    auto_tok._tokenizer = tok
+    auto_tok._special_token_map = {
+        "bos_token": "<|begin_of_sentence|>",
+        "eos_token": "<|end_of_sentence|>",
+        "pad_token": "<|_pad_|>",
+        "unk_token": "<|_unk_|>",
+    }
+    auto_tok.set_chat_template(_CHAT_TEMPLATE)
+    return auto_tok
+
+
+@pytest.fixture(scope="session")
+def chat_tokenizer():
+    return _build_chat_tokenizer()
+
+
+@pytest.fixture
+def temp_dir():
+    d = tempfile.mkdtemp()
+    yield d
+    import shutil
+
+    shutil.rmtree(d, ignore_errors=True)
+
+
+def make_chat_config():
+    return PipelineConfig(
+        input=InputConfig(sections=_CHAT_SECTIONS),
+        mask={"system": "mask", "user": "mask", "assistant": "train"},
+        mask_default="mask",
+        preprocessing=ProcessingConfig(max_seq_len=2048),
+    )
+
+
+def make_instruction_config():
+    return PipelineConfig(
+        input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+        mask={"prompt": "mask", "response": "train"},
+        mask_default="mask",
+        preprocessing=ProcessingConfig(max_seq_len=2048),
+    )
+
+
+def make_text_config():
+    return PipelineConfig(
+        input=InputConfig(sections=_TEXT_SECTIONS),
+        preprocessing=ProcessingConfig(
+            max_seq_len=2048, min_chars=1, max_chars=2_000_000
+        ),
+    )
+
+
+def make_dpo_chat_config():
+    return PipelineConfig(
+        input=InputConfig(
+            sources={
+                "chosen": {
+                    "sections": [
+                        {"field": "chosen", "action": "$role", "template": True}
+                    ]
+                },
+                "rejected": {
+                    "sections": [
+                        {"field": "rejected", "action": "$role", "template": True}
+                    ]
+                },
+            }
+        ),
+        mask={"user": "mask", "assistant": "train"},
+        mask_default="mask",
+        preprocessing=ProcessingConfig(max_seq_len=2048),
+    )
+
+
+def make_grpo_config():
+    return PipelineConfig(
+        input=InputConfig(
+            sources={
+                "prompts": {
+                    "sections": [
+                        {"field": "prompt", "action": "mask", "template": True}
+                    ]
+                },
+                "responses": {
+                    "sections": _GRPO_RESPONSE_SECTIONS,
+                    "list_field": True,
+                    "mask_key": "masks",
+                },
+                "rewards": {
+                    "sections": [{"field": "rewards", "action": "value"}],
+                },
+            }
+        ),
+        mask={"user": "mask", "assistant": "train"},
+        mask_default="mask",
+        preprocessing=ProcessingConfig(max_seq_len=2048),
+    )
+
+
+def make_grpo_no_template_config():
+    return PipelineConfig(
+        input=InputConfig(
+            sources={
+                "prompts": {
+                    "sections": [
+                        {
+                            "field": "prompt",
+                            "action": "mask",
+                            "add_special_tokens": True,
+                        }
+                    ]
+                },
+                "responses": {
+                    "sections": _GRPO_RESPONSE_SECTIONS,
+                    "list_field": True,
+                    "mask_key": "masks",
+                },
+                "rewards": {
+                    "sections": [{"field": "rewards", "action": "value"}],
+                },
+            }
+        ),
+        mask={"user": "mask", "assistant": "train"},
+        mask_default="mask",
+        preprocessing=ProcessingConfig(max_seq_len=2048),
+    )
+
+
+@pytest.fixture
+def builder():
+    return SectionedMaskBuilder()
+
+
+@pytest.fixture
+def tokenizer_dir(temp_dir, test_tokenizer):
+    d = os.path.join(temp_dir, "tok")
+    os.makedirs(d, exist_ok=True)
+    test_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
+    with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
+        json.dump(
+            {"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}}, f
+        )
+    return d
+
+
+@pytest.fixture
+def chat_tokenizer_dir(temp_dir, chat_tokenizer):
+    d = os.path.join(temp_dir, "tok")
+    os.makedirs(d, exist_ok=True)
+    chat_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
+    with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
+        json.dump(
+            {"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE},
+            f,
+        )
+    return d

tests/data/test_checkpoint.py

@@ -1,3 +1,4 @@
+import os
 import tempfile
 
 import torch
@@ -24,7 +25,9 @@ def test_single_process():
 
         scheduler.step()
 
-    checkpoint = Checkpoint(state_dict=model.state_dict(), epoch=3, iteration=30)
+    checkpoint = Checkpoint(
+        state_dict=model.state_dict(), epoch=3, consumed_samples=120
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         checkpoint.save(tmpdir)
@@ -32,7 +35,34 @@ def test_single_process():
         loaded_checkpoint = Checkpoint.load(tmpdir)
 
         assert loaded_checkpoint.epoch == 3
-        assert loaded_checkpoint.iteration == 30
+        assert loaded_checkpoint.consumed_samples == 120
+
+
+def test_checkpoint_with_extra():
+    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,
+        consumed_samples=40,
+        extra=extra,
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        checkpoint.save(tmpdir)
+
+        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():
@@ -52,7 +82,7 @@ def simple_training():
     checkpoint = Checkpoint(
         state_dict=model.state_dict(),
         epoch=2,
-        iteration=10,
+        consumed_samples=40,
     )
 
     rank = get_rank()