feat: 服务化基础设施 - 有界队列/超时/优雅关闭/metrics

- astrai/inference/scheduler.py: 有界队列 (max_queue_size) 拒绝满时入队抛 RuntimeError -> 请求超时检测 (deadline + _abort_expired_tasks)，超时任务 abort 释放页并通知回调 -> stop() 改为 drain 模式：等待活跃任务自然结束再强制清理 -> get_stats() 扩展 latency P50/P95/P99 + cache hit rate - astrai/inference/engine.py: generate/generate_async 新增 timeout 参数 -> _generate_streaming/_generate_non_streaming 捕获 add_task 异常并清理 - astrai/inference/server.py: 新增 /metrics 端点 (Prometheus 格式) -> chat completions 端点捕获 RuntimeError 返回 503 -> configure_server 传递 max_queue_size/request_timeout - astrai/inference/cache.py: 新增 lookup_hits/lookup_misses 计数器 - tests/: fix stats key total_tasks -> total_requests
2026-05-10 18:16:51 +08:00
152 changed files with 5449 additions and 22633 deletions
@@ -2,7 +2,7 @@
 name: Bug report
 about: Create a report to help us improve
 title: "[BUG]"
-labels: bug
+labels: enhancement
 assignees: ''

 ---
@@ -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 . --select I`)
+- [ ] My code follows the style guidelines of this project (run `ruff format .` and `ruff check --fix .`)
 - [ ] I have performed a self-review of my own code
- [ ] Code is self-documenting (no unnecessary comments)
+- [ ] I have commented my code, particularly in hard-to-understand areas
 - [ ] 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
@@ -1,71 +0,0 @@
-name: Release
-
-on:
-  push:
-    tags:
-      - "v*"
-
-jobs:
-  build-pure:
-    name: Build pure-Python wheel
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
-
-      - name: Build wheel (no CUDA)
-        run: |
-          pip wheel . --no-deps -w dist/
-
-      - uses: actions/upload-artifact@v4
-        with:
-          name: pure-wheel
-          path: dist/*.whl
-
-  build-cuda-linux:
-    name: Build CUDA wheel (Linux)
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
-
-      - name: Install torch (CUDA 12.8)
-        run: |
-          pip install torch --index-url https://download.pytorch.org/whl/cu128
-
-      - name: Setup CUDA
-        uses: Jimver/cuda-toolkit@v0.2.35
-        with:
-          cuda: "12.8.0"
-
-      - name: Build wheel (with CUDA kernels)
-        run: |
-          CSRC_KERNELS=true pip wheel . --no-deps --no-build-isolation -w dist/
-
-      - uses: actions/upload-artifact@v4
-        with:
-          name: cuda-wheel-linux
-          path: dist/*.whl
-
-  release:
-    name: Attach wheels to release
-    needs: [build-pure, build-cuda-linux]
-    runs-on: ubuntu-latest
-    permissions:
-      contents: write
-    steps:
-      - uses: actions/download-artifact@v4
-        with:
-          pattern: "*-wheel"
-          merge-multiple: true
-
-      - name: Create release & upload assets
-        uses: softprops/action-gh-release@v2
-        with:
-          files: ./*.whl
-          tag_name: ${{ github.ref_name }}
-          generate_release_notes: true
@@ -5,16 +5,8 @@
 !*/

 # Allow specific file types and root files
-!astrai/**/*.py
-!scripts/**/*.py
-!tests/**/*.py
-!csrc/**/*.py
-
-!csrc/**/*.cu
-!csrc/**/*.h
-!csrc/**/*.cuh
-
-!scripts/**/*.sh
+!*.py
+!*.sh

 # Allow GitHub files
 !/.github/**
@@ -29,8 +21,3 @@
 !/LICENSE
 !/pyproject.toml
 !/README.md
-# Allow extension modules (only source .py)
-!/astrai/extension/**/*.py
-
-# Allow build files
-!/setup.py
@@ -1,100 +1,68 @@
 # Contributing to AstrAI

-Thank you for your interest in contributing! This document provides step-by-step guidelines.
+Thank you for your interest in contributing to AstrAI! This document provides guidelines and steps for contributing.

-## Quick Start
+## How to Contribute

-```bash
-git clone https://github.com/ViperEkura/AstrAI.git
-cd AstrAI
-pip install -e ".[dev]"     # install with dev dependencies (pytest, ruff)
-```
+### 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.).

-## Before You Commit
+### 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.

-Run the following checks **in order** — CI will reject if any fail.
+## Code Style

-### 1. Format
+AstrAI uses [Ruff](https://docs.astral.sh/ruff/) for code formatting and linting. Please ensure your code is formatted before submitting.

-```bash
-ruff 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`).

-> **Note**: `ruff format` may rename parameters (e.g. `mask` → `attn_mask`).  
-> Always review the diff after formatting.
+## Testing

-### 2. Import sorting
+If you add or modify functionality, please include appropriate tests.

-```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`.
+- Run the test suite with:
+  ```bash
+  pytest
+  ```
+- Ensure all tests pass before submitting your PR.

 ## Code Review

- 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.
+All submissions will be reviewed. We may request changes or discuss alternatives. Please be responsive to feedback.

 ## License

-By contributing, you agree that your contributions will be licensed under the [GPL-3.0 License](LICENSE).
+By contributing, you agree that your contributions will be licensed under the same [GPL-3.0 License](LICENSE) that covers the project.

 ---

-Questions? Ask in [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions) or open an issue.
+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!
@@ -1,7 +1,7 @@
 # AstrAI Dockerfile - Multi-stage Build (Optimized)

 # Build stage - use base image with minimal build tools
-FROM ubuntu:24.04 AS builder
+FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS builder

 WORKDIR /app

@@ -18,22 +18,21 @@ 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 (deps read from pyproject.toml)
+# Copy source code and install dependencies
 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/cu128
+    --extra-index-url https://download.pytorch.org/whl/cu126

 # Production stage
-FROM ubuntu:24.04 AS production
+FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production

 WORKDIR /app

-# Install Python 3.12 runtime and healthcheck dependency
+# Install Python 3.12 runtime
 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
@@ -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/tag/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">
+  <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">
 </div>
 <br>

@@ -20,7 +20,7 @@
  <a href="assets/docs/README-zh-CN.md">中文</a> •
  <a href="https://github.com/ViperEkura/AstrAI/issues">Issue Tracker</a> •
  <a href="https://github.com/ViperEkura/AstrAI/discussions">Discussions</a> •
-  <a href="https://huggingface.co/ViperEkura">HuggingFace</a>
+  <a href="https://huggingface.co/ViperEk/">HuggingFace</a>
 </div>

 <br>
@@ -28,8 +28,7 @@
 ## 📖 Table of Contents

 - [Features](#features)
- [Getting Started](#getting-started)
- [Demo](#demo)
+- [Quick Start](#quick-start)
 - [Documentation](#documentation)
 - [Contributing](#contributing)
 - [Community](#community)
@@ -50,116 +49,45 @@
 - 🤗 **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.

-### Getting Started
+### Quick Start

-End-to-end walkthrough in 5 steps:
-
-**1. Install**
+#### Installation

 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .                                          # pure PyTorch (no CUDA kernels)
-# CSRC_KERNELS=true pip install -e . --no-build-isolation  # optional: fused CUDA kernels
-# pip install -e ".[dev]"                                  # dev dependencies (pytest, ruff)
+pip install -e .
 ```

-**2. Download model**
+For development dependencies:

 ```bash
-python scripts/demo/download.py    # downloads 1B checkpoint to params/
+pip install -e ".[dev]"
 ```

-**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"}
-}
-```
+#### Train a Model

 ```bash
-python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
+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
 ```

-**4. Train**
+Full reference at [Parameter Guide](assets/docs/params.md).

-```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 \
-    --weight_decay=0.1 \
-    --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:
+#### Generate Text

 ```bash
 python scripts/tools/generate.py \
-    --param_path ./params \
-    --input_json_file input.jsonl \
-    --output_json_file output.jsonl
+    --param_path /path/to/model \
+    --input_json_file /path/to/input.json \
+    --output_json_file /path/to/output.json
 ```

 #### Docker
@@ -173,6 +101,9 @@ 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
@@ -189,42 +120,87 @@ docker compose --profile cpu up -d

 > **Note**: `--gpus all` is required for CUDA support. Without it, `torch.cuda.is_available()` will return `False`.

-#### HTTP API Examples
+#### Start HTTP Server

-Additional request examples beyond the [Getting Started](#getting-started) flow:
+Start the inference server with OpenAI and Anthropic-compatible HTTP API:

 ```bash
+python -m scripts.tools.server --port 8000 --device cuda
+```
+
+Make requests:
+
+```bash
+# OpenAI-compatible
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "Hello"}],
+    "max_tokens": 512
+  }'
+
 # OpenAI-compatible streaming
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
-  -d '{"messages":[{"role":"user","content":"Tell a story"}],"stream":true,"max_tokens":500}'
+  -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
 ```

-See [Inference Guide](assets/docs/inference.md) for SSE streaming format, error codes, and stats endpoint.
+#### 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).

 ### Documentation

 | Document | Description |
 |----------|-------------|
-| [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 |
+| [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 |

 ### Contributing

@@ -241,7 +217,7 @@ For major changes, please open an issue first to discuss what you would like to

 - **GitHub Issues**: [Issue Tracker](https://github.com/ViperEkura/AstrAI/issues)
 - **Discussions**: [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions)
- **HuggingFace**: [Model Hub](https://huggingface.co/ViperEkura)
+- **HuggingFace**: [Model Hub](https://huggingface.co/ViperEk)

 ### License

@@ -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/tag/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">
+  <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">
 </div>

 <br>
@@ -27,15 +27,14 @@
  <a href="#chinese">中文</a> •
  <a href="https://github.com/ViperEkura/AstrAI/issues">问题追踪</a> •
  <a href="https://github.com/ViperEkura/AstrAI/discussions">讨论区</a> •
-  <a href="https://huggingface.co/ViperEkura">HuggingFace</a>
+  <a href="https://huggingface.co/ViperEk">HuggingFace</a>
 </div>
 <br>

 ## 📖 目录

 - [特性](#特性)
- [快速上手](#快速上手)
- [演示](#演示)
+- [快速开始](#快速开始)
 - [文档](#文档)
 - [贡献](#贡献)
 - [社区](#社区)
@@ -56,116 +55,45 @@
 - 🤗 **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 .                                          # 纯 PyTorch（不含 CUDA 内核）
-# CSRC_KERNELS=true pip install -e . --no-build-isolation  # 可选：融合 CUDA 内核加速
-# pip install -e ".[dev]"                                  # 可选：开发依赖（pytest, ruff）
+pip install -e .
 ```

-**2. 下载模型**
+安装开发依赖：

 ```bash
-python scripts/demo/download.py    # 下载 1B 检查点到 params/
+pip install -e ".[dev]"
 ```

-**3. 预处理数据**
-
-创建 `pretrain.json`（`seq` 策略的预处理配置）：
-
-```json
-{
-    "version": 1,
-    "input": {"sections": [{"field": "text", "action": "train"}]},
-    "preprocessing": {"max_seq_len": 2048},
-    "output": {"storage_format": "bin"}
-}
-```
+#### 训练模型

 ```bash
-python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
+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
 ```

-**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 \
-    --weight_decay=0.1 \
-    --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) 上的视频演示。
-
---
-
-更多选项请参考[文档](#文档)。
+完整参数列表见[参数说明](./params.md)。

 #### 文本生成

-从 JSONL 文件批量生成：
-
 ```bash
 python scripts/tools/generate.py \
-    --param_path ./params \
-    --input_json_file input.jsonl \
-    --output_json_file output.jsonl
+    --param_path /path/to/model \
+    --input_json_file /path/to/input.json \
+    --output_json_file /path/to/output.json
 ```

 #### Docker
@@ -179,6 +107,9 @@ 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
@@ -195,42 +126,87 @@ docker compose --profile cpu up -d

 > **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。

-#### HTTP API 示例
+#### 启动 HTTP 服务

-除[快速上手](#快速上手)流程外，更多请求示例：
+启动推理服务器，支持 OpenAI 和 Anthropic 兼容的 HTTP API：

 ```bash
+python -m scripts.tools.server --port 8000 --device cuda
+```
+
+发起请求：
+
+```bash
+# OpenAI 兼容
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "你好"}],
+    "max_tokens": 512
+  }'
+
 # OpenAI 兼容流式
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
-  -d '{"messages":[{"role":"user","content":"讲个故事"}],"stream":true,"max_tokens":500}'
+  -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
 ```

-SSE 流式格式、错误码和统计端点详见[推理文档](./inference.md)。
+#### 演示
+
+查看 `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) 上的视频演示。

 ### 文档

 | 文档 | 说明 |
 |------|------|
-| [CLI 参考](./params.md) | 所有 CLI 工具参数（训练、服务、生成、预处理） |
-| [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
-| [训练文档](./training.md) | 训练循环、策略与公式 |
-| [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
-| [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 |
-| [数据预处理](./preprocessing.md) | 声明式 JSON 驱动数据预处理 |
+| [参数说明](./params.md) | 训练与推理参数配置 |
+| [设计文档](./design.md) | 系统架构与模块设计 |
+| [数据流程](./dataflow.md) | 数据处理管道详解 |
+| [模型介绍](./introduction.md) | 模型架构与技术细节 |

 ### 贡献

@@ -247,7 +223,7 @@ SSE 流式格式、错误码和统计端点详见[推理文档](./inference.md)

 - **GitHub Issues**: [问题追踪](https://github.com/ViperEkura/AstrAI/issues)
 - **Discussions**: [GitHub 讨论区](https://github.com/ViperEkura/AstrAI/discussions)
- **HuggingFace**: [模型中心](https://huggingface.co/ViperEkura)
+- **HuggingFace**: [模型中心](https://huggingface.co/ViperEk)

 ### 许可证

@@ -1,130 +1,237 @@
-# Data Flow
+# AstrAI Data Flow Documentation

-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)
+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.

 ## Overview

-```
-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
+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
 ```

-## Data Preparation
+## Detailed Module Descriptions

-Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
+### 1. Serialization (`astrai/serialization.py`)

-### Tokenization
+- **`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

-The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](preprocessing.md)), and produces flat token sequences:
+### 2. Dataset Module

-```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
-```
+#### 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)

-The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
+#### 2.2 Sampler (`sampler.py`)
+- **`ResumableDistributedSampler`**: Tracks `epoch` and `iter` for breakpoint resume; supports shuffle and drop_last

-### Format Detection
+### 3. Model Module

-`detect_format(load_path)` inspects the path:
+#### 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

- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, `.jsonl` → `"jsonl"`, unknown suffix raises `ValueError`
- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, `*.bin` + `**/meta.json` → `"bin"`, or `*.jsonl` + `dataset_config.json` → `"jsonl"`
+#### 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

-### Store Backends
+### 4. Training Module

-Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
+#### 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):

 ```
-StoreFactory.create("h5")    → H5Store
-StoreFactory.create("bin")   → MmapStore
-StoreFactory.create("jsonl") → JsonlStore
+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
 ```

-All three inherit `Store` (base, owns `_data`/`_cum`/`_offsets`/`_normalize`) plus the `Streamable` and `Recordable` mixins, so every backend supports both `fetch(begin, end, keys)` (stream) and `fetch_record(index, keys)` (record) APIs.
+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`

-**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage. `segments_are_records=True` — each `data_i` dataset is one record.
+#### 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

-**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. `segments_are_records=False` — bin segments are contiguous streams; record access is driven by `_offsets` (written when `save_bin(..., record_keys=...)` was used at preprocessing time).
+#### 4.4 Scheduler (`schedule.py`)
+- **`CosineScheduler`**: Cosine decay + linear warmup
+- **`SGDRScheduler`**: Cosine annealing with warm restarts
+- Created by `SchedulerFactory` and bound to optimizer

-**JsonlStore**: On-the-fly tokenization of raw JSONL files at load time. Requires a `dataset_config.json` alongside the `.jsonl` files following the same `PipelineConfig` schema with an additional `tokenizer_path` field. Two modes: eager (default, applies `TokenizeTransform` to all records at load) and lazy (`processor=fn` given, defers tokenisation to `fetch_record` — used by DPO/GRPO).
+#### 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`

-All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record-mode indexing). Nested keys (GRPO `responses`/`masks` as `List[List[Tensor]]`) are stored as-is and excluded from both bookkeepings — they are only accessed record-by-record.
+### 5. Inference Module

-## Data Keys by Training Type
+#### 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)

-| Type | Storage Keys | Access Mode |
-|------|-------------|-------------|
-| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | stream (`fetch`) |
-| `sft` | `sequence`, `loss_mask`, `position_ids` | stream (`fetch`) |
-| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | record (`fetch_record`) |
-| `grpo` | `prompts`, `responses`, `masks`, `rewards` | record (`fetch_record`) |
+#### 5.2 Scheduler 4-Phase Loop (`scheduler.py`)

-## Dataset Architecture
+Background thread runs continuously:

 ```
-DatasetFactory.load(train_type, load_path, window_size, stride=None,
-                    storage_type=None, tokenizer_path=None,
-                    max_len=2048, store=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]]]   (stream mode)
-          + _offsets[Dict[str, List[int]]]  (record mode)
-
-Stream datasets (SEQ/SFT):
-  BaseDataset.__getitem__(idx)
-    → get_index(idx) → [begin, end)
-    → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
-
-Record datasets (DPO/GRPO via RecordDataset):
-  RecordDataset.__getitem__(idx)
-    → Store.fetch_record(idx, keys) → Tensor / Dict[str, Tensor]
+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
 ```

-Class hierarchy: `BaseDataset` ← `SEQDataset` / `SFTDataset` (stream); `BaseDataset` ← `RecordDataset` ← `DPODataset` / `GRPODataset` (record).
+- **`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

-`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). Only meaningful for stream datasets — record datasets ignore both. `storage_type` defaults to `None` (auto-detect via `detect_format`).
+#### 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`

-`tokenizer_path` triggers lazy on-the-fly tokenisation for record datasets on raw JSONL (DPO builds a `dpo_processor`; SEQ/SFT/pre-tokenised backends ignore it). `store` (pre-built `Store`) bypasses `load_path`/`storage_type`/`tokenizer_path` entirely — the caller controls Store construction.
+### 6. Tokenizer Module

-`Store.fetch(begin, end, keys)` (stream mode, on `Streamable`): 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()`.
+- **`AutoTokenizer`**: Wraps HuggingFace tokenizers (BBPE); `encode`/`decode`/`apply_chat_template`
+- **`ChatTemplate`**: Jinja2-based template rendering for multi-turn chat

-`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present (bin layout with per-record offsets), otherwise indexes `_data[key]` directly (H5/JSONL where each segment is one record).
+### 7. Factory & Parallel

-## Sampler
+- **`Registry` / `BaseFactory`**: Decorator-based component registration
+- **`spawn_parallel_fn`**: Multi-process DDP launcher with NCCL backend
+- **`ParallelModel` / `ColumnParallelLinear` / `RowParallelLinear`**: Tensor model parallelism

-`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
+## Training Data Flow — Detailed Steps

- Tracks `start_epoch` / `start_iter` for resume
- Shuffle via `torch.Generator(seed + epoch)`
- Per-replica index slicing for DDP
+1. **Data Preparation**
+   - Raw text → token IDs via `AutoTokenizer.encode()`
+   - Save as `.h5` files (groups of tensor lists per data key)

-## DataLoader
+2. **Dataset Loading**
+   - `BaseDataset.load()` calls `load_h5()`, builds `MultiSegmentFetcher`
+   - Sliding window of `window_size` with `stride` determines sample boundaries

-Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
+3. **Sampling & Batching**
+   - `ResumableDistributedSampler` produces shuffled index sequences
+   - `DataLoader` fetches `[batch_size, window_size]` tensors via `__getitem__`

-> Document Update Time: 2026-07-19
+4. **Strategy Forward**
+   - Strategy receives batch, calls `Transformer.forward()` for logits
+   - Computes task-specific loss (cross-entropy, DPO, GRPO)
+
+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
+
+6. **Checkpoint**
+   - `CheckpointCallback` saves `model.state_dict()` + metadata to safetensors at `ckpt_interval` iterations
+   - Does NOT save optimizer/scheduler state (resume resets those)
+
+## Inference Data Flow — Detailed Steps
+
+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
@@ -0,0 +1,719 @@
+## 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
@@ -1,252 +0,0 @@
-# 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, with two concrete cache implementations:
-
-### ContiguousCache (default)
-
-```
-ContiguousCache (simple contiguous per-slot cache)
-  ├── ContiguousCacheView  bundles k/v tensors + slot indices for attention layers
-```
-
-Created by default when no cache is passed to `InferenceScheduler`. Each task occupies a fixed slot of `[max_seq_len, n_kv_heads, head_dim]`. Simple and efficient for small-to-medium batch sizes.
-
-### PageCache (paged with prefix sharing)
-
-```
-PageCache (paged KV cache with prefix sharing, alternative)
-  ├── PagePool               orchestrates page allocation + prefix matching
-  │     ├── Allocator         bitmask-based page allocator + ref-count + LRU
-  │     └── PrefixCache       hash-based prefix matching (page_hash via polynomial hash)
-  ├── 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)
-  └── PageCacheView          bundles Storage + page_table + total_len for attention layers
-```
-
-`isinstance(cache, KVCache)` checks dispatch to the correct view. Both implement the abstract `KVCache` interface used by `Executor` and `InferenceScheduler`.
-
-## Continuous Batching
-
-`InferenceScheduler` runs a daemon thread with a 4-phase loop:
-
-```
-1. Cleanup → Remove finished tasks, free KV cache slots/pages
-2. Refill  → Pop from waiting_queue, task_alloc resources, activate
-3. Prefill → Group by (prompt_len, start_pos), run full forward
-4. Decode  → Run single-token forward for each same-position group
-```
-
-## 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-07-09
@@ -0,0 +1,334 @@
+## 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
@@ -1,11 +1,4 @@
-# CLI Parameter Reference
-
-## Contents
-
- [Training Parameters](#training-parameters)
- [Inference Server](#inference-server-serverpy)
- [Generate](#generate-generatepy)
- [Preprocess](#preprocess-preprocesspy)
+# Parameter Documentation

 ## Training Parameters

@@ -17,28 +10,24 @@
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
 | `--n_epoch` | Total training epochs | 1 |
-| `--batch_per_device` | Batch size per device | 1 |
-| `--grad_accum_steps` | Gradient accumulation steps between optimizer steps | 1 |
+| `--batch_size` | Batch size | 1 |
+| `--accumulation_steps` | Gradient accumulation steps between optimizer steps | 1 |

 ### Learning Rate Scheduling

 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
+| `--warmup_steps` | Warmup steps | 1000 |
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
-| `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | None |
+| `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |

-### Optimizer (MuonMix)
-
-Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
+### Optimizer (AdamW)

 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
-| `--muon_momentum` | Muon momentum factor | 0.95 |
-| `--muon_nesterov` | Enable Nesterov momentum for Muon | True |
-| `--muon_ns_steps` | Newton-Schulz iteration steps for Muon | 5 |
-| `--muon_adjust_lr` | Muon LR adjustment strategy (`original`, `match_rms_adamw`) | `match_rms_adamw` |
+| `--adamw_beta1` | AdamW beta1 | 0.9 |
+| `--adamw_beta2` | AdamW beta2 | 0.95 |
+| `--adamw_weight_decay` | AdamW weight decay | 0.01 |

 ### Data Loading

@@ -57,148 +46,113 @@ Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`f
 | `--ckpt_interval` | Iterations between checkpoints | 5000 |
 | `--ckpt_dir` | Checkpoint save directory | checkpoint |
 | `--start_epoch` | Resume from epoch (0 = from scratch) | 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 |
-| `--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 |
+| `--start_batch` | Resume from batch iteration | 0 |

 ### 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.0 | `seq`, `sft` |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.1 | `seq`, `sft` |
 | `--group_size` | GRPO group size | 4 | `grpo` |
 | `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
 | `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
 | `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` |
-| `--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: 0.01) |
-| `--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
-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 \
-    --weight_decay=0.1 \
-    --window_size=2048 \
-    --ckpt_interval=10000 \
-    --ckpt_dir=./checkpoint \
-    --random_seed=3407 \
-    --label_smoothing=0.05 \
-    > out.log 2> err.log &
+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
 ```

 ---

-## Inference Server (`server.py`)
+## Generation 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 |
+### GenerationRequest Parameters

-Usage:
-```bash
-python scripts/tools/server.py --param_path ./params --device cuda --dtype bfloat16
+| 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,
+)
 ```

-See [Inference Guide](inference.md) for HTTP API documentation.
+### Generation Modes

-## Generate (`generate.py`)
+| Mode | Description |
+|------|-------------|
+| `stream=True` | Streaming output, yields token by token |
+| `stream=False` | Non-streaming output, returns complete result |

-| 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-07-19
+> Document Update Time: 2026-04-09
@@ -1,364 +0,0 @@
-# Preprocessing Pipeline
-
-Declarative JSON-driven data preprocessing. `MaskBuilderFactory` supports three registered builders: `"single"` (single-output via `input.sections`), `"multi"` (multi-output via `input.sources`), and `"sectioned"` (façade dispatching to `single` or `multi` based on config).
-
-## 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-07-09
@@ -1,225 +0,0 @@
-# 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`, `MetricCallback`, `ProgressBarCallback` |
-| `on_batch_end` | Every batch | `CheckpointCallback` |
-| `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
-| `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
-| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricCallback`, `GradientCheckpointingCallback` |
-
-Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
-
-## 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`, `position_ids`. 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="sum"`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
-
-### GRPO (Group Relative Policy Optimization)
-
-Token-level PPO with group-normalized advantages. Advantages are derived from
-scalar per-response rewards, group-normalized, and broadcast across all response
-tokens. Only response tokens contribute to the loss (prompt tokens are masked
-out):
-
-$$
-\text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
-$$
-
-$$
-L_{\text{GRPO}} = -\mathbb{E}_t\left[\min\left(\rho_t A,\; \text{clip}\left(\rho_t, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}_t\left[\frac{\pi_{\text{ref}}}{\pi_\theta} - \log\frac{\pi_{\text{ref}}}{\pi_\theta} - 1\right]
-$$
-
-where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the
-per-token importance sampling ratio against the behaviour policy
-(`old_model`, synced externally between data-generation rounds) and the
-expectations are over valid response tokens. The KL term regularises
-$\pi_\theta$ towards a frozen reference model (`ref_model`, typically
-the SFT checkpoint).
-
-Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. External sync of `old_model` weights via `sync_old_model()` between data-generation rounds.
-
-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 \
-    --weight_decay=0.1 \
-    --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-07-19
@@ -1,98 +1,32 @@
-__version__ = "1.3.10"
+__version__ = "1.3.4"
 __author__ = "ViperEkura"

 from astrai.config import (
-    AutoRegressiveLMConfig,
-    BaseModelConfig,
-    ConfigFactory,
-    EncoderConfig,
-    PipelineConfig,
+    ModelConfig,
    TrainConfig,
 )
-from astrai.dataset import (
-    BaseDataset,
-    DatasetFactory,
-    RDSampler,
-    Store,
-    StoreFactory,
-)
+from astrai.dataset import DatasetFactory
 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__ = [
-    "AutoRegressiveLM",
-    "AutoRegressiveLMConfig",
-    "AutoModel",
-    "AutoTokenizer",
-    "BaseDataset",
-    "BaseFactory",
-    "BaseModelConfig",
-    "BaseScheduler",
-    "BaseStrategy",
-    "CallbackFactory",
-    "ChatTemplate",
-    "Checkpoint",
-    "ConfigFactory",
+    "Transformer",
+    "ModelConfig",
+    "TrainConfig",
    "DatasetFactory",
-    "EmbeddingEncoder",
-    "EncoderConfig",
-    "ExecutorFactory",
+    "AutoTokenizer",
    "GenerationRequest",
    "InferenceEngine",
-    "LoRAConfig",
-    "Pipeline",
-    "PipelineConfig",
-    "ProtocolHandler",
-    "RDSampler",
-    "SamplingPipeline",
-    "SchedulerFactory",
-    "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",
+    "CallbackFactory",
+    "StrategyFactory",
+    "SchedulerFactory",
+    "BaseFactory",
+    "AutoModel",
 ]
@@ -1,25 +1,8 @@
-from astrai.config.model_config import (
-    AutoRegressiveLMConfig,
-    BaseModelConfig,
-    ConfigFactory,
-    EncoderConfig,
-)
-from astrai.config.preprocess_config import (
-    InputConfig,
-    OutputConfig,
-    PipelineConfig,
-    ProcessingConfig,
-)
+from astrai.config.model_config import ModelConfig
 from astrai.config.train_config import TrainConfig

 __all__ = [
-    "BaseModelConfig",
-    "AutoRegressiveLMConfig",
-    "EncoderConfig",
-    "ConfigFactory",
+    # Model configuration
+    "ModelConfig",
    "TrainConfig",
-    "InputConfig",
-    "OutputConfig",
-    "PipelineConfig",
-    "ProcessingConfig",
 ]
@@ -1,98 +0,0 @@
-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)
@@ -1,82 +1,42 @@
-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)
+import json
+from dataclasses import asdict, dataclass
+from typing import Optional, Self


@dataclass
-class BaseModelConfig(BaseConfig):
-    """Base config with ``model_type`` dispatch and file I/O."""
-
+class ModelConfig:
+    # basic config
    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

-    attn_type: str = "gqa"
+    # GQA
    n_heads: Optional[int] = None
    n_kv_heads: Optional[int] = None
    use_qk_norm: Optional[bool] = None
    use_gated_attention: Optional[bool] = None

-    kv_lora_rank: Optional[int] = None
-    qk_nope_head_dim: Optional[int] = None
-    qk_rope_head_dim: Optional[int] = None
+    def load(self, config_path: str) -> Self:
+        config = {}
+        with open(config_path, "r") as f:
+            config.update(json.load(f))

-    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
+        for key, value in config.items():
+            if hasattr(self, key):
+                setattr(self, key, value)

+        return self

-@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
+    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)
@@ -1,109 +0,0 @@
-"""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)
@@ -1,80 +1,43 @@
-from dataclasses import dataclass, field, fields
-from typing import Any, Callable, Dict, List, Optional
+from dataclasses import dataclass, field
+from typing import Callable, 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(BaseConfig):
+class TrainConfig:
    # basic setting
-    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.")
-    )
+    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."})
    optimizer_fn: Callable[[nn.Module], Optimizer] = field(
-        default=None, metadata=required(help="Optimizer factory for training.")
+        default=None, metadata={"help": "Optimizer factory for training."}
    )
    scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
-        default=None, metadata=required(help="Scheduler factory for training.")
+        default=None, metadata={"help": "Scheduler factory for training."}
    )
    n_epoch: int = field(default=1, metadata={"help": "Number of epochs for training."})
-    batch_per_device: int = field(
-        default=4, metadata={"help": "Batch size per device."}
-    )
-    grad_accum_steps: int = field(
+    batch_size: int = field(default=4, metadata={"help": "Batch size for training."})
+    accumulation_steps: int = field(
        default=1, metadata={"help": "Number of iterations between steps."}
    )
-    max_grad_norm: Optional[float] = field(
-        default=None,
-        metadata={"help": "Maximum gradient norm. None disables clipping."},
-    )
-    gradient_checkpointing_modules: List[str] = field(
-        default_factory=list,
-        metadata={"help": "Module types to enable activation checkpointing for."},
+    max_grad_norm: float = field(
+        default=1.0, metadata={"help": "Maximum gradient norm."}
    )

    # checkpoint setting
    start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
-    start_samples: int = field(
-        default=0,
-        metadata={
-            "help": "Start samples count (per rank). Superseded by checkpoint consumed_samples."
-        },
+    start_batch: int = field(
+        default=0, metadata={"help": "Start batch iteration for training."}
    )
    ckpt_dir: str = field(
        default="./checkpoint", metadata={"help": "Checkpoint directory."}
    )
    ckpt_interval: int = field(
-        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."},
+        default=5000, metadata={"help": "Number of iterations between checkpoints."}
    )

    # dataloader setting
@@ -88,10 +51,6 @@ class TrainConfig(BaseConfig):
    pin_memory: bool = field(
        default=False, metadata={"help": "Pin memory for dataloader."}
    )
-    collate_fn: Optional[Callable[[List[Any]], Any]] = field(
-        default=None,
-        metadata={"help": "Collate function for dataloader (e.g. dpo_collate_fn)."},
-    )

    # distributed training
    nprocs: int = field(
@@ -107,42 +66,18 @@ class TrainConfig(BaseConfig):
    master_port: str = field(
        default="29500", metadata={"help": "Master port for distributed training."}
    )
-    parallel_mode: str = field(
-        default="none",
-        metadata={"help": "Parallel strategy: none, ddp, fsdp."},
+    parallel_wrapper: Optional[Callable] = field(
+        default=None, metadata={"help": "Parallel function for training."}
    )
-    start_method: str = field(
-        default="spawn",
-        metadata={"help": "Multiprocessing start method (spawn/fork/forkserver)."},
+    state_dict_fn: Optional[Callable] = field(
+        default=None, metadata={"help": "Parallel function for state  dict saving."}
    )

    # others
    device_type: str = field(
        default="cuda", metadata={"help": "Device type for distributed training."}
    )
-    val_dataset: Optional[Dataset] = field(
-        default=None, metadata={"help": "Dataset for validation."}
-    )
-    val_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(
+    extra_kwargs: dict = field(
        default_factory=dict, metadata={"help": "Other arguments."}
    )

@@ -150,6 +85,14 @@ class TrainConfig(BaseConfig):
        self.validate()

    def validate(self):
-        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.")
+        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.")
@@ -1,43 +1,19 @@
 from astrai.dataset.dataset import (
    BaseDataset,
+    BaseSegmentFetcher,
    DatasetFactory,
-    dpo_collate_fn,
-    grpo_collate_fn,
-)
-from astrai.dataset.sampler import RDSampler
-from astrai.dataset.storage import (
-    H5Store,
-    JsonlStore,
-    MmapStore,
-    Recordable,
-    Store,
-    StoreFactory,
-    Streamable,
-    detect_format,
-)
-from astrai.serialization import (
-    load_bin,
-    load_h5,
-    save_bin,
-    save_h5,
+    MultiSegmentFetcher,
 )
+from astrai.dataset.sampler import ResumableDistributedSampler

 __all__ = [
+    # Base classes
    "BaseDataset",
+    # Factory
    "DatasetFactory",
-    "dpo_collate_fn",
-    "grpo_collate_fn",
-    "Store",
-    "Streamable",
-    "Recordable",
-    "StoreFactory",
-    "H5Store",
-    "MmapStore",
-    "JsonlStore",
-    "detect_format",
-    "save_h5",
-    "load_h5",
-    "save_bin",
-    "load_bin",
-    "RDSampler",
+    # Fetchers
+    "BaseSegmentFetcher",
+    "MultiSegmentFetcher",
+    # Sampler
+    "ResumableDistributedSampler",
 ]
@@ -1,176 +1,308 @@
-"""Dataset implementations for training.
-
-Composition over inheritance — every dataset is a thin wrapper that
-binds a :class:`Store` to a particular train-type's key mapping.  All
-sample-id → token/record indexing lives on the Store; datasets never
-know about window/stride math or segment layouts.
-
-Class hierarchy:
-
-    BaseDataset (ABC)         — holds a Store, exposes __len__/keys,
-                                 overrides __getitem__
-    ├── SEQDataset            — next-token prediction (stream)
-    ├── SFTDataset            — loss-mask + position_ids (stream)
-    ├── DPODataset            — chosen/rejected pairs (record)
-    └── GRPODataset           — prompt + response group (record)
-
-``DatasetFactory.load(train_type, load_path, window_size, stride, …)``
-builds the Store (auto-detecting format) before constructing the
-matching dataset.  Passing ``store=`` skips Store construction.
-
-When a record dataset (DPO) reads from raw JSONL, a *processor*
-function (pure ``record -> Dict[str, Tensor]``) is forwarded to
-:class:`JsonlStore` so tokenisation happens on the fly.
-"""
+"""Dataset implementations with factory pattern for training."""

+import bisect
 from abc import ABC, abstractmethod
-from functools import partial
-from typing import Callable, Dict, List, Optional
+from typing import Dict, List, Optional, Union

 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.tokenize import AutoTokenizer
+from astrai.serialization import load_h5


-def dpo_tokenize(
-    record: dict,
-    tokenizer,
-    max_len: int = 2048,
-) -> Optional[dict]:
-    """Tokenize one DPO record into chosen/rejected + masks.
+class BaseSegmentFetcher:
+    """Fetches data segments across multiple tensor segments.

-    Applies the tokenizer's chat template so token sequences match the
-    SFT checkpoint's format.  Prompt is rendered with
-    ``add_generation_prompt=True``; chosen/rejected are appended as a
-    single assistant turn.
-
-    Accepts:
-
-    - Flat:   ``{"prompt": str, "chosen": str, "rejected": str}``
-    - Conv:   ``{"prompt": [{role, content}, ...], "chosen": [...], ...}``
-    - Legacy: ``{"input": str, "chosen": str, "rejected": str}``
-
-    No packing, no ``position_ids`` — DPO sequences are independent.
+    Maintains cumulative lengths for efficient range queries across
+    multiple discontinuous segments.
    """
-    prompt = record.get("prompt") or record.get("input")
-    chosen = record.get("chosen")
-    rejected = record.get("rejected")
-    if prompt is None or chosen is None or rejected is None:
-        return None

-    prompt_messages = _to_messages(prompt)
-    chosen_text = _extract_text(chosen)
-    rejected_text = _extract_text(rejected)
-    if chosen_text is None or rejected_text is None:
-        return None
-    chosen_messages = prompt_messages + [{"role": "assistant", "content": chosen_text}]
-    rejected_messages = prompt_messages + [
-        {"role": "assistant", "content": rejected_text}
-    ]
+    def __init__(self, segments: List[Tensor]):
+        self.segments = segments
+        self.cum_lengths = []

-    prompt_ids = tokenizer.apply_chat_template(
-        prompt_messages, tokenize=True, add_generation_prompt=True
-    )
-    ch_ids = tokenizer.apply_chat_template(
-        chosen_messages, tokenize=True, add_generation_prompt=False
-    )
-    re_ids = tokenizer.apply_chat_template(
-        rejected_messages, tokenize=True, add_generation_prompt=False
-    )
+        total = 0
+        for seg in segments:
+            total += torch.numel(seg)
+            self.cum_lengths.append(total)

-    full_ch = ch_ids[:max_len]
-    full_re = re_ids[:max_len]
+        self.total_length = total

-    prompt_len = min(len(prompt_ids), max_len)
-    ch_mask = [0] * prompt_len + [1] * max(0, len(full_ch) - prompt_len)
-    ch_mask = ch_mask[:max_len]
-    re_mask = [0] * prompt_len + [1] * max(0, len(full_re) - prompt_len)
-    re_mask = re_mask[:max_len]
+    def __len__(self) -> int:
+        return self.total_length

-    return {
-        "chosen": full_ch,
-        "rejected": full_re,
-        "chosen_mask": ch_mask,
-        "rejected_mask": re_mask,
+    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 _to_messages(value) -> list:
-    """Accept str or conversation list; return message list."""
-    if isinstance(value, str):
-        return [{"role": "user", "content": value}]
-    if isinstance(value, list):
-        return value
-    return [{"role": "user", "content": str(value)}]
+    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

-def _extract_text(value) -> Optional[str]:
-    """Accept str or conversation list; return plain text."""
-    if value is None:
-        return None
-    if isinstance(value, str):
-        return value
-    if isinstance(value, list):
-        return "".join(m.get("content", "") for m in value if isinstance(m, dict))
-    return None
-
-
-def dpo_processor(
-    record: dict,
-    tokenizer,
-    max_len: int = 2048,
-) -> Dict[str, Tensor]:
-    """DPO processor: wraps :func:`dpo_tokenize` and returns tensors."""
-    result = dpo_tokenize(record, tokenizer, max_len=max_len)
-    if result is None:
-        raise ValueError(f"Malformed DPO record: {list(record.keys())}")
-    return {
-        "chosen": torch.tensor(result["chosen"], dtype=torch.int32),
-        "rejected": torch.tensor(result["rejected"], dtype=torch.int32),
-        "chosen_mask": torch.tensor(result["chosen_mask"], dtype=torch.bool),
-        "rejected_mask": torch.tensor(result["rejected_mask"], dtype=torch.bool),
-    }
-
-
-def dpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
-    """Collate variable-length DPO samples into padded 2-D tensors.
-
-    Input: list of dicts, each with:
-      - chosen:        [C_i]
-      - rejected:      [R_i]
-      - chosen_mask:   [C_i]
-      - rejected_mask: [R_i]
-
-    Output (padded to the max length across chosen/rejected within the batch):
-      - chosen:        [B, S_max]
-      - rejected:      [B, S_max]
-      - chosen_mask:   [B, S_max]
-      - rejected_mask: [B, S_max]
+        Returns:
+            Dictionary of tensors if multiple keys, single tensor if one key
        """
-    B = len(batch)
-    S_max = max(b["chosen"].size(0) for b in batch)
-    S_max = max(S_max, max(b["rejected"].size(0) for b in batch))
+        fetch_dict = {}
+        keys = [keys] if isinstance(keys, str) else keys

-    chosen = torch.zeros(B, S_max, dtype=torch.long)
-    rejected = torch.zeros(B, S_max, dtype=torch.long)
-    chosen_mask = torch.zeros(B, S_max, dtype=torch.bool)
-    rejected_mask = torch.zeros(B, S_max, dtype=torch.bool)
+        for key in keys:
+            fetcher = self.multi_fetchers[key]
+            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
+            fetch_dict[key] = fetch_tensor

-    for i, b in enumerate(batch):
-        c_len = b["chosen"].size(0)
-        r_len = b["rejected"].size(0)
-        chosen[i, :c_len] = b["chosen"]
-        rejected[i, :r_len] = b["rejected"]
-        chosen_mask[i, :c_len] = b["chosen_mask"]
-        rejected_mask[i, :r_len] = b["rejected_mask"]
+        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.
+    """
+
+    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
+
+    def load(self, load_path: str):
+        """Load dataset from HDF5 file.
+
+        Args:
+            load_path: Path to the HDF5 data file
+        """
+        self.segments = load_h5(load_path)
+        self.fetcher = MultiSegmentFetcher(self.segments)
+        self.total_samples = len(self.fetcher)
+
+    def get_index(self, index: int) -> tuple:
+        """Calculate begin and end indices for a sample.
+
+        Args:
+            index: Sample index
+
+        Returns:
+            Tuple of (begin_idx, end_idx)
+        """
+        assert self.total_samples > 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)
+
+        return begin_idx, end_idx
+
+    @abstractmethod
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        """Get a single sample by index.
+
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def __len__(self) -> int:
+        assert self.total_samples is not None
+        if self.total_samples <= self.window_size:
+            return 0
+        return (self.total_samples - 1 - self.window_size) // self.stride + 1
+
+
+class DatasetFactory(BaseFactory["BaseDataset"]):
+    """Factory class for creating dataset instances.
+
+    Supports decorator-based registration for extensible dataset types.
+    All default dataset types (seq, sft, dpo, grpo) are registered automatically
+    when their classes are defined with the decorator.
+
+    Example usage:
+        @DatasetFactory.register("custom")
+        class CustomDataset(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,
+        train_type: str,
+        load_path: str,
+        window_size: int,
+        stride: Optional[int] = None,
+    ) -> "BaseDataset":
+        """Create and load a dataset in one step.
+
+        Args:
+            train_type: Type of training dataset
+            load_path: Path to the data file
+            window_size: Window size for data sampling
+            stride: Stride between consecutive samples (default: same as window_size)
+
+        Returns:
+            Loaded dataset instance
+        """
+        if stride is None:
+            stride = window_size
+
+        dataset = cls.create(train_type, window_size, stride)
+        dataset.load(load_path)
+
+        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)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
+
+    def __getitem__(self, index):
+        begin_idx, end_idx = self.get_index(index)
+
+        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
+
+        return {"input_ids": x, "target_ids": y}
+
+
+@DatasetFactory.register("sft")
+class SFTDataset(BaseDataset):
+    """Dataset for supervised fine-tuning with loss masking."""
+
+    def __init__(self, window_size: int, stride: int):
+        super().__init__(window_size, stride)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_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
+        )
+
+        return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask}
+
+
+@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)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+
+    def __getitem__(self, index: int):
+        begin_idx, end_idx = self.get_index(index)
+
+        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
+        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
+        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(
+            dtype=torch.bool
+        )
+        rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(
+            dtype=torch.bool
+        )

        return {
            "chosen": chosen,
@@ -180,40 +312,23 @@ def dpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
        }


-def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
-    """Collate variable-length GRPO samples into padded 3-D tensors.
+@DatasetFactory.register("grpo")
+class GRPODataset(BaseDataset):
+    """Dataset for Group Relative Policy Optimization training."""

-    Input: list of dicts, each with:
-      - prompts:  [P_i]
-      - responses: list of G tensors, each [R_ij]
-      - masks:     list of G tensors, each [R_ij]
-      - rewards:  [G]
+    def __init__(self, window_size: int, stride: int):
+        super().__init__(window_size, stride)

-    Output:
-      - prompts:   [B, P_max]
-      - responses: [B, G, R_max]
-      - masks:     [B, G, R_max]
-      - rewards:   [B, G]
-    """
-    B = len(batch)
-    G = len(batch[0]["responses"])
-    P_max = max(b["prompts"].size(0) for b in batch)
-    R_max = max(r.size(0) for b in batch for r in b["responses"])
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, key)

-    prompts = torch.zeros(B, P_max, dtype=torch.long)
-    responses = torch.zeros(B, G, R_max, dtype=torch.long)
-    masks = torch.zeros(B, G, R_max, dtype=torch.bool)
-    rewards = torch.zeros(B, G, dtype=torch.float32)
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        begin_idx, end_idx = self.get_index(index)

-    for i, b in enumerate(batch):
-        p_len = b["prompts"].size(0)
-        prompts[i, :p_len] = b["prompts"]
-        rewards[i, : b["rewards"].size(0)] = b["rewards"]
-        for g in range(min(G, len(b["responses"]))):
-            r_len = b["responses"][g].size(0)
-            responses[i, g, :r_len] = b["responses"][g]
-            if g < len(b["masks"]):
-                masks[i, g, :r_len] = b["masks"][g]
+        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")
+        rewards = self._fetch_data(begin_idx, end_idx, "rewards")

        return {
            "prompts": prompts,
@@ -221,274 +336,3 @@ def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
            "masks": masks,
            "rewards": rewards,
        }
-
-
-def validate_keys(store: Store, required: List[str]) -> None:
-    """Raise ``KeyError`` if *store* is missing any *required* key."""
-    if not required:
-        return
-    actual = set(store.keys)
-    missing = [k for k in required if k not in actual]
-    if missing:
-        raise KeyError(
-            f"Store at {getattr(store, '_load_path', '?')} is missing required "
-            f"keys {missing}; available keys are {sorted(actual)}."
-        )
-
-
-class BaseDataset(Dataset, ABC):
-    """Abstract base class for dataset types.
-
-    Holds a :class:`Store`.  All sample-id indexing is delegated to the
-    store — this class exposes ``__len__`` as ``len(store)`` and the
-    ``keys`` property as ``store.keys``.  Subclasses implement
-    ``__getitem__`` with the train-type-specific key mapping and any
-    training-only index arithmetic (e.g. the next-token ``+1`` shift).
-    """
-
-    required_keys: List[str] = []
-
-    def __init__(self, store: Store):
-        super().__init__()
-        self.store: Store = store
-        validate_keys(store, self.required_keys)
-
-    def __len__(self) -> int:
-        return len(self.store)
-
-    @property
-    def keys(self) -> List[str]:
-        return self.store.keys
-
-    @property
-    def token_count(self) -> int:
-        return self.store.token_count
-
-    @abstractmethod
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        raise NotImplementedError
-
-
-class DatasetFactory(BaseFactory["BaseDataset"]):
-    """Factory for creating dataset instances by train-type.
-
-    Use :meth:`DatasetFactory.register("custom")` to register new
-    dataset classes; they must inherit from :class:`BaseDataset`.
-    """
-
-    @classmethod
-    def load(
-        cls,
-        train_type: str,
-        load_path: Optional[str] = None,
-        window_size: int = 0,
-        stride: Optional[int] = None,
-        storage_type: Optional[str] = None,
-        tokenizer_path: Optional[str] = None,
-        max_len: int = 2048,
-        store: Optional[Store] = None,
-        **kwargs,
-    ) -> "BaseDataset":
-        """Create and load a dataset in one step.
-
-        Two entry points:
-
-        - **store given**: bind it directly — the caller fully controls
-          Store construction and processor setup.  *load_path*,
-          *storage_type*, *tokenizer_path*, *window_size*, *stride* are
-          ignored.
-        - **store is None**: build a Store from *load_path*, auto-detecting
-          format and constructing a processor when *tokenizer_path* is
-          given for a record dataset on JSONL.
-
-        Args:
-            train_type: Registered dataset name ("seq", "sft", "dpo",
-                "grpo", …).
-            load_path: Path to the data file or directory (ignored if
-                *store* is given).
-            window_size: Stream window length — only meaningful for
-                stream datasets (SEQ/SFT).  Record datasets ignore it.
-            stride: Stride between consecutive stream samples
-                (default: same as *window_size*).
-            storage_type: Storage backend ("h5", "bin", "jsonl") or
-                None for auto-detection.
-            tokenizer_path: Path to tokenizer for lazy JSONL
-                tokenisation (record datasets only).
-            max_len: Max sequence length forwarded to processors.
-            store: Pre-built, already-loaded Store instance.
-            **kwargs: Extra arguments forwarded to ``store.load()``.
-
-        Returns:
-            Loaded dataset instance.
-        """
-        if store is not None:
-            return cls.create(train_type, store=store)
-
-        if load_path is None:
-            raise ValueError("Either load_path or store must be provided")
-
-        if storage_type is None:
-            storage_type = detect_format(load_path)
-
-        if stride is None:
-            stride = window_size
-
-        processor = cls._maybe_build_processor(
-            train_type, storage_type, tokenizer_path, max_len
-        )
-
-        store_window = cls._store_window_for(train_type, window_size)
-        store = StoreFactory.create(
-            storage_type,
-            window_size=store_window,
-            stride=stride if stride else store_window,
-        )
-        if processor is not None:
-            store.load(load_path, processor=processor, **kwargs)
-        else:
-            store.load(load_path, **kwargs)
-
-        return cls.create(train_type, store=store)
-
-    @staticmethod
-    def _store_window_for(train_type: str, window_size: int) -> int:
-        """Stream datasets consume ``window_size``; record datasets ignore it.
-
-        Record datasets (dpo/grpo) treat each record as an independent
-        training unit and never window, so the store is built with
-        ``window_size=0`` and ``len(store)`` returns the record count.
-        """
-        if train_type in ("seq", "sft"):
-            return window_size
-        return 0
-
-    @staticmethod
-    def _maybe_build_processor(
-        train_type: str,
-        storage_type: str,
-        tokenizer_path: Optional[str],
-        max_len: int,
-    ) -> Optional[Callable[[dict], Dict[str, Tensor]]]:
-        """Build an on-the-fly tokenisation processor if applicable.
-
-        Only raw JSONL + record datasets (DPO/GRPO) need a processor;
-        pre-tokenised backends (H5/bin) and stream datasets (SEQ/SFT)
-        return ``None`` so no tokenizer is loaded.
-        """
-        if tokenizer_path is None or storage_type != "jsonl":
-            return None
-        if train_type == "dpo":
-            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
-            return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
-        return None
-
-
-@DatasetFactory.register("seq")
-class SEQDataset(BaseDataset):
-    """Dataset for sequential next-token prediction training.
-
-    Stream mode: ``store.fetch(begin, end, "sequence")`` returns the
-    input window; the +1 shifted call returns the next-token target.
-    """
-
-    required_keys = ["sequence"]
-
-    def __getitem__(self, index: int):
-        begin, end = self.store.sample_window(index)
-        x = self.store.fetch(begin, end, "sequence")
-        y = self.store.fetch(begin + 1, end + 1, "sequence")
-        return {
-            "input_ids": x.to(dtype=torch.long),
-            "target_ids": y.to(dtype=torch.long),
-        }
-
-
-@DatasetFactory.register("sft")
-class SFTDataset(BaseDataset):
-    """Dataset for supervised fine-tuning with loss masking.
-
-    Stream mode: ``sequence``/``loss_mask``/``position_ids`` are sliced
-    to the window.  ``loss_mask`` and ``target_ids`` use the +1 shifted
-    slice so they align with the predicted positions.
-    """
-
-    required_keys = ["sequence", "loss_mask", "position_ids"]
-
-    def __getitem__(self, index: int):
-        begin, end = self.store.sample_window(index)
-        x = self.store.fetch(begin, end, "sequence")
-        y = self.store.fetch(begin + 1, end + 1, "sequence")
-        position_ids = self.store.fetch(begin, end, "position_ids")
-        loss_mask = self.store.fetch(begin + 1, end + 1, "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):
-    """Record-structured dataset for Direct Preference Optimization.
-
-    Each sample is one preference pair (chosen + rejected) and is an
-    independent training unit — no windowing, stride, or cross-record
-    concatenation.  This keeps each sequence self-contained so attention
-    never leaks across preference pairs.
-
-    Two loading paths (handled by :class:`DatasetFactory`):
-
-    - **Pre-tokenized** (H5/bin): ``store.load(path)`` reads per-record
-      tensors; ``__getitem__`` returns them directly.
-    - **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor
-      via :func:`dpo_processor` that tokenises on the fly — no packing,
-      no ``position_ids``.
-    """
-
-    required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"]
-
-    def make_processor(self, tokenizer, max_len: int):
-        return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
-
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        return {
-            "chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long),
-            "rejected": self.store.fetch_record(index, "rejected").to(dtype=torch.long),
-            "chosen_mask": self.store.fetch_record(index, "chosen_mask").to(
-                dtype=torch.bool
-            ),
-            "rejected_mask": self.store.fetch_record(index, "rejected_mask").to(
-                dtype=torch.bool
-            ),
-        }
-
-
-@DatasetFactory.register("grpo")
-class GRPODataset(BaseDataset):
-    """Dataset for offline Group Relative Policy Optimization.
-
-    Each sample is one prompt with its group of responses and scalar
-    rewards — an independent training unit with no windowing or stride.
-
-    Expected storage layout (produced by JsonlStore or pre-tokenized):
-
-    - ``prompts``:   List[Tensor]  — one 1-D token tensor per record
-    - ``responses``: List[List[Tensor]] — G response tensors per record
-    - ``masks``:     List[List[Tensor]] — G mask tensors per record
-    - ``rewards``:   List[Tensor]  — one 1-D float tensor (len G) per record
-    """
-
-    required_keys = ["prompts", "responses", "masks", "rewards"]
-
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        prompts = self.store.fetch_record(index, "prompts")
-        responses = self.store.fetch_record(index, "responses")
-        masks = self.store.fetch_record(index, "masks")
-        rewards = self.store.fetch_record(index, "rewards")
-        return {
-            "prompts": prompts.to(dtype=torch.long),
-            "responses": [r.to(dtype=torch.long) for r in responses],
-            "masks": [m.to(dtype=torch.bool) for m in masks],
-            "rewards": rewards.to(dtype=torch.float32),
-        }
@@ -5,15 +5,7 @@ import torch.distributed as dist
 from torch.utils.data import Dataset, Sampler


-class RDSampler(Sampler[int]):
-    """Resumable Distributed Sampler.
-
-    A distributed sampler that supports checkpoint-based resume: iteration
-    state (epoch, position) is tracked so training can continue from the
-    exact sample after a restart.  Shards the dataset across
-    ``dist.world_size`` replicas with optional shuffling.
-    """
-
+class ResumableDistributedSampler(Sampler[int]):
    def __init__(
        self,
        data_source: Dataset,
@@ -51,7 +43,6 @@ class RDSampler(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

@@ -82,12 +73,6 @@ class RDSampler(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._remaining
+        return self.num_samples_per_replica
@@ -1,636 +0,0 @@
-"""Storage backends for different data formats.
-
-Architecture (composition over inheritance):
-
-    Store (ABC)               — owns _data/_cum/_offsets bookkeeping
-                                  + window_size/stride for sample-id
-                                  indexing.  __getitem__/__len__ produce
-                                  the smallest iterable unit so Dataset
-                                  classes are pure delegators.
-    Streamable (mixin)        — raw token slice  fetch(begin, end, keys)
-    Recordable (mixin)        — raw record slice fetch_record(idx, keys)
-
-    H5Store(Store, Streamable, Recordable)
-    MmapStore(Store, Streamable, Recordable)
-    JsonlStore(Store, Streamable, Recordable)
-
-Each mixin is a stateless trait that relies on ``self._data`` etc.
-provided by :class:`Store`.  Concrete stores mix in whichever access
-primitives they support — ``Store`` is the sole base class, so there is
-no diamond inheritance or MRO ambiguity.
-
-Sample-id indexing lives on :class:`Store`, not on the dataset:
-
- **Stream mode** (``window_size > 0``):  ``len(store)`` returns the number
-  of ``(window_size, stride)`` windows that fit in the token river;
-  ``store[i]`` returns the *i*-th window as a dict of per-key tensors;
-  ``store.sample_window(i)`` exposes the underlying ``(begin, end)``
-  token slice for callers (e.g. next-token trainers) that need a +1
-  shifted companion window.
- **Record mode** (``num_records > 0``):  ``len(store)`` returns the
-  record count; ``store[i]`` returns the *i*-th record dict.
-
-Raw token/record access via :meth:`fetch` / :meth:`fetch_record`
-remains available for low-level callers that want explicit index
-control.  ``store.token_count`` is the total stream token count (what
-``len(store)`` used to mean in the legacy stream-only API).
-
-``segments_are_records`` (class attribute on each Store subclass)
-tells ``_normalize`` whether segments are inherently per-record (H5/
-JSONL) or opaque shards (bin).  Record access for bin relies on
-``_offsets`` instead.
-
-:class:`JsonlStore` supports a lazy mode (``processor=fn``) that keeps
-raw records and defers tokenisation to ``fetch_record`` — used by DPO
-to train directly from a ``.jsonl`` file without a pre-tokenised copy.
-"""
-
-import bisect
-import glob
-import json
-import logging
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Callable, Dict, List, Optional, Tuple, Union
-
-import torch
-from torch import Tensor
-
-from astrai.factory import BaseFactory
-from astrai.preprocessing.transform import TokenizeTransform
-from astrai.serialization import (
-    load_bin,
-    load_bin_offsets,
-    load_h5,
-)
-
-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", "jsonl", or "processed")
-
-    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):
-    """Common base for all storage backends.
-
-    A Store owns both its data layout AND its sample-id → token/record
-    index translation.  Datasets are thin wrappers that bind a Store
-    to a particular train-type's key mapping; they never know about
-    window/stride math.
-
-    Two iteration modes:
-
-    - **Stream** (``window_size > 0``): data is treated as one long
-      token river.  ``len(store)`` returns the number of windows;
-      ``store[i]`` slices every stream-compatible key to window ``i``;
-      ``store.sample_window(i)`` returns the ``(begin, end)`` token
-      slice for callers needing a +1 shifted companion window.
-    - **Record** (``num_records > 0``): data is per-record.
-      ``len(store)`` returns ``num_records``; ``store[i]`` returns
-      the *i*-th record as a dict.
-
-    Raw token slicing is still available via :meth:`fetch` (mixed in
-    by :class:`Streamable`) when a store has stream support configured.
-    Raw record slicing via :meth:`fetch_record` (mixed in by
-    :class:`Recordable`) when a store has record support.
-
-    ``token_count`` exposes the raw total stream length — this is what
-    ``len(store)`` returned in the legacy stream-only API and what
-    stream-bound ``fetch`` uses for its bounds check.
-    """
-
-    segments_are_records: bool = False
-
-    def __init__(
-        self,
-        window_size: int = 0,
-        stride: Optional[int] = None,
-    ):
-        self._data: Dict[str, List[Tensor]] = {}
-        self._cum: Dict[str, List[int]] = {}
-        self._offsets: Dict[str, List[int]] = {}
-        self._length: int = 0
-        self._num_records: int = 0
-        self._window_size: int = int(window_size)
-        self._stride: int = int(stride) if stride is not None else int(window_size)
-
-    @abstractmethod
-    def load(self, path: str, **kwargs) -> None:
-        raise NotImplementedError
-
-    @property
-    def keys(self) -> List[str]:
-        return list(self._data.keys())
-
-    @property
-    def window_size(self) -> int:
-        return self._window_size
-
-    @property
-    def stride(self) -> int:
-        return self._stride
-
-    @property
-    def token_count(self) -> int:
-        """Total tokens across all stream segments.
-
-        Useful for the bounds-checked raw :meth:`fetch` and as the
-        legacy ``len(store)`` value.
-        """
-        return self._length
-
-    @property
-    def num_records(self) -> int:
-        """Number of records available via :meth:`fetch_record`.
-
-        Non-zero only when the backing layout provides per-record
-        indexing (H5/JSONL segments or bin ``_offsets``).
-        """
-        return self._num_records
-
-    @property
-    def num_samples(self) -> int:
-        """Number of items produced by ``__getitem__``.
-
-        Stream-mode wins when ``window_size > 0`` and there are tokens
-        to slice; otherwise falls back to ``num_records``.
-        """
-        if self._window_size > 0 and self._length > 0:
-            total = self._length
-            w = self._window_size
-            if total <= w:
-                return 0
-            return (total - 1 - w) // self._stride + 1
-        return self._num_records
-
-    def __len__(self) -> int:
-        return self.num_samples
-
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        if index < 0:
-            index += self.num_samples
-        if not 0 <= index < self.num_samples:
-            raise IndexError(
-                f"Store index out of range: {index}, num_samples={self.num_samples}"
-            )
-        if self._window_size > 0 and self._length > 0:
-            begin, end = self.sample_window(index)
-            keys = self._stream_keys()
-            return {k: self.fetch(begin, end, k) for k in keys}
-        return self.fetch_record(index, self._record_keys())
-
-    def sample_window(self, index: int) -> Tuple[int, int]:
-        """Return ``(begin, end)`` token positions for stream sample *index*.
-
-        The clipped tail keeps the last reachable window inside the
-        token river instead of overshooting.  Caller is responsible
-        for staying within :attr:`num_samples`: an out-of-range index
-        raises ``IndexError``.
-        """
-        if self._window_size <= 0:
-            raise RuntimeError("sample_window() requires window_size > 0 (stream mode)")
-        if self._window_size <= 0 or self._length <= self._window_size:
-            raise IndexError(
-                f"Data too short for window: token_count={self._length}, "
-                f"window_size={self._window_size}"
-            )
-        if not 0 <= index < self.num_samples:
-            raise IndexError(
-                f"Sample index out of range: {index}, num_samples={self.num_samples}"
-            )
-        total = self._length
-        begin = min(index * self._stride, total - 1 - self._window_size)
-        end = min(begin + self._window_size, total - 1)
-        return begin, end
-
-    def _stream_keys(self) -> List[str]:
-        out: List[str] = []
-        for k, tensors in self._data.items():
-            if tensors and isinstance(tensors[0], list):
-                continue
-            out.append(k)
-        return out
-
-    def _record_keys(self) -> List[str]:
-        return list(self._data.keys())
-
-    def _normalize(
-        self,
-        raw: Dict[str, list],
-        offsets: Optional[Dict[str, List[int]]] = None,
-    ):
-        """Register segments and pre-compute indices for both access modes.
-
-        Stream mode: ``_cum[key]`` accumulates per-segment lengths so
-        ``Streamable._fetch_stream_key`` can bisect across segments
-        without concatenation.
-
-        Record mode: if *offsets* is provided (bin layout),
-        ``_offsets[key]`` stores cumulative per-record offsets into the
-        single concatenated segment.  Otherwise, when
-        ``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is
-        a per-record list and ``fetch_record`` indexes it directly.
-
-        Nested keys (GRPO ``responses``/``masks`` as
-        ``List[List[Tensor]]``) are stored as-is and excluded from both
-        cumulative bookkeepings — they are only accessed record-by-record.
-        """
-        flat_lengths = []
-        for key, tensors in raw.items():
-            self._data[key] = tensors
-            if not tensors:
-                self._cum[key] = []
-                flat_lengths.append(0)
-                continue
-            if isinstance(tensors[0], list):
-                self._cum[key] = []
-                continue
-            cum = []
-            total = 0
-            for t in tensors:
-                total += t.shape[0]
-                cum.append(total)
-            self._cum[key] = cum
-            flat_lengths.append(cum[-1] if cum else 0)
-        self._length = min(flat_lengths) if flat_lengths else 0
-
-        valid_offsets: Dict[str, List[int]] = {}
-        if offsets:
-            for key, off in offsets.items():
-                segs = self._data.get(key, [])
-                if len(segs) == 1 and len(off) > 1:
-                    valid_offsets[key] = off
-                elif len(segs) > 1:
-                    logger.warning(
-                        "Key '%s' has %d segments with offsets — record mode "
-                        "disabled for this key (multi-shard bin+offsets not "
-                        "supported). Merge shards or use H5/JSONL.",
-                        key,
-                        len(segs),
-                    )
-        self._offsets = valid_offsets
-        if valid_offsets:
-            record_counts = [len(v) - 1 for v in valid_offsets.values()]
-            self._num_records = min(record_counts) if record_counts else 0
-        elif self.segments_are_records:
-            per_record_counts = []
-            for key, tensors in self._data.items():
-                if tensors and isinstance(tensors[0], list):
-                    continue
-                per_record_counts.append(len(tensors))
-            self._num_records = min(per_record_counts) if per_record_counts else 0
-        else:
-            self._num_records = 0
-
-
-class Streamable:
-    """Mixin granting raw token-stream access via :meth:`fetch`.
-
-    Stateless trait relying on ``self._data``, ``self._cum``,
-    ``self._length`` maintained by :class:`Store`.  Stream mode is
-    active when the owning store has ``window_size > 0``; for stores
-    that can also serve record access (H5/JSONL/bin+offsets), the
-    ``fetch_record`` API from :class:`Recordable` is used instead.
-    """
-
-    def fetch(
-        self,
-        begin: int,
-        end: int,
-        keys: Union[str, List[str]],
-    ):
-        return _stream_fetch(self, begin, end, keys)
-
-
-def _stream_fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
-    if not getattr(self, "_data", None):
-        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 _fetch_stream_key(self, keys, begin, end)
-    return {k: _fetch_stream_key(self, k, begin, end) for k in keys}
-
-
-def _fetch_stream_key(self, key: str, begin: int, end: int) -> Tensor:
-    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)
-
-
-class Recordable:
-    """Mixin granting raw record access via :meth:`fetch_record`.
-
-    Stateless trait relying on ``self._data``, ``self._offsets``,
-    ``self._num_records`` maintained by :class:`Store`.
-    """
-
-    def fetch_record(
-        self,
-        index: int,
-        keys: Union[str, List[str]],
-    ):
-        return _record_fetch(self, index, keys)
-
-
-def _record_fetch(self, index: int, keys: Union[str, List[str]]):
-    if not getattr(self, "_data", None) and self._num_records == 0:
-        raise RuntimeError("Store not loaded")
-    if not 0 <= index < self._num_records:
-        raise ValueError(
-            f"Record index out of bounds: {index}, num_records={self._num_records}"
-        )
-    if isinstance(keys, str):
-        return _fetch_record_key(self, keys, index)
-    return {k: _fetch_record_key(self, k, index) for k in keys}
-
-
-def _fetch_record_key(self, key: str, index: int):
-    offsets = self._offsets.get(key)
-    if offsets:
-        start = offsets[index]
-        end = (
-            offsets[index + 1]
-            if index + 1 < len(offsets)
-            else self._data[key][0].shape[0]
-        )
-        return self._data[key][0][start:end]
-    return self._data[key][index]
-
-
-class StoreFactory(BaseFactory["Store"]):
-    """Factory for creating Store instances by type name."""
-
-
-@StoreFactory.register("h5")
-class H5Store(Store, Streamable, Recordable):
-    """HDF5-based storage backend (pre-tokenized data).
-
-    Each key is stored as a group of per-record datasets (``data_0``,
-    ``data_1``, …).  Supports both access modes:
-
-    - **Stream**: ``fetch(begin, end, key)`` and ``store[i]`` slice
-      across concatenated records via ``_cum`` — used by SEQ/SFT.
-    - **Record**: ``fetch_record(i, key)`` and ``store[i]`` (when
-      ``window_size == 0``) index ``_data[key]`` directly — used by
-      DPO/GRPO.
-    """
-
-    segments_are_records = True
-
-    def __init__(
-        self,
-        window_size: int = 0,
-        stride: Optional[int] = None,
-    ):
-        super().__init__(window_size=window_size, stride=stride)
-
-    def load(self, path: str, **kwargs):
-        self._normalize(load_h5(path))
-
-
-@StoreFactory.register("bin")
-class MmapStore(Store, Streamable, Recordable):
-    """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.
-
-    Supports both access modes:
-
-    - **Stream**: always available via :meth:`fetch`.
-    - **Record** (``fetch_record(i, key)``): only when ``meta.json``
-      contains per-record ``offsets`` (written via
-      ``save_bin(..., record_keys=...)``).  Legacy bin files without
-      offsets have ``num_records == 0`` and ``len(store)`` reflects the
-      windowed sample count when ``window_size > 0``.
-
-    ``segments_are_records`` is ``False`` here (bin segments are
-    contiguous streams, not per-record) — record access is driven
-    purely by ``_offsets``.
-    """
-
-    segments_are_records = False
-
-    def __init__(
-        self,
-        window_size: int = 0,
-        stride: Optional[int] = None,
-    ):
-        super().__init__(window_size=window_size, stride=stride)
-        self._mmap_refs: List[Tensor] = []
-
-    def load(self, path: str, **kwargs):
-        self._mmap_refs = []
-        root = Path(path)
-        all_raw: Dict[str, List[Tensor]] = {}
-        all_offsets: Dict[str, List[int]] = {}
-        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))
-            off = load_bin_offsets(str(meta_path.parent))
-            for key, tensors in raw.items():
-                if key not in all_raw:
-                    all_raw[key] = []
-                all_raw[key].extend(tensors)
-            for key, o in off.items():
-                if key not in all_offsets:
-                    all_offsets[key] = []
-                all_offsets[key].extend(o)
-        if not meta_paths:
-            raise FileNotFoundError(f"No meta.json found under {path}")
-        self._normalize(all_raw, offsets=all_offsets or None)
-        for tensors in self._data.values():
-            self._mmap_refs.extend(tensors)
-
-
-class JsonlSource:
-    """Read raw JSON records from a ``.jsonl`` file or directory.
-
-    A thin reader used by :class:`JsonlStore` in processor mode — holds
-    no tokenizer, performs no tokenisation, just yields dicts.
-    """
-
-    def __init__(self, path: str):
-        self.path = Path(path)
-        self._records: Optional[List[dict]] = None
-
-    def load(self) -> List[dict]:
-        if self._records is None:
-            self._records = self._read(self.path)
-        return self._records
-
-    @staticmethod
-    def _read(root: Path) -> List[dict]:
-        if root.is_file():
-            return JsonlSource._read_file(root)
-        return JsonlSource._read_dir(root)
-
-    @staticmethod
-    def _read_file(path: Path) -> List[dict]:
-        records: List[dict] = []
-        with open(path, "r", encoding="utf-8") as f:
-            for line in f:
-                line = line.strip()
-                if not line:
-                    continue
-                try:
-                    records.append(json.loads(line))
-                except json.JSONDecodeError:
-                    logger.warning("Failed to parse JSON line in %s, skipping", path)
-        return records
-
-    @staticmethod
-    def _read_dir(root: Path) -> List[dict]:
-        records: List[dict] = []
-        for jsonl_path in sorted(root.glob("*.jsonl")):
-            records.extend(JsonlSource._read_file(jsonl_path))
-        return records
-
-
-@StoreFactory.register("jsonl")
-class JsonlStore(Store, Streamable, Recordable):
-    """JSONL reader with two tokenisation modes.
-
-    A JSONL dataset is a ``.jsonl`` file or a directory of ``*.jsonl``
-    files plus (optionally) a ``dataset_config.json`` describing the
-    tokenization pipeline.
-
-    Two modes, selected at :meth:`load` time:
-
-    - **Eager** (default): applies a :class:`TokenizeTransform` to every
-      record at load time and registers per-key tensors via
-      ``_normalize``.  Both ``fetch`` (stream) and ``fetch_record``
-      (record) work.
-    - **Lazy** (``processor=fn`` passed): keeps raw records and defers
-      tokenisation to ``fetch_record``.  Only record access works —
-      ``len(store)`` returns ``num_records``; stream primitives raise.
-    """
-
-    CONFIG_NAME = "dataset_config.json"
-    segments_are_records = True
-
-    def __init__(
-        self,
-        window_size: int = 0,
-        stride: Optional[int] = None,
-    ):
-        super().__init__(window_size=window_size, stride=stride)
-        self._source: Optional[JsonlSource] = None
-        self._processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None
-        self._keys_cache: Optional[List[str]] = None
-
-    def load(self, path: str, transform=None, processor=None, **kwargs):
-        self._source = JsonlSource(path)
-        records = self._source.load()
-
-        if processor is not None:
-            self._processor = processor
-            self._num_records = len(records)
-            return
-
-        if transform is None:
-            root = Path(path)
-            config_path = root / self.CONFIG_NAME if root.is_dir() else None
-            if config_path is None or not config_path.exists():
-                raise FileNotFoundError(
-                    f"JSONL dataset config not found. Expected "
-                    f"{self.CONFIG_NAME} alongside *.jsonl files, pass an "
-                    f"explicit transform, or pass processor= for lazy "
-                    f"on-the-fly tokenisation."
-                )
-            transform = TokenizeTransform.from_config_file(str(config_path))
-
-        transformed = transform.apply(records)
-        self._normalize(transformed)
-
-    @property
-    def keys(self) -> List[str]:
-        if self._processor is not None:
-            if self._keys_cache is None and self._num_records > 0:
-                sample = self._processor(self._source.load()[0])
-                self._keys_cache = list(sample.keys())
-            return self._keys_cache or []
-        return list(self._data.keys())
-
-    def fetch_record(self, index: int, keys: Union[str, List[str]]):
-        if self._processor is not None:
-            if not 0 <= index < self._num_records:
-                raise ValueError(
-                    f"Record index out of bounds: {index}, "
-                    f"num_records={self._num_records}"
-                )
-            record = self._source.load()[index]
-            data = self._processor(record)
-            if isinstance(keys, str):
-                return data[keys]
-            return {k: data[k] for k in keys}
-        return _record_fetch(self, index, keys)
-
-    def fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
-        if self._processor is not None:
-            raise RuntimeError(
-                "JsonlStore in lazy (processor) mode does not support "
-                "stream fetch(); use fetch_record() instead."
-            )
-        return _stream_fetch(self, begin, end, keys)
-
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        if self._processor is not None:
-            return self.fetch_record(index, self._record_keys())
-        return super().__getitem__(index)
@@ -1,29 +0,0 @@
-"""CUDA attention kernel wrappers with torch fallback.
-
-Public API:
-    - ``attn_decode`` — single-query decode attention
-    - ``attn_prefill`` — multi-query prefill attention
-    - ``attn_paged_decode`` — paged decode attention (direct page-table access)
-
-Interface (shared by all wrappers):
-    causal_offset: -1 = non-causal; >=0 = absolute position of first Q token
-    mask:       2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True = keep)
-    scale:      0.0 = auto (1/sqrt(head_dim)); >0 = explicit
-    layout:     "bhld" (default) or "blhd"
-
-Causal and mask can coexist — both are applied simultaneously.
-
-Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``)
-when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
-"""
-
-from astrai.extension.loader import KERNEL_NAMES, is_available
-from astrai.extension.ops import attn_decode, attn_paged_decode, attn_prefill
-
-__all__ = [
-    "attn_decode",
-    "attn_paged_decode",
-    "attn_prefill",
-    "is_available",
-    "KERNEL_NAMES",
-]
@@ -1,36 +0,0 @@
-"""Dynamic discovery and loading of compiled CUDA kernel modules.
-
-Each kernel is registered in ``csrc/build.py`` and built into a ``.so`` placed
-in this package directory. On import we try to load each one; kernels that
-failed to build (or are running on a CPU-only machine) are marked unavailable
-so the wrapper functions can fall back to ``torch`` SDPA.
-"""
-
-import importlib
-import logging
-
-logger = logging.getLogger(__name__)
-
-KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode"]
-
-_available: dict[str, bool] = {}
-_modules: dict[str, object] = {}
-
-for _name in KERNEL_NAMES:
-    try:
-        _mod = importlib.import_module(f".{_name}", package=__package__)
-        _available[_name] = True
-        _modules[_name] = _mod
-    except ImportError:
-        _available[_name] = False
-        _modules[_name] = None
-
-
-def is_available(name: str) -> bool:
-    """Return ``True`` if the compiled kernel ``name`` was loaded."""
-    return _available.get(name, False)
-
-
-def get_module(name: str) -> object:
-    """Return the loaded kernel module for ``name``, or ``None`` if unavailable."""
-    return _modules.get(name)
@@ -1,246 +0,0 @@
-"""GQA attention wrapper functions — one entry point per compiled kernel.
-
-Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when
-available, otherwise falls back to ``torch`` SDPA.
-
-Interface (all functions):
-    causal_offset: -1 = non-causal; >=0 = absolute position of first Q token
-    mask:       2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool)
-    scale:      0.0 = auto (1/sqrt(head_dim)); >0 = explicit
-    layout:     "bhld" (default) or "blhd"
-
-Add new kernel wrappers here; split into per-variant files only if this file
-grows large.
-"""
-
-import math
-
-import torch
-import torch.nn.functional as F
-
-from astrai.extension.loader import _available, _modules
-
-_LAYOUT_CODES: dict[str, int] = {"bhld": 0, "blhd": 1}
-
-
-def _parse_layout(layout: str | int) -> int:
-    if isinstance(layout, int):
-        return layout
-    code = _LAYOUT_CODES.get(layout.lower())
-    if code is None:
-        raise ValueError(
-            f"unknown layout '{layout}', expected one of {list(_LAYOUT_CODES)}"
-        )
-    return code
-
-
-def _to_bhld(t: torch.Tensor, layout: int) -> torch.Tensor:
-    """Normalize to b h l d view. Zero-copy transpose if layout==1 (b l h d)."""
-    if layout == 1:
-        return t.transpose(1, 2)
-    return t
-
-
-def _expand_kv_heads(
-    k: torch.Tensor, v: torch.Tensor, q_head: int
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Expand K/V heads to match Q heads for GQA fallback."""
-    kv_head = k.size(1)
-    if kv_head == q_head:
-        return k, v
-    group = q_head // kv_head
-    k = k.repeat_interleave(group, dim=1)
-    v = v.repeat_interleave(group, dim=1)
-    return k, v
-
-
-def _build_attn_mask(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    mask: torch.Tensor | None,
-    causal_offset: int,
-    scale: float,
-) -> tuple[torch.Tensor | None, float]:
-    """Build SDPA-compatible attn_mask + resolved scale.
-
-    q and k must already be in b h l d layout.
-    Causal and mask can coexist: causal sets -inf above the diagonal, mask
-    sets -inf for padded positions. Both are OR'd into a single bool mask.
-    """
-    q_len = q.size(2)
-    kv_len = k.size(2)
-    head_dim = q.size(3)
-    resolved_scale = scale if scale and scale > 0 else 1.0 / math.sqrt(head_dim)
-
-    attn_mask = None
-
-    if mask is not None:
-        if mask.dim() == 2:
-            # [batch, kv_len] → [batch, 1, 1, kv_len]
-            attn_mask = mask[:, None, None, :]
-        elif mask.dim() == 3:
-            # [batch, q_len, kv_len] → [batch, 1, q_len, kv_len]
-            attn_mask = mask[:, None, :, :]
-        else:
-            raise ValueError(f"mask must be 2D or 3D, got {mask.dim()}D")
-
-    if causal_offset >= 0:
-        batch = q.size(0)
-        # q row i attends to kv cols 0..(causal_offset + i)
-        q_idx = torch.arange(q_len, device=q.device).unsqueeze(1)  # [q_len, 1]
-        kv_idx = torch.arange(kv_len, device=q.device).unsqueeze(0)  # [1, kv_len]
-        causal_bool = kv_idx > (causal_offset + q_idx)  # True = masked out
-        causal_mask = causal_bool.unsqueeze(0).expand(
-            batch, -1, -1
-        )  # [batch, q_len, kv_len]
-        causal_mask = causal_mask[:, None, :, :]  # [batch, 1, q_len, kv_len]
-
-        if attn_mask is not None:
-            attn_mask = attn_mask | causal_mask
-        else:
-            attn_mask = causal_mask
-
-    return attn_mask, resolved_scale
-
-
-def _torch_fallback(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None,
-    causal_offset: int,
-    scale: float,
-    q_layout: int,
-    kv_layout: int | None = None,
-) -> torch.Tensor:
-    """Reference attention via ``scaled_dot_product_attention``.
-
-    q_layout / kv_layout: 0 = b h l d, 1 = b l h d.
-    If kv_layout is None, uses q_layout (Q and K/V share the same layout).
-    """
-    if kv_layout is None:
-        kv_layout = q_layout
-    q = _to_bhld(q, q_layout)
-    k = _to_bhld(k, kv_layout)
-    v = _to_bhld(v, kv_layout)
-    k, v = _expand_kv_heads(k, v, q.size(1))
-    attn_mask, resolved_scale = _build_attn_mask(q, k, mask, causal_offset, scale)
-    out = F.scaled_dot_product_attention(
-        q, k, v, attn_mask=attn_mask, is_causal=False, scale=resolved_scale
-    )
-    # Restore Q's original layout
-    if q_layout == 1:
-        out = out.transpose(1, 2)
-    return out
-
-
-def _gather_kv_from_pages(
-    page_table: torch.Tensor,
-    k_cache: torch.Tensor,
-    v_cache: torch.Tensor,
-    page_size: int,
-    kv_len: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Gather contiguous K/V from paged cache for torch SDPA fallback.
-
-    Shapes:
-        page_table : [batch, max_pages] (int64)
-        k_cache    : [n_pages, page_size, n_kv_heads, head_dim]
-        v_cache    : same as k_cache
-    Returns:
-        k, v : [batch, kv_len, n_kv_heads, head_dim]  (b l h d)
-    """
-    batch, max_pages = page_table.shape
-    _, ps, n_kv_heads, head_dim = k_cache.shape
-    if ps != page_size:
-        raise ValueError(f"k_cache page_size mismatch: {ps} vs {page_size}")
-
-    # Vectorized gather: build physical page + offset indices, then advanced-index
-    positions = torch.arange(kv_len, device=page_table.device)
-    logical_pages = positions // page_size  # [kv_len]
-    page_offsets = positions % page_size  # [kv_len]
-
-    phys_pages = page_table[:, logical_pages]  # [batch, kv_len]
-    # k_cache[phys_pages, page_offsets] → [batch, kv_len, n_kv_heads, head_dim] (b l h d)
-    k = k_cache[phys_pages, page_offsets]
-    v = v_cache[phys_pages, page_offsets]
-    return k, v
-
-
-def attn_decode(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None = None,
-    causal_offset: int = -1,
-    scale: float = 0.0,
-    layout: str = "bhld",
-) -> torch.Tensor:
-    li = _parse_layout(layout)
-    if _available["attn_decode"]:
-        return _modules["attn_decode"].attn_decode(
-            q,
-            k,
-            v,
-            mask=mask,
-            causal_offset=causal_offset,
-            scale=scale,
-            layout=li,
-        )
-    return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li)
-
-
-def attn_prefill(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None = None,
-    causal_offset: int = -1,
-    scale: float = 0.0,
-    layout: str = "bhld",
-) -> torch.Tensor:
-    li = _parse_layout(layout)
-    if _available["attn_prefill"]:
-        return _modules["attn_prefill"].attn_prefill(
-            q,
-            k,
-            v,
-            mask=mask,
-            causal_offset=causal_offset,
-            scale=scale,
-            layout=li,
-        )
-    return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li)
-
-
-def attn_paged_decode(
-    q: torch.Tensor,
-    page_table: torch.Tensor,
-    k_cache: torch.Tensor,
-    v_cache: torch.Tensor,
-    page_size: int,
-    kv_len: int,
-    mask: torch.Tensor | None = None,
-    causal_offset: int = -1,
-    scale: float = 0.0,
-    layout: str = "bhld",
-) -> torch.Tensor:
-    li = _parse_layout(layout)
-    if _available["attn_paged_decode"]:
-        return _modules["attn_paged_decode"].attn_paged_decode(
-            q,
-            page_table,
-            k_cache,
-            v_cache,
-            page_size,
-            kv_len,
-            mask=mask,
-            causal_offset=causal_offset,
-            scale=scale,
-            layout=li,
-        )
-    # Gathered K/V are always b l h d
-    k, v = _gather_kv_from_pages(page_table, k_cache, v_cache, page_size, kv_len)
-    return _torch_fallback(
-        q, k, v, mask, causal_offset, scale, q_layout=li, kv_layout=1
-    )
@@ -1,144 +1,190 @@
-"""Base factory with decorator-based registration and kwarg-filtered instantiation."""
+"""Base factory class for extensible component registration."""

-import inspect
-import sys
 from abc import ABC
-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
+from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar

 T = TypeVar("T")


-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
+class Registry:
+    """Flexible registry for component classes with category and priority support.

-    name = arg if isinstance(arg, str) else arg.__forward_arg__
-    if name == factory_cls.__name__:
-        return factory_cls
+    This registry stores component classes with optional metadata (category, priority).
+    It provides methods for registration, retrieval, and listing with filtering.
+    """

-    mod = sys.modules.get(factory_cls.__module__)
-    if mod is None:
-        return None
-    ns = vars(mod)
+    def __init__(self):
+        self._entries = {}  # name -> (component_cls, category, priority)

-    if isinstance(arg, ForwardRef):
-        return arg._evaluate(ns, None, recursive_guard=frozenset())
+    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)

-    return ns.get(name)
+    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()


 class BaseFactory(ABC, Generic[T]):
-    """Generic factory with decorator-based component registration.
+    """Generic factory class for component registration and creation.

-        class MyFactory(BaseFactory[MyBase]):
+    This base class provides a decorator-based registration pattern
+    for creating extensible component factories.
+
+    Example usage:
+        class MyFactory(BaseFactory[MyBaseClass]):
            pass

        @MyFactory.register("custom")
-        class CustomComponent(MyBase):
+        class CustomComponent(MyBaseClass):
            ...

-        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.
+        component = MyFactory.create("custom", *args, **kwargs)
    """

-    _entries: Dict[str, Type[T]]
+    _registry: Registry

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
-        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
+        cls._registry = Registry()

    @classmethod
-    def register(cls, name: str) -> Callable[[Type[T]], Type[T]]:
-        """Decorator to register a component class.
+    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.

-        Validates that the decorated class inherits from the generic
-        type parameter ``T`` declared on the factory.
+        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
        """

        def decorator(component_cls: Type[T]) -> Type[T]:
            cls._validate_component(component_cls)
-            if name in cls._entries:
-                raise ValueError(f"Component '{name}' is already registered")
-            cls._entries[name] = component_cls
+            cls._registry.register(
+                name, component_cls, category=category, priority=priority
+            )
            return component_cls

        return decorator

    @classmethod
    def create(cls, name: str, *args, **kwargs) -> T:
-        """Create a component instance by name, filtering kwargs to match
-        the component's ``__init__`` signature.
+        """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
        """
-        entry = cls._entries.get(name)
-        if entry is None:
+        if not cls._registry.contains(name):
            raise ValueError(
-                f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
+                f"Unknown component: '{name}'. "
+                f"Supported types: {sorted(cls._registry.list_names())}"
            )
-        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}
+        component_cls = cls._registry.get(name)
        return component_cls(*args, **kwargs)

    @classmethod
-    def _validate_component(cls, component_cls: Type[T]):
-        """Validate the decorated class inherits from the factory's base type.
+    def _validate_component(cls, component_cls: Type[T]) -> None:
+        """Validate that the component class is valid for this factory.

-        Override for custom validation beyond ``issubclass``.
+        Override this method in subclasses to add custom validation.
+
+        Args:
+            component_cls: Component class to validate
+
+        Raises:
+            TypeError: If the component class is invalid
        """
-        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__}"
-            )
+        pass

    @classmethod
-    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
+    def list_registered(cls) -> list:
+        """List all registered component names.

-    @classmethod
-    def list_registered(cls) -> List[str]:
-        """List all registered component names."""
-        return sorted(cls._entries)
+        Returns:
+            List of registered component names
+        """
+        return cls._registry.list_names()

    @classmethod
    def is_registered(cls, name: str) -> bool:
-        """Check if a component name is registered."""
-        return name in cls._entries
+        """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"]
@@ -1,105 +1,46 @@
 """Inference module for continuous batching.

 Layers:
-  - 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, FrequencyPenaltyStrategy)
+  - engine.py:    Facade (InferenceEngine), Value Object (GenerationParams, GenerationRequest)
+  - scheduler.py: Continuous-batching loop, Task state machine, TaskStatus enum
+  - cache.py:     PagedCache (page-table-indirected KV cache with alloc/free)
+  - sampling.py:  Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
+  - server.py:    FastAPI HTTP server (OpenAI-compatible endpoints)
 """

-from astrai.inference.api import (
-    AnthropicMessage,
-    BaseToolParser,
-    ChatCompletionRequest,
-    ChatMessage,
-    FunctionDef,
-    GenContext,
-    MessagesRequest,
-    ProtocolHandler,
-    SimpleJsonToolParser,
-    StopChecker,
-    ToolDef,
-    ToolParserFactory,
-    get_app,
-    run_server,
+from astrai.inference.engine import (
+    GenerationParams,
+    GenerationRequest,
+    InferenceEngine,
 )
-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 (
+from astrai.inference.sampling import (
    BaseSamplingStrategy,
-    FrequencyPenaltyStrategy,
    SamplingPipeline,
    TemperatureStrategy,
    TopKStrategy,
    TopPStrategy,
    sample,
 )
+from astrai.inference.scheduler import (
+    InferenceScheduler,
+    Task,
+    TaskStatus,
+)

 __all__ = [
+    # Engine / Requests
    "InferenceEngine",
    "GenerationRequest",
+    "GenerationParams",
+    # Scheduler
    "InferenceScheduler",
-    "Executor",
-    "STOP",
    "Task",
-    "TaskManager",
    "TaskStatus",
-    "Allocator",
-    "CacheView",
-    "KVCache",
-    "ContiguousCache",
-    "ContiguousCacheView",
-    "PageCache",
-    "PageCacheView",
-    "PagePool",
-    "PrefixCache",
-    "Storage",
-    "TaskTable",
-    "page_hash",
+    # Sampling (Strategy pattern)
    "sample",
    "BaseSamplingStrategy",
    "TemperatureStrategy",
    "TopKStrategy",
    "TopPStrategy",
-    "FrequencyPenaltyStrategy",
    "SamplingPipeline",
-    "ProtocolHandler",
-    "StopChecker",
-    "GenContext",
-    "BaseToolParser",
-    "SimpleJsonToolParser",
-    "ToolParserFactory",
-    "OpenAIResponseBuilder",
-    "AnthropicResponseBuilder",
-    "ChatMessage",
-    "ChatCompletionRequest",
-    "FunctionDef",
-    "ToolDef",
-    "AnthropicMessage",
-    "MessagesRequest",
-    "get_app",
-    "run_server",
 ]
@@ -1,39 +0,0 @@
-"""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",
-]
@@ -1,142 +0,0 @@
-"""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,
-            },
-        }
@@ -1,277 +0,0 @@
-"""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",
-    "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,
-            },
-        }
@@ -1,200 +0,0 @@
-"""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,
-            frequency_penalty=getattr(self.request, "frequency_penalty", 0.0),
-        )
-
-        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)
@@ -1,202 +0,0 @@
-"""
-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,
-    )
@@ -1,344 +0,0 @@
-"""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 json
-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)``.
-    """
-    if complete:
-        try:
-            obj = json.loads(json_str)
-        except json.JSONDecodeError:
-            return None, "", False
-        name = obj.get("name")
-        if not isinstance(name, str) or not name:
-            return None, "", False
-        args = obj.get("arguments")
-        if isinstance(args, dict):
-            if not args:
-                args = ""
-            else:
-                args = json.dumps(args, ensure_ascii=False)
-                args = args[1:-1].rstrip()
-        elif isinstance(args, list):
-            args = json.dumps(args, ensure_ascii=False) if args else ""
-        elif isinstance(args, str):
-            pass
-        else:
-            args = str(args) if args is not None else ""
-        return name, args, True
-
-    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 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]
-
-        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 '"name"' not in 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
@@ -0,0 +1,211 @@
+"""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 and persistent prefix caching.
+
+    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)
+      - LRU eviction for persistent cross-batch prefix caching
+
+    Pages with recorded hashes persist after refcount reaches 0 (pinned).
+    They are evicted via LRU only when alloc() finds no free pages.
+    """
+
+    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] = {}
+        self._lru: List[int] = []
+        self._pin: List[bool] = [False] * n_pages
+        self.lookup_hits: int = 0
+        self.lookup_misses: int = 0
+
+    def _touch(self, idx: int) -> None:
+        if self._refs[idx] == 0 and idx in self._lru:
+            self._lru.remove(idx)
+            self._lru.append(idx)
+
+    def _evict_one(self) -> int:
+        while self._lru:
+            idx = self._lru.pop(0)
+            h = self._page_to_hash.pop(idx, None)
+            if h is not None:
+                self._hash_to_page.pop(h, None)
+            self._pin[idx] = False
+            self._refs[idx] = 1
+            return idx
+        return -1
+
+    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
+        self._pin[page_idx] = True
+        if page_idx in self._lru:
+            self._lru.remove(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:
+                self.lookup_misses += 1
+                break
+            self.lookup_hits += 1
+            self._touch(p)
+            hits.append(p)
+        return hits
+
+    def inc_ref(self, idx: int) -> None:
+        self._refs[idx] += 1
+        if self._refs[idx] == 1 and idx in self._lru:
+            self._lru.remove(idx)
+
+    def alloc(self) -> int:
+        if self._free_mask:
+            lsb = self._free_mask & -self._free_mask
+            idx = lsb.bit_length() - 1
+            self._free_mask ^= lsb
+            self._refs[idx] = 1
+            if idx in self._lru:
+                self._lru.remove(idx)
+            return idx
+        return self._evict_one()
+
+    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:
+            h = self._page_to_hash.get(idx)
+            if h is not None and self._pin[idx]:
+                self._lru.append(idx)
+            else:
+                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)
+                self._pin[idx] = False
+
+    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
@@ -1,40 +0,0 @@
-"""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",
-]
@@ -1,533 +0,0 @@
-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,
-        write_positions: Optional[Tensor] = None,
-    ) -> 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,
-        write_positions: Optional[Tensor] = None,
-    ) -> 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,
-        write_positions: Optional[Tensor] = None,
-    ):
-        self._cache = cache
-        self._batch_indices = batch_indices
-        self._total_len = total_len
-        self._write_positions = write_positions
-
-    def write(self, layer_id: int, k: Tensor, v: Tensor):
-        seq_len = k.size(1)
-        indices = self._batch_indices
-        if self._write_positions is not None and seq_len == 1:
-            pos = self._write_positions
-            self._cache.k[layer_id, indices, pos] = k.squeeze(1)
-            self._cache.v[layer_id, indices, pos] = v.squeeze(1)
-            for s, p in zip(indices.tolist(), pos.tolist()):
-                cur = self._cache._slot_len.get(s, 0)
-                if p + 1 > cur:
-                    self._cache._slot_len[s] = p + 1
-        else:
-            start_pos = self._total_len - seq_len
-            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 task_cached(self, task_id: str) -> int:
-        slot = self._task_slot.get(task_id)
-        if slot is None:
-            return 0
-        return self._slot_len.get(slot, 0)
-
-    def bind_tasks(
-        self,
-        task_ids: List[str],
-        total_len: int,
-        device: torch.device,
-        write_positions: Optional[Tensor] = None,
-    ) -> 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, write_positions=write_positions
-        )
@@ -1,127 +0,0 @@
-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)
-        freq_penalties = torch.tensor(
-            [t.frequency_penalty for t in tasks], device=self.device
-        )
-
-        history_lists = []
-        mask_lists = []
-        for t in tasks:
-            window = t.rep_window
-            prompt_part = t.prompt_ids[-window:]
-            ids = prompt_part + t.output_ids
-            history_lists.append(ids)
-            mask_lists.append([True] * len(ids))
-
-        max_len = max(len(h) for h in history_lists)
-        padded_ids = torch.zeros(
-            len(tasks), max_len, dtype=torch.long, device=self.device
-        )
-        padded_mask = torch.zeros(
-            len(tasks), max_len, dtype=torch.bool, device=self.device
-        )
-        for i, (h, m) in enumerate(zip(history_lists, mask_lists)):
-            padded_ids[i, : len(h)] = torch.tensor(
-                h, dtype=torch.long, device=self.device
-            )
-            padded_mask[i, : len(m)] = torch.tensor(
-                m, dtype=torch.bool, 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,
-                    write_positions=position_ids,
-                ),
-                position_ids=position_ids.unsqueeze(1),
-            )
-            logits = outputs["logits"][:, -1, :]
-
-        return sample(
-            logits,
-            temperature=temperatures,
-            top_k=top_ks,
-            top_p=top_ps,
-            frequency_penalty=freq_penalties,
-            input_ids=padded_ids,
-            input_mask=padded_mask,
-        ).tolist()
@@ -1,199 +0,0 @@
-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
-                            )
-
-                decode_tasks = self._task_mgr.get_active_tasks()
-
-                valid: List[Task] = []
-                for t in sorted(decode_tasks, key=lambda t: t.task_id):
-                    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
-                        new_text = t.decode_new_token(self._task_mgr.tokenizer)
-                        if new_text:
-                            self._task_mgr.invoke_callback(t.task_id, new_text)
-
-                    for t in valid:
-                        if t.is_finished(stop_ids):
-                            remaining = t.flush_remaining(self._task_mgr.tokenizer)
-                            if remaining:
-                                self._task_mgr.invoke_callback(t.task_id, remaining)
-                            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()
@@ -1,286 +0,0 @@
-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 StreamDecoder:
-    """Incremental decoder for byte-level BPE streaming.
-
-    Byte-level BPE may split a single Unicode character (e.g. em-dash,
-    smart quotes) across multiple tokens. Decoding such a token in
-    isolation produces U+FFFD (replacement char). This decoder
-    accumulates token IDs and only emits text once the trailing
-    characters are complete, buffering incomplete multi-byte sequences
-    until the next token arrives.
-    """
-
-    __slots__ = ("_tokenizer", "_ids", "_emitted")
-
-    def __init__(self, tokenizer: AutoTokenizer):
-        self._tokenizer = tokenizer
-        self._ids: List[int] = []
-        self._emitted: str = ""
-
-    def push(self, token_id: int) -> str:
-        """Append a token ID and return newly completed text.
-
-        Returns "" while a multi-byte character is still incomplete.
-        """
-        self._ids.append(token_id)
-        full = self._tokenizer.decode(self._ids, skip_special_tokens=True)
-        if full.endswith("\ufffd"):
-            return ""
-        if len(full) > len(self._emitted):
-            diff = full[len(self._emitted) :]
-            self._emitted = full
-            return diff
-        return ""
-
-
-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,
-        frequency_penalty: float = 0.0,
-        rep_window: int = 64,
-    ):
-        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.frequency_penalty = frequency_penalty
-        self.rep_window = rep_window
-
-        self.status = TaskStatus.PENDING
-        self.output_ids: List[int] = []
-        self.input_tokens: int = 0
-        self.output_tokens: int = 0
-        self.arrival_time = time.time()
-        self.finish_time: Optional[float] = None
-        self._decoder: Optional[StreamDecoder] = None
-
-    def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
-        """Decode the last appended output token, buffering incomplete
-        multi-byte sequences across calls.
-
-        Lazily creates a :class:`StreamDecoder` on first use.
-        """
-        if self._decoder is None:
-            self._decoder = StreamDecoder(tokenizer)
-        return self._decoder.push(self.output_ids[-1])
-
-    def flush_remaining(self, tokenizer: AutoTokenizer) -> str:
-        """Emit any text still buffered in the decoder.
-
-        Called when generation terminates (max_tokens reached, stop
-        sequence, or external removal) to avoid dropping a final
-        incomplete-looking fragment that is actually complete when
-        adjacent to the stop token.
-        """
-        if self._decoder is None or not self.output_ids:
-            return ""
-        full = tokenizer.decode(self.output_ids, skip_special_tokens=True)
-        if len(full) > len(self._decoder._emitted):
-            diff = full[len(self._decoder._emitted) :]
-            self._decoder._emitted = full
-            return diff
-        return ""
-
-    @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,
-        frequency_penalty: float = 0.0,
-        rep_window: int = 64,
-        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,
-            frequency_penalty=frequency_penalty,
-            rep_window=rep_window,
-        )
-
-        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()
@@ -1,32 +1,130 @@
-"""Unified inference engine for continuous batching."""
+"""Unified inference engine for continuous batching.
+
+Layers:
+  - GenerationParams:    Immutable value object for sampling parameters.
+  - GenerationRequest:   User-facing request DTO with validation.
+  - _Result:             Thread-safe token accumulator (Observer pattern).
+  - InferenceEngine:     Facade over InferenceScheduler + async wrapper.
+"""

 import asyncio
 import gc
 import threading
+from dataclasses import dataclass
 from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, Union

 import torch
 import torch.nn as nn

-from astrai.inference.core.cache import KVCache
-from astrai.inference.core.scheduler import InferenceScheduler
-from astrai.inference.core.task import STOP
+from astrai.inference.cache import STOP
+from astrai.inference.scheduler import InferenceScheduler
 from astrai.tokenize import AutoTokenizer


-class GenerateResult:
-    """Thread-safe token accumulator for streaming and non-streaming modes."""
+@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.
+    """

    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[Tuple[int, str]] = []
+        self.tokens: List[str] = []
        self.results: List[str] = [""] * count
        self._done: List[bool] = [False] * count
        self._completed = 0
        self._total = count

    def append(self, token: str, 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((idx, token))
            if token is not STOP:
@@ -39,6 +137,11 @@ class GenerateResult:
            self._event.set()

    def pop_all(self) -> List[Tuple[int, str]]:
+        """Returns and clears all accumulated (idx, token) pairs.
+
+        Returns:
+            List of (index, token_string) tuples since the last call.
+        """
        with self._cond:
            out = self.tokens.copy()
            self.tokens.clear()
@@ -47,83 +150,78 @@ class GenerateResult:
            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, timeout: float = 300.0):
+    def wait_completion(self) -> None:
+        """Blocks until all tasks complete (non-streaming).
+
+        Uses a Condition to sleep efficiently instead of busy-waiting.
+        The calling thread is parked until a STOP signal arrives.
+        """
        with self._cond:
-            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)"
-                )
+            self._cond.wait_for(lambda: self._completed >= self._total)

    def get_results(self) -> List[str]:
+        """Returns all accumulated results for non-streaming mode.
+
+        Returns:
+            List of complete generated strings, one per task index.
+        """
        with self._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,
-        frequency_penalty: float = 0.0,
-        rep_window: int = 64,
-        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 non-negative number")
-        if not (
-            isinstance(frequency_penalty, (int, float))
-            and -2.0 <= frequency_penalty <= 2.0
-        ):
-            raise ValueError("frequency_penalty must be between -2.0 and 2.0")
-        if not (isinstance(rep_window, int) and rep_window > 0):
-            raise ValueError("rep_window must be a positive integer")
-
-        self.messages = messages
-        self.top_k = top_k
-        self.top_p = top_p
-        self.temperature = temperature
-        self.max_tokens = max_tokens
-        self.frequency_penalty = frequency_penalty
-        self.rep_window = rep_window
-        self.stream = stream
-
-
 class InferenceEngine:
-    """Unified inference engine backed by continuous-batching scheduler."""
+    """Unified inference engine backed by continuous-batching scheduler.
+
+    Usage:
+        with InferenceEngine(model, tokenizer) as engine:
+            for token in engine.generate("hello", stream=True):
+                print(token, end="")
+
+            text = engine.generate("hello")
+    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: AutoTokenizer,
        max_batch_size: int = 1,
+        max_queue_size: int = 64,
+        request_timeout: float = 60.0,
        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.
+            page_size: Number of tokens per KV cache page.
+        """
        self.model = model
        self.tokenizer = tokenizer
        self.scheduler = InferenceScheduler(
            model=self.model,
            tokenizer=self.tokenizer,
            max_batch_size=max_batch_size,
+            max_queue_size=max_queue_size,
+            request_timeout=request_timeout,
            max_seq_len=max_seq_len,
            max_prompt_len=max_prompt_len,
-            cache=cache,
+            page_size=page_size,
        )

        self.scheduler.start()
@@ -139,58 +237,68 @@ class InferenceEngine:
        self,
        prompt: Union[str, List[str]],
        stream: bool = False,
-        max_tokens: Optional[int] = None,
+        max_tokens: int = 1024,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
-        frequency_penalty: float = 0.0,
-        rep_window: int = 64,
+        timeout: Optional[float] = None,
    ) -> Union[Generator, 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.
+            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).
+            timeout: Per-request timeout in seconds (None = use scheduler default).
+
+        Returns:
+            stream=False, single prompt: str
+            stream=False, batch: List[str]
+            stream=True, single prompt: Generator[str, None, None]
+            stream=True, batch: Generator[Tuple[int, str], None, None]
+        """
        is_batch = isinstance(prompt, list)
        prompts = prompt if is_batch else [prompt]

        if stream:
            return self._generate_streaming(
-                prompts,
-                is_batch,
-                max_tokens,
-                temperature,
-                top_p,
-                top_k,
-                frequency_penalty,
-                rep_window,
+                prompts, is_batch, max_tokens, temperature, top_p, top_k, timeout
            )
        else:
            return self._generate_non_streaming(
-                prompts,
-                is_batch,
-                max_tokens,
-                temperature,
-                top_p,
-                top_k,
-                frequency_penalty,
-                rep_window,
+                prompts, is_batch, max_tokens, temperature, top_p, top_k, timeout
            )

    def generate_async(
        self,
        prompt: str,
-        max_tokens: Optional[int] = None,
+        max_tokens: int = 1024,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
-        frequency_penalty: float = 0.0,
-        rep_window: int = 64,
+        timeout: Optional[float] = None,
    ) -> 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.
+            timeout: Per-request timeout in seconds.
+
+        Yields:
+            Decoded token strings as they are generated.
+        """
        sync_gen = self._generate_streaming(
-            [prompt],
-            False,
-            max_tokens,
-            temperature,
-            top_p,
-            top_k,
-            frequency_penalty,
-            rep_window,
+            [prompt], False, max_tokens, temperature, top_p, top_k, timeout
        )

        async def _agen():
@@ -205,6 +313,14 @@ 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:
@@ -213,74 +329,75 @@ 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.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            top_k=request.top_k,
-            frequency_penalty=request.frequency_penalty,
-            rep_window=request.rep_window,
+            max_tokens=request.params.max_tokens,
+            temperature=request.params.temperature,
+            top_p=request.params.top_p,
+            top_k=request.params.top_k,
        )

-    def _submit_tasks(
+    def _generate_streaming(
        self,
        prompts: List[str],
-        max_tokens: Optional[int],
+        is_batch: bool,
+        max_tokens: int,
        temperature: float,
        top_p: float,
        top_k: int,
-        frequency_penalty: float,
-        rep_window: int,
-    ) -> Tuple[GenerateResult, List[str]]:
+        timeout: Optional[float] = None,
+    ) -> Generator:
+        """Internal streaming generator.
+
+        Polls the _Result accumulator in a loop, yielding tokens as they arrive.
+        Single prompt yields raw token strings; batch yields (idx, token) tuples.
+
+        Args:
+            prompts: List of prompts.
+            is_batch: If True, yields (idx, token) tuples; else yields raw tokens.
+            max_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature.
+            top_p: Nucleus sampling threshold.
+            top_k: Top-k sampling count.
+            timeout: Per-request timeout in seconds.
+
+        Yields:
+            Single prompt: decoded token strings.
+            Batch: (sequence_index, token_string) tuples.
+        """
        n = len(prompts)
-        result = GenerateResult(count=n)
+        result = _Result(count=n)
        task_ids = []
+
+        try:
            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,
-                frequency_penalty=frequency_penalty,
-                rep_window=rep_window,
-                stream_callback=cb,
+                    stream_callback=lambda tok, idx=i: result.append(tok, idx),
+                    timeout=timeout,
                )
                task_ids.append(task_id)
-        return result, task_ids
+        except RuntimeError:
+            for tid in task_ids:
+                self.scheduler.remove_task(tid)
+            raise

-    @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: Optional[int],
-        temperature: float,
-        top_p: float,
-        top_k: int,
-        frequency_penalty: float,
-        rep_window: int,
-    ) -> Generator:
-        result, task_ids = self._submit_tasks(
-            prompts,
-            max_tokens,
-            temperature,
-            top_p,
-            top_k,
-            frequency_penalty,
-            rep_window,
-        )
-        n = len(prompts)
        remaining = n
        finished = [False] * n

@@ -297,7 +414,8 @@ class InferenceEngine:
                        else:
                            yield (idx, token) if is_batch else token
                    if remaining > 0:
-                        result.wait(timeout=0.05)
+                        if not result.wait(timeout=0.05):
+                            pass
            finally:
                for tid in task_ids:
                    self.scheduler.remove_task(tid)
@@ -308,40 +426,70 @@ class InferenceEngine:
        self,
        prompts: List[str],
        is_batch: bool,
-        max_tokens: Optional[int],
+        max_tokens: int,
        temperature: float,
        top_p: float,
        top_k: int,
-        frequency_penalty: float,
-        rep_window: int,
+        timeout: Optional[float] = None,
    ) -> Union[str, List[str]]:
-        result, task_ids = self._submit_tasks(
-            prompts,
-            max_tokens,
-            temperature,
-            top_p,
-            top_k,
-            frequency_penalty,
-            rep_window,
-        )
+        """Internal non-streaming generator.
+
+        Submits all prompts to the scheduler and waits for all to complete.
+
+        Args:
+            prompts: List of prompt strings.
+            is_batch: Whether multiple prompts were provided.
+            max_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature.
+            top_p: Nucleus sampling threshold.
+            top_k: Top-k sampling count.
+            timeout: Per-request timeout in seconds.
+
+        Returns:
+            Single string for one prompt, list of strings for batch.
+        """
+        result = _Result(count=len(prompts))
+        task_ids = []

        try:
-            result.wait_completion()
-        except TimeoutError:
+            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),
+                    timeout=timeout,
+                )
+                task_ids.append(task_id)
+        except RuntimeError:
            for tid in task_ids:
                self.scheduler.remove_task(tid)
            raise

-        for tid in task_ids:
-            self.scheduler.remove_task(tid)
+        result.wait_completion()
+
+        for task_id in task_ids:
+            self.scheduler.remove_task(task_id)

        res = result.get_results()
        return res if is_batch else res[0]

    def get_stats(self) -> Dict[str, Any]:
+        """Returns current engine statistics.
+
+        Returns:
+            Dict with total_tasks, total_tokens, active_tasks, waiting_queue.
+        """
        return self.scheduler.get_stats()

-    def shutdown(self):
+    def shutdown(self) -> None:
+        """Shuts down the engine, stops the scheduler, and frees GPU memory."""
        self.scheduler.stop()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
@@ -1,343 +0,0 @@
-"""Composable sampling strategies for logit transformation.
-
-Implements the Strategy pattern: each sampling technique
-(temperature, top-k, top-p, frequency penalty) is a pluggable
-strategy that can be composed into a pipeline.
-
-All strategies accept both scalar and per-sample tensor
-parameters, so a single pipeline works for any batch size.
-"""
-
-from abc import ABC, abstractmethod
-from typing import List, Optional, Union
-
-import torch
-from torch import Tensor
-
-
-class BaseSamplingStrategy(ABC):
-    """Abstract base for a logit transformation strategy."""
-
-    @abstractmethod
-    def apply(
-        self,
-        logits: Tensor,
-        filter_value: float = -float("inf"),
-        input_ids: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        """Applies the strategy to logits.
-
-        Args:
-            logits: Raw logits tensor (batch, vocab_size).
-            filter_value: Value assigned to filtered-out positions.
-            input_ids: Previously generated token IDs ``[batch, seq_len]``,
-                padded with 0. Used by frequency penalty.
-            input_mask: Boolean mask ``[batch, seq_len]``, True for real
-                tokens, False for padding. Used to exclude padding from
-                penalty computation.
-
-        Returns:
-            Transformed logits tensor.
-        """
-        raise NotImplementedError
-
-
-class TemperatureStrategy(BaseSamplingStrategy):
-    """Divides logits by temperature to control randomness.
-
-    Args:
-        temperature: Scalar or ``[batch]`` tensor.
-    """
-
-    def __init__(self, temperature: Union[float, Tensor] = 1.0):
-        self.temperature = temperature
-
-    def apply(
-        self,
-        logits: Tensor,
-        filter_value: float = -float("inf"),
-        input_ids: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-    ) -> 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
-        elif t != 1.0:
-            logits = logits / max(t, 1e-8)
-        return logits
-
-
-class TopKStrategy(BaseSamplingStrategy):
-    """Keeps only the top-k logits, setting the rest to filter_value.
-
-    Args:
-        top_k: Scalar or ``[batch]`` tensor (0 disables).
-    """
-
-    def __init__(self, top_k: Union[int, Tensor] = 0):
-        self.top_k = top_k
-
-    def apply(
-        self,
-        logits: Tensor,
-        filter_value: float = -float("inf"),
-        input_ids: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-    ) -> 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
-            max_k = min(max_k, logits.size(-1))
-            values, _ = torch.topk(logits, max_k, dim=-1)
-            per_row_k = tk.clamp(max=max_k)
-            thresholds = torch.full_like(logits[..., -1:], -float("inf"))
-            positive = per_row_k > 0
-            if positive.any():
-                row_idx = torch.arange(logits.size(0), device=logits.device)[positive]
-                thresholds[positive] = values[
-                    row_idx, per_row_k[positive] - 1
-                ].unsqueeze(-1)
-            logits[logits < thresholds] = filter_value
-            return logits
-        if tk > 0:
-            k = min(tk, logits.size(-1))
-            thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
-            logits[logits < thresholds] = filter_value
-        return logits
-
-
-class TopPStrategy(BaseSamplingStrategy):
-    """Nucleus (top-p) filtering: keeps the smallest set of tokens whose
-    cumulative probability exceeds top_p.
-
-    Args:
-        top_p: Scalar or ``[batch]`` tensor (1.0 disables).
-    """
-
-    def __init__(self, top_p: Union[float, Tensor] = 1.0):
-        self.top_p = top_p
-
-    def _apply(
-        self, logits: 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
-        remove[..., 1:] = remove[..., :-1].clone()
-        remove[..., 0] = False
-        mask = torch.zeros_like(logits, dtype=torch.bool)
-        mask.scatter_(1, sorted_indices, remove)
-        logits[mask] = filter_value
-        return logits
-
-    def apply(
-        self,
-        logits: Tensor,
-        filter_value: float = -float("inf"),
-        input_ids: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        tp = self.top_p
-        if isinstance(tp, Tensor):
-            tp = tp.to(logits.device, non_blocking=True)
-            if (tp < 1.0).any():
-                logits = self._apply(logits, tp.view(-1, 1), filter_value)
-        elif tp < 1.0:
-            logits = self._apply(logits, tp, filter_value)
-        return logits
-
-
-class FrequencyPenaltyStrategy(BaseSamplingStrategy):
-    """Penalizes tokens based on how many times they appeared in history.
-
-    Subtracts ``penalty * count(token)`` from each token's logit, where
-    ``count(token)`` is the number of occurrences in the generation history
-    (prompt + output). A penalty of ``0.0`` disables the strategy.
-
-    Unlike repetition penalty (which only checks *presence*), frequency
-    penalty scales linearly with occurrence count: the first use is
-    penalized once, the third use three times. This allows natural
-    repetition of common words while suppressing degenerate loops.
-
-    Reference: OpenAI API ``frequency_penalty`` parameter.
-
-    Args:
-        penalty: Scalar or ``[batch]`` tensor (0.0 disables, range -2.0~2.0).
-    """
-
-    def __init__(self, penalty: Union[float, Tensor] = 0.0):
-        self.penalty = penalty
-
-    def apply(
-        self,
-        logits: Tensor,
-        filter_value: float = -float("inf"),
-        input_ids: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        if input_ids is None:
-            return logits
-
-        p = self.penalty
-        if isinstance(p, Tensor):
-            p = p.to(logits.device, non_blocking=True).view(-1, 1)
-            if (p == 0.0).all():
-                return logits
-        elif p == 0.0:
-            return logits
-
-        input_ids = input_ids.to(logits.device, non_blocking=True)
-
-        if input_mask is not None:
-            input_mask = input_mask.to(logits.device, non_blocking=True)
-            masked_ids = input_ids.clone()
-            masked_ids[~input_mask] = -1
-        else:
-            masked_ids = input_ids
-
-        batch_sz, seq_len = masked_ids.shape
-        vocab_size = logits.size(-1)
-
-        if isinstance(p, Tensor):
-            penalty_per_row = p.expand(batch_sz, 1)
-        else:
-            penalty_per_row = torch.full(
-                (batch_sz, 1), float(p), device=logits.device, dtype=logits.dtype
-            )
-
-        counts = torch.zeros(
-            batch_sz, vocab_size, device=logits.device, dtype=logits.dtype
-        )
-        valid_mask = masked_ids >= 0
-        if valid_mask.any():
-            valid_ids = masked_ids[valid_mask]
-            row_indices = (
-                torch.arange(batch_sz, device=logits.device)
-                .unsqueeze(1)
-                .expand_as(masked_ids)[valid_mask]
-            )
-            counts.index_put_(
-                (row_indices, valid_ids),
-                torch.ones_like(valid_ids, dtype=logits.dtype),
-                accumulate=True,
-            )
-
-        return logits - penalty_per_row * counts
-
-
-class SamplingPipeline(BaseSamplingStrategy):
-    """Composes multiple sampling strategies into a single transformation.
-
-    Strategies are applied sequentially in the order they are provided,
-    matching the original temperature -> top-k -> top-p ordering.
-
-    Usage::
-
-        pipeline = SamplingPipeline([
-            TemperatureStrategy(0.8),
-            TopKStrategy(50),
-            TopPStrategy(0.95),
-        ])
-        logits = pipeline.apply(logits)
-        token = pipeline.sample(logits)       # softmax + multinomial
-    """
-
-    def __init__(self, strategies: List[BaseSamplingStrategy]):
-        self.strategies = strategies
-
-    def apply(
-        self,
-        logits: Tensor,
-        filter_value: float = -float("inf"),
-        input_ids: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        for strategy in self.strategies:
-            logits = strategy.apply(logits, filter_value, input_ids, input_mask)
-        return logits
-
-    @staticmethod
-    def _is_greedy(temperature: Union[float, Tensor]) -> bool:
-        if isinstance(temperature, Tensor):
-            return temperature.numel() == 1 and temperature.item() == 0
-        return temperature == 0
-
-    @torch.inference_mode()
-    def sample(
-        self,
-        logits: Tensor,
-        filter_value: float = -float("inf"),
-        input_ids: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        """Apply strategies then sample (softmax + multinomial).
-
-        Short-circuits to ``argmax`` when temperature is exactly 0
-        (deterministic / greedy decode).
-
-        Args:
-            logits: Raw logits ``[batch, vocab_size]``.
-            input_ids: Previously generated token IDs ``[batch, seq_len]``.
-            input_mask: Boolean mask for ``input_ids`` padding.
-
-        Returns:
-            Sampled token IDs ``[batch]``.
-        """
-        for s in self.strategies:
-            if isinstance(s, TemperatureStrategy) and self._is_greedy(s.temperature):
-                return logits.argmax(dim=-1)
-            break
-
-        return torch.multinomial(
-            torch.softmax(
-                self.apply(logits, filter_value, input_ids, input_mask), dim=-1
-            ),
-            num_samples=1,
-        ).squeeze(-1)
-
-
-@torch.inference_mode()
-def sample(
-    logits: Tensor,
-    temperature: Union[float, Tensor] = 1.0,
-    top_k: Union[int, Tensor] = 0,
-    top_p: Union[float, Tensor] = 1.0,
-    frequency_penalty: Union[float, Tensor] = 0.0,
-    input_ids: Optional[Tensor] = None,
-    input_mask: Optional[Tensor] = None,
-    filter_value: float = -float("inf"),
-) -> Tensor:
-    """Apply sampling strategies then sample (softmax + multinomial).
-
-    Shortcut for ``SamplingPipeline(...).sample(logits)``.
-
-    When **temperature** is exactly 0 (scalar or single-element tensor)
-    the function short-circuits to ``argmax`` for deterministic decode.
-
-    Args:
-        logits: Raw logits ``[batch, vocab_size]``.
-        frequency_penalty: Penalty per occurrence for repeated tokens
-            (0.0 disables, range -2.0~2.0).
-        input_ids: Previously generated token IDs ``[batch, seq_len]``.
-        input_mask: Boolean mask for ``input_ids`` padding.
-
-    Returns:
-        Sampled token IDs ``[batch]``.
-    """
-    if SamplingPipeline._is_greedy(temperature):
-        return logits.argmax(dim=-1)
-    return SamplingPipeline(
-        [
-            TemperatureStrategy(temperature),
-            TopKStrategy(top_k),
-            TopPStrategy(top_p),
-            FrequencyPenaltyStrategy(frequency_penalty),
-        ]
-    ).sample(logits, filter_value, input_ids, input_mask)
@@ -0,0 +1,178 @@
+"""Composable sampling strategies for logit transformation.
+
+Implements the Strategy pattern: each sampling technique
+(temperature, top-k, top-p) is a pluggable strategy that
+can be composed into a pipeline.
+
+All strategies accept both scalar and per-sample tensor
+parameters, so a single pipeline works for any batch size.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Union
+
+import torch
+from torch import Tensor
+
+
+class BaseSamplingStrategy(ABC):
+    """Abstract base for a logit transformation strategy."""
+
+    @abstractmethod
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
+        """Applies the strategy to logits.
+
+        Args:
+            logits: Raw logits tensor (batch, vocab_size).
+            filter_value: Value assigned to filtered-out positions.
+
+        Returns:
+            Transformed logits tensor.
+        """
+
+
+class TemperatureStrategy(BaseSamplingStrategy):
+    """Divides logits by temperature to control randomness.
+
+    Args:
+        temperature: Scalar or ``[batch]`` tensor.
+    """
+
+    def __init__(self, temperature: Union[float, Tensor] = 1.0):
+        self.temperature = temperature
+
+    def apply(self, logits, filter_value=-float("inf")):
+        t = self.temperature
+        if isinstance(t, Tensor):
+            if (t != 1.0).any():
+                logits = logits / t.to(logits.device, non_blocking=True).view(-1, 1)
+        elif t != 1.0:
+            logits = logits / t
+        return logits
+
+
+class TopKStrategy(BaseSamplingStrategy):
+    """Keeps only the top-k logits, setting the rest to filter_value.
+
+    Args:
+        top_k: Scalar or ``[batch]`` tensor (0 disables).
+    """
+
+    def __init__(self, top_k: Union[int, Tensor] = 0):
+        self.top_k = top_k
+
+    def apply(self, logits, filter_value=-float("inf")):
+        tk = self.top_k
+        if isinstance(tk, Tensor):
+            max_k = int(tk.max().item())
+            if max_k <= 0:
+                return logits
+            k = min(max_k, logits.size(-1))
+        elif tk > 0:
+            k = min(tk, logits.size(-1))
+        else:
+            return logits
+        thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
+        logits[logits < thresholds] = filter_value
+        return logits
+
+
+class TopPStrategy(BaseSamplingStrategy):
+    """Nucleus (top-p) filtering: keeps the smallest set of tokens whose
+    cumulative probability exceeds top_p.
+
+    Args:
+        top_p: Scalar or ``[batch]`` tensor (1.0 disables).
+    """
+
+    def __init__(self, top_p: Union[float, Tensor] = 1.0):
+        self.top_p = top_p
+
+    def _apply(self, logits, top_p, filter_value):
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+        remove = cum_probs > top_p
+        remove[..., 1:] = remove[..., :-1].clone()
+        remove[..., 0] = False
+        mask = torch.zeros_like(logits, dtype=torch.bool)
+        mask.scatter_(1, sorted_indices, remove)
+        logits[mask] = filter_value
+        return logits
+
+    def apply(self, logits, filter_value=-float("inf")):
+        tp = self.top_p
+        if isinstance(tp, Tensor):
+            tp = tp.to(logits.device, non_blocking=True)
+            if (tp < 1.0).any():
+                logits = self._apply(logits, tp.view(-1, 1), filter_value)
+        elif tp < 1.0:
+            logits = self._apply(logits, tp, filter_value)
+        return logits
+
+
+class SamplingPipeline(BaseSamplingStrategy):
+    """Composes multiple sampling strategies into a single transformation.
+
+    Strategies are applied sequentially in the order they are provided,
+    matching the original temperature -> top-k -> top-p ordering.
+
+    Usage::
+
+        pipeline = SamplingPipeline([
+            TemperatureStrategy(0.8),
+            TopKStrategy(50),
+            TopPStrategy(0.95),
+        ])
+        logits = pipeline.apply(logits)
+        token = pipeline.sample(logits)       # softmax + multinomial
+    """
+
+    def __init__(self, strategies: List[BaseSamplingStrategy]):
+        self.strategies = strategies
+
+    def apply(self, logits, filter_value=-float("inf")):
+        for strategy in self.strategies:
+            logits = strategy.apply(logits, filter_value)
+        return logits
+
+    @torch.no_grad()
+    def sample(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
+        """Apply strategies then sample (softmax + multinomial).
+
+        Args:
+            logits: Raw logits ``[batch, vocab_size]``.
+
+        Returns:
+            Sampled token IDs ``[batch]``.
+        """
+        return torch.multinomial(
+            torch.softmax(self.apply(logits, filter_value), dim=-1),
+            num_samples=1,
+        ).squeeze(-1)
+
+
+@torch.inference_mode()
+def sample(
+    logits: Tensor,
+    temperature: Union[float, Tensor] = 1.0,
+    top_k: Union[int, Tensor] = 0,
+    top_p: Union[float, Tensor] = 1.0,
+    filter_value: float = -float("inf"),
+) -> Tensor:
+    """Apply sampling strategies then sample (softmax + multinomial).
+
+    Shortcut for ``SamplingPipeline(...).sample(logits)``.
+
+    Args:
+        logits: Raw logits ``[batch, vocab_size]``.
+
+    Returns:
+        Sampled token IDs ``[batch]``.
+    """
+    return SamplingPipeline(
+        [
+            TemperatureStrategy(temperature),
+            TopKStrategy(top_k),
+            TopPStrategy(top_p),
+        ]
+    ).sample(logits, filter_value)
@@ -0,0 +1,505 @@
+"""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.deadline: float = 0.0
+        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_queue_size: int = 64,
+        request_timeout: float = 60.0,
+        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_queue_size = max_queue_size
+        self.request_timeout = request_timeout
+        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._draining = False
+        self._task_event = threading.Event()
+        self._lock = threading.Lock()
+
+        self._total_tasks = 0
+        self._total_tokens = 0
+        self._total_requests = 0
+        self._total_rejected = 0
+        self._total_timeouts = 0
+        self._request_latencies: List[float] = []
+
+    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,
+        timeout: Optional[float] = 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,
+        )
+        task.deadline = time.time() + (
+            timeout if timeout is not None else self.request_timeout
+        )
+
+        with self._lock:
+            if len(self.waiting_queue) >= self.max_queue_size:
+                self._total_rejected += 1
+                raise RuntimeError("Request queue is full")
+            self.waiting_queue.append(task)
+            self._total_requests += 1
+            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 _abort_task(self, task: Task) -> None:
+        task.status = TaskStatus.ABORTED
+        task.finish_time = time.time()
+        if not task._pages_freed:
+            self._free_pages(task.page_table)
+            task.page_table.clear()
+            task.n_pages = 0
+            task._pages_freed = True
+        if task.stream_callback:
+            task.stream_callback(STOP)
+
+    def _abort_expired_tasks(self) -> None:
+        now = time.time()
+        alive = []
+        for t in self.active_tasks:
+            if now > t.deadline:
+                self._abort_task(t)
+                self._total_timeouts += 1
+            else:
+                alive.append(t)
+        self.active_tasks = alive
+
+        with self._lock:
+            keep = []
+            for t in self.waiting_queue:
+                if now > t.deadline:
+                    t.status = TaskStatus.ABORTED
+                    if t.stream_callback:
+                        t.stream_callback(STOP)
+                    self._total_timeouts += 1
+                else:
+                    keep.append(t)
+            self.waiting_queue = keep
+
+    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
+                self._request_latencies.append(task.finish_time - task.arrival_time)
+                if len(self._request_latencies) > 1000:
+                    self._request_latencies.pop(0)
+
+        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 or (self._draining and self.active_tasks):
+                self._abort_expired_tasks()
+                self._remove_finished_tasks()
+                if not self._draining:
+                    self._refill_active_batch()
+
+                if not self.active_tasks:
+                    if self._draining:
+                        break
+                    if 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, timeout: float = 30.0) -> None:
+        self._draining = True
+        self._running = False
+        self._task_event.set()
+        if hasattr(self, "_loop_thread"):
+            self._loop_thread.join(timeout=timeout)
+
+        for task in self.active_tasks:
+            if not task._pages_freed:
+                self._free_pages(task.page_table)
+                task._pages_freed = True
+            if task.stream_callback:
+                task.stream_callback(STOP)
+
+        with self._lock:
+            for task in self.waiting_queue:
+                task.status = TaskStatus.ABORTED
+                if task.stream_callback:
+                    task.stream_callback(STOP)
+            self.waiting_queue.clear()
+            self.active_tasks.clear()
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def get_stats(self) -> Dict[str, Any]:
+        latencies = self._request_latencies
+        sorted_lat = sorted(latencies) if latencies else []
+        n = len(sorted_lat)
+        p50 = sorted_lat[n // 2] if n > 0 else 0.0
+        p95 = sorted_lat[int(n * 0.95)] if n > 0 else 0.0
+        p99 = sorted_lat[int(n * 0.99)] if n > 0 else 0.0
+
+        cache = self.page_cache
+        total_lookups = cache.lookup_hits + cache.lookup_misses
+        hit_rate = cache.lookup_hits / total_lookups if total_lookups > 0 else 0.0
+
+        return {
+            "total_requests": self._total_requests,
+            "total_rejected": self._total_rejected,
+            "total_timeouts": self._total_timeouts,
+            "total_tokens": self._total_tokens,
+            "active_tasks": len(self.active_tasks),
+            "waiting_queue": len(self.waiting_queue),
+            "latency_p50": p50,
+            "latency_p95": p95,
+            "latency_p99": p99,
+            "cache_hit_rate": hit_rate,
+            "cache_hits": cache.lookup_hits,
+            "cache_misses": cache.lookup_misses,
+        }
@@ -0,0 +1,528 @@
+"""
+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 PlainTextResponse, 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,
+        max_queue_size=64,
+        request_timeout=60.0,
+    )
+
+
+@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.get("/metrics")
+async def metrics():
+    s = _get_engine().get_stats()
+    lines = [
+        "# HELP astrai_requests_total Total requests received",
+        "# TYPE astrai_requests_total counter",
+        f'astrai_requests_total{{status="accepted"}} {s["total_requests"]}',
+        f'astrai_requests_total{{status="rejected"}} {s["total_rejected"]}',
+        f'astrai_requests_total{{status="timeout"}} {s["total_timeouts"]}',
+        "# HELP astrai_tokens_generated Total generated tokens",
+        "# TYPE astrai_tokens_generated counter",
+        f"astrai_tokens_generated {s['total_tokens']}",
+        "# HELP astrai_active_tasks Currently active tasks",
+        "# TYPE astrai_active_tasks gauge",
+        f"astrai_active_tasks {s['active_tasks']}",
+        "# HELP astrai_queue_depth Waiting queue depth",
+        "# TYPE astrai_queue_depth gauge",
+        f"astrai_queue_depth {s['waiting_queue']}",
+        "# HELP astrai_request_latency_seconds Request latency quantiles",
+        "# TYPE astrai_request_latency_seconds gauge",
+        f'astrai_request_latency_seconds{{quantile="0.5"}} {s["latency_p50"]:.3f}',
+        f'astrai_request_latency_seconds{{quantile="0.95"}} {s["latency_p95"]:.3f}',
+        f'astrai_request_latency_seconds{{quantile="0.99"}} {s["latency_p99"]:.3f}',
+        "# HELP astrai_cache_hit_rate Prefix cache hit ratio",
+        "# TYPE astrai_cache_hit_rate gauge",
+        f"astrai_cache_hit_rate {s['cache_hit_rate']:.3f}",
+        "# HELP astrai_cache_lookups_total Prefix cache page lookups",
+        "# TYPE astrai_cache_lookups_total counter",
+        f'astrai_cache_lookups_total{{result="hit"}} {s["cache_hits"]}',
+        f'astrai_cache_lookups_total{{result="miss"}} {s["cache_misses"]}',
+    ]
+    return PlainTextResponse("\n".join(lines) + "\n")
+
+
+@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:
+        try:
+            agen = engine.generate_async(
+                prompt=prompt,
+                max_tokens=request.max_tokens,
+                temperature=request.temperature,
+                top_p=request.top_p,
+                top_k=request.top_k,
+            )
+        except RuntimeError as e:
+            raise HTTPException(status_code=503, detail=str(e))
+
+        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] = []
+    try:
+        agen = engine.generate_async(
+            prompt=prompt,
+            max_tokens=request.max_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            top_k=request.top_k,
+        )
+    except RuntimeError as e:
+        raise HTTPException(status_code=503, detail=str(e))
+    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,
+    )
@@ -1,18 +1,12 @@
 from astrai.model.automodel import AutoModel
-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.module import (
+    GQA,
+    MLP,
+    DecoderBlock,
+    Linear,
+    RMSNorm,
 )
-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
+from astrai.model.transformer import Transformer

 __all__ = [
    # Modules
@@ -22,13 +16,6 @@ __all__ = [
    "GQA",
    "DecoderBlock",
    # Models
-    "AutoRegressiveLM",
-    "EmbeddingEncoder",
+    "Transformer",
    "AutoModel",
-    # LoRA
-    "LoRAConfig",
-    "inject_lora",
-    "merge_lora",
-    "save_lora",
-    "load_lora",
 ]
@@ -4,22 +4,18 @@ AutoModel base class for model loading and saving.

 from contextlib import contextmanager
 from pathlib import Path
-from typing import Self, Union
+from typing import Self, Type, Union

+import safetensors.torch as st
 import torch.nn as nn

-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
+from astrai.config import ModelConfig
+from astrai.factory import Registry


@contextmanager
 def _disable_random_init(enable: bool = True):
-    if not enable:
-        yield
-        return
-
-    names = (
+    init_functions = [
        "xavier_normal_",
        "xavier_uniform_",
        "kaiming_normal_",
@@ -29,27 +25,60 @@ def _disable_random_init(enable: bool = True):
        "constant_",
        "normal_",
        "uniform_",
-    )
-    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)
+    ]
+    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)
    try:
        yield
    finally:
-        for n, fn in orig.items():
-            setattr(nn.init, n, fn)
+        if enable:
+            for name, orig_func in original_funcs.items():
+                setattr(nn.init, name, orig_func)


-class AutoModel(BaseFactory["AutoModel"], nn.Module):
+class AutoModel(nn.Module):
    """
    Autoregressive language model base class.
-    Provides model loading/saving, registration, and generation.
+    Provides model loading/saving and generation capabilities.
    """

-    def __init__(self, config: BaseModelConfig):
+    _registry = Registry()
+
+    def __init__(self, config: ModelConfig):
        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,
@@ -60,22 +89,24 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):

        model_path = Path(path)

+        # Load config
+        config = ModelConfig()
        config_path = model_path / "config.json"
-        if not config_path.exists():
+        if config_path.exists():
+            config.load(str(config_path))
+        else:
            raise FileNotFoundError(f"Config file not found: {config_path}")

-        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)
+        model_type = config.model_type or "transformer"
+        actual_cls = cls.get_model_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 = load_model_weights(str(model_path))
+            state_dict = st.load_file(str(weights_path))
            model.load_state_dict(state_dict, strict=strict)

        return model
@@ -83,12 +114,15 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
    def save_pretrained(
        self,
        save_directory: Union[str, Path],
-    ):
-        save_model(
-            config=self.config.to_dict(),
-            state_dict=self.state_dict(),
-            save_directory=str(save_directory),
-        )
+    ) -> 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"))

    def to(self, *args, **kwargs) -> Self:
        """Move model to device/dtype."""
@@ -1,25 +0,0 @@
-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",
-]
@@ -1,214 +0,0 @@
-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
@@ -1,39 +0,0 @@
-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
@@ -1,26 +0,0 @@
-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
@@ -1,24 +0,0 @@
-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)
@@ -1,194 +0,0 @@
-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
@@ -1,100 +0,0 @@
-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
@@ -1,15 +0,0 @@
-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)
@@ -1,71 +0,0 @@
-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)
@@ -1,89 +0,0 @@
-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
@@ -0,0 +1,337 @@
+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)
@@ -1,83 +1,98 @@
-from typing import Any, Dict, Mapping, Optional
+from typing import Any, Mapping, Optional

 import torch
 import torch.nn as nn
 from torch import Tensor

-from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.core.cache import CacheView
+from astrai.config.model_config import ModelConfig
+from astrai.inference.cache import CacheView
 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.linear import Linear
-from astrai.model.components.norm import RMSNorm
-from astrai.model.components.rope import RotaryEmbedding
+from astrai.model.module import (
+    DecoderBlock,
+    Embedding,
+    Linear,
+    RMSNorm,
+    RotaryEmbedding,
+)


 def process_attention_mask(
+    seq_mask: Tensor,
    input_tensor: Tensor,
-    position_ids: Optional[Tensor],
-    input_mask: Optional[Tensor] = None,
+    start_pos: int = 0,
    is_causal: bool = False,
-) -> Optional[Tensor]:
-    if position_ids is None:
-        return None
-    if input_mask is not None and input_mask.dim() > 2:
-        return input_mask
-
+) -> Tensor:
+    """Build 4D attention mask from 2D seq_mask, with optional causal masking."""
    device = input_tensor.device
-    B = input_tensor.size(0)
-    T = position_ids.max().item() + 1
+    dtype = input_tensor.dtype
+    seq_len = input_tensor.size(1)

-    if input_mask is None:
-        if position_ids.min().item() == 0 and is_causal:
-            return None
-        attend = torch.ones(B, 1, T, dtype=torch.bool, device=device)
+    if seq_mask is None:
+        if start_pos != 0:
+            seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
        else:
-        attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1)
+            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
+    )

    if is_causal:
-        causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
-        attend = attend & causal
+        expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)

-    return attend.unsqueeze(1)
+    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


-@AutoModel.register("autoregressive_lm")
-class AutoRegressiveLM(AutoModel):
-    """Autoregressive language model with paged KV cache."""
+@AutoModel.register("transformer")
+class Transformer(AutoModel):
+    """Transformer language model with paged KV cache."""

-    def __init__(self, config: AutoRegressiveLMConfig):
+    def __init__(self, config: ModelConfig):
        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(
-            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
+            config.dim // config.n_heads, config.max_len
        )
+        self.embed_tokens = Embedding(config.vocab_size, config.dim)

        self.layers = nn.ModuleList(
-            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
+            [
+                DecoderBlock(
+                    config.dim,
+                    config.n_heads,
+                    config.dim_ffn,
+                    config.n_kv_heads,
+                    config.norm_eps,
+                    config.use_qk_norm,
+                    config.use_gated_attention,
+                    layer_id,
+                )
+                for layer_id in range(config.n_layers)
+            ]
        )

        self.norm = RMSNorm(config.dim, config.norm_eps)
        self.lm_head = Linear(config.dim, config.vocab_size)

-        if self.config.tie_weight is True:
+        if self.config.tie_weight:
            self.lm_head.weight = self.embed_tokens.weight

-        self.apply(self._init_weights)
+        self._init_weights()

-    def _init_weights(self, module):
-        if hasattr(module, "reset_parameters"):
-            module.reset_parameters()
+    def _init_weights(self):
+        for param in self.parameters():
+            if param.dim() > 1:
+                nn.init.normal_(param, mean=0.0, std=0.006)

    def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
        lm_head_key = "lm_head.weight"
@@ -85,7 +100,7 @@ class AutoRegressiveLM(AutoModel):

        state_dict = dict(state_dict)

-        if self.config.tie_weight is True:
+        if self.config.tie_weight:
            # same tensor for embed and lm_head
            if embed_key in state_dict:
                state_dict[lm_head_key] = state_dict[embed_key]
@@ -101,7 +116,7 @@ class AutoRegressiveLM(AutoModel):
            destination=destination, prefix=prefix, keep_vars=keep_vars
        )

-        if self.config.tie_weight is True:
+        if self.config.tie_weight:
            lm_head_key = prefix + "lm_head.weight"
            if lm_head_key in state_dict:
                del state_dict[lm_head_key]
@@ -113,16 +128,17 @@ class AutoRegressiveLM(AutoModel):
        input_ids: Tensor,
        input_mask: Optional[Tensor] = None,
        paged_cache: Optional[CacheView] = None,
-        position_ids: Optional[Tensor] = None,
-    ) -> Dict[str, Tensor]:
+        start_pos: int = 0,
+    ) -> Tensor:
        assert input_ids.ndim == 2

        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=True)
+        rotary_emb = self.rotary_embedding(x, start_pos)
+
+        attn_mask = process_attention_mask(input_mask, x, start_pos, is_causal=True)

        for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, paged_cache)
+            x = layer(x, rotary_emb, attn_mask, paged_cache, start_pos)

        hidden_states = self.norm(x)
        logits = self.lm_head(hidden_states)
@@ -1,13 +1,3 @@
-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,
@@ -27,12 +17,4 @@ __all__ = [
    "spawn_parallel_fn",
    "RowParallelLinear",
    "ColumnParallelLinear",
-    "ExecutorFactory",
-    "BaseExecutor",
-    "GradientState",
-    "AccumOptimizer",
-    "AccumScheduler",
-    "NoneExecutor",
-    "DDPExecutor",
-    "FSDPExecutor",
 ]
@@ -1,311 +0,0 @@
-"""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.distributed as dist
-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()
-
-    @contextmanager
-    def checkpoint_context(self, model: nn.Module):
-        if self.use_distributed:
-            dist.barrier()
-        state_dict = self._gather_state_dict(model)
-        yield state_dict
-        if self.use_distributed:
-            dist.barrier()
-
-    def _gather_state_dict(self, model: nn.Module):
-        state_dict = self.unwrap_model(model)
-        if self.use_distributed and get_rank() != 0:
-            return None
-        return 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: Optional[float]) -> float:
-        if max_norm is None:
-            total_norm = torch.norm(
-                torch.stack(
-                    [p.grad.norm(2) for p in model.parameters() if p.grad is not None]
-                )
-            )
-            return total_norm.item()
-        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: Optional[float]) -> float:
-        if max_norm is None:
-            return super().clip_grad_norm(model, max_norm)
-        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=True),
-            ):
-                return model.state_dict()
-
-        return model.state_dict()
@@ -1,21 +1,13 @@
 import os
-import socket
-from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from functools import wraps
-from typing import Callable, Optional
+from typing import Callable

 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp


-def find_free_port() -> str:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("", 0))
-        return str(s.getsockname()[1])
-
-
 def get_current_device():
    return os.environ["LOCAL_DEVICE"]

@@ -38,7 +30,6 @@ 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",
@@ -50,26 +41,20 @@ 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, local_rank)
+    device_id = torch.device(device_type, rank)

    os.environ["MASTER_ADDR"] = master_addr
    os.environ["MASTER_PORT"] = master_port
-    os.environ["LOCAL_RANK"] = str(local_rank)
+    os.environ["LOCAL_RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["LOCAL_DEVICE"] = str(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)
+    dist.init_process_group(
+        rank=rank, world_size=world_size, backend=backend, device_id=device_id
+    )

    try:
        if backend == "nccl" and torch.cuda.is_available():
@@ -105,7 +90,7 @@ def only_on_rank(rank, sync=False):
    return decorator


-def _run_single_rank(
+def wrapper_spawn_func(
    rank: int,
    world_size: int,
    backend: str,
@@ -115,10 +100,10 @@ def _run_single_rank(
    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,
@@ -126,118 +111,51 @@ def _run_single_rank(
        ):
            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,
-            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)
-
-
-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()
+    except Exception as e:
+        print(f"Error in rank {rank}: {e}")
        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(
    func: Callable,
    world_size: int,
    backend: str = "nccl",
    master_addr: str = "localhost",
-    master_port: Optional[str] = None,
+    master_port: str = "29500",
    device_type: str = "cuda",
-    start_method: str = "spawn",
    **kwargs,
 ):
-    if master_port is None:
-        master_port = find_free_port()
-    launcher = _detect_launcher()
-    if launcher in ("torchelastic", "torchrun", "external"):
-        strategy = TorchrunStrategy(
-            world_size, backend, master_addr, master_port, device_type, start_method
+    # 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,
    )
-    else:
-        strategy = LocalStrategy(
-            world_size, backend, master_addr, master_port, device_type, start_method
+
+    mp.spawn(
+        wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
    )
-    strategy.launch(func, **kwargs)
@@ -1,40 +0,0 @@
-from astrai.preprocessing.builder import (
-    BaseMaskBuilder,
-    MaskBuilderFactory,
-    MultiOutputMaskBuilder,
-    SectionedMaskBuilder,
-    SingleOutputMaskBuilder,
-)
-from astrai.preprocessing.packing import (
-    PackingStrategy,
-    PackingStrategyFactory,
-    plan_bfd,
-)
-from astrai.preprocessing.pipeline import Pipeline, filter_by_length
-from astrai.preprocessing.position_id import (
-    PositionIdStrategy,
-    PositionIdStrategyFactory,
-)
-from astrai.preprocessing.transform import TokenizeTransform
-from astrai.preprocessing.writer import (
-    StoreWriter,
-    StoreWriterFactory,
-)
-
-__all__ = [
-    "BaseMaskBuilder",
-    "MaskBuilderFactory",
-    "MultiOutputMaskBuilder",
-    "PackingStrategy",
-    "PackingStrategyFactory",
-    "Pipeline",
-    "PositionIdStrategy",
-    "PositionIdStrategyFactory",
-    "SectionedMaskBuilder",
-    "SingleOutputMaskBuilder",
-    "StoreWriter",
-    "StoreWriterFactory",
-    "TokenizeTransform",
-    "filter_by_length",
-    "plan_bfd",
-]
@@ -1,337 +0,0 @@
-"""Mask building for preprocessing pipeline.
-
-:class:`SectionRenderer` converts section specs into token ids and loss
-masks (template / text / value extraction).  :class:`SingleOutputMaskBuilder`
-handles single-output (SFT / pretrain), :class:`MultiOutputMaskBuilder`
-handles multi-output (DPO / GRPO), and :class:`SectionedMaskBuilder`
-orchestrates both modes as a façade.
-"""
-
-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):
-        """Tokenize a list-valued field, preserving per-element boundaries.
-
-        Returns ``(list_of_id_lists, list_of_mask_lists)`` where each
-        inner list corresponds to one element of the source list.  This
-        is critical for GRPO where each response must stay a separate
-        sequence so the strategy can form a ``[G, R]`` tensor.
-        """
-        per_item_ids: list[list[int]] = []
-        per_item_masks: list[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:
-                ids: list[int] = []
-                mask: list[int] = []
-                if use_template:
-                    if isinstance(val, list):
-                        wrapper = {field: val}
-                        self._append_template(
-                            wrapper, field, action, tokenizer, config, ids, mask
-                        )
-                else:
-                    wrapper = {field: str(val)}
-                    self._append_text(
-                        wrapper,
-                        field,
-                        action,
-                        tokenizer,
-                        False,
-                        False,
-                        config,
-                        ids,
-                        mask,
-                    )
-                if ids:
-                    max_len = config.preprocessing.max_seq_len
-                    ids = ids[:max_len]
-                    mask = mask[: len(ids)]
-                    per_item_ids.append(ids)
-                    per_item_masks.append(mask)
-
-        if not per_item_ids:
-            return None, None
-        return per_item_ids, per_item_masks
-
-    @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("single")
-class SingleOutputMaskBuilder(BaseMaskBuilder):
-    """Build a single output sequence with optional loss mask.
-
-    Expects ``config.input.sections`` (list of section specs).
-    """
-
-    def __init__(self, renderer: Optional[SectionRenderer] = None):
-        self.renderer = renderer or SectionRenderer()
-
-    def build(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
-
-
-@MaskBuilderFactory.register("multi")
-class MultiOutputMaskBuilder(BaseMaskBuilder):
-    """Build multiple output sequences (DPO / GRPO).
-
-    Expects ``config.input.sources`` (dict of output_key → spec).
-    """
-
-    def __init__(self, renderer: Optional[SectionRenderer] = None):
-        self.renderer = renderer or SectionRenderer()
-
-    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
-        sources_spec = getattr(config.input, "sources", None)
-        if not sources_spec:
-            return None
-
-        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
-                )
-                if ids is None:
-                    continue
-                # ids is List[List[int]] — preserve per-response structure
-                result[output_key] = ids
-                if mask is not None:
-                    result[mask_key] = mask
-                any_output = True
-                continue
-
-            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
-
-
-@MaskBuilderFactory.register("sectioned")
-class SectionedMaskBuilder(BaseMaskBuilder):
-    """Façade that dispatches to SingleOutputMaskBuilder or MultiOutputMaskBuilder.
-
-    Preserves backward compatibility for existing configs and code that rely
-    on the ``"sectioned"`` factory name.
-    """
-
-    def __init__(self):
-        self._single = SingleOutputMaskBuilder()
-        self._multi = MultiOutputMaskBuilder()
-
-    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
-        sources_spec = getattr(config.input, "sources", None)
-        if sources_spec:
-            return self._multi.build(item, config, tokenizer)
-        return self._single.build(item, config, tokenizer)
@@ -1,124 +0,0 @@
-"""Shared preprocessing kernel used by both :class:`Pipeline` and
-:class:`TokenizeTransform`.
-
-The two entry points previously duplicated ~60 % of their logic:
-record iteration, mask-builder invocation, primary-id extraction,
-per-key accumulation, dtype inference and position-id generation.
-This module factors out the common core as pure functions so that
-the online (``TokenizeTransform``) and offline (``Pipeline``) paths
-stay in lockstep.
-"""
-
-from itertools import chain
-from typing import Dict, Iterator, List, Optional
-
-import torch
-
-from astrai.config.preprocess_config import PipelineConfig
-from astrai.preprocessing.builder import MaskBuilderFactory
-from astrai.preprocessing.position_id import PositionIdStrategyFactory
-from astrai.tokenize import AutoTokenizer
-
-
-def build_preprocessing_components(config: PipelineConfig, tokenizer_path: str):
-    """Load tokenizer, mask builder and position-id strategy together.
-
-    Both ``Pipeline`` and ``TokenizeTransform`` need the same triple;
-    centralising the construction avoids drift (e.g. one path forgetting
-    to create the position-id strategy).
-    """
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
-    mask_builder = MaskBuilderFactory.create("sectioned")
-    position_strategy = PositionIdStrategyFactory.create(
-        config.output.position_ids_mode
-    )
-    return tokenizer, mask_builder, position_strategy
-
-
-def primary_ids(result: dict) -> List[int]:
-    """Return the first flat int-list value in *result*.
-
-    Used for token counting and position-id generation when the
-    primary key name is not known (DPO uses ``chosen``, GRPO uses
-    ``prompts``, SFT uses ``sequence``).
-    """
-    for val in result.values():
-        if isinstance(val, list) and val and isinstance(val[0], int):
-            return val
-    return []
-
-
-def infer_dtype(ids: List) -> torch.dtype:
-    """Float values become float32, everything else int32."""
-    if ids and isinstance(ids[0], float):
-        return torch.float32
-    return torch.int32
-
-
-def iter_raw_records(
-    records: List[dict],
-    mask_builder,
-    config: PipelineConfig,
-    tokenizer,
-) -> Iterator[dict]:
-    """Yield mask-builder output dicts for each record, skipping failures.
-
-    Drops ``domain`` from the result (callers that need it should read
-    it before calling this).  Each yielded dict maps a key
-    (``sequence``, ``chosen``, ``responses``…) to either a flat
-    ``List[int]`` or a nested ``List[List[int]]`` (GRPO responses/masks).
-    """
-    for item in records:
-        result = mask_builder.build(item, config, tokenizer)
-        if result is None:
-            continue
-        result.pop("domain", None)
-        if not primary_ids(result):
-            continue
-        yield result
-
-
-def to_per_record_tensors(
-    raw: Dict[str, list],
-) -> Dict[str, List[torch.Tensor]]:
-    """Convert an accumulated ``{key: [per-record ids]}`` dict to tensors.
-
-    Handles three shapes transparently:
-
-    - ``List[int]`` per record (``sequence``, ``chosen``…) → one tensor per record.
-    - ``List[List[int]]`` per record (GRPO ``responses``/``masks``) → one
-      ``List[Tensor]`` per record (nested), preserving the per-response
-      boundary so downstream code can index responses individually.
-    - ``List[int]`` for the whole shard (pre-packed keys) → single tensor.
-
-    The detection mirrors the previous inline logic in
-    ``Pipeline._flush`` and ``TokenizeTransform.apply``.
-    """
-    tensors: Dict[str, List[torch.Tensor]] = {}
-    for key, ids_list in raw.items():
-        if ids_list and isinstance(ids_list[0], list):
-            tensors[key] = [
-                [torch.tensor(sub, dtype=infer_dtype(sub)) for sub in ids]
-                if ids and isinstance(ids[0], list)
-                else torch.tensor(ids, dtype=infer_dtype(ids))
-                for ids in ids_list
-            ]
-        else:
-            tensors[key] = [
-                torch.tensor(list(chain.from_iterable(ids_list)), dtype=torch.int32)
-            ]
-    return tensors
-
-
-def build_position_ids(
-    sequences: List[List[int]],
-    strategy,
-) -> Optional[List[int]]:
-    """Generate position ids for *sequences* using *strategy*.
-
-    Returns ``None`` when the strategy produces no ids (e.g. ``none``
-    mode), so callers can skip attaching the key instead of storing
-    an empty list.
-    """
-    pos_ids = strategy.generate(sequences)
-    return pos_ids or None
@@ -1,176 +0,0 @@
-"""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]
-
-
-def plan_bfd(
-    sequences: List[List[int]], max_packed_len: int, truncation_mode: str = "keep_start"
-) -> List[List[int]]:
-    """Best-Fit Decreasing bin packing of *sequences* into bins.
-
-    Returns a list of bins, each bin a list of original indices into
-    *sequences*.  Bin capacities are respected on the *truncated*
-    length of each sequence (so a sequence longer than
-    *max_packed_len* counts at *max_packed_len*).
-
-    Pure index-based so callers can apply the same plan to any
-    aligned key (``loss_mask``, ``position_ids``…).
-    """
-    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
-
-
-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 = plan_bfd(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
-
-
-@PackingStrategyFactory.register("bfd_split")
-class BFDSplitPacking(BFDPacking):
-    """BFD packing with over-length sequences split into chunks.
-
-    Sequences longer than *max_packed_len* are split into consecutive
-    chunks of at most *max_packed_len* tokens instead of being
-    truncated.  Each chunk becomes an independent sequence that enters
-    BFD planning.  All keys (``loss_mask``, ``position_ids``, …) are
-    split in lockstep so per-token alignment is preserved.
-
-    Note: because each chunk is treated as a separate document, the
-    second chunk of a split sequence loses the preceding context.
-    """
-
-    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
-        if max_packed_len <= 0:
-            return super().apply(keys, max_packed_len, truncation_mode)
-
-        split_keys = self._split_all(keys, max_packed_len)
-        return super().apply(split_keys, max_packed_len, truncation_mode)
-
-    @staticmethod
-    def _split_all(
-        keys: Dict[str, List[List[int]]], max_packed_len: int
-    ) -> Dict[str, List[List[int]]]:
-        """Split every sequence exceeding *max_packed_len* into chunks,
-        applying the same chunk boundaries to all keys."""
-        sequences = keys["sequence"]
-        chunk_bounds = [list(range(0, len(s), max_packed_len)) for s in sequences]
-        result: Dict[str, List[List[int]]] = {}
-        for key, vals in keys.items():
-            split_vals: List[List[int]] = []
-            for val, starts in zip(vals, chunk_bounds):
-                for start in starts:
-                    split_vals.append(val[start : start + max_packed_len])
-            result[key] = split_vals
-        return result
@@ -1,248 +0,0 @@
-"""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.
-
-Record iteration, mask building, primary-id extraction and per-key
-accumulation are shared with :class:`TokenizeTransform` via the
-:mod:`astrai.preprocessing.core` helpers.
-"""
-
-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.core import (
-    build_preprocessing_components,
-    iter_raw_records,
-    primary_ids,
-)
-from astrai.preprocessing.packing import PackingStrategyFactory
-from astrai.preprocessing.writer import StoreWriterFactory
-
-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.tokenizer, self.mask_builder, self._position_id = (
-            build_preprocessing_components(config, tokenizer_path)
-        )
-        self._packer = PackingStrategyFactory.create(
-            config.preprocessing.packing_strategy
-        )
-        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):
-        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__")
-            ids = primary_ids(result)
-            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 _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:
-                if path.endswith(".json"):
-                    data = json.load(f)
-                    if isinstance(data, dict):
-                        yield data
-                    elif isinstance(data, list):
-                        yield from data
-                else:
-                    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
-
-            keys = self._inject_doc_reset_position_ids(keys, mode, original_sequences)
-            keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode)
-            tensors = self._to_tensors(keys)
-            tensors = self._inject_continuous_position_ids(
-                tensors, mode, keys.get("sequence", [])
-            )
-
-            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)"
-            )
-
-    def _inject_doc_reset_position_ids(
-        self,
-        keys: Dict[str, list],
-        mode: str,
-        original_sequences: List[List[int]],
-    ) -> Dict[str, list]:
-        """Attach per-document position_ids before packing (``doc_reset``).
-
-        ``doc_reset`` position ids must enter the packer so that each
-        packed bin concatenates the per-doc ranges in bin order.  The
-        per-record structure ``[range(len(s)) for s in seqs]`` is required
-        by the packer (it concatenates per-record lists per bin); the
-        ``PositionIdStrategy.generate`` flattens, so it cannot be used
-        directly here — it is only consulted for the ``continuous``
-        post-packing path.
-        """
-        if mode != "doc_reset" or not original_sequences:
-            return keys
-        keys["position_ids"] = [list(range(len(s))) for s in original_sequences]
-        return keys
-
-    def _inject_continuous_position_ids(
-        self,
-        tensors: Dict[str, List[torch.Tensor]],
-        mode: str,
-        packed_sequences: List[List[int]],
-    ) -> Dict[str, List[torch.Tensor]]:
-        """Attach a single continuous position_ids tensor after packing.
-
-        ``continuous`` mode spans the whole shard (post-packing), so it
-        cannot participate in bin packing — it is computed from the
-        packed sequences and appended directly to the tensor dict.
-        """
-        if mode != "continuous" or not packed_sequences:
-            return tensors
-        pos_ids = self._position_id.generate(packed_sequences)
-        if pos_ids:
-            tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
-        return tensors
-
-    def _to_tensors(self, keys: Dict[str, list]) -> Dict[str, List[torch.Tensor]]:
-        """Convert packed per-key id lists to tensors.
-
-        Honours ``config.output.dtype`` overrides per key; falls back to
-        ``int32``.  Handles three shapes (see
-        :func:`astrai.preprocessing.core.to_per_record_tensors` for the
-        equivalent online-path helper):
-        - ``List[int]`` per record → one tensor per record.
-        - ``List[List[int]]`` per record (GRPO responses/masks) → one tensor
-          per record, inner lists flattened.
-        - ``List[int]`` for the whole shard (pre-packed keys) → single tensor.
-        """
-        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
-            )
-            if ids_list and isinstance(ids_list[0], list):
-                tensors[key] = [
-                    torch.tensor(
-                        list(chain.from_iterable(ids))
-                        if ids and isinstance(ids[0], list)
-                        else ids,
-                        dtype=dt,
-                    )
-                    for ids in ids_list
-                ]
-            else:
-                tensors[key] = [
-                    torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
-                ]
-        return tensors
@@ -1,46 +0,0 @@
-"""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))
@@ -1,92 +0,0 @@
-"""Tokenization transform for JSONL record streams.
-
-Bridges the Reader layer (``JsonlStore`` reads raw JSON records) and the
-Dataset layer (expects per-record tensors).  Holds the tokenizer,
-mask-builder and position-id strategy together so that I/O code stays
-free of model dependencies.
-
-The record-processing core (mask building, primary-id extraction,
-per-key tensorisation, position-id generation) is shared with
-:class:`astrai.preprocessing.pipeline.Pipeline` via the
-:mod:`astrai.preprocessing.core` helpers.
-"""
-
-import json
-from pathlib import Path
-from typing import Dict, List
-
-import torch
-
-from astrai.config.preprocess_config import PipelineConfig
-from astrai.preprocessing.core import (
-    build_position_ids,
-    build_preprocessing_components,
-    iter_raw_records,
-    to_per_record_tensors,
-)
-
-
-class TokenizeTransform:
-    """Tokenize raw JSONL record dicts into per-key tensor lists.
-
-    Owns the three preprocessing concerns that were previously inlined in
-    ``JsonlStore``: tokenization, loss-mask construction and position-id
-    generation.  Constructing it loads the tokenizer, so it is intentionally
-    cheap to pass around once built.
-
-    Args:
-        config: Pipeline config describing sections / masks / position mode.
-        tokenizer_path: Path passed to ``AutoTokenizer.from_pretrained``.
-    """
-
-    def __init__(self, config: PipelineConfig, tokenizer_path: str):
-        self.config = config
-        self.tokenizer, self.mask_builder, self.position_strategy = (
-            build_preprocessing_components(config, tokenizer_path)
-        )
-
-    @classmethod
-    def from_config_file(cls, config_path: str) -> "TokenizeTransform":
-        """Build from a ``dataset_config.json`` file path.
-
-        The config file follows :class:`PipelineConfig` schema with an
-        extra ``tokenizer_path`` field.  When omitted, the config's
-        parent directory is used as the tokenizer path.
-        """
-        root = Path(config_path).parent
-        with open(config_path, "r", encoding="utf-8") as f:
-            raw_config = json.load(f)
-        tokenizer_path = raw_config.pop("tokenizer_path", None) or str(root)
-        config = PipelineConfig.from_dict(raw_config)
-        return cls(config, tokenizer_path)
-
-    def apply(self, records: List[dict]) -> Dict[str, list]:
-        """Tokenize a list of raw record dicts.
-
-        Returns a dict mapping key (``sequence``, ``chosen``, ``responses``,
-        …) to a list of per-record tensors (or nested tensor lists for
-        multi-response keys such as GRPO ``responses``).
-        """
-        raw: Dict[str, list] = {}
-        doc_sequences: List[List[int]] = []
-
-        for result in iter_raw_records(
-            records, self.mask_builder, self.config, self.tokenizer
-        ):
-            primary = None
-            for val in result.values():
-                if isinstance(val, list) and val and isinstance(val[0], int):
-                    primary = val
-                    break
-            if primary is not None:
-                doc_sequences.append(primary)
-            for key, ids in result.items():
-                raw.setdefault(key, []).append(ids)
-
-        tensors = to_per_record_tensors(raw)
-
-        pos_ids = build_position_ids(doc_sequences, self.position_strategy)
-        if pos_ids is not None:
-            tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
-
-        return tensors
@@ -1,75 +0,0 @@
-"""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
@@ -1,21 +0,0 @@
-"""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): ...
@@ -0,0 +1,116 @@
+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,
+        )
@@ -1,45 +0,0 @@
-"""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_bin_offsets,
-    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_bin_offsets",
-    "load_h5",
-    "save_bin",
-    "save_h5",
-]
@@ -1,201 +0,0 @@
-"""Model checkpoint serialization helpers."""
-
-import io
-import json
-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)
-
-        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,
-            meta=meta,
-            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
@@ -1,123 +0,0 @@
-"""Dataset storage serialization helpers (HDF5 / memory-mapped binary)."""
-
-import json
-import os
-from pathlib import Path
-from typing import Any, Dict, List, Optional
-
-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)
-    if root_path.is_file() and root_path.suffix in (".h5", ".hdf5"):
-        h5_files = [root_path]
-    else:
-        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]],
-    record_keys: Optional[List[str]] = None,
-):
-    """Save tensors as memory-mapped binary files.
-
-    When *record_keys* is provided, those keys are written with per-record
-    cumulative offsets in ``meta.json`` so that ``MmapStore.fetch_record``
-    can slice individual records from the concatenated binary without
-    cross-record concatenation.  Keys not in *record_keys* (e.g. SEQ
-    ``sequence``) are written as a single contiguous stream without
-    offsets, preserving backward compatibility.
-
-    Nested keys (``List[List[Tensor]]`` such as GRPO ``responses``) are
-    not supported in bin format — use H5 for those.
-    """
-    os.makedirs(file_path, exist_ok=True)
-    record_keys = set(record_keys or [])
-    meta = {}
-    for key, tensors in tensor_group.items():
-        if tensors and isinstance(tensors[0], list):
-            raise ValueError(
-                f"Nested key '{key}' (List[List[Tensor]]) is not supported "
-                f"in bin format. Use H5 or JSONL storage instead."
-            )
-        cat = torch.cat(tensors, dim=0)
-        entry: Dict[str, Any] = {
-            "shape": list(cat.shape),
-            "dtype": str(cat.dtype).split(".")[-1],
-        }
-        if key in record_keys:
-            offsets = [0]
-            for t in tensors:
-                offsets.append(offsets[-1] + t.shape[0])
-            entry["offsets"] = offsets
-        meta[key] = entry
-        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
-
-
-def load_bin_offsets(file_path: str) -> Dict[str, List[int]]:
-    """Read per-record cumulative offsets from ``meta.json``.
-
-    Returns an empty dict when no key has offsets (legacy bin files),
-    in which case record-mode access falls back to per-record segment
-    indexing (H5/JSONL layout).
-    """
-    with open(os.path.join(file_path, "meta.json"), "r") as f:
-        meta = json.load(f)
-    offsets: Dict[str, List[int]] = {}
-    for key, info in meta.items():
-        if "offsets" in info:
-            offsets[key] = info["offsets"]
-    return offsets
@@ -1,11 +1,13 @@
-from functools import cached_property
+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.

@@ -13,36 +15,23 @@ class ChatTemplate:
        name: Unique identifier for the template.
        template_str: Jinja2 template string.
        description: Optional description.
-        default_variables: Optional dictionary of default variable values.
+        default_variables: Optional dictionary of default variable values
+            that will be passed to the template if not overridden during rendering.
        special_tokens: Optional dictionary mapping token names to their string values.
+            These tokens are automatically added to the template variables.
    """

-    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 {}
+    name: str
+    template_str: str
+    description: str = ""
+    default_variables: Dict[str, Any] = None
+    special_tokens: Dict[str, str] = None

-    @cached_property
-    def _compiled(self) -> Template:
-        """Lazy-compiled Jinja2 template, cached on first access.
-
-        The compiled :class:`~jinja2.Template` holds a dynamically-generated
-        ``root`` render function whose ``__module__`` is ``None``; under
-        ``pickle`` it falls back to ``__main__`` and breaks ``spawn``-based
-        multiprocessing.  By deferring compilation to first access, the
-        default pickle protocol serialises only ``template_str``; each
-        worker rebuilds the cache on first render.
-        """
-        return Template(self.template_str)
+    def __post_init__(self):
+        if self.default_variables is None:
+            self.default_variables = {}
+        if self.special_tokens is None:
+            self.special_tokens = {}

    @classmethod
    def from_string(
@@ -54,7 +43,7 @@ class ChatTemplate:
    ) -> "ChatTemplate":
        """Create a ChatTemplate instance directly from a template string."""
        return cls(
-            name="",
+            name="",  # empty name for ad‑hoc templates
            template_str=template_str,
            description=description,
            default_variables=default_variables,
@@ -84,4 +73,5 @@ class ChatTemplate:
        if system_prompt is not None:
            variables["system_prompt"] = system_prompt

-        return self._compiled.render(**variables)
+        jinja_template = Template(self.template_str)
+        return jinja_template.render(**variables)
@@ -51,26 +51,9 @@ class AutoTokenizer:
                self.set_chat_template(config["chat_template"])

    @classmethod
-    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."
-            )
+    def from_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoTokenizer":
+        """Load tokenizer from pretrained directory."""
        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):
@@ -164,14 +147,7 @@ class AutoTokenizer:
          - tokenizer.bos_token   → returns string
          - tokenizer.bos_token_id → returns corresponding integer ID
          - tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens
-
-        Internal/private attrs are not intercepted: during unpickle
-        ``__dict__`` is empty, so probing ``self._special_token_map``
-        would recurse infinitely.
        """
-        if key.startswith("_"):
-            raise AttributeError(key)
-
        # Handle stop_ids - return IDs for all special tokens
        if key == "stop_ids":
            stop_ids = []
@@ -1,25 +1,75 @@
 from typing import Dict

-import torch
 import torch.nn as nn


-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
-
-    total_sq = torch.stack([g.pow(2).sum() for g in grads]).sum()
-    if per_param:
+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():
-            if param.grad is not None:
-                norms[name] = param.grad.norm(2).item()
-            else:
        norms[name] = 0.0
-        norms["total"] = total_sq.sqrt().item()
+        if param.grad:
+            norm = param.grad.data.norm(norm_type).item()
+            norms[name] = norm
    return norms
-    return total_sq.sqrt().item()
+
+
+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


 def ctx_get_loss(ctx):
@@ -30,9 +80,25 @@ 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 ctx.grad_norm
+    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)
@@ -2,7 +2,7 @@

 import math
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Type

 from torch.optim.lr_scheduler import LRScheduler

@@ -31,6 +31,7 @@ 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")
@@ -40,6 +41,33 @@ 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 -----------

@@ -53,7 +81,7 @@ class CosineScheduler(BaseScheduler):
        optimizer,
        warmup_steps: int,
        lr_decay_steps: int,
-        min_rate: float = 0.01,
+        min_rate: float = 0.05,
        last_epoch: int = -1,
    ):
        self.warmup_steps = warmup_steps
@@ -65,15 +93,11 @@ 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 / max(self.warmup_steps, 1)
-            )
+            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
            return [base_lr * warmup_factor for base_lr in self.base_lrs]

        # cosine decay
-        decay_progress = (self.last_epoch - self.warmup_steps) / max(
-            self.lr_decay_steps, 1
-        )
+        decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
        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)
@@ -108,7 +132,7 @@ class SGDRScheduler(BaseScheduler):
        optimizer,
        warmup_steps: int,
        cycle_length: int,
-        min_rate: float = 0.01,
+        min_rate: float = 0.05,
        t_mult: int = 2,
        last_epoch: int = -1,
    ):
@@ -122,9 +146,7 @@ class SGDRScheduler(BaseScheduler):
    def get_lr(self):
        # warmup
        if self.last_epoch < self.warmup_steps:
-            warmup_factor = max(
-                self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
-            )
+            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
            return [base_lr * warmup_factor for base_lr in self.base_lrs]

        # SGDR
@@ -170,66 +192,3 @@ 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)
@@ -1,28 +1,39 @@
 """Training strategy implementations with factory pattern."""

+import copy
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, Union
+from typing import Any, 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 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)
+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)
    ref_model.requires_grad_(False)
    ref_model.eval()
    return ref_model


-def move_to_device(batch: Dict[str, Tensor], device: str) -> Dict[str, Tensor]:
+def move_to_device(batch: Dict[str, Tensor], device: str) -> Any:
    """Move batch tensors to specified device with non-blocking transfer."""
    return {key: value.to(device, non_blocking=True) for key, value in batch.items()}

@@ -32,7 +43,7 @@ def get_logprobs(
    input_ids: Tensor,
    mask: Tensor,
    reduction: str,
-) -> Tensor:
+):
    """Compute token-wise log probabilities from model outputs.

    Args:
@@ -70,34 +81,14 @@ 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[nn.Module, Callable[..., Dict[str, Tensor]]],
-        device: str,
-        **kwargs,
+        self, model: Union[Callable[..., Dict[str, Tensor]]], device: str, **kwargs
    ):
        self.model = model
        self.device = device
-        self.executor = kwargs.pop("executor", None)
        self.extra_kwargs = kwargs

    @abstractmethod
@@ -131,6 +122,32 @@ 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
@@ -143,13 +160,7 @@ class SEQStrategy(BaseStrategy):
    Computes cross-entropy loss for next token prediction.
    """

-    def __init__(
-        self,
-        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
-        device: str,
-        label_smoothing: float = 0.0,
-        **kwargs,
-    ):
+    def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
        super().__init__(model, device, **kwargs)
        self.label_smoothing = label_smoothing

@@ -174,31 +185,21 @@ class SFTStrategy(BaseStrategy):
    Applies cross-entropy loss only to tokens where loss_mask is True.
    """

-    def __init__(
-        self,
-        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
-        device: str,
-        label_smoothing: float = 0.0,
-        **kwargs,
-    ):
+    def __init__(self, model, device, 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, position_ids, loss_mask = (
+        input_ids, target_ids, loss_mask = (
            batch["input_ids"],
            batch["target_ids"],
-            batch["position_ids"],
            batch["loss_mask"],
        )

        ignore_index = -100
-        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"]
+        logits = self.model(input_ids=input_ids)["logits"]
+        target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)

        loss = F.cross_entropy(
            input=logits.flatten(0, 1).float(),
@@ -222,13 +223,12 @@ class DPOStrategy(BaseStrategy):
        self,
        model: nn.Module,
        device: str,
-        ref_model: nn.Module,
        beta: float = 0.1,
-        reduction: str = "sum",
+        reduction: str = "mean",
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
-        self.ref_model = ref_model
+        self.ref_model = create_ref_model(model)
        self.beta = beta
        self.reduction = reduction

@@ -265,45 +265,41 @@ class DPOStrategy(BaseStrategy):
 class GRPOStrategy(BaseStrategy):
    """Group Relative Policy Optimization strategy.

-    Implements GRPO following DeepSeek-R1 with token-level PPO clipping.
-    Advantages are group-normalized from scalar per-response rewards and
-    broadcast across all response tokens.  The loss is computed **only on
-    response tokens** — prompt tokens are masked out.
-
-    Three model roles are distinguished:
-
-    * **Policy** ``self.model`` — the model being trained.
-    * **Old policy** ``self.old_model`` — the behaviour policy that generated
-      the responses.  Used for the importance sampling ratio
-      ``ρ = π_θ / π_old``.  Synced externally after each data-generation round.
-    * **Reference model** ``self.ref_model`` — a frozen copy of the initial
-      policy (typically the SFT checkpoint) used **only** for the KL
-      regularisation term.  It is never updated during training.
+    On-policy GRPO following DeepSeek-R1: the policy model is updated while
+    a frozen ref_model stores the old-policy log-probs.  ratio = exp(logπ_θ - logπ_ref),
+    clipped PPO objective.  Call ``sync_ref_model()`` after each data-generation round.
    """

    def __init__(
        self,
        model: nn.Module,
        device: str,
-        old_model: nn.Module,
-        ref_model: nn.Module,
        clip_eps: float = 0.2,
        kl_coef: float = 0.01,
        group_size: int = 4,
+        reduction: str = "mean",
+        sync_interval: int = 200,
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
-        self.old_model = old_model
-        self.ref_model = ref_model
+        self.ref_model = create_ref_model(model)
        self.clip_eps = clip_eps
        self.kl_coef = kl_coef
        self.group_size = group_size
+        self.reduction = reduction
+        self.sync_interval = sync_interval
+        self._step = 0

-    def sync_old_model(self):
-        """Copy current policy weights to old model."""
-        self.old_model.load_state_dict(self.executor.unwrap_model(self.model))
+    def sync_ref_model(self):
+        """Copy current model weights to ref model."""
+        ref_state = self.model.state_dict()
+        self.ref_model.load_state_dict(ref_state)

    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        self._step += 1
+        if self._step % self.sync_interval == 0:
+            self.sync_ref_model()
+
        batch = move_to_device(batch, self.device)
        prompts = batch["prompts"]
        responses = batch["responses"]
@@ -314,60 +310,33 @@ class GRPOStrategy(BaseStrategy):
        responses_flat = responses.view(-1, response_len)
        masks_flat = masks.view(-1, response_len)
        prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
-        prompt_len = prompt_expanded.size(1)

        full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
-        # Prompt tokens are masked out (0) so logprobs are computed only for
-        # response tokens.  get_logprobs shifts the mask by one position, so
-        # the first response token's logprob (predicted from the last prompt
-        # token) is correctly included.
-        full_masks = torch.cat([torch.zeros_like(prompt_expanded), masks_flat], dim=-1)
+        full_masks = torch.cat([torch.ones_like(prompt_expanded), masks_flat], dim=-1)
+
+        log_probs_policy = get_logprobs(
+            self.model, full_sequences, full_masks, self.reduction
+        )
+        log_probs_policy = log_probs_policy.view(batch_size, group_size)

-        # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
-        # Response token logprobs occupy the last ``response_len`` positions
-        # (the first response token is predicted from the last prompt token).
-        token_log_probs_policy = get_logprobs(
-            self.model, full_sequences, full_masks, "none"
-        )[:, prompt_len - 1 :]
        with torch.no_grad():
-            token_log_probs_old = get_logprobs(
-                self.old_model, full_sequences, full_masks, "none"
-            )[:, prompt_len - 1 :]
-            token_log_probs_ref = get_logprobs(
-                self.ref_model, full_sequences, full_masks, "none"
-            )[:, prompt_len - 1 :]
+            log_probs_ref = get_logprobs(
+                self.ref_model, full_sequences, full_masks, self.reduction
+            )
+            log_probs_ref = log_probs_ref.view(batch_size, group_size)

-        # Reshape to [B, G, response_len]
-        token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
-        token_log_probs_old = token_log_probs_old.view(batch_size, group_size, -1)
-        token_log_probs_ref = token_log_probs_ref.view(batch_size, group_size, -1)
-        token_masks = masks_flat.view(batch_size, group_size, -1).float()
-
-        # Group-normalized advantages from scalar per-response rewards.
-        eps = 1e-8
+        eps = torch.finfo(log_probs_policy.dtype).eps
        mean = rewards.mean(dim=-1, keepdim=True)
-        std = rewards.std(dim=-1, keepdim=True, unbiased=False)
+        std = rewards.std(dim=-1, keepdim=True)
        advantages = (rewards - mean) / (std + eps)
-        # Broadcast scalar advantage to every response token: [B, G, 1]
-        advantages = advantages.unsqueeze(-1)

-        # Token-level ratio (π_θ / π_old) and PPO clipping.
-        log_ratio = token_log_probs_policy - token_log_probs_old
-        ratio = torch.exp(log_ratio)
+        ratio = torch.exp(log_probs_policy - log_probs_ref)

        surr1 = ratio * advantages
        surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
-        per_token_policy_loss = -torch.min(surr1, surr2)
-        token_count = token_masks.sum().clamp(min=1.0)
-        policy_loss = (per_token_policy_loss * token_masks).sum() / token_count
-
-        # KL penalty to frozen reference model with k1 estimator (non-negative):
-        # k1 = π_ref / π_θ - log(π_ref / π_θ) - 1, where π_ref / π_θ = exp(log_ref - log_policy).
-        log_ref_ratio = token_log_probs_ref - token_log_probs_policy
-        r = torch.exp(log_ref_ratio)
-        kl_per_token = r - torch.log(r + eps) - 1.0
-        kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count

+        policy_loss = -torch.min(surr1, surr2).mean()
+        kl_penalty = self.kl_coef * (log_probs_policy - log_probs_ref).square().mean()
        total_loss = policy_loss + kl_penalty

        return total_loss
@@ -1,31 +1,28 @@
 import json
-import logging
 import os
-import sys
 import time
 from pathlib import Path
-from typing import IO, Callable, List, Optional, Protocol, runtime_checkable
+from typing import Callable, List, Optional, Protocol, runtime_checkable

-import torch
-import torch.distributed as dist
 import torch.nn as nn
-from torch.utils.checkpoint import checkpoint as torch_checkpoint
+from torch.nn.utils import clip_grad_norm_
 from tqdm import tqdm

 from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
-from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
+    ctx_get_grad_max,
+    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):
@@ -45,15 +42,18 @@ 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,6 +69,12 @@ 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):
@@ -79,45 +85,28 @@ class GradientClippingCallback(TrainCallback):
    def __init__(self, max_grad_norm: float):
        self.max_grad_norm = max_grad_norm

-    def on_optimizer_step(self, context: TrainContext):
-        context.grad_norm = context.executor.clip_grad_norm(
-            context.model, self.max_grad_norm
-        )
+    def on_step_begin(self, context: TrainContext):
+        _ = context
+        clip_grad_norm_(context.model.parameters(), self.max_grad_norm)


-@CallbackFactory.register("gradient_checkpointing")
-class GradientCheckpointingCallback(TrainCallback):
+@CallbackFactory.register("scheduler")
+class SchedulerCallback(TrainCallback):
    """
-    Activation checkpointing callback — trades compute for memory
-    by recomputing specified module activations during the backward pass.
-
-    Args:
-        modules: Module types to apply checkpointing to.
+    Scheduler callback for trainer.
    """

-    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 __init__(self):
+        pass

    def on_train_begin(self, context: TrainContext):
-        context.model.apply(self._enable)
-        logger.info("Gradient checkpointing enabled")
+        for group in context.optimizer.param_groups:
+            if "initial_lr" not in group:
+                group["initial_lr"] = group["lr"]

-    def on_train_end(self, context: TrainContext):
-        context.model.apply(self._disable)
+    def on_batch_end(self, context: TrainContext):
+        if context.scheduler:
+            context.scheduler.step()


@CallbackFactory.register("checkpoint")
@@ -126,65 +115,54 @@ 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.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
-        self.last_ckpt_step = None
-
-    def on_train_begin(self, context: TrainContext):
-        self.last_ckpt_step = context.optimizer_step
+        self.state_dict_fn = state_dict_fn
+        self.save_extra_fn = save_extra_fn
+        self.last_ckpt_iter = 0

+    @only_on_rank(0)
    def _save_checkpoint(self, context: TrainContext):
-        self.last_ckpt_step = context.optimizer_step
-
-        with context.executor.checkpoint_context(context.model) as state_dict:
-            if state_dict is not None:
        save_path = os.path.join(
-                    self.save_dir,
-                    f"epoch_{context.epoch}_step_{context.optimizer_step}",
+            self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
        )
-                extra = self.save_extra_fn(context)
-                meta = context.config.to_dict()
+        state_dict = (
+            self.state_dict_fn(context.model)
+            if self.state_dict_fn
+            else context.model.state_dict()
+        )
+
+        extra = self.save_extra_fn(context) if self.save_extra_fn else None
        context.checkpoint = Checkpoint(
            state_dict=state_dict,
            epoch=context.epoch,
-                    consumed_samples=context.consumed_samples,
-                    config=context.model_config,
+            iteration=context.iteration,
            extra=extra,
-                    meta=meta,
        )
+
        context.checkpoint.save(save_path)
+        self.last_ckpt_iter = context.iteration

    def on_batch_end(self, context: TrainContext):
-        if context.optimizer_step - self.last_ckpt_step >= self.interval:
+        if context.iteration - self.last_ckpt_iter >= self.interval:
            self._save_checkpoint(context)

    def on_train_end(self, context: TrainContext):
-        if context.optimizer_step != self.last_ckpt_step:
+        if context.iteration != self.last_ckpt_iter:
            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):
@@ -192,36 +170,26 @@ class ProgressBarCallback(TrainCallback):
    Progress bar callback for trainer.
    """

-    def __init__(
-        self, num_epoch: int, log_interval: int = 100, file: Optional[IO[str]] = None
-    ):
+    def __init__(self, num_epoch: int):
        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(
-            total=total_steps,
+            context.dataloader,
            desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
            dynamic_ncols=True,
-            file=self.file or sys.stdout,
        )

    @only_on_rank(0)
-    def on_optimizer_step(self, context: TrainContext):
-        postfix = {
-            "step": f"{context.optimizer_step:d}",
+    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}",
            }
-        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)
+        )
        self.progress_bar.update(1)

    @only_on_rank(0)
@@ -231,20 +199,19 @@ class ProgressBarCallback(TrainCallback):
            self.progress_bar.close()


-@CallbackFactory.register("metric")
-class MetricCallback(TrainCallback):
+@CallbackFactory.register("metric_logger")
+class MetricLoggerCallback(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_flush_step = None
+        self.last_log_iter = 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)
@@ -254,91 +221,46 @@ class MetricCallback(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 _metrics(self, context: TrainContext, names):
+    def _get_log_data(self, context: TrainContext):
        return {
-            m: self._metric_funcs[m](context)
-            for m in names
-            if self._metric_funcs[m](context) is not None
-        }
-
-    @only_on_rank(0)
-    def _append(self, event_type: str, context: TrainContext, **extra):
-        entry = {
-            "type": event_type,
-            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "epoch": context.epoch,
-            "step": context.optimizer_step,
-            "consumed_samples": context.consumed_samples,
-            **extra,
+            "iter": context.iteration,
+            **{m: self._metric_funcs[m](context) for m in self.metrics},
        }
-        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
-
-    def on_train_begin(self, context: TrainContext):
-        self.last_log_flush_step = context.optimizer_step

    @only_on_rank(0)
-    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)
+    def _add_log(self, log_data):
+        self.log_cache.append(log_data)
+
+    @only_on_rank(0)
+    def _save_log(self, epoch, iter):
+        log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
+
        with open(log_file, "w") as f:
            for log in self.log_cache:
                f.write(json.dumps(log) + "\n")

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

-        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)
+        if context.iteration - self.last_log_iter >= self.save_interval:
+            self._save_log(context.epoch, context.iteration)
+            self.last_log_iter = context.iteration

    def on_train_end(self, context):
-        if (
-            self.last_log_flush_step is None
-            or context.optimizer_step != self.last_log_flush_step
-        ):
-            self._flush(context.epoch, context.optimizer_step)
-            self.last_log_flush_step = context.optimizer_step
+        if context.iteration != self.last_log_iter:
+            self._save_log(context.epoch, context.iteration)

    def on_error(self, context):
-        self._flush(context.epoch, context.optimizer_step)
+        self._save_log(context.epoch, context.iteration)
@@ -1,19 +1,16 @@
 from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any, Dict, Optional, Self
+from typing import Callable, Optional, Self

-import torch
 import torch.nn as nn
-from torch.utils.data import DataLoader, random_split
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from torch.utils.data import DataLoader

 from astrai.config.train_config import TrainConfig
-from astrai.dataset import RDSampler
-from astrai.model.components.lora import inject_lora
-from astrai.parallel.executor import BaseExecutor, ExecutorFactory
+from astrai.dataset import ResumableDistributedSampler
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
-from astrai.protocols import OptimizerProtocol, SchedulerProtocol
-from astrai.serialization import Checkpoint, load_json
-from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
+from astrai.serialization import Checkpoint
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory


@dataclass
@@ -21,197 +18,84 @@ class TrainContext:
    model: nn.Module = field(default=None)
    strategy: BaseStrategy = field(default=None)
    dataloader: DataLoader = field(default=None)
-    optimizer: OptimizerProtocol = field(default=None)
-    scheduler: SchedulerProtocol = field(default=None)
+    optimizer: Optimizer = field(default=None)
+    scheduler: LRScheduler = 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)
-    consumed_samples: int = field(default=0)
+    iteration: 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[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
-        )
+    kwargs: dict = field(default_factory=dict)


 class TrainContextBuilder:
    def __init__(
        self,
        config: TrainConfig,
+        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
    ):
        self.config = config
-        self._param_path: Optional[str] = None
-        self._resume: bool = False
+        self._checkpoint: Optional[Checkpoint] = None
+        self._load_extra_fn = load_extra_fn

-    def with_param_path(self, param_path: Optional[str], resume: bool = False) -> Self:
-        self._param_path = param_path
-        self._resume = resume
+    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+        self._checkpoint = checkpoint
        return self

    def build(self) -> TrainContext:
-        cfg = self.config
-        device = get_current_device()
-
-        executor = ExecutorFactory.create(
-            cfg.parallel_mode,
-            grad_accum_steps=cfg.grad_accum_steps,
-            **cfg.executor_kwargs,
-        )
-
-        model = cfg.model_fn()
-        model = model.to(device=device)
-
-        model_config = {}
-        if self._param_path:
-            config_path = Path(self._param_path) / "config.json"
-            if config_path.exists():
-                model_config = load_json(config_path)
-
-        if not model_config and hasattr(model, "config"):
-            model_config = model.config.to_dict()
-
        context = TrainContext(
-            model=model,
+            model=self.config.model,
            world_size=get_world_size(),
            rank=get_rank(),
-            config=cfg,
-            model_config=model_config,
-            executor=executor,
        )

-        if self._param_path:
-            checkpoint = Checkpoint.load_any(self._param_path)
-            if checkpoint is not None:
-                model.load_state_dict(checkpoint.state_dict, strict=False)
-                if checkpoint.config:
-                    context.model_config = checkpoint.config
+        device = get_current_device()
+        context.model = context.model.to(device=device)

-                if self._resume:
-                    context.epoch = checkpoint.epoch or cfg.start_epoch
-                    if checkpoint.consumed_samples > 0:
-                        per_step = (
-                            cfg.batch_per_device
-                            * context.world_size
-                            * cfg.grad_accum_steps
-                        )
-                        context.consumed_samples = (
-                            checkpoint.consumed_samples // per_step
-                        ) * per_step
+        if self.config.nprocs > 1 and self.config.parallel_wrapper:
+            context.model = self.config.parallel_wrapper(context.model)
+
+        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.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.checkpoint = Checkpoint(
+                state_dict=context.model.state_dict(),
            )

-        context.optimizer = cfg.optimizer_fn(model)
-        context.scheduler = cfg.scheduler_fn(context.optimizer)
+        context.optimizer = self.config.optimizer_fn(context.model)
+        context.scheduler = self.config.scheduler_fn(context.optimizer)

-        train_dataset = cfg.dataset
-        val_dataset = cfg.val_dataset
+        if self._checkpoint and self._checkpoint.extra and self._load_extra_fn:
+            self._load_extra_fn(self._checkpoint.extra, context)

-        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 = RDSampler(
-            data_source=train_dataset,
+        cfg = self.config
+        sampler_offset = context.iteration * cfg.batch_size
+        sampler = ResumableDistributedSampler(
+            data_source=cfg.dataset,
            start_epoch=context.epoch,
            start_iter=sampler_offset,
            seed=cfg.random_seed,
        )
        context.dataloader = DataLoader(
-            train_dataset,
-            batch_size=cfg.batch_per_device,
+            cfg.dataset,
+            batch_size=cfg.batch_size,
            sampler=sampler,
            num_workers=cfg.num_workers,
            pin_memory=cfg.pin_memory,
            prefetch_factor=cfg.prefetch_factor,
-            collate_fn=cfg.collate_fn,
        )

-        if val_dataset is not None:
-            val_sampler = RDSampler(
-                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,
-                collate_fn=cfg.collate_fn,
-            )
-
-        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])
-
-        strategy_kwargs = dict(cfg.extra_kwargs)
-
-        if cfg.strategy in ("dpo", "grpo"):
-            ref_model = create_ref_model(
-                cfg.model_fn, executor.unwrap_model(context.model)
-            ).to(device=device)
-            strategy_kwargs["ref_model"] = ref_model
-
-        if cfg.strategy == "grpo":
-            old_model = create_ref_model(
-                cfg.model_fn, executor.unwrap_model(context.model)
-            ).to(device=device)
-            strategy_kwargs["old_model"] = old_model
-
        context.strategy = StrategyFactory.create(
-            cfg.strategy,
            model=context.model,
+            train_type=self.config.strategy,
            device=device,
-            executor=executor,
-            **strategy_kwargs,
+            **self.config.extra_kwargs,
        )

        return context
@@ -3,6 +3,7 @@ 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,
@@ -24,27 +25,18 @@ class Trainer:

    def _get_default_callbacks(self) -> List[TrainCallback]:
        cfg = self.train_config
-        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,
-            ),
+        return [
            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"),
        ]
-        return callbacks
+
+    def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
+        return (
+            TrainContextBuilder(self.train_config).with_checkpoint(checkpoint).build()
+        )

    def _call_callbacks(self, method_name: str, context: TrainContext):
        for callback in self.callbacks:
@@ -52,61 +44,55 @@ class Trainer:
            if method:
                method(context)

-    def _trainer_loop(self, param_path: Optional[str] = None, resume: bool = False):
-        context = (
-            TrainContextBuilder(self.train_config)
-            .with_param_path(param_path, resume=resume)
-            .build()
+    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,
        )
-        executor = context.executor
+
+    def _train_impl(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
+        context = self._build_context(checkpoint)
        self._call_callbacks("on_train_begin", context)

        try:
            context.model.train()
-
-            for epoch in range(context.epoch, context.config.n_epoch):
+            # 1.epoch
+            for epoch in range(context.epoch, self.train_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:
-                    with executor.accumulate(context.model):
+                    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)
+
+                    # 3. batch
                    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
-                        )
+                    context.iteration += 1
+
+                    # to make the loss normalized by accumulation steps
+                    stand_loss = loss / accumulation_steps
+                    stand_loss.backward()
+
                    self._call_callbacks("on_batch_end", context)

-                        if executor.sync_gradients:
-                            self._call_callbacks("on_optimizer_step", context)
-                            context.optimizer.step()
-                            context.optimizer.zero_grad()
-
-                            if context.scheduler:
-                                context.scheduler.step()
-
                self._call_callbacks("on_epoch_end", context)

        except Exception as e:
-            logger.error("Training failed: %s", str(e), exc_info=True)
+            logger.error(f"Training failed: {str(e)}", exc_info=True)
            self._call_callbacks("on_error", context)
            raise
        finally:
            self._call_callbacks("on_train_end", context)
-
-    def train(self, param_path: Optional[str] = None, resume: bool = False):
-        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,
-            param_path=param_path,
-            resume=resume,
-        )
@@ -1,2 +0,0 @@
-# Source directory for CUDA kernels — build-time only.
-# Compiled .so files live in astrAI/_ext/.
@@ -1,48 +0,0 @@
-from pathlib import Path
-
-
-def _arch_flags() -> list[str]:
-    import torch
-
-    if torch.cuda.is_available():
-        cap = torch.cuda.get_device_capability()
-    else:
-        cap = (8, 0)
-    ver = f"{cap[0]}{cap[1]}"
-    flags = [f"-gencode=arch=compute_{ver},code=sm_{ver}"]
-    # tensor-core mma path (mma.sync.m16n8k16.bf16) requires sm_80+; decide the
-    # kernel dispatch at build time via this define rather than at runtime.
-    if cap[0] < 8:
-        flags.append("-DASTRAI_NO_MMA")
-    return flags
-
-
-_kernels_dir = Path("csrc/kernels")
-REGISTRY: dict[str, dict] = {}
-
-CXX_FLAGS = ["-O3", "-funroll-loops"]
-NVCC_FLAGS = [
-    "-O3",
-    "--expt-relaxed-constexpr",
-    "--use_fast_math",
-    "--ptxas-options=-O3,-v",
-    "--extra-device-vectorization",
-    "--threads=8",
-]
-
-
-def register(name: str, sources: list[str] | None = None, **kwargs):
-    if sources is None:
-        sources = [str(_kernels_dir / f"{name}.cu")]
-    REGISTRY[name] = {
-        "sources": sources,
-        "cxx_flags": [*CXX_FLAGS],
-        "nvcc_flags": [*NVCC_FLAGS, *_arch_flags()],
-        "extra_link_args": kwargs.pop("extra_link_args", []),
-        **kwargs,
-    }
-
-
-register("attn_decode")
-register("attn_prefill")
-register("attn_paged_decode")
@@ -1,68 +0,0 @@
-#pragma once
-
-
-template<typename T, typename AT = float>
-struct AttentionParams {
-    int batch;
-    int q_head;
-    int kv_head;
-    int q_len;
-    int kv_len;
-    int head_dim;
-    int use_mask;
-    int causal_offset;   // -1 = non-causal; >=0 = absolute position of first Q token
-    int num_splits;
-    float scale;
-
-    // Q strides (element offsets for each dim — layout-agnostic)
-    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
-    // KV strides (K and V share the same layout — only base pointers differ)
-    int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
-
-    // Mask: 2D [batch, kv_len] (mask_q_stride=0) or 3D [batch, q_len, kv_len]
-    int mask_b_stride;   // = kv_len (both 2D and 3D)
-    int mask_q_stride;   // 2D: 0 (all q rows share); 3D: kv_len
-
-    const T* __restrict__ q;
-    const T* __restrict__ k;
-    const T* __restrict__ v;
-    const bool* __restrict__ mask;
-
-    T* __restrict__ o;
-    AT* __restrict__ o_part;
-    AT* __restrict__ ml_part;
-};
-
-template<typename T, typename AT = float>
-struct PagedAttentionParams {
-    int batch;
-    int q_head;
-    int kv_head;
-    int q_len;
-    int kv_len;
-    int head_dim;
-    int use_mask;
-    int causal_offset;
-    float scale;
-
-    int num_splits;
-    int page_size;
-    int max_pages;
-
-    // Q strides (layout-agnostic)
-    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
-
-    // Mask strides (2D or 3D)
-    int mask_b_stride;
-    int mask_q_stride;
-
-    const T* __restrict__ q;
-    const T* __restrict__ k_cache;
-    const T* __restrict__ v_cache;
-    const bool* __restrict__ mask;
-    const int64_t* __restrict__ page_table;
-
-    T* __restrict__ o;
-    AT* __restrict__ o_part;
-    AT* __restrict__ ml_part;
-};
@@ -1,82 +0,0 @@
-#include "attn_decode_split_kv.cuh"
-#include "attn_entry_utils.cuh"
-
-#ifndef ASTRAI_NO_MMA
-#include "attn_decode_split_kv_mma.cuh"
-#endif
-
-// Scalar fallback: one warp per query head, split-KV across grid.z.
-static void launch_scalar_decode(AttentionParams<bf16>& p) {
-    int group_size = p.q_head / p.kv_head;
-    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
-    alloc_split_partials(p);
-
-    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
-    attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, group_size), smem>>>(p);
-    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
-}
-
-#ifndef ASTRAI_NO_MMA
-// MMA head-packing requires G <= 16 (BR=16 rows). sm_80+ tensor-core
-// + cp.async wins even at G=1 (decode is memory-bound, not compute-bound).
-// STAGES=2 (double-buffer) for D<=128 (smem 16 KB); STAGES=1 for D=256
-// (double-buffer would be 32 KB, near the 48 KB static cap — keep single
-// to preserve occupancy).
-template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
-static void launch_mma_decode(AttentionParams<bf16>& p) {
-    int tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
-    alloc_split_partials(p);
-
-    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
-    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
-}
-#endif
-
-template <int HEAD_DIM>
-static void dispatch_decode(AttentionParams<bf16>& p) {
-#ifndef ASTRAI_NO_MMA
-    int G = p.q_head / p.kv_head;
-    if (G >= 1 && G <= 16) {
-        launch_mma_decode<HEAD_DIM, 32>(p);
-        return;
-    }
-#endif
-    launch_scalar_decode(p);
-}
-
-torch::Tensor attn_decode(
-    torch::Tensor q,
-    torch::Tensor k,
-    torch::Tensor v,
-    c10::optional<torch::Tensor> mask,
-    int64_t causal_offset,
-    double scale,
-    int64_t layout
-) {
-    AttentionParams<bf16> p;
-    attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
-    TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
-    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
-
-    // O matches Q's original layout
-    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
-    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
-
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
-    return O;
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("attn_decode", &attn_decode,
-        py::arg("q"),
-        py::arg("k"),
-        py::arg("v"),
-        py::arg("mask") = py::none(),
-        py::arg("causal_offset") = -1,
-        py::arg("scale") = 0.0,
-        py::arg("layout") = 0,
-        "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
-}
@@ -1,132 +0,0 @@
-#pragma once
-#include <cuda_bf16.h>
-#include <float.h>
-#include "attn_common.h"
-
-using bf16 = __nv_bfloat16;
-constexpr int DC_CHUNK = 64;
-
-__device__ inline float warp_reduce_sum(float val) {
-    for (int offset = 16; offset > 0; offset >>= 1)
-        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
-    return val;
-}
-
-__global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
-    int batch = blockIdx.x / p.kv_head;
-    int kv_head = blockIdx.x % p.kv_head;
-    int split = blockIdx.z;
-    int group_size = blockDim.y;
-    int q_head = kv_head * group_size + threadIdx.y;
-    int lane = threadIdx.x;
-    int hd_per_thread = p.head_dim / 32;
-
-    // Q: [batch, q_head, q_len=1, head_dim] — stride-based
-    float q_reg[8];
-    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
-              + lane * hd_per_thread * p.q_stride_d;
-    for (int i = 0; i < hd_per_thread; i++)
-        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
-
-    // KV: [batch, kv_head, kv_len, head_dim] — stride-based base
-    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    int mask_base = batch * p.mask_b_stride;
-
-    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
-
-    extern __shared__ __align__(16) bf16 k_smem[];
-
-    // Split-KV: each split processes a contiguous subset of chunks
-    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
-    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
-    int ch_begin = split * chunks_per_split;
-    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
-
-    for (int ci = ch_begin; ci < ch_end; ci++) {
-        int chunk_start = ci * DC_CHUNK;
-        int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start);
-
-        // Load K into shared memory (gather from strided global)
-        int total = this_chunk * p.head_dim;
-        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) {
-            int s = i / p.head_dim;
-            int d_dim = i % p.head_dim;
-            int kv_idx = chunk_start + s;
-            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
-            k_smem[i] = p.k[g_off];
-        }
-        __syncthreads();
-
-        for (int s = 0; s < this_chunk; s++) {
-            float partial = 0.0f;
-            for (int i = 0; i < hd_per_thread; i++)
-                partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
-            partial = warp_reduce_sum(partial) * p.scale;
-
-            int kv_idx = chunk_start + s;
-            if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx])
-                partial = -FLT_MAX;
-            if (p.causal_offset >= 0 && kv_idx > p.causal_offset)
-                partial = -FLT_MAX;
-
-            float new_m = fmaxf(m, partial);
-            float alpha = expf(m - new_m);
-            float beta  = expf(partial - new_m);
-            d = d * alpha + beta;
-
-            // V: stride-based read
-            int v_off = kv_base + kv_idx * p.kv_stride_l + lane * hd_per_thread * p.kv_stride_d;
-            for (int i = 0; i < hd_per_thread; i++)
-                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i * p.kv_stride_d]) * beta;
-            m = new_m;
-        }
-        __syncthreads();
-    }
-
-    // ---- write UN-normalised partials for this split ----
-    size_t bh = (size_t)batch * p.q_head + q_head;
-    size_t slot = bh * p.num_splits + split;
-    int d0 = lane * hd_per_thread;
-    for (int i = 0; i < hd_per_thread; i++) {
-        int dd = d0 + i;
-        p.o_part[slot * p.head_dim + dd] = acc_reg[i];
-    }
-    if (lane == 0) {
-        p.ml_part[slot * 2] = m;
-        p.ml_part[slot * 2 + 1] = d;
-    }
-}
-
-// Reduce split-K partials into the final bf16 output. One block per (batch,
-// q_head); each thread folds across all splits with a single-pass
-// online-rescale reduction (expf + FMA counts halved vs 3-pass original).
-__global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
-    int bh = blockIdx.x;
-    int d = threadIdx.x;
-    if (d >= p.head_dim) return;
-
-    int batch = bh / p.q_head;
-    int q_head = bh % p.q_head;
-
-    size_t split_base = (size_t)bh * p.num_splits;
-    const float* mlp = p.ml_part + split_base * 2;
-    const float* op = p.o_part + split_base * p.head_dim;
-
-    float m = -FLT_MAX, l = 0.0f, acc = 0.0f;
-    for (int s = 0; s < p.num_splits; s++) {
-        float mi = mlp[s * 2];
-        if (mi <= -FLT_MAX) continue;
-        float li = mlp[s * 2 + 1];
-        float nm = fmaxf(m, mi);
-        float corr = __expf(m - nm);
-        float e = __expf(mi - nm);
-        acc = acc * corr + op[s * p.head_dim + d] * e;
-        l = l * corr + li * e;
-        m = nm;
-    }
-
-    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
-    // Stride-based output write (q_len=1 for decode, so stride_l not needed)
-    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
-    p.o[o_off] = __float2bfloat16(acc * inv);
-}
@@ -1,176 +0,0 @@
-#pragma once
-#include <cfloat>
-#include <cuda_bf16.h>
-#include "attn_common.h"
-#include "attn_mma_utils.cuh"
-
-using bf16 = __nv_bfloat16;
-
-// Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
-//
-// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core
-// work on its own. But GQA gives us G = q_head / kv_head query heads that all
-// share one kv_head. We pack those G heads into the M=16 rows of
-// mma.sync.m16n8k16, turning G independent GEMVs into a single GEMM that
-// reuses each loaded K/V tile across all G heads (K/V load is the decode
-// bottleneck, so the reuse is the win, not the flops). The KV sequence is
-// partitioned across gridDim.z blocks so that a decode with only
-// batch*kv_head independent tasks can fill all SMs. Each (batch, kv_head,
-// split) block computes an UN-normalised partial (Oacc, m, l) over its KV
-// slice; the combine kernel below reduces across splits. Fixes the "grid too
-// small" bottleneck (0.04 waves/SM → many blocks) for long-context,
-// small-batch decode.
-
-template <int HEAD_DIM, int BC, int STAGES = 2>
-__global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
-    constexpr int KD = HEAD_DIM / 16;
-    constexpr int NC8 = BC / 8;
-    constexpr int KT2 = BC / 16;
-    constexpr int DN8 = HEAD_DIM / 8;
-    constexpr int LD = HEAD_DIM;
-    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
-    constexpr int VEC = 8;
-    constexpr int TOTAL = BC * HEAD_DIM;
-
-    const int lane = threadIdx.x;
-    const int gid = lane >> 2;
-    const int tid4 = lane & 3;
-
-    const int kv_head = blockIdx.x;
-    const int batch = blockIdx.y;
-    const int split = blockIdx.z;
-    const int G = p.q_head / p.kv_head;
-    const int q_head0 = kv_head * G;
-
-    // Double-buffered shared memory for K/V (no sQ needed — Q goes direct
-    // from global to registers).
-    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
-    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
-
-    // ---- Load Q directly from global into mma A-operand registers ----
-    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
-    const int qra = gid;
-    const int qrb = gid + 8;
-    const bool va = qra < G, vb = qrb < G;
-    unsigned Qa[KD][4];
-    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
-                          qra, qrb, va, vb, tid4, Qa);
-
-    float Oacc[DN8][4];
-#pragma unroll
-    for (int j = 0; j < DN8; j++)
-        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
-    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
-
-    // KV: stride-based base — [batch, kv_head, kv_len, head_dim]
-    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    const int tiles_total = (p.kv_len + BC - 1) / BC;
-    const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
-    const int ti_begin = split * tiles_per_split;
-    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
-    const int has_mask = p.use_mask && p.mask;
-
-    // ---- Load tile lambda: predicated cp.async, unified full/partial ----
-    auto load_tile = [&](int ti, int buf) {
-        int kv0 = ti * BC;
-        bf16* dK = sK + buf * BC * LD;
-        bf16* dV = sV + buf * BC * LD;
-#pragma unroll
-        for (int i = lane * VEC; i < TOTAL; i += 32 * VEC) {
-            int r = i / HEAD_DIM, d = i % HEAD_DIM;
-            int kc = kv0 + r;
-            bool valid = kc < p.kv_len;
-            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
-            // KV stride-based: contiguous within head_dim (stride_d == 1 typically)
-            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
-            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
-            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
-        }
-        cp_async_commit();
-    };
-
-    // ---- Prologue: issue first tile load ----
-    if (ti_begin < ti_end) {
-        load_tile(ti_begin, 0);
-    }
-
-    for (int ti = ti_begin; ti < ti_end; ti++) {
-        constexpr int BUF_MASK = (STAGES > 1) ? (STAGES - 1) : 0;
-        int buf = (ti - ti_begin) & BUF_MASK;
-
-        // Wait for current tile, then issue next tile's prefetch (overlaps
-        // with this tile's compute). Single syncwarp covers both hazards.
-        // When STAGES==1, no prefetch — load happens at end of prior iter.
-        cp_async_wait_group<0>();
-        __syncwarp();
-        if constexpr (STAGES > 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
-        }
-
-        const bf16* bK = sK + buf * BC * LD;
-        const bf16* bV = sV + buf * BC * LD;
-        int kv0 = ti * BC;
-
-        float Sacc[NC8][4];
-        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
-
-        #pragma unroll
-        for (int n8 = 0; n8 < NC8; n8++)
-            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
-            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
-
-        // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant
-        int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
-        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
-                                    0, 0,
-                                    p.mask_b_stride, 0,
-                                    batch,
-                                    p.mask, has_mask,
-                                    Sacc, Oacc, m0, m1, l0, l1, lane);
-
-        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
-        __syncwarp();
-
-        if constexpr (STAGES == 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, 0);
-        }
-    }
-
-    // ---- write UN-normalised partials for this split ----
-    auto split_slot = [&](int h) -> size_t {
-        size_t bh = (size_t)batch * p.q_head + h;
-        return bh * p.num_splits + split;
-    };
-#pragma unroll
-    for (int dn8 = 0; dn8 < DN8; dn8++) {
-        int d = dn8 * 8 + 2 * tid4;
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            int h = q_head0 + r0;
-            float* op = p.o_part + split_slot(h) * HEAD_DIM;
-            op[d] = Oacc[dn8][0];
-            op[d + 1] = Oacc[dn8][1];
-        }
-        if (r1 < G) {
-            int h = q_head0 + r1;
-            float* op = p.o_part + split_slot(h) * HEAD_DIM;
-            op[d] = Oacc[dn8][2];
-            op[d + 1] = Oacc[dn8][3];
-        }
-    }
-    if (tid4 == 0) {
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            int h = q_head0 + r0;
-            float* mp = p.ml_part + split_slot(h) * 2;
-            mp[0] = m0; mp[1] = l0;
-        }
-        if (r1 < G) {
-            int h = q_head0 + r1;
-            float* mp = p.ml_part + split_slot(h) * 2;
-            mp[0] = m1; mp[1] = l1;
-        }
-    }
-}
@@ -1,177 +0,0 @@
-#pragma once
-#include <torch/extension.h>
-#include <c10/cuda/CUDAGuard.h>
-#include "attn_common.h"
-
-using bf16 = __nv_bfloat16;
-
-inline int compute_num_splits(int base_blocks, int tiles_total) {
-    int sm_count = 0;
-    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
-    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
-    return std::max(1, std::min(n, std::min(tiles_total, 32)));
-}
-
-// Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
-// Usage: DISPATCH_HEAD_DIM(hd, fn, arg)
-//   Expands to: fn<32>(arg); fn<64>(arg); etc.
-#define DISPATCH_HEAD_DIM(hd, fn, arg) \
-    switch (hd) { \
-        case 32:  fn<32>(arg); break; \
-        case 64:  fn<64>(arg); break; \
-        case 128: fn<128>(arg); break; \
-        case 256: fn<256>(arg); break; \
-        default: \
-            TORCH_CHECK(false, "unsupported head_dim ", hd, \
-                         " (supported: 32, 64, 128, 256)"); \
-    }
-
-template<typename P>
-inline void alloc_split_partials(P& p) {
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
-    p.o_part = (float*)o_part.data_ptr();
-    p.ml_part = (float*)ml_part.data_ptr();
-}
-
-// ---- Shared Q-dims + strides extraction ----
-template <typename P>
-inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
-    if (layout == 1) q = q.transpose(1, 2);
-    p.batch = (int)q.size(0);
-    p.q_head = (int)q.size(1);
-    p.q_len = (int)q.size(2);
-    p.head_dim = (int)q.size(3);
-    p.q_stride_b = (int)q.stride(0);
-    p.q_stride_h = (int)q.stride(1);
-    p.q_stride_l = (int)q.stride(2);
-    p.q_stride_d = (int)q.stride(3);
-}
-
-// ---- Shared mask packing ----
-template <typename P>
-inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
-    if (p.use_mask) {
-        auto m = mask.value();
-        TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
-        TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
-        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
-        TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
-        if (m.dim() == 2) {
-            p.mask_b_stride = (int)m.stride(0);
-            p.mask_q_stride = 0;
-        } else if (m.dim() == 3) {
-            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
-            p.mask_b_stride = (int)m.stride(0);
-            p.mask_q_stride = (int)m.stride(1);
-        } else {
-            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
-        }
-        p.mask = m.data_ptr<bool>();
-    } else {
-        p.mask = nullptr;
-        p.mask_b_stride = 0;
-        p.mask_q_stride = 0;
-    }
-}
-
-// ---- attn_pack_params (contiguous KV) ----
-template<typename T>
-inline void attn_pack_params(
-    torch::Tensor q,
-    torch::Tensor k,
-    torch::Tensor v,
-    c10::optional<torch::Tensor> mask,
-    int64_t causal_offset,
-    double scale,
-    int64_t layout,
-    AttentionParams<T>& p
-) {
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
-
-    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
-    TORCH_CHECK(q.dtype() == torch::kBFloat16);
-    TORCH_CHECK(k.dtype() == torch::kBFloat16);
-    TORCH_CHECK(v.dtype() == torch::kBFloat16);
-    TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
-    TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
-
-    extract_q_dims_and_strides(q, layout, p);
-
-    if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2);
-
-    p.kv_head = (int)k.size(1);
-    p.kv_len = (int)k.size(2);
-    TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
-
-    p.kv_stride_b = (int)k.stride(0);
-    p.kv_stride_h = (int)k.stride(1);
-    p.kv_stride_l = (int)k.stride(2);
-    p.kv_stride_d = (int)k.stride(3);
-
-    p.causal_offset = (int)causal_offset;
-    p.use_mask = mask.has_value() ? 1 : 0;
-    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
-
-    p.q = (const T*)q.data_ptr();
-    p.k = (const T*)k.data_ptr();
-    p.v = (const T*)v.data_ptr();
-    p.o = nullptr;
-    p.o_part = nullptr;
-    p.ml_part = nullptr;
-
-    pack_mask(mask, p);
-}
-
-// ---- attn_pack_paged_params ----
-template<typename T>
-inline void attn_pack_paged_params(
-    torch::Tensor q,
-    torch::Tensor page_table,
-    torch::Tensor k_cache,
-    torch::Tensor v_cache,
-    int64_t page_size,
-    int64_t kv_len,
-    c10::optional<torch::Tensor> mask,
-    int64_t causal_offset,
-    double scale,
-    int64_t layout,
-    PagedAttentionParams<T>& p
-) {
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
-
-    TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
-    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16");
-    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16");
-    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16");
-    TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
-    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
-
-    extract_q_dims_and_strides(q, layout, p);
-
-    p.kv_head = (int)k_cache.size(2);
-    p.kv_len = (int)kv_len;
-    p.page_size = (int)page_size;
-    p.max_pages = (int)page_table.size(1);
-
-    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
-    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
-    TORCH_CHECK(k_cache.size(1) == page_size,
-                "k_cache dim 1 must equal page_size, got ",
-                k_cache.size(1), " vs ", page_size);
-
-    p.causal_offset = (int)causal_offset;
-    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
-    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
-
-    p.page_table = page_table.data_ptr<int64_t>();
-    p.k_cache = (const T*)k_cache.data_ptr();
-    p.v_cache = (const T*)v_cache.data_ptr();
-    p.q = (const T*)q.data_ptr();
-    p.o = nullptr;
-    p.o_part = nullptr;
-    p.ml_part = nullptr;
-
-    pack_mask(mask, p);
-}
@@ -1,293 +0,0 @@
-#pragma once
-#include <cfloat>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-
-// Shared MMA utilities for tensor-core GQA kernels.
-// mma.sync.m16n8k16 PTX wrappers, ldmatrix helpers, and bf16 packing.
-
-// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
-__device__ __forceinline__ void mma16816(float* d, const unsigned* a,
-                                          const unsigned* b, const float* c) {
-    asm volatile(
-        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
-        : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
-        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
-          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
-}
-
-// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
-__device__ __forceinline__ unsigned ld2(const bf16* p) {
-    return *reinterpret_cast<const unsigned*>(p);
-}
-
-// pack two floats into one bf16x2 as .b32
-__device__ __forceinline__ unsigned pk2(float a, float b) {
-    __nv_bfloat162 v = __floats2bfloat162_rn(a, b);
-    return *reinterpret_cast<unsigned*>(&v);
-}
-
-// pack two (non-contiguous) bf16 into one .b32
-__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
-    __nv_bfloat162 v;
-    v.x = a;
-    v.y = b;
-    return *reinterpret_cast<unsigned*>(&v);
-}
-
-// ldmatrix: cooperatively load mma fragments from smem (one instruction per
-// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
-// scalar per-thread fragment packing, cutting shared-load instructions and bank
-// conflicts. Each lane supplies the shared address of one 8-wide row.
-__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
-                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
-                 : "r"(a));
-}
-__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
-                 : "=r"(r[0]), "=r"(r[1])
-                 : "r"(a));
-}
-__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
-                 : "=r"(r[0]), "=r"(r[1])
-                 : "r"(a));
-}
-
-// XOR swizzle for shared-memory column at 8-bf16 chunk granularity.
-// Eliminates ldmatrix bank conflicts without LD padding: consecutive rows
-// land in distinct bank groups. swiz_col(d, r, mask) = ((d>>3)^(r&mask))<<3 | (d&7).
-// mask must cover log2(HEAD_DIM/8) chunk bits but stay within LD: use 7 for
-// HEAD_DIM>=64 (8+ chunks), 3 for HEAD_DIM=32 (4 chunks). Default 7 keeps
-// existing HEAD_DIM>=64 call sites working unchanged.
-__device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
-    return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
-}
-
-// cp.async: copy 16 bytes (8 bf16) from global to shared memory directly,
-// bypassing registers. Eliminates shared-store bank conflicts and cuts
-// load-loop instruction count in half (1 cp.async vs 1 LDG + 1 STS).
-// Requires sm_80+.
-__device__ __forceinline__ void cp_async_16(bf16* smem_ptr, const void* gmem_ptr) {
-    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
-    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;"
-                 :: "r"(smem_addr), "l"(gmem_ptr));
-}
-
-// Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill the
-// destination (src-size operand = 0 → no bytes read from src, so an
-// out-of-bounds src address is never dereferenced). Lets full and partial
-// tiles share one uniform async load path — no scalar fallback branch.
-__device__ __forceinline__ void cp_async_16_pred(bf16* smem_ptr,
-                                                  const void* gmem_ptr,
-                                                  bool pred) {
-    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
-    int src_size = pred ? 16 : 0;
-    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
-                 :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
-}
-
-__device__ __forceinline__ void cp_async_commit() {
-    asm volatile("cp.async.commit_group;");
-}
-
-__device__ __forceinline__ void cp_async_wait_all() {
-    asm volatile("cp.async.wait_all;");
-}
-
-// Wait until at most N commit groups are still in flight. Used for
-// double-buffered pipelining: wait_group<1> lets the next tile's cp.async
-// continue while ensuring the current tile's data is ready.
-template <int N>
-__device__ __forceinline__ void cp_async_wait_group() {
-    asm volatile("cp.async.wait_group %0;" :: "n"(N));
-}
-
-// ---------------------------------------------------------------------------
-// Q-load: load query rows directly from global memory into mma A-operand
-// register layout. One call replaces ~15 duplicated lines in each MMA kernel.
-// stride_row is p.q_stride_h for decode (q_len=1, G heads) or
-//              p.q_stride_l for prefill (multi-q rows).
-// ---------------------------------------------------------------------------
-template <int KD>
-__device__ inline void load_q_mma_frags(
-    const bf16* __restrict__ q,
-    int stride_row,
-    int stride_d,
-    int qra, int qrb,
-    bool va, bool vb,
-    int tid4,
-    unsigned Qa[KD][4])
-{
-    #pragma unroll
-    for (int kt = 0; kt < KD; kt++) {
-        int c = kt * 16 + tid4 * 2;
-        const unsigned* pau = reinterpret_cast<const unsigned*>(
-            &q[qra * stride_row + c * stride_d]);
-        const unsigned* pbu = reinterpret_cast<const unsigned*>(
-            &q[qrb * stride_row + c * stride_d]);
-        Qa[kt][0] = va ? pau[0] : 0u;
-        Qa[kt][1] = vb ? pbu[0] : 0u;
-        Qa[kt][2] = va ? pau[4] : 0u;
-        Qa[kt][3] = vb ? pbu[4] : 0u;
-    }
-}
-
-// ---------------------------------------------------------------------------
-// Shared MMA compute functions — used by both decode and prefill MMA kernels.
-// Extracted because S=Q@K^T, online softmax, and P@V are structurally identical
-// between the two kernels; only the per-row causal/mask bounds differ.
-// ---------------------------------------------------------------------------
-
-// S = Q @ K^T  (Qa pre-loaded by the caller; scale applied post-mma in the
-// caller to avoid bf16 precision loss).
-// LD and SWIZ_MASK are constexpr in the calling kernel — passing them as
-// runtime ints lets the compiler fold them while keeping the signature clean.
-template <int KD, int NC8>
-__device__ inline void mma_compute_scores(
-    const unsigned Qa[KD][4],
-    const bf16* __restrict__ sK,
-    int LD, 
-    int SWIZ_MASK, 
-    int lane,
-    float Sacc[NC8][4])
-{
-    #pragma unroll
-    for (int n8 = 0; n8 < NC8; n8++) {
-        Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
-        int krow_l = n8 * 8 + (lane & 7);
-        int kcol_h = (lane & 8) ? 8 : 0;
-        #pragma unroll
-        for (int kt = 0; kt < KD; kt++) {
-            unsigned b[2];
-            ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
-            mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
-        }
-    }
-}
-
-// Online softmax + Oacc rescale for one K/V tile.
-//   maxc0/maxc1: per-row KV column bounds (prefill: per-query-row causal limits;
-//     decode: same value for both rows since q_len==1).
-//   qrow0/qrow1: query row indices (for 3D mask indexing; decode passes 0).
-//   mask_b_stride/mask_q_stride: mask layout (2D: mask_q_stride=0; 3D: =kv_len).
-// Reads Sacc (Q@K^T scores), applies causal/mask, computes P = exp(S - nm),
-// rescales Oacc by exp(m_old - nm), and updates m/l — all in place.
-template <int NC8, int DN8>
-__device__ inline void mma_softmax_tile(
-    int kv0,
-    int maxc0, 
-    int maxc1,
-    int qrow0, 
-    int qrow1,
-    int mask_b_stride, 
-    int mask_q_stride,
-    int mask_batch,
-    const bool* __restrict__ mask,
-    bool has_mask,
-    float Sacc[NC8][4],
-    float Oacc[DN8][4],
-    float& m0, float& m1,
-    float& l0, float& l1,
-    int lane)
-{
-    int tid4 = lane & 3;
-
-    // Mask out-of-bounds / masked columns: set -FLT_MAX so expf → 0 downstream
-    // without per-element sentinel checks. Compute tile-local row maxima.
-    float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
-    int mask_base0 = mask_batch * mask_b_stride + qrow0 * mask_q_stride;
-    int mask_base1 = mask_batch * mask_b_stride + qrow1 * mask_q_stride;
-    #pragma unroll
-    for (int n8 = 0; n8 < NC8; n8++) {
-        int cc = kv0 + n8 * 8 + 2 * tid4;
-        int c1 = cc + 1;
-        bool b0 = (cc >= maxc0) || (has_mask && !mask[mask_base0 + cc]);
-        bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base0 + c1]);
-        bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base1 + cc]);
-        bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base1 + c1]);
-        float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
-        float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
-        float s2 = b2 ? -FLT_MAX : Sacc[n8][2];
-        float s3 = b3 ? -FLT_MAX : Sacc[n8][3];
-        Sacc[n8][0] = s0; Sacc[n8][1] = s1;
-        Sacc[n8][2] = s2; Sacc[n8][3] = s3;
-        rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
-        rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
-    }
-    // Warp-reduce row maxima across the 4-lane thread group (xor 1, xor 2).
-    rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
-    rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
-    rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
-    rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2));
-
-    // nm = max(running max m, tile-local max rmax) — updated running maximum.
-    float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1);
-    // corr rescales Oacc and l by exp(m_old - nm). When all-masked (m == nm ==
-    // -FLT_MAX), exp(0) = 1 — correct, no guard needed.
-    float corr0 = __expf(m0 - nm0);
-    float corr1 = __expf(m1 - nm1);
-    // pn guards only the all-masked-row edge: if nm == -FLT_MAX, exp(S - nm)
-    // gives 1 not 0 for masked entries. Two scalar masks replace 4*NC8
-    // per-element comparisons.
-    float pn0 = (nm0 == -FLT_MAX) ? 0.0f : 1.0f;
-    float pn1 = (nm1 == -FLT_MAX) ? 0.0f : 1.0f;
-
-    // P = exp(S - nm) for each element. Masked entries (Sacc = -FLT_MAX) give
-    // exp(-inf) ≈ 0 naturally; pn zero-fills the all-masked-row edge.
-    float rsum0 = 0.0f, rsum1 = 0.0f;
-    #pragma unroll
-    for (int n8 = 0; n8 < NC8; n8++) {
-        float p0 = pn0 * __expf(Sacc[n8][0] - nm0);
-        float p1 = pn0 * __expf(Sacc[n8][1] - nm0);
-        float p2 = pn1 * __expf(Sacc[n8][2] - nm1);
-        float p3 = pn1 * __expf(Sacc[n8][3] - nm1);
-        Sacc[n8][0] = p0; Sacc[n8][1] = p1;
-        Sacc[n8][2] = p2; Sacc[n8][3] = p3;
-        rsum0 += p0 + p1;
-        rsum1 += p2 + p3;
-    }
-    rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1);
-    rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2);
-    rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1);
-    rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2);
-    l0 = l0 * corr0 + rsum0;
-    l1 = l1 * corr1 + rsum1;
-    m0 = nm0; m1 = nm1;
-
-    #pragma unroll
-    for (int j = 0; j < DN8; j++) {
-        Oacc[j][0] *= corr0; Oacc[j][1] *= corr0;
-        Oacc[j][2] *= corr1; Oacc[j][3] *= corr1;
-    }
-}
-
-// O += P @ V  (Sacc must contain P = attention weights after softmax).
-template <int DN8, int KT2>
-__device__ inline void mma_pv_accumulate(
-    float Sacc[][4],
-    const bf16* __restrict__ sV,
-    int LD, int SWIZ_MASK, int lane,
-    float Oacc[DN8][4])
-{
-    #pragma unroll
-    for (int kt2 = 0; kt2 < KT2; kt2++) {
-        unsigned Pa[4];
-        Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]);
-        Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
-        Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
-        Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
-        int vrow_l = kt2 * 16 + (lane & 15);
-        #pragma unroll
-        for (int dn8 = 0; dn8 < DN8; dn8++) {
-            unsigned b[2];
-            ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
-            mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
-        }
-    }
-}
@@ -1,82 +0,0 @@
-#include "attn_paged_decode_split_kv.cuh"
-#ifndef ASTRAI_NO_MMA
-#include "attn_paged_decode_split_kv_mma.cuh"
-#endif
-
-#include "attn_entry_utils.cuh"
-
-static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
-    int group_size = p.q_head / p.kv_head;
-    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
-    alloc_split_partials(p);
-
-    size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
-    dim3 grid = dim3(p.batch * p.kv_head, 1, p.num_splits);
-    dim3 block = dim3(32, group_size);
-    paged_attn_decode_split_kv_kernel<<<grid, block, smem>>>(p);
-    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
-}
-
-#ifndef ASTRAI_NO_MMA
-template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
-static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
-    int tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
-    alloc_split_partials(p);
-
-    paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
-    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
-}
-#endif
-
-template <int HEAD_DIM>
-static void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
-#ifndef ASTRAI_NO_MMA
-    int G = p.q_head / p.kv_head;
-    if (G >= 1 && G <= 16 && p.page_size >= 32) {
-        launch_paged_mma_decode<HEAD_DIM, 32>(p);
-        return;
-    }
-#endif
-    launch_paged_scalar_decode(p);
-}
-
-torch::Tensor attn_paged_decode(
-    torch::Tensor q,
-    torch::Tensor page_table,
-    torch::Tensor k_cache,
-    torch::Tensor v_cache,
-    int64_t page_size,
-    int64_t kv_len,
-    c10::optional<torch::Tensor> mask,
-    int64_t causal_offset,
-    double scale,
-    int64_t layout
-) {
-    PagedAttentionParams<bf16> p;
-    attn_pack_paged_params(q, page_table, k_cache, v_cache,
-                           page_size, kv_len, mask, causal_offset, scale, layout, p);
-
-    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
-    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
-
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
-    return O;
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("attn_paged_decode", &attn_paged_decode,
-        py::arg("q"),
-        py::arg("page_table"),
-        py::arg("k_cache"),
-        py::arg("v_cache"),
-        py::arg("page_size"),
-        py::arg("kv_len"),
-        py::arg("mask") = py::none(),
-        py::arg("causal_offset") = -1,
-        py::arg("scale") = 0.0,
-        py::arg("layout") = 0,
-        "Paged GQA decode — split-KV with direct page-table access.");
-}
@@ -1,147 +0,0 @@
-#pragma once
-#include <cuda_bf16.h>
-#include <float.h>
-#include "attn_common.h"
-
-using bf16 = __nv_bfloat16;
-constexpr int PDC_CHUNK = 64;
-
-__device__ inline float paged_warp_reduce_sum(float val) {
-    for (int offset = 16; offset > 0; offset >>= 1)
-        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
-    return val;
-}
-
-// Split-KV scalar decode: one warp per query head, grid.z partitions KV.
-__global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p) {
-    int batch = blockIdx.x / p.kv_head;
-    int kv_head = blockIdx.x % p.kv_head;
-    int split = blockIdx.z;
-    int group_size = blockDim.y;
-    int q_head = kv_head * group_size + threadIdx.y;
-    int lane = threadIdx.x;
-    int hd_per_thread = p.head_dim / 32;
-
-    // Q: stride-based [batch, q_head, q_len=1, head_dim]
-    float q_reg[8];
-    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
-              + lane * hd_per_thread * p.q_stride_d;
-    #pragma unroll
-    for (int i = 0; i < hd_per_thread; i++)
-        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
-
-    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
-
-    extern __shared__ __align__(16) bf16 k_smem[];
-
-    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
-    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
-    int ch_begin = split * chunks_per_split;
-    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
-
-    const int mask_base = batch * p.mask_b_stride;
-
-    for (int ci = ch_begin; ci < ch_end; ci++) {
-        int chunk_start = ci * PDC_CHUNK;
-        int this_chunk = min(PDC_CHUNK, p.kv_len - chunk_start);
-
-        int total = this_chunk * p.head_dim;
-        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) {
-            int s = i / p.head_dim;
-            int d_dim = i % p.head_dim;
-            int pos = chunk_start + s;
-            int logical_page = pos / p.page_size;
-            int page_offset = pos % p.page_size;
-            int phys_page = p.page_table[batch * p.max_pages + logical_page];
-            if (phys_page >= 0) {
-                int64_t off = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
-                            + (int64_t)page_offset * p.kv_head * p.head_dim
-                            + (int64_t)kv_head * p.head_dim
-                            + d_dim;
-                k_smem[i] = p.k_cache[off];
-            } else {
-                k_smem[i] = __float2bfloat16(0.0f);
-            }
-        }
-        __syncthreads();
-
-        for (int s = 0; s < this_chunk; s++) {
-            float partial = 0.0f;
-            #pragma unroll
-            for (int i = 0; i < hd_per_thread; i++)
-                partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
-            partial = paged_warp_reduce_sum(partial) * p.scale;
-
-            int kv_idx = chunk_start + s;
-            if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx])
-                partial = -FLT_MAX;
-            if (p.causal_offset >= 0 && kv_idx > p.causal_offset)
-                partial = -FLT_MAX;
-
-            float new_m = fmaxf(m, partial);
-            float alpha = expf(m - new_m);
-            float beta  = expf(partial - new_m);
-            d = d * alpha + beta;
-
-            int pos = chunk_start + s;
-            int logical_page = pos / p.page_size;
-            int page_offset = pos % p.page_size;
-            int phys_page = p.page_table[batch * p.max_pages + logical_page];
-            if (phys_page >= 0) {
-                int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
-                               + (int64_t)page_offset * p.kv_head * p.head_dim
-                               + (int64_t)kv_head * p.head_dim;
-                #pragma unroll
-                for (int i = 0; i < hd_per_thread; i++)
-                    acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v_cache[v_base + lane * hd_per_thread + i]) * beta;
-            } else {
-                #pragma unroll
-                for (int i = 0; i < hd_per_thread; i++)
-                    acc_reg[i] = acc_reg[i] * alpha + 0.0f * beta;
-            }
-            m = new_m;
-        }
-        __syncthreads();
-    }
-
-    size_t bh = (size_t)batch * p.q_head + q_head;
-    size_t slot = bh * p.num_splits + split;
-    int d0 = lane * hd_per_thread;
-    #pragma unroll
-    for (int i = 0; i < hd_per_thread; i++)
-        p.o_part[slot * p.head_dim + (d0 + i)] = acc_reg[i];
-    if (lane == 0) {
-        p.ml_part[slot * 2] = m;
-        p.ml_part[slot * 2 + 1] = d;
-    }
-}
-
-__global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
-    int bh = blockIdx.x;
-    int d = threadIdx.x;
-    if (d >= p.head_dim) return;
-
-    int batch = bh / p.q_head;
-    int q_head = bh % p.q_head;
-
-    size_t split_base = (size_t)bh * p.num_splits;
-    const float* mlp = p.ml_part + split_base * 2;
-    const float* op = p.o_part + split_base * p.head_dim;
-
-    float m = -FLT_MAX, l = 0.0f, acc = 0.0f;
-    for (int s = 0; s < p.num_splits; s++) {
-        float mi = mlp[s * 2];
-        if (mi <= -FLT_MAX) continue;
-        float li = mlp[s * 2 + 1];
-        float nm = fmaxf(m, mi);
-        float corr = __expf(m - nm);
-        float e = __expf(mi - nm);
-        acc = acc * corr + op[s * p.head_dim + d] * e;
-        l = l * corr + li * e;
-        m = nm;
-    }
-
-    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
-    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
-    p.o[o_off] = __float2bfloat16(acc * inv);
-}
@@ -1,170 +0,0 @@
-#pragma once
-#include <cfloat>
-#include <cuda_bf16.h>
-#include "attn_common.h"
-#include "attn_mma_utils.cuh"
-
-using bf16 = __nv_bfloat16;
-
-// Paged split-KV tensor-core decode via GQA head-packing.
-// Identical algorithm to attn_decode_split_kv_mma_kernel but reads K/V
-// directly from the page pool through a page table, eliminating the gather
-// copy.  Each tile (BC=32) fits within a single page (page_size >= 32), so
-// the page-table lookup happens once per tile for cp.async.
-
-template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
-__global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
-    constexpr int KD = HEAD_DIM / 16;
-    constexpr int NC8 = BC / 8;
-    constexpr int KT2 = BC / 16;
-    constexpr int DN8 = HEAD_DIM / 8;
-    constexpr int LD = HEAD_DIM;
-    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
-    constexpr int VEC = 8;
-    constexpr int TOTAL = BC * HEAD_DIM;
-
-    const int lane = threadIdx.x;
-    const int gid = lane >> 2;
-    const int tid4 = lane & 3;
-
-    const int kv_head = blockIdx.x;
-    const int batch = blockIdx.y;
-    const int split = blockIdx.z;
-    const int G = p.q_head / p.kv_head;
-    const int q_head0 = kv_head * G;
-
-    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
-    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
-
-    // ---- Load Q directly from global into mma A-operand registers ----
-    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
-    const int qra = gid;
-    const int qrb = gid + 8;
-    const bool va = qra < G, vb = qrb < G;
-    unsigned Qa[KD][4];
-    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
-                          qra, qrb, va, vb, tid4, Qa);
-
-    float Oacc[DN8][4];
-#pragma unroll
-    for (int j = 0; j < DN8; j++)
-        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
-    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
-
-    const int tiles_total = (p.kv_len + BC - 1) / BC;
-    const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
-    const int ti_begin = split * tiles_per_split;
-    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
-    const int has_mask = p.use_mask && p.mask;
-
-    // Paged strides (constant for the block)
-    const int64_t page_stride = (int64_t)p.page_size * p.kv_head * HEAD_DIM;
-    const int64_t pos_stride  = (int64_t)p.kv_head * HEAD_DIM;
-    const int64_t head_off    = (int64_t)kv_head * HEAD_DIM;
-
-    // ---- Load tile lambda: predicated cp.async, paged addressing ----
-    auto load_tile = [&](int ti, int buf) {
-        int kv0 = ti * BC;
-        bf16* dK = sK + buf * BC * LD;
-        bf16* dV = sV + buf * BC * LD;
-        int logical_page = kv0 / p.page_size;
-        int phys_page = p.page_table[batch * p.max_pages + logical_page];
-        bool page_valid = (phys_page >= 0);
-#pragma unroll
-        for (int i = lane * VEC; i < TOTAL; i += 32 * VEC) {
-            int r = i / HEAD_DIM, d = i % HEAD_DIM;
-            int kc = kv0 + r;
-            bool valid = (kc < p.kv_len) && page_valid;
-            int page_off = kc % p.page_size;
-            int64_t gmem_base = (int64_t)phys_page * page_stride
-                              + (int64_t)page_off * pos_stride
-                              + head_off;
-            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
-            cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
-            cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
-        }
-        cp_async_commit();
-    };
-
-    // ---- Prologue: issue first tile load ----
-    if (ti_begin < ti_end) {
-        load_tile(ti_begin, 0);
-    }
-
-    for (int ti = ti_begin; ti < ti_end; ti++) {
-        constexpr int BUF_MASK = (STAGES > 1) ? (STAGES - 1) : 0;
-        int buf = (ti - ti_begin) & BUF_MASK;
-
-        cp_async_wait_group<0>();
-        __syncwarp();
-        if constexpr (STAGES > 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
-        }
-
-        const bf16* bK = sK + buf * BC * LD;
-        const bf16* bV = sV + buf * BC * LD;
-        int kv0 = ti * BC;
-
-        float Sacc[NC8][4];
-        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
-
-#pragma unroll
-        for (int n8 = 0; n8 < NC8; n8++)
-            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
-            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
-
-        // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant
-        int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
-        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
-                                    0, 0,
-                                    p.mask_b_stride, 0,
-                                    batch,
-                                    p.mask, has_mask,
-                                    Sacc, Oacc, m0, m1, l0, l1, lane);
-
-        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
-        __syncwarp();
-
-        if constexpr (STAGES == 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, 0);
-        }
-    }
-
-    // ---- write UN-normalised partials for this split ----
-    auto split_slot = [&](int h) -> size_t {
-        size_t bh = (size_t)batch * p.q_head + h;
-        return bh * p.num_splits + split;
-    };
-#pragma unroll
-    for (int dn8 = 0; dn8 < DN8; dn8++) {
-        int d = dn8 * 8 + 2 * tid4;
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            int h = q_head0 + r0;
-            float* op = p.o_part + split_slot(h) * HEAD_DIM;
-            op[d] = Oacc[dn8][0];
-            op[d + 1] = Oacc[dn8][1];
-        }
-        if (r1 < G) {
-            int h = q_head0 + r1;
-            float* op = p.o_part + split_slot(h) * HEAD_DIM;
-            op[d] = Oacc[dn8][2];
-            op[d + 1] = Oacc[dn8][3];
-        }
-    }
-    if (tid4 == 0) {
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            int h = q_head0 + r0;
-            float* mp = p.ml_part + split_slot(h) * 2;
-            mp[0] = m0; mp[1] = l0;
-        }
-        if (r1 < G) {
-            int h = q_head0 + r1;
-            float* mp = p.ml_part + split_slot(h) * 2;
-            mp[0] = m1; mp[1] = l1;
-        }
-    }
-}
@@ -1,64 +0,0 @@
-#include "attn_prefill_split_q.cuh"
-#include "attn_entry_utils.cuh"
-
-#ifndef ASTRAI_NO_MMA
-#include "attn_prefill_split_q_mma.cuh"
-#endif
-
-template <int HEAD_DIM>
-static void dispatch_prefill(AttentionParams<bf16>& p) {
-#ifndef ASTRAI_NO_MMA
-    constexpr int WARPS = 4, BR = 16;
-    // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier +
-    // loop overhead over more tensor-core work (this kernel is latency-bound,
-    // not compute/bandwidth-bound), so BC=32 wins ~6-8% over BC=16 for
-    // D<=128. D=256 stays at 16: BC=32 double-buffered would need 64KB smem,
-    // over the 48KB static cap.
-    constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
-    dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
-    dim3 block(WARPS * 32, 1, 1);
-    // Static shared memory — double-buffered K/V only (no sQ: Q goes direct
-    // to registers). 2*BC*LD bf16 each for sK and sV → 4*BC*HEAD_DIM*2 bytes.
-    // Occupancy is smem-capped: D=64→3 blocks/SM (16KB), D=128→1 (32KB),
-    // D=256→1 (32KB, BC=16).
-    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC><<<grid, block>>>(p);
-#else
-    constexpr int G = 8, ROWS = 32, P_BC = 32;
-    dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
-    dim3 block(G, ROWS, 1);
-    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
-#endif
-}
-
-torch::Tensor attn_prefill(
-    torch::Tensor q,
-    torch::Tensor k,
-    torch::Tensor v,
-    c10::optional<torch::Tensor> mask,
-    int64_t causal_offset,
-    double scale,
-    int64_t layout
-) {
-    AttentionParams<bf16> p;
-    attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
-    TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
-
-    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
-    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
-
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
-    return O;
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("attn_prefill", &attn_prefill,
-        py::arg("q"),
-        py::arg("k"),
-        py::arg("v"),
-        py::arg("mask") = py::none(),
-        py::arg("causal_offset") = -1,
-        py::arg("scale") = 0.0,
-        py::arg("layout") = 0,
-        "GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
-}
@@ -1,152 +0,0 @@
-#pragma once
-#include <cfloat>
-#include <cuda_bf16.h>
-#include "attn_common.h"
-
-using bf16 = __nv_bfloat16;
-
-// v9: group-split register blocking. G threads cooperate on one query row,
-// each owning HEAD_DIM/G dims of qreg[]/acc[]. Small per-thread footprint keeps
-// occupancy high; the S dot product is reduced across the G-lane group with a
-// short shuffle chain (log2(G) shuffles) instead of a full 32-lane warp reduce.
-// Online (per-kv) softmax — cheap because acc[] is only HEAD_DIM/G long.
-// Templated on <HEAD_DIM, G, ROWS, P_BC>. Block = (G, ROWS). G power-of-two,
-// G*ROWS a multiple of 32 with groups warp-aligned.
-
-template <int G>
-__device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) {
-#pragma unroll
-    for (int o = G / 2; o > 0; o >>= 1)
-        v += __shfl_xor_sync(mask, v, o);
-    return v;
-}
-
-// load 8 contiguous bf16 from (16-byte aligned) smem as one float4, unpack to
-// 8 floats — cuts shared-load instructions 8x vs scalar bf16 loads.
-__device__ __forceinline__ void ld8(const bf16* p, float* o) {
-    float4 raw = *reinterpret_cast<const float4*>(p);
-    const __nv_bfloat162* h = reinterpret_cast<const __nv_bfloat162*>(&raw);
-#pragma unroll
-    for (int j = 0; j < 4; j++) {
-        float2 f = __bfloat1622float2(h[j]);
-        o[2 * j] = f.x;
-        o[2 * j + 1] = f.y;
-    }
-}
-
-template <int HEAD_DIM, int G, int ROWS, int P_BC>
-__global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
-    constexpr int DPT = HEAD_DIM / G;
-
-    int q_tile = blockIdx.x;
-    int q_head = blockIdx.y;
-    int batch  = blockIdx.z;
-    int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
-    int row    = threadIdx.y;  // 0..ROWS-1
-    int q_row  = q_tile * ROWS + row;
-
-    int kv_head = q_head / (p.q_head / p.kv_head);
-
-    __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM];
-    __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM];
-
-    // Q: stride-based load [batch, q_head, q_len, head_dim]
-    float qreg[DPT];
-    if (q_row < p.q_len) {
-        int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
-                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
-#pragma unroll
-        for (int i = 0; i < DPT; i++)
-            qreg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]) * p.scale;
-    }
-
-    float m = -FLT_MAX, l = 0.0f;
-    float acc[DPT];
-#pragma unroll
-    for (int i = 0; i < DPT; i++)
-        acc[i] = 0.0f;
-
-    // KV: stride-based base
-    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    int mask_batch_base = batch * p.mask_b_stride;
-    int tiles   = (p.kv_len + P_BC - 1) / P_BC;
-    int tt      = G * ROWS;
-    int lid     = row * G + gpos;
-
-    // per-group shuffle mask: only the G lanes of this row's group participate,
-    // so causal masking (differing loop bounds across rows in a warp) is safe.
-    int lane_in_warp = lid & 31;
-    unsigned gmask = (G == 32) ? 0xFFFFFFFFu
-                               : (((1u << G) - 1u) << (lane_in_warp & ~(G - 1)));
-
-    for (int ti = 0; ti < tiles; ti++) {
-        int kv0  = ti * P_BC;
-        int tlen = min(P_BC, p.kv_len - kv0);
-
-        // Load K/V into shared memory from strided global
-        for (int i = lid; i < tlen * HEAD_DIM; i += tt) {
-            int s = i / HEAD_DIM;
-            int d_dim = i % HEAD_DIM;
-            int kv_idx = kv0 + s;
-            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
-            sK[i] = p.k[g_off];
-            sV[i] = p.v[g_off];
-        }
-        __syncthreads();
-
-        int lim = tlen;
-        if (p.causal_offset >= 0 && q_row < p.q_len) {
-            int ep = q_row + p.causal_offset + 1;
-            if (kv0 >= ep)
-                lim = 0;
-            else if (kv0 + tlen > ep)
-                lim = ep - kv0;
-        }
-
-        int mask_row_base = mask_batch_base + q_row * p.mask_q_stride;
-        for (int s = 0; s < lim; s++) {
-            const bf16* kr = sK + s * HEAD_DIM + gpos * DPT;
-            float part = 0.0f;
-#pragma unroll
-            for (int i = 0; i < DPT; i += 8) {
-                float k8[8];
-                ld8(kr + i, k8);
-#pragma unroll
-                for (int j = 0; j < 8; j++)
-                    part = fmaf(qreg[i + j], k8[j], part);
-            }
-            float dot = group_reduce_sum<G>(part, gmask);
-
-            int kv_idx = kv0 + s;
-            if (p.use_mask && p.mask && !p.mask[mask_row_base + kv_idx])
-                dot = -FLT_MAX;
-
-            float nm = fmaxf(m, dot);
-            float al = __expf(m - nm);
-            float be = __expf(dot - nm);
-            l = l * al + be;
-
-            const bf16* vr = sV + s * HEAD_DIM + gpos * DPT;
-#pragma unroll
-            for (int i = 0; i < DPT; i += 8) {
-                float v8[8];
-                ld8(vr + i, v8);
-#pragma unroll
-                for (int j = 0; j < 8; j++)
-                    acc[i + j] = fmaf(v8[j], be, acc[i + j] * al);
-            }
-            m = nm;
-        }
-        __syncthreads();
-    }
-
-    if (q_row < p.q_len) {
-        // O: stride-based write
-        int o_off = batch * p.q_stride_b + q_head * p.q_stride_h
-                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
-        float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
-#pragma unroll
-        for (int i = 0; i < DPT; i++)
-            p.o[o_off + i * p.q_stride_d] = __float2bfloat16(acc[i] * rl);
-    }
-}
@@ -1,196 +0,0 @@
-#pragma once
-#include <cfloat>
-#include <cuda_bf16.h>
-#include "attn_common.h"
-#include "attn_mma_utils.cuh"
-
-using bf16 = __nv_bfloat16;
-
-// Tensor-core prefill flash attention (raw mma.sync PTX).
-// One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
-// cores via mma.sync.m16n8k16 (f32 accumulate). Q fragments are loaded once
-// straight from global into the mma A-operand layout (no smem staging) and
-// kept resident in registers across the tile loop. S, O, and the online-softmax
-// stats (m, l) also live in registers.
-// Shared memory is statically sized via template parameters — no dynamic
-// allocation. The mma fragment layout is used directly: the S accumulator
-// (f32) maps element-for-element onto the P matrix_a (bf16) operand, so
-// softmax needs no shuffle repack; row reductions fold across the 4-lane
-// thread group. Templated on <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
-//
-// Software pipeline: K/V are double-buffered and loaded via cp.async one tile
-// ahead, so the next tile streams from global memory while the current tile's
-// tensor-core math runs — hiding load latency (long_scoreboard). A single
-// __syncthreads per tile both publishes the freshly loaded tile cross-warp and
-// (because it runs before the next prefetch) guards the buffer being refilled,
-// so no second barrier is needed. Predicated cp.async (cp_async_16_pred)
-// zero-fills rows past kv_len, unifying full and partial tiles on one path.
-// BC=32 (D<=128) amortizes the per-tile wait+barrier+loop overhead over more
-// tensor-core work — this kernel is latency-bound (low occupancy from high
-// register pressure), so fewer, larger tiles beat many tiny ones.
-//
-// Optimizations: load Q fragments directly from global in mma A-operand layout
-// (no sQ staging, no prologue barriers); post-multiply scale in float after
-// S=Q@K^T to avoid bf16 precision loss; packed bf16x2 output stores;
-// causal tile skipping (block-level prefetch bound + warp-level compute skip);
-// XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without LD
-// padding (LD=HEAD_DIM).
-
-template <int HEAD_DIM, int WARPS, int BC>
-__global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
-    constexpr int BR = 16;
-    constexpr int KD = HEAD_DIM / 16;  // Q/K k-tiles
-    constexpr int NC8 = BC / 8;        // S n-tiles (N=8 each)
-    constexpr int KT2 = BC / 16;       // P k-tiles (K=16 each)
-    constexpr int DN8 = HEAD_DIM / 8;  // O n-tiles (N=8 each)
-    constexpr int LD = HEAD_DIM;   // XOR swizzle (swiz_col) handles bank conflicts
-    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);  // chunk bits, stay within LD
-
-    const int warp = threadIdx.x / 32;
-    const int lane = threadIdx.x % 32;
-    const int gid = lane >> 2;   // 0..7  → rows gid, gid+8
-    const int tid4 = lane & 3;   // 0..3
-    const int nthreads = WARPS * 32;
-
-    const int q_head = blockIdx.y;
-    const int batch = blockIdx.z;
-    const int kv_head = q_head / (p.q_head / p.kv_head);
-    const int qrow0 = (blockIdx.x * WARPS + warp) * BR;
-
-    // ---- Static shared memory: double-buffered K/V ----
-    // K/V are double-buffered (STAGES=2): the next tile's cp.async load runs
-    // while the current tile's tensor-core math executes, hiding global-load
-    // latency (FA2-style software pipeline). No dynamic smem / carveout opt-in.
-    constexpr int STAGES = 2;
-    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
-    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
-
-    // Load Q fragments straight from global into mma A-operand layout.
-    // stride_row = p.q_stride_l for prefill (multi-q rows across q_len).
-    // See attn_mma_utils.cuh for the shared template.
-    const int q_base = batch * p.q_stride_b + q_head * p.q_stride_h;
-    const int qra = qrow0 + gid;
-    const int qrb = qrow0 + gid + 8;
-    const bool va = qra < p.q_len, vb = qrb < p.q_len;
-    unsigned Qa[KD][4];
-    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
-                          qra, qrb, va, vb, tid4, Qa);
-
-    float Oacc[DN8][4];
-#pragma unroll
-    for (int j = 0; j < DN8; j++)
-        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
-    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
-
-    // KV: stride-based base
-    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    const int tiles = (p.kv_len + BC - 1) / BC;
-    const int qr0 = qrow0 + gid;      // row for c0/c1
-    const int qr1 = qrow0 + gid + 8;  // row for c2/c3
-
-    // Causal tile-skip bounds (no-op when causal_offset < 0)
-    const int use_skip = (p.causal_offset >= 0) ? 1 : 0;
-    const int max_kv = qrow0 + BR - 1 + p.causal_offset;
-    const int block_max_kv =
-        blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset;
-    const int has_mask = p.use_mask && p.mask;
-
-    // Last active tile: block-level causal bound (all warps in the block share
-    // the K/V load, so the prefetch range is the block max, not per-warp).
-    int t_end = tiles - 1;
-    if (use_skip) {
-        int bt = block_max_kv / BC;
-        if (bt < t_end) t_end = bt;
-    }
-
-    constexpr int VEC = 8;  // bf16 per cp.async unit (16 bytes)
-    constexpr int TOTAL = BC * HEAD_DIM;
-
-    // ---- Load tile lambda: predicated cp.async ----
-    // Issue cp.async loads for tile `ti` into shared buffer `buf`. Predicated
-    // loads zero-fill rows past kv_len, so partial tiles need no scalar path.
-    auto load_tile = [&](int ti, int buf) {
-        int kv0 = ti * BC;
-        bf16* dK = sK + buf * BC * LD;
-        bf16* dV = sV + buf * BC * LD;
-#pragma unroll
-        for (int i = threadIdx.x * VEC; i < TOTAL; i += nthreads * VEC) {
-            int r = i / HEAD_DIM, d = i % HEAD_DIM;
-            int kc = kv0 + r;
-            bool valid = kc < p.kv_len;
-            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
-            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
-            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
-            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
-        }
-        cp_async_commit();
-    };
-
-    // ---- Prologue: issue first tile load ----
-    load_tile(0, 0);
-
-    for (int ti = 0; ti <= t_end; ti++) {
-        int buf = ti & 1;
-
-        // Wait for the current tile's async copies, then a single barrier: it
-        // both publishes this tile's data cross-warp AND guarantees the prior
-        // compute on the buffer we are about to refill has finished. Issuing
-        // the next tile's load *after* this barrier lets one barrier cover both
-        // hazards (vs two), while the load still overlaps this tile's math.
-        cp_async_wait_group<0>();
-        __syncthreads();
-        if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
-
-        const bf16* bK = sK + buf * BC * LD;
-        const bf16* bV = sV + buf * BC * LD;
-        int kv0 = ti * BC;
-
-        // Warp-level causal skip
-        if (!use_skip || kv0 <= max_kv) {
-
-        // S = Q @ K^T + scale + online softmax + O += P @ V
-        float Sacc[NC8][4];
-        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
-
-        // post-multiply scale in float (no bf16 precision loss from pre-scaling Q)
-        #pragma unroll
-        for (int n8 = 0; n8 < NC8; n8++)
-            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
-            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
-
-        int maxc0 = (p.causal_offset >= 0) ? min(p.kv_len, qr0 + p.causal_offset + 1)
-                                        : p.kv_len;
-        int maxc1 = (p.causal_offset >= 0) ? min(p.kv_len, qr1 + p.causal_offset + 1)
-                                        : p.kv_len;
-        mma_softmax_tile<NC8, DN8>(kv0, maxc0, maxc1,
-                                    qr0, qr1,
-                                    p.mask_b_stride, p.mask_q_stride,
-                                    batch,
-                                    p.mask, has_mask,
-                                    Sacc, Oacc, m0, m1, l0, l1, lane);
-
-        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
-        }  // if active (warp-level causal skip)
-    }
-
-    // ---- write output ---- (packed bf16x2 stores: one 32-bit STG per pair,
-    // halves store count and removes the uncoalesced scalar-store penalty)
-    float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
-    float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
-    // O: stride-based write
-    const int o_base = batch * p.q_stride_b + q_head * p.q_stride_h;
-#pragma unroll
-    for (int dn8 = 0; dn8 < DN8; dn8++) {
-        int d = dn8 * 8 + 2 * tid4;
-        if (qr0 < p.q_len) {
-            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
-                                                     Oacc[dn8][1] * rl0);
-            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v;
-        }
-        if (qr1 < p.q_len) {
-            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
-                                                     Oacc[dn8][3] * rl1);
-            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v;
-        }
-    }
-}
--- a/Show More
+++ b/Show More