chore: bump version to 1.3.9

refactor: use actual q_len from input, remove dead num_splits init
fix: auto-assign free port in spawn_parallel_fn to avoid EADDRINUSE
2026-07-12 21:04:20 +08:00 · 2026-07-12 19:48:17 +08:00 · 2026-07-12 19:21:56 +08:00 · 2026-07-12 15:45:08 +08:00 · 2026-07-12 15:00:53 +08:00 · 2026-07-12 14:14:54 +08:00
137 changed files with 16535 additions and 4006 deletions
@@ -2,7 +2,7 @@
 name: Bug report
 about: Create a report to help us improve
 title: "[BUG]"
-labels: enhancement
+labels: bug
 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 --fix .`)
+- [ ] My code follows the style guidelines of this project (run `ruff format .` and `ruff check . --select I`)
 - [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
+- [ ] Code is self-documenting (no unnecessary comments)
 - [ ] I have made corresponding changes to the documentation
 - [ ] My changes generate no new warnings
 - [ ] I have added tests that prove my fix is effective or that my feature works
@@ -0,0 +1,71 @@
 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,8 +5,16 @@
 !*/
 # Allow specific file types and root files
-!*.py
+!astrai/**/*.py
-!*.sh
+!scripts/**/*.py
 !tests/**/*.py
 !csrc/**/*.py
 !csrc/**/*.cu
 !csrc/**/*.h
 !csrc/**/*.cuh
 !scripts/**/*.sh
 # Allow GitHub files
 !/.github/**
@@ -20,4 +28,9 @@
 !/CONTRIBUTING.md
 !/LICENSE
 !/pyproject.toml
-!/README.md
+!/README.md
 # Allow extension modules (only source .py)
 !/astrai/extension/**/*.py
 # Allow build files
 !/setup.py
@@ -1,68 +1,100 @@
 # Contributing to AstrAI
-Thank you for your interest in contributing to AstrAI! This document provides guidelines and steps for contributing.
+Thank you for your interest in contributing! This document provides step-by-step guidelines.
-## How to Contribute
+## Quick Start
-### Reporting Issues
+```bash
-If you encounter a bug or have a feature request, please open an issue on GitHub. Include as much detail as possible:
+git clone https://github.com/ViperEkura/AstrAI.git
- A clear description of the problem or request.
+cd AstrAI
- Steps to reproduce (for bugs).
+pip install -e ".[dev]"     # install with dev dependencies (pytest, ruff)
- Your environment (Python version, OS, etc.).
+```
-### Submitting Changes
+## Before You Commit
 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.
-## Code Style
+Run the following checks **in order** — CI will reject if any fail.
-AstrAI uses [Ruff](https://docs.astral.sh/ruff/) for code formatting and linting. Please ensure your code is formatted before submitting.
+### 1. Format
- Run Ruff to format and lint (requires conda environment `nlp`):
+```bash
-  ```bash
+ruff format .
-  conda run -n nlp ruff format .
+```
  conda run -n nlp ruff check --fix .
  ```
 - The project uses **double quotes** for strings and **4‑space indentation** (as configured in `pyproject.toml`).
-## Testing
+> **Note**: `ruff format` may rename parameters (e.g. `mask` → `attn_mask`).  
 > Always review the diff after formatting.
-If you add or modify functionality, please include appropriate tests.
+### 2. Import sorting
- Run the test suite with:
+```bash
-  ```bash
+ruff check . --select I
-  conda run -n nlp python -u -m pytest
+```
-  ```
+
- Ensure all tests pass before submitting your PR.
+If this fails, **manually fix** import ordering (ruff does not auto-fix in this project's CI):
 ```bash
 ruff check . --select I --fix .
 ruff format .    # re-format after fix
 ```
 ### 3. Run tests
 ```bash
 python -u -m pytest tests/ -v
 ```
 > Failed tests may leave orphan tempdirs under `%TEMP%`. Clean them manually if needed.
 ### 4. (Optional) Full pre-commit check
 If you have Git Bash available:
 ```bash
 bash scripts/pre_commit.sh
 ```
 This runs format check, import sort check, and tests in one go.
 ## Commit Style
 ```
 fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description (~50 chars)
 - bullet point body (each ~60 chars)
 ```
 - **Type** must be one of: `fix`, `feat`, `chore`, `docs`, `refactor`, `perf`, `test`, `style`, `ci`, `build`, `revert`.
 - **Subject line** ends with no period.
 - **Body** uses bullet points starting with `-`.
 - No `(scope)` parentheses.
 ## Common Issues
 | Problem | Cause | Fix |
 |---------|-------|-----|
 | `ruff check --select I` fails | Wrong import order | `ruff check . --select I --fix .` then `ruff format .` |
 | `ruff format` changed many files | Not formatted before commit | Review diff carefully before staging |
 | Pre-commit hook rejects | Tests or lint failed | Fix individually, do not `--no-verify` |
 | Tests fail with tempdir left | Test crash | Clean `%TEMP%` manually |
 ## Submitting Changes
 1. Fork the repo.
 2. Create a feature branch: `git checkout -b feat/my-feature`
 3. Make changes following the steps above.
 4. Commit with the commit style above.
 5. Push: `git push origin feat/my-feature`
 6. Open a Pull Request against `main`.
 ## Code Review
-All submissions will be reviewed. We may request changes or discuss alternatives. Please be responsive to feedback.
+- All PRs are reviewed. We may request changes.
 - CI runs `ruff format --check .` then `ruff check . --select I` (no `--fix` in CI).
 - Ensure all tests pass.
 ## License
-By contributing, you agree that your contributions will be licensed under the same [GPL-3.0 License](LICENSE) that covers the project.
+By contributing, you agree that your contributions will be licensed under the [GPL-3.0 License](LICENSE).
 ---
-If you have any questions, feel free to ask in the [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions) or open an issue.
+Questions? Ask in [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 nvidia/cuda:12.6.0-base-ubuntu24.04 AS builder
+FROM ubuntu:24.04 AS builder
 WORKDIR /app
@@ -18,21 +18,22 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
 RUN python3.12 -m venv --copies /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
-# Copy source code and install dependencies
+# Copy source code and install (deps read from pyproject.toml)
 COPY astrai/ ./astrai/
 COPY pyproject.toml .
 RUN pip install --no-cache-dir --upgrade pip \
    && pip install --no-cache-dir . \
-    --extra-index-url https://download.pytorch.org/whl/cu126
+    --extra-index-url https://download.pytorch.org/whl/cu128
 # Production stage
-FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production
+FROM ubuntu:24.04 AS production
 WORKDIR /app
-# Install Python 3.12 runtime
+# Install Python 3.12 runtime and healthcheck dependency
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    python3.12 \
    curl \
    && rm -rf /var/lib/apt/lists/*
 # Copy virtual environment from builder
@@ -9,9 +9,9 @@
 <div align="center">
  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/tag/ViperEkura/AstrAI?label=Release&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/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
-  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.forks_count&label=forks&suffix=%20forks&color=76bad9" alt="forks">
+  <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
 <br>
@@ -28,7 +28,8 @@
 ## 📖 Table of Contents
 - [Features](#features)
- [Quick Start](#quick-start)
+- [Getting Started](#getting-started)
 - [Demo](#demo)
 - [Documentation](#documentation)
 - [Contributing](#contributing)
 - [Community](#community)
@@ -49,55 +50,116 @@
 - 🤗 **HuggingFace-Style API**: AutoModel/AutoTokenizer APIs inspired by HuggingFace for easy model and tokenizer loading.
 - 🔌 **Dual API Compatibility**: Supports both OpenAI and Anthropic chat completion APIs out of the box.
-### Quick Start
+### Getting Started
-#### Installation
+End-to-end walkthrough in 5 steps:
 **1. Install**
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .
+pip install -e .                                          # 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)
 ```
-For development dependencies:
+**2. Download model**
 ```bash
-pip install -e ".[dev]"
+python scripts/demo/download.py    # downloads 1B checkpoint to params/
 ```
-#### Download Pre-trained Model
+**3. Preprocess data**
-Download pre-trained model weights (1B bilingual checkpoint) to `params/`:
+Create `pretrain.json` (preprocessing config for `seq` strategy):
 ```json
 {
    "version": 1,
    "input": {"sections": [{"field": "text", "action": "train"}]},
    "preprocessing": {"max_seq_len": 2048},
    "output": {"storage_format": "bin"}
 }
 ```
 ```bash
-python scripts/demo/download.py
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
 ```
-Or download manually from [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) into `params/`.
+**4. Train**
 #### Train a Model
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
+export CUDA_VISIBLE_DEVICES=0,1,2,3
-    --train_type seq \
+
-    --data_root_path /path/to/dataset \
+nohup python scripts/tools/train.py \
-    --param_path /path/to/model \
+    --nprocs=4 \
-    --batch_size 4 \
+    --parallel_mode=ddp \
-    --accumulation_steps 8 \
+    --train_type=seq \
-    --max_lr 3e-4 \
+    --data_root_path=/path/to/dataset \
-    --warmup_steps 1000 \
+    --param_path=/path/to/model \
-    --n_epoch 1
+    --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 reference at [Parameter Guide](assets/docs/params.md).
+**5. Serve & query**
-#### Generate Text
+```bash
 # Terminal 1: start server
 python scripts/tools/server.py --param_path ./params --device cuda
 # Terminal 2: query
 curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
 ```
 ### Demo
 Check out the demos in the `scripts/demo/` folder:
 ```bash
 # Download model weights (required before running demos)
 python scripts/demo/download.py                      # model → params/
 # Interactive streaming chat (multi-turn, maintains history)
 python scripts/demo/stream_chat.py
 # Type your message after >>, type !exit to quit
 # Batch generation (5 hardcoded prompts, non-streaming)
 python scripts/demo/generate_batch.py
 # Single-prompt autoregressive streaming
 python scripts/demo/generate_ar.py
 ```
 All generation demos use `temperature=0.8`, `top_p=0.95`, `top_k=50`, `max_tokens=2048` by default and require `params/` to contain model weights (run `download.py` first).
 Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).
 ---
 See [Documentation](#documentation) for full references beyond the examples above.
 #### Text Generation
 Batch generation from a JSONL file:
 ```bash
 python scripts/tools/generate.py \
-    --param_path /path/to/model \
+    --param_path ./params \
-    --input_json_file /path/to/input.json \
+    --input_json_file input.jsonl \
-    --output_json_file /path/to/output.json
+    --output_json_file output.jsonl
 ```
 #### Docker
@@ -111,9 +173,6 @@ docker build -t astrai:latest .
 # Run with GPU support
 docker run --gpus all -it astrai:latest
 # Run with specific GPUs
 docker run --gpus '"device=0,1"' -it astrai:latest
 # Run inference server
 docker run --gpus all -p 8000:8000 astrai:latest \
  python -m scripts.tools.server --port 8000 --device cuda
@@ -130,87 +189,42 @@ docker compose --profile cpu up -d
 > **Note**: `--gpus all` is required for CUDA support. Without it, `torch.cuda.is_available()` will return `False`.
-#### Start HTTP Server
+#### HTTP API Examples
-Start the inference server with OpenAI and Anthropic-compatible HTTP API:
+Additional request examples beyond the [Getting Started](#getting-started) flow:
 ```bash
 python -m scripts.tools.server --port 8000 --device cuda
 ```
 Make requests:
 ```bash
 # OpenAI-compatible
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "Hello"}],
    "max_tokens": 512
  }'
 # OpenAI-compatible streaming
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
-  -d '{
+  -d '{"messages":[{"role":"user","content":"Tell a story"}],"stream":true,"max_tokens":500}'
    "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 '{
+  -d '{"model":"astrai","system":"You are a helpful assistant.","messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
    "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 '{
+  -d '{"model":"astrai","messages":[{"role":"user","content":"Write a story"}],"max_tokens":500,"stream":true,"stop_sequences":["The end"]}'
    "model": "astrai",
    "messages": [{"role": "user", "content": "Write a story"}],
    "max_tokens": 500,
    "stream": true,
    "stop_sequences": ["The end"]
  }'
 # Health check
 curl http://localhost:8000/health
 ```
-#### Demo
+See [Inference Guide](assets/docs/inference.md) for SSE streaming format, error codes, and stats endpoint.
 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 |
 |----------|-------------|
-| [Parameter Guide](./assets/docs/params.md) | Training & inference parameters |
+| [CLI Reference](./assets/docs/params.md) | Parameters for all CLI tools (train, server, generate, preprocess) |
-| [Design Document](./assets/docs/design.md) | Framework architecture & module design |
+| [Architecture](./assets/docs/architecture.md) | System architecture, class diagram & design patterns |
-| [Data Flow](./assets/docs/dataflow.md) | Data processing pipeline details |
+| [Training](./assets/docs/training.md) | Training loop, strategies & formulas |
-| [Model Introduction](./assets/docs/introduction.md) | Model architecture & technical details |
+| [Inference](./assets/docs/inference.md) | KVCache, continuous batching, sampling & HTTP API |
 | [Data Flow](./assets/docs/dataflow.md) | Data pipeline, storage backends & dataset architecture |
 | [Preprocessing](./assets/docs/preprocessing.md) | Declarative JSON-driven data preprocessing |
 ### Contributing
@@ -15,9 +15,9 @@
 <div align="center">
  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/tag/ViperEkura/AstrAI?label=Release&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/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
-  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.forks_count&label=forks&suffix=%20forks&color=76bad9" alt="forks">
+  <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
 <br>
@@ -34,7 +34,8 @@
 ## 📖 目录
 - [特性](#特性)
- [快速开始](#快速开始)
+- [快速上手](#快速上手)
 - [演示](#演示)
 - [文档](#文档)
 - [贡献](#贡献)
 - [社区](#社区)
@@ -55,55 +56,116 @@
 - 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
 - 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
-### 快速开始
+### 快速上手
-#### 安装
+端到端演示，只需 5 步：
 **1. 安装**
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .
+pip install -e .                                          # 纯 PyTorch（不含 CUDA 内核）
 # CSRC_KERNELS=true pip install -e . --no-build-isolation  # 可选：融合 CUDA 内核加速
 # pip install -e ".[dev]"                                  # 可选：开发依赖（pytest, ruff）
 ```
-安装开发依赖：
+**2. 下载模型**
 ```bash
-pip install -e ".[dev]"
+python scripts/demo/download.py    # 下载 1B 检查点到 params/
 ```
-#### 下载预训练模型
+**3. 预处理数据**
-下载预训练模型权重（1B 双语检查点）到 `params/` 目录：
+创建 `pretrain.json`（`seq` 策略的预处理配置）：
 ```json
 {
    "version": 1,
    "input": {"sections": [{"field": "text", "action": "train"}]},
    "preprocessing": {"max_seq_len": 2048},
    "output": {"storage_format": "bin"}
 }
 ```
 ```bash
-python scripts/demo/download.py
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
 ```
-或从 [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) 手动下载放入 `params/`。
+**4. 训练**
 #### 训练模型
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
+export CUDA_VISIBLE_DEVICES=0,1,2,3
-    --train_type seq \
+
-    --data_root_path /path/to/dataset \
+nohup python scripts/tools/train.py \
-    --param_path /path/to/model \
+    --nprocs=4 \
-    --batch_size 4 \
+    --parallel_mode=ddp \
-    --accumulation_steps 8 \
+    --train_type=seq \
-    --max_lr 3e-4 \
+    --data_root_path=/path/to/dataset \
-    --warmup_steps 1000 \
+    --param_path=/path/to/model \
-    --n_epoch 1
+    --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 &
 ```
-完整参数列表见[参数说明](./params.md)。
+**5. 启动服务并调用**
 ```bash
 # 终端 1：启动服务
 python scripts/tools/server.py --param_path ./params --device cuda
 # 终端 2：发起请求
 curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"你好"}],"max_tokens":512}'
 ```
 ### 演示
 查看 `scripts/demo/` 文件夹中的演示：
 ```bash
 # 下载模型权重（运行演示前必需）
 python scripts/demo/download.py                      # model → params/
 # 交互式流式聊天（多轮对话，保持历史记录）
 python scripts/demo/stream_chat.py
 # 在 >> 后输入消息，输入 !exit 退出
 # 批量生成（5 条硬编码提示词，非流式）
 python scripts/demo/generate_batch.py
 # 单条提示词自回归流式生成
 python scripts/demo/generate_ar.py
 ```
 所有生成演示默认使用 `temperature=0.8`、`top_p=0.95`、`top_k=50`、`max_tokens=2048`，需要 `params/` 目录包含模型权重（请先运行 `download.py`）。
 观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
 ---
 更多选项请参考[文档](#文档)。
 #### 文本生成
 从 JSONL 文件批量生成：
 ```bash
 python scripts/tools/generate.py \
-    --param_path /path/to/model \
+    --param_path ./params \
-    --input_json_file /path/to/input.json \
+    --input_json_file input.jsonl \
-    --output_json_file /path/to/output.json
+    --output_json_file output.jsonl
 ```
 #### Docker
@@ -117,9 +179,6 @@ docker build -t astrai:latest .
 # 启用 GPU 运行
 docker run --gpus all -it astrai:latest
 # 指定特定 GPU
 docker run --gpus '"device=0,1"' -it astrai:latest
 # 运行推理服务
 docker run --gpus all -p 8000:8000 astrai:latest \
  python -m scripts.tools.server --port 8000 --device cuda
@@ -136,87 +195,42 @@ docker compose --profile cpu up -d
 > **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。
-#### 启动 HTTP 服务
+#### HTTP API 示例
-启动推理服务器，支持 OpenAI 和 Anthropic 兼容的 HTTP API：
+除[快速上手](#快速上手)流程外，更多请求示例：
 ```bash
 python -m scripts.tools.server --port 8000 --device cuda
 ```
 发起请求：
 ```bash
 # OpenAI 兼容
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "你好"}],
    "max_tokens": 512
  }'
 # OpenAI 兼容流式
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
-  -d '{
+  -d '{"messages":[{"role":"user","content":"讲个故事"}],"stream":true,"max_tokens":500}'
    "messages": [{"role": "user", "content": "讲个故事"}],
    "stream": true,
    "max_tokens": 500
  }'
 # Anthropic 兼容
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
-  -d '{
+  -d '{"model":"astrai","system":"你是一个乐于助人的助手。","messages":[{"role":"user","content":"你好"}],"max_tokens":512}'
    "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 '{
+  -d '{"model":"astrai","messages":[{"role":"user","content":"写个故事"}],"max_tokens":500,"stream":true,"stop_sequences":["结束"]}'
    "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) 上的视频演示。
 ### 文档
 | 文档 | 说明 |
 |------|------|
-| [参数说明](./params.md) | 训练与推理参数配置 |
+| [CLI 参考](./params.md) | 所有 CLI 工具参数（训练、服务、生成、预处理） |
-| [设计文档](./design.md) | 系统架构与模块设计 |
+| [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
-| [数据流程](./dataflow.md) | 数据处理管道详解 |
+| [训练文档](./training.md) | 训练循环、策略与公式 |
-| [模型介绍](./introduction.md) | 模型架构与技术细节 |
+| [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
 | [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 |
 | [数据预处理](./preprocessing.md) | 声明式 JSON 驱动数据预处理 |
 ### 贡献
@@ -1,237 +1,112 @@
-# AstrAI Data Flow Documentation
+# Data Flow
-This document describes the data flow of the AstrAI project (a training and inference framework for autoregressive Transformer language models). It covers the complete flow from raw data to model training and inference.
+This document describes the data pipeline: from raw text to model input tensors. For creating preprocessing configs, see [Preprocessing Guide](preprocessing.md).
 ## Contents
 - [Overview](#overview)
 - [Data Preparation](#data-preparation) — tokenization, format detection, backends
 - [Data Keys by Training Type](#data-keys-by-training-type)
 - [Dataset Architecture](#dataset-architecture)
 - [Sampler](#sampler)
 - [DataLoader](#dataloader)
 ## Overview
-AstrAI adopts a modular design with the following main components:
+```
- **Dataset Module** (`astrai/dataset/`): Dataset, sampler, serialization tools
+JSONL Lines → Pipeline (mask builder) → Tokenized Tensors
- **Model Module** (`astrai/model/`): AutoModel, Transformer model and its submodules
+                                              ↓
- **Training Module** (`astrai/trainer/`): Trainer, training context, strategies, schedulers, callbacks, metric utilities
+                                      .h5 or .bin storage
- **Inference Module** (`astrai/inference/`): Inference engine with continuous batching, streaming generation
+                                              ↓
- **Config Module** (`astrai/config/`): ModelConfig, TrainConfig
+                                      Store.load()
- **Factory Module** (`astrai/factory/`): Registry, BaseFactory for component registration
+                                              ↓
- **Parallel Module** (`astrai/parallel/`): Distributed training support
+                                      Store.fetch(begin, end, keys)
- **Serialization** (`astrai/serialization.py`): Checkpoint management with safetensors
+                                              ↓
-
+                                      BaseDataset.__getitem__(idx)
-## Data Flow Diagram
+                                              ↓
-
+                                      Sampler → DataLoader → Training / Inference
 ```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
 ```
-## Detailed Module Descriptions
+## Data Preparation
-### 1. Data Serialization (`astrai/dataset/storage.py` & `astrai/serialization.py`)
+Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
- **`save_h5`**: Saves tensors by groups as HDF5 files (`.h5`), each key maps to a list of tensors
+### Tokenization
 - **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory
 - **`Checkpoint`**: Encapsulates model state dict + epoch + iteration; uses safetensors
-### 2. Dataset Module
+The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](preprocessing.md)), and produces flat token sequences:
-#### 2.1 Dataset (`dataset.py`)
+```python
- **`BaseDataset`**: Abstract base class for windowed sequence sampling
+# Per JSONL line: messages → chat template → token IDs + loss mask
- **`BaseSegmentFetcher` / `MultiSegmentFetcher`**: Fetch tensor segments by index range
+tokens = tokenizer.encode(rendered_text)        # List[int]
- **`DatasetFactory`**: Creates dataset instances by `train_type` (`seq`, `sft`, `dpo`, `grpo`)
+loss_mask = [0, 0, 0, 1, 1, 1, 1, 1, 1]        # 0=masked, 1=train
- Data keys: `"sequence"` (SEQ), `"loss_mask"` (SFT), `"chosen_mask"/"rejected_mask"` (DPO), `"masks"` (GRPO)
+# Stored as flat tensors, packed with other lines by packing strategy
 #### 2.2 Sampler (`sampler.py`)
 - **`ResumableDistributedSampler`**: Tracks `epoch` and `iter` for breakpoint resume; supports shuffle and drop_last
 ### 3. Model Module
 #### 3.1 Transformer / AutoModel
 - **`AutoModel`**: Base class with `from_pretrained()` / `save_pretrained()`
 - **`Transformer`**: Decoder-only architecture, registered via `@AutoModel.register('transformer')`
 - Embedding → N×DecoderBlock → RMSNorm → Linear lm_head
 - RoPE position encoding, optional weight tying
 #### 3.2 Submodules (`module.py`)
 - **`DecoderBlock`**: GQA attention + residual + MLP + RMSNorm
 - **`GQA`**: Grouped Query Attention (also `MLA` for multi-latent attention)
 - **`MLP`**: `SiLU(gate(x)) * up(x)` → down projection
 - **`RotaryEmbedding`**: RoPE complex cache (freqs_cis)
 - **`RMSNorm`**: Layer normalization
 ### 4. Training Module
 #### 4.1 Training Context (`train_context.py`)
 - **`TrainContext`**: Dataclass holding model, optimizer, dataloader, strategy, scheduler, checkpoint state
 - **`TrainContextBuilder`**: Builder pattern — takes checkpoint for resume, builds all components
 #### 4.2 Trainer (`trainer.py`)
 The training loop is nested: **epoch** → **batch** (with step phase interspersed):
 ```
 on_train_begin
  on_epoch_begin
    for each accumulation window of batches:        ← step phase
      on_step_begin
        for each batch in window:                    ← batch phase
          on_batch_begin → strategy(batch) → loss → backward → on_batch_end
          iteration += 1
      on_step_end
      optimizer.step() → zero_grad
    on_epoch_end
 on_train_end
 ```
-Key points:
+The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
 - `on_step_*` fires every `accumulation_steps` batches, wrapping optimizer step AFTER the hook
 - `on_batch_*` fires every batch, wrapping loss computation
 - `GradientClippingCallback` fires on `on_step_end`
 - LR scheduler steps inline (no `SchedulerCallback` class)
-#### 4.3 Strategy (`strategy.py`)
+### Format Detection
 - **`SEQStrategy`**: Next-token prediction, cross-entropy with label smoothing
 - **`SFTStrategy`**: Supervised fine-tuning with loss masking
 - **`DPOStrategy`**: Direct Preference Optimization with reference model
 - **`GRPOStrategy`**: Group Relative Policy Optimization with clipped ratio
-#### 4.4 Scheduler (`schedule.py`)
+`detect_format(load_path)` inspects the path:
 - **`CosineScheduler`**: Cosine decay + linear warmup
 - **`SGDRScheduler`**: Cosine annealing with warm restarts
 - Created by `SchedulerFactory` and bound to optimizer
-#### 4.5 Callbacks
+- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, `.jsonl` → `"jsonl"`, unknown suffix raises `ValueError`
- **`CheckpointCallback`**: Saves safetensors at `ckpt_interval` iterations
+- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, `*.bin` + `**/meta.json` → `"bin"`, or `*.jsonl` + `dataset_config.json` → `"jsonl"`
 - **`ProgressBarCallback`**: tqdm progress display
 - **`MetricLoggerCallback`**: Writes JSONL metrics to `{ckpt_dir}/logs/`
 - **`GradientClippingCallback`**: `clip_grad_norm_` on `on_step_end`
-### 5. Inference Module
+### Store Backends
-#### 5.1 Inference Engine (`engine.py`)
+Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
 - **`InferenceEngine`**: Facade over scheduler; provides `generate()`, `generate_with_request()`, `generate_async()`
 - Accepts `prompt: str | List[str]`, returns generator (stream) or string (non-stream)
 #### 5.2 Scheduler 4-Phase Loop (`scheduler.py`)
 Background thread runs continuously:
 ```
-1. Cleanup → Remove finished tasks, free KV cache pages
+StoreFactory.create("h5")    → H5Store
-2. Refill  → Pop from waiting_queue, alloc pages, add to active
+StoreFactory.create("bin")   → MmapStore
-3. Prefill → Group active tasks by prompt_len, run full forward pass
+StoreFactory.create("jsonl") → JsonlStore
 4. Decode  → Pick largest same-position group, run single-token forward
 ```
- **`Task`**: Tracks prompt_ids, output_ids, status (PENDING/RUNNING/FINISHED/ABORTED)
+**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage.
 - **`KVCache`**: Facade over `Allocator` + `PrefixCache` + `PagePool` + `Storage` for paged KV cache
 - **`KvcacheView`**: Batch view bundling cache + page table for attention layers
 - **`sample()`**: Temperature → top-k → top-p → multinomial
-#### 5.3 Server (`server.py`)
+**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
 - FastAPI with OpenAI `/v1/chat/completions` and Anthropic `/v1/messages` endpoints
 - Streaming via SSE, health check at `/health`, stats at `/stats`
-### 6. Tokenizer Module
+**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.
- **`AutoTokenizer`**: Wraps HuggingFace tokenizers (BBPE); `encode`/`decode`/`apply_chat_template`
+All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based indexing).
 - **`ChatTemplate`**: Jinja2-based template rendering for multi-turn chat
-### 7. Factory & Parallel
+## Data Keys by Training Type
- **`Registry` / `BaseFactory`**: Decorator-based component registration
+| Type | Storage Keys |
- **`spawn_parallel_fn`**: Multi-process DDP launcher with NCCL backend
+|------|-------------|
- **`ParallelModel` / `ColumnParallelLinear` / `RowParallelLinear`**: Tensor model parallelism
+| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) |
 | `sft` | `sequence`, `loss_mask`, `position_ids` |
 | `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` |
 | `grpo` | `prompts`, `responses`, `masks`, `rewards` |
-## Training Data Flow — Detailed Steps
+## Dataset Architecture
-1. **Data Preparation**
+```
-   - Raw text → token IDs via `AutoTokenizer.encode()`
+DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_type=None)
-   - Save as `.h5` files (groups of tensor lists per data key)
+  → BaseDataset.load(load_path, storage_type=None)
    → detect_format(load_path)
    → StoreFactory.create(storage_type)
    → Store.load(load_path)
      → _normalize(raw)  # base Store, shared by both backends
        → Store._data[Dict[str, List[Tensor]]] + _cum[Dict[str, List[int]]]
          → BaseDataset.__getitem__(idx)
            → get_index(idx) → [begin, end)
            → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
 ```
-2. **Dataset Loading**
+`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). `storage_type` defaults to `None` (auto-detect via `detect_format`).
   - `BaseDataset.load()` calls `load_h5()`, builds `MultiSegmentFetcher`
   - Sliding window of `window_size` with `stride` determines sample boundaries
-3. **Sampling & Batching**
+`Store.fetch(begin, end, keys)` accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
   - `ResumableDistributedSampler` produces shuffled index sequences
   - `DataLoader` fetches `[batch_size, window_size]` tensors via `__getitem__`
-4. **Strategy Forward**
+## Sampler
   - Strategy receives batch, calls `Transformer.forward()` for logits
   - Computes task-specific loss (cross-entropy, DPO, GRPO)
-5. **Backward & Accumulation**
+`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
   - `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**
+- Tracks `start_epoch` / `start_iter` for resume
-   - `CheckpointCallback` saves `model.state_dict()` + metadata to safetensors at `ckpt_interval` iterations
+- Shuffle via `torch.Generator(seed + epoch)`
-   - Does NOT save optimizer/scheduler state (resume resets those)
+- Per-replica index slicing for DDP
-## Inference Data Flow — Detailed Steps
+## DataLoader
-1. **Model Loading**
+Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
   - `AutoModel.from_pretrained(path)` loads weights from safetensors
   - `torch.inference_mode()` wraps generation
-2. **Prompt Construction**
+> Document Update Time: 2026-07-09
   - 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, `PagePool.task_alloc()` 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`
      - `PagePool.task_alloc()` allocates pages 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-14
@@ -1,779 +0,0 @@
 ## 1. Why I Created This Project
 There are many large language models on the market today, such as GPT, LLaMA, and others, with tens of billions or even hundreds of billions of parameters. But honestly, these models have extremely high hardware requirements, making them inaccessible for ordinary developers. I thought: **Can we create a model that is both useful and can run on ordinary computers?** This is also what most people currently hope for - a locally deployable AI project that achieves complete privatization while maintaining some level of intelligence.
 Thus, the AstrAI project was born - 1B parameters, Chinese-English bilingual, supporting dialogue, text generation, and the training code is open source!
 ## 2. System Architecture
 ```mermaid
 classDiagram
    namespace config {
        class ModelConfig {
            +int vocab_size
            +int dim
            +int n_layers
            +float norm_eps
            +int dim_ffn
            +bool tie_weight
            +int max_len
            +float rope_theta
            +int n_heads
            +int n_kv_heads
            +bool use_qk_norm
            +bool use_gated_attention
            +load(config_path) ModelConfig
            +save(config_path)
        }
        class TrainConfig {
            +nn.Module model
            +str strategy
            +Dataset dataset
            +Callable optimizer_fn
            +Callable scheduler_fn
            +int n_epoch
            +int batch_size
            +int accumulation_steps
            +float max_grad_norm
            +int start_epoch
            +int start_batch
            +str ckpt_dir
            +int ckpt_interval
            +int random_seed
            +int num_workers
            +int prefetch_factor
            +bool pin_memory
            +int nprocs
            +str backend
            +str master_addr
            +str master_port
            +Callable parallel_wrapper
            +Callable state_dict_fn
            +str device_type
            +dict extra_kwargs
            +validate()
        }
    }
    namespace dataset {
        class BaseDataset {
            +int window_size
            +int stride
            +BaseStorage storage
            +load(load_path, storage_type, tokenizer)
            +__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 BaseStorage {
            +MultiSegmentFetcher _fetcher
            +keys (property)
            +load(load_path, tokenizer)
            +fetch(begin, end, keys)
            +__len__()
        }
        class H5Storage {
            +load(load_path, tokenizer)
            +fetch(begin, end, keys) Dict
            +keys() List
        }
        class JSONStorage {
            +load(load_path, tokenizer)
            +fetch(begin, end, keys) Dict
            +keys() List
        }
        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_component_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, position_ids) Tensor
            +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) 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, attn_mask, paged_cache) 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, attn_mask, paged_cache) 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, position_ids=None) Tensor
        }
        class Embedding {
            +Parameter weight
            +forward(x) Tensor
        }
    }
    namespace tokenize {
        class AutoTokenizer {
            +vocab_size int
            +encode(tokens, out_ids, add_special_tokens) List[int]
            +decode(tokens, skip_special_tokens) str
            +__getattr__(name) Any (bos_id, eos_id, pad_id, stop_ids)
            +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
            +str reduction
            +int sync_interval
            +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 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
            +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
            +KVCache _page_cache
            +int max_batch_size
            +int max_seq_len
            +int max_prompt_len
            +int page_size
            +TaskManager _task_mgr
            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
            +remove_task(task_id)
            +start()
            +stop()
            +get_stats() Dict
        }
        class Allocator {
            +int _free_mask
            +int refs_count
            +LRU _lru
            +alloc() int
            +free(idx, keep_cached)
            +inc_ref(idx)
            +touch(idx)
            +ref_count(idx) int
        }
        class PrefixCache {
            +int _page_size
            +evict(page_idx)
            +has_page(idx) bool
            +lookup(token_ids) List[int]
            +record(page_idx, token_ids, logical_page_idx)
        }
        class PagePool {
            -Allocator _alloc
            -PrefixCache _prefix
            +alloc() int
            +free(idx)
            +inc_ref(idx)
            +lookup(token_ids) List[int]
            +record(page_idx, token_ids, logical_page_idx)
        }
        class Storage {
            +int n_layers
            +int page_size
            +int head_dim
            +int n_kv_heads
            +Tensor k_cache
            +Tensor v_cache
            +write(layer_id, page_table, start_pos, k, v)
            +gather(layer_id, page_table, total_len) Tuple[Tensor, Tensor]
        }
        class KVCache {
            -PagePool _pool
            -Storage _storage
            -TaskTable _table
            +int page_size
            +task_alloc(task_id, prompt_ids) bool
            +task_free(task_id)
            +task_extend(task_id, pos) bool
            +task_cached(task_id) int
            +task_record_hashes(task_id, prompt_ids, start_logical_page)
            +make_table_tensor(task_ids, device) Tensor
            +bind(page_table, total_len) KvcacheView
        }
        class KvcacheView {
            -Storage _storage
            +Tensor _page_table
            +int _total_len
            +write(layer_id, k, v)
            +gather(layer_id) Tuple[Tensor, Tensor]
        }
        class TaskTable {
            +set(task_id, page_table, cached)
            +get(task_id) List[int]
            +get_cached(task_id) int
            +get_ref(task_id) List[int]
            +pop(task_id) Tuple[List[int], int]
            +table_tensor(task_ids, device) 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
            +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
            +int top_k
            +float top_p
            +float temperature
            +Optional[int] max_tokens
            +bool stream
        }
        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 GenerateResult {
            +List[Tuple[int, str]] tokens
            +List[str] results
            +List[bool] _done
            +append(token, idx)
            +get_results() List[str]
            +pop_all() List[str]
            +wait(timeout) bool
            +wait_completion()
        }
        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 Functions {
            +spawn_parallel_fn(fn, nprocs)
            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type)
            +get_current_device() str
            +get_world_size() int
            +get_rank() int
        }
        class ParallelModel {
            +dist.ProcessGroup process_group
            +int rank
            +int world_size
        }
        class ColumnParallelLinear {
            +forward(x) Tensor
        }
        class RowParallelLinear {
            +forward(x) Tensor
        }
    }
    %% Relationships
    TrainConfig --> BaseDataset : uses
    TrainConfig ..> BaseStrategy : selects
    StrategyFactory ..> BaseStrategy : creates
    BaseStrategy <|-- SEQStrategy
    BaseStrategy <|-- SFTStrategy
    BaseStrategy <|-- DPOStrategy
    BaseStrategy <|-- GRPOStrategy
    DPOStrategy --> Transformer : uses
    GRPOStrategy --> Transformer : uses
    Trainer --> TrainConfig : uses
    Trainer --> TrainContextBuilder : uses
    Trainer --> TrainCallback : manages
    TrainContextBuilder --> TrainContext : creates
    TrainContextBuilder --> StrategyFactory : uses
    Checkpoint ..> Checkpoint : serializes
    TrainContext --> Checkpoint : manages
    TrainContext --> BaseStrategy : uses
    TrainContext --> BaseScheduler : uses
    SchedulerFactory ..> BaseScheduler : creates
    BaseScheduler <|-- CosineScheduler
    BaseScheduler <|-- SGDRScheduler
    CallbackFactory ..> TrainCallback : creates
    TrainCallback <|-- GradientClippingCallback
    TrainCallback <|-- CheckpointCallback
    TrainCallback <|-- ProgressBarCallback
    TrainCallback <|-- MetricLoggerCallback
    PagePool --> Allocator : composes
    PagePool --> PrefixCache : composes
    KVCache --> PagePool : composes
    KVCache --> Storage : composes
    KVCache --> TaskTable : composes
    KvcacheView --> Storage : wraps
    InferenceEngine --> InferenceScheduler : uses
    InferenceEngine --> GenerationRequest : uses
    InferenceEngine --> GenerateResult : creates
    InferenceScheduler --> Task : manages
    InferenceScheduler --> TaskStatus : uses
    InferenceScheduler --> KVCache : uses
    InferenceScheduler --> Transformer : uses
    Task --> TaskStatus : uses
    InferenceEngine --> Transformer : uses
    BaseSamplingStrategy <|-- TemperatureStrategy
    BaseSamplingStrategy <|-- TopKStrategy
    BaseSamplingStrategy <|-- TopPStrategy
    SamplingPipeline --> BaseSamplingStrategy : composes
    BaseDataset <|-- SEQDataset
    BaseDataset <|-- SFTDataset
    BaseDataset <|-- DPODataset
    BaseDataset <|-- GRPODataset
    DatasetFactory ..> BaseDataset : creates
    BaseStorage <|-- H5Storage
    BaseStorage <|-- JSONStorage
    BaseDataset --> BaseStorage : uses
    MultiSegmentFetcher --> BaseSegmentFetcher : 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
    BaseFactory <|-- AutoModel
    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, BaseStorage, H5Storage, JSONStorage, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory, save_h5, load_h5 | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | 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, KVCache, KvcacheView, Allocator, PrefixCache, PagePool, Storage, TaskTable, Task, TaskStatus, GenerationRequest, BaseSamplingStrategy, TemperatureStrategy, TopKStrategy, TopPStrategy, SamplingPipeline, ChatMessage, ChatCompletionRequest | Inference service with continuous batching and paged KV cache |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank, get_world_size, get_current_device, 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** | `Allocator`, `PagePool` | Page-based KV cache with O(1) alloc/free via bitmask + LRU eviction |
 | **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** | `GenerateResult`, `GenerationRequest` | Event-based result notification for streaming/non-streaming generation |
 ### Core Relationships
 1. **Configuration → Training**: `TrainConfig` holds model, dataset, optimizer_fn, scheduler_fn and other training configuration 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 `KVCache` (backed by `Allocator` + `PrefixCache` + `PagePool` + `Storage`) 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-05-14
@@ -0,0 +1,252 @@
 # 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
@@ -1,334 +0,0 @@
 ## Model Introduction
 ### 1. Model Architecture
 This model uses the Transformer architecture with GQA mechanism (q_head=24, kv_head=4), which saves KV cache memory compared to traditional MHA. The model is built by stacking multiple 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 ...]
        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_tokens=None,
    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-05-14
@@ -1,4 +1,11 @@
-# Parameter Documentation
+# CLI Parameter Reference
 ## Contents
 - [Training Parameters](#training-parameters)
 - [Inference Server](#inference-server-serverpy)
 - [Generate](#generate-generatepy)
 - [Preprocess](#preprocess-preprocesspy)
 ## Training Parameters
@@ -10,24 +17,28 @@
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
 | `--n_epoch` | Total training epochs | 1 |
-| `--batch_size` | Batch size | 1 |
+| `--batch_per_device` | Batch size per device | 1 |
-| `--accumulation_steps` | Gradient accumulation steps between optimizer steps | 1 |
+| `--grad_accum_steps` | Gradient accumulation steps between optimizer steps | 1 |
 ### Learning Rate Scheduling
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--warmup_steps` | Warmup steps | 1000 |
+| `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
 | `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
-### Optimizer (AdamW)
+### Optimizer (MuonMix)
 Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--adamw_beta1` | AdamW beta1 | 0.9 |
+| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
-| `--adamw_beta2` | AdamW beta2 | 0.95 |
+| `--muon_momentum` | Muon momentum factor | 0.95 |
-| `--adamw_weight_decay` | AdamW weight decay | 0.01 |
+| `--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` |
 ### Data Loading
@@ -46,113 +57,148 @@
 | `--ckpt_interval` | Iterations between checkpoints | 5000 |
 | `--ckpt_dir` | Checkpoint save directory | checkpoint |
 | `--start_epoch` | Resume from epoch (0 = from scratch) | 0 |
-| `--start_batch` | Resume from batch iteration | 0 |
+| `--start_samples` | Resume from sample count per rank | 0 |
 ### Validation
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--val_split` | Ratio to split from training dataset for validation (e.g. 0.05) | None |
 | `--val_step` | Number of optimizer steps between validation runs | 1000 |
 ### Logging
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--log_dir` | Directory for metric logs | checkpoint/logs |
 | `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr", "grad_norm"] |
 ### Gradient Checkpointing
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--gradient_checkpointing` | Enable activation checkpointing for DecoderBlock modules | False |
 ### Distributed Training
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--nprocs` | Number of GPUs / processes | 1 |
 | `--parallel_mode` | Parallel strategy (`none`, `ddp`, or `fsdp`) | none |
 | `--device_type` | Device type | cuda |
 | `--start_method` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) | spawn |
 | `--backend` | Distributed training backend | nccl |
 | `--master_addr` | Master node address | localhost |
 | `--master_port` | Master node port | 29500 |
 ### Strategy-specific
 | Parameter | Description | Default | Used by |
 |-----------|-------------|---------|---------|
 | `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
-| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.1 | `seq`, `sft` |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` |
 | `--group_size` | GRPO group size | 4 | `grpo` |
 | `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
 | `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
 | `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` |
 | `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` |
 ### Scheduler
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine |
 | `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default: 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
-python scripts/tools/train.py \
+export CUDA_VISIBLE_DEVICES=0,1,2,3
-  --train_type seq \
+
-  --data_root_path /path/to/dataset \
+nohup python scripts/tools/train.py \
-  --param_path /path/to/model \
+    --nprocs=4 \
-  --n_epoch 3 \
+    --parallel_mode=ddp \
-  --batch_size 4 \
+    --train_type=seq \
-  --accumulation_steps 8 \
+    --data_root_path=/path/to/dataset \
-  --max_lr 3e-4 \
+    --param_path=/path/to/model \
-  --warmup_steps 2000 \
+    --batch_per_device=4 \
-  --max_grad_norm 1.0 \
+    --grad_accum_steps=8 \
-  --ckpt_interval 5000 \
+    --warmup_ratio=0.05 \
-  --ckpt_dir ./checkpoints \
+    --max_lr=1e-4 \
-  --num_workers 4 \
+    --max_grad_norm=1.0 \
-  --nprocs 1 \
+    --weight_decay=0.1 \
-  --device_type cuda
+    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &
 ```
 ---
-## Generation Parameters
+## Inference Server (`server.py`)
-### GenerationRequest Parameters
+| Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `--host` | str | `0.0.0.0` | Host address |
 | `--port` | int | `8000` | Port number |
 | `--param_path` | path | `project_root/params` | Path to model parameters |
 | `--device` | str | `cuda` | Device to load model on |
 | `--dtype` | str | `bfloat16` | Model weights dtype (`bfloat16`, `float16`, `float32`) |
 | `--max_batch_size` | int | `16` | Maximum batch size for continuous batching |
 | `--reload` | flag | `False` | Enable auto-reload for development |
-| Parameter | Description | Default Value |
+Usage:
-|-----------|-------------|---------------|
+```bash
-| `messages` | List of message dictionaries (role, content) | required |
+python scripts/tools/server.py --param_path ./params --device cuda --dtype bfloat16
 | `temperature` | Sampling temperature (higher = more random) | 1.0 |
 | `top_p` | Nucleus sampling threshold | 1.0 |
 | `top_k` | Top-k sampling count | 50 |
 | `max_tokens` | Maximum generation length | None (unlimited) |
 | `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_tokens=None,
 )
 # 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,
 )
 ```
-### Generation Modes
+See [Inference Guide](inference.md) for HTTP API documentation.
-| Mode | Description |
+## Generate (`generate.py`)
 |------|-------------|
 | `stream=True` | Streaming output, yields token by token |
 | `stream=False` | Non-streaming output, returns complete result |
-> Document Update Time: 2026-05-14
+| 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-09
@@ -0,0 +1,364 @@
 # Preprocessing Pipeline
 Declarative JSON-driven data preprocessing. One `SectionedMaskBuilder` handles all formats via `input.sections` (single-output) or `input.sources` (multi-output).
 ## Contents
 - [Philosophy](#philosophy)
 - [Config Structure](#config-structure)
 - [Quick Start](#quick-start) — SFT Chat, SFT Instruction, Pretrain, DPO, GRPO examples
 - [Configuration Reference](#configuration-reference) — all fields
 - [Mask Algorithm](#mask-algorithm)
 - [Output Layout](#output-layout)
 - [CLI](#cli)
 - [Python API](#python-api)
 ## Philosophy
 | Component | Responsibility |
 |-----------|---------------|
 | `tokenizer_config.json` (`chat_template`) | Formatting -- how roles become tokens |
 | `pipeline.json` (`mask`) | Masking -- which roles participate in training |
 A single config file captures the entire pipeline, reusable and version-controllable.
 ## Config Structure
 ```json
 {
  "version":       1,
  "input":         {},   // sections (single) or sources (multi)
  "mask":          {},   // role -> "train" | "mask"
  "mask_default":  "mask",
  "preprocessing": {},
  "output":        {}
 }
 ```
 ### Section Fields
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `field` | str | -- | JSONL key to read |
 | `action` | str | -- | `"train"` / `"mask"` / `"$role"` |
 | `template` | bool | `false` | Apply `chat_template` per message |
 | `add_special_tokens` | bool | `true` for first non-template section | Add special tokens during encode |
 ### Source Fields (multi-output mode)
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `sections` | list[dict] | -- | Same as single-output section list |
 | `list_field` | bool | `false` | JSONL field holds a list; tokenise each element |
 | `mask_key` | str | `"{key}_mask"` | Explicit output key for loss mask |
 ---
 ## Quick Start
 ### SFT Chat
 Input JSONL:
 ```json
 {"messages": [{"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]}
 ```
 Config:
 ```json
 {
  "input": {
    "sections": [
      {"field": "messages", "action": "$role", "template": true}
    ]
  },
  "mask": {
    "system": "mask",
    "user": "mask",
    "assistant": "train"
  },
  "mask_default": "mask",
  "preprocessing": {
    "max_seq_len": 2048
  },
  "output": {
    "storage_format": "bin",
    "dtype": {"loss_mask": "bool"}
  }
 }
 ```
 Output keys: `sequence` (int32), `loss_mask` (bool)
 ### SFT Instruction
 Input JSONL:
 ```json
 {"prompt": "Translate to French: Hello", "response": "Bonjour"}
 ```
 Config:
 ```json
 {
  "input": {
    "sections": [
      {"field": "prompt",   "action": "mask", "add_special_tokens": true},
      {"field": "response", "action": "train"}
    ]
  },
  "mask_default": "mask",
  "preprocessing": {
    "max_seq_len": 2048
  }
 }
 ```
 Output keys: `sequence`, `loss_mask`
 ### Pretrain
 Input JSONL:
 ```json
 {"text": "Artificial Intelligence is a field of computer science..."}
 ```
 Config:
 ```json
 {
  "input": {
    "sections": [
      {"field": "text", "action": "train"}
    ]
  },
  "preprocessing": {
    "max_seq_len": 8192,
    "min_chars": 100
  }
 }
 ```
 Output keys: `sequence` (no `loss_mask` — all tokens trained)
 ### DPO
 Input JSONL:
 ```json
 {"chosen": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"}], "rejected": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "5"}]}
 ```
 Config:
 ```json
 {
  "input": {
    "sources": {
      "chosen": {
        "sections": [
          {"field": "chosen", "action": "$role", "template": true}
        ]
      },
      "rejected": {
        "sections": [
          {"field": "rejected", "action": "$role", "template": true}
        ]
      }
    }
  },
  "mask": {
    "user": "mask",
    "assistant": "train"
  },
  "mask_default": "mask"
 }
 ```
 Output keys: `chosen`, `chosen_mask`, `rejected`, `rejected_mask`
 ### GRPO
 Input JSONL:
 ```json
 {"prompt": [{"role": "user", "content": "What is 2+2?"}], "responses": ["4", "Five", "Four"], "rewards": [1.0, 0.3, 0.8]}
 ```
 Config:
 ```json
 {
  "input": {
    "sources": {
      "prompts": {
        "sections": [
          {"field": "prompt", "action": "mask", "template": true}
        ]
      },
      "responses": {
        "sections": [
          {"field": "responses", "action": "train"}
        ],
        "list_field": true,
        "mask_key": "masks"
      },
      "rewards": {
        "sections": [
          {"field": "rewards", "action": "value"}
        ]
      }
    }
  },
  "mask": {
    "user": "mask",
    "assistant": "train"
  },
  "mask_default": "mask"
 }
 ```
 Output keys: `prompts`, `prompts_mask`, `responses`, `masks`, `rewards` (float32)
 - `action: "value"` — extract raw values from JSONL without tokenisation
 - `list_field: true` — tokenise each list element independently, then concatenate
 - `mask_key: "masks"` — rename the auto-generated mask key (default: `responses_mask`)
 - `prompts_mask` is auto-generated (all masked) and unused by GRPOStrategy
 ---
 ## Configuration Reference
 ### `input`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `sections` | list[dict] or null | `null` | Section specs for single-output mode |
 | `sources` | dict[str, dict] or null | `null` | Source specs for multi-output mode (DPO/GRPO) |
 When `sources` is set, `sections` is ignored.
 ### `mask`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `mask` | dict | `{}` | `{role: "train" \| "mask"}` |
 | `mask_default` | str | `"mask"` | Default action for unlisted roles |
 ### `preprocessing`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `max_seq_len` | int | `2048` | Truncate sequences to this length |
 | `min_chars` | int | `50` | Skip text-mode items shorter than this |
 | `max_chars` | int | `2000000` | Skip text-mode items longer than this |
 | `max_items` | int or null | `null` | Stop after N documents |
 | `packing_strategy` | str | `"simple"` | Packing strategy: `"simple"`, `"bfd"`, `"bfd_split"` |
 | `max_packed_len` | int | `8192` | Maximum length of a packed bin |
 | `truncation_mode` | str | `"keep_start"` | How to truncate sequences: `"keep_start"` or `"keep_end"` |
 ### `output`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `domain_key` | str or null | `null` | JSONL key for domain grouping |
 | `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` |
 | `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens |
 | `dtype` | dict[str, str] | `{}` | Per-key tensor dtype override (e.g. `{"loss_mask": "bool"}`) |
 | `position_ids_mode` | str | `"doc_reset"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` |
 ---
 ## Mask Algorithm
 ### Template mode (`template: true`)
 1. Prepend BOS token (masked)
 2. For each message in the field's array:
   1. Render through `chat_template` for that single message
   2. Encode rendered text
   3. Apply mask rule for the message's role
 ### Non-template mode
 Encode the field value as text. Mask value is 1 (train) or 0 (mask) per the section's `action`.
 ### Text config detection
 When no section uses `template` and all sections have `action: "train"`, the builder omits `loss_mask` from the output — all tokens are trained.
 ---
 ## Output Layout
 ### Single-Shard (`bin`)
 ```
 output/
  __default__/
    shard_0000/
      meta.json
      sequence.bin
      loss_mask.bin
  wiki/
    shard_0000/
      meta.json
      sequence.bin
      loss_mask.bin
 ```
 ### Multi-Shard (`bin`)
 When `max_tokens_per_shard` is exceeded:
 ```
 output/
  __default__/
    shard_0000/
      meta.json
      sequence.bin
      loss_mask.bin
    shard_0001/
      meta.json
      sequence.bin
      loss_mask.bin
 ```
 For `bin` format, `MmapStore` discovers all shards under the domain directory via `rglob("meta.json")`. For `h5` format, `H5Store` discovers `.h5`/`.hdf5` files via recursive glob.
 ---
 ## CLI
 ```bash
 # SFT
 python scripts/tools/preprocess.py data/sft/*.jsonl -o output/sft/ -c configs/sft_chat.json
 # DPO
 python scripts/tools/preprocess.py data/dpo/*.jsonl -o output/dpo/ -c configs/dpo.json --tokenizer_path params
 # GRPO
 python scripts/tools/preprocess.py data/grpo/*.jsonl -o output/grpo/ -c configs/grpo.json
 ```
 ---
 ## Python API
 ```python
 from astrai.preprocessing.pipeline import Pipeline
 from astrai.config.preprocess_config import PipelineConfig
 config = PipelineConfig.from_file("sft.json")
 Pipeline(
    config,
    ["data_part1.jsonl", "data_part2.jsonl"],
    output_dir="output/",
    tokenizer_path="params",
 ).run()
 ```
 > Document Update Time: 2026-07-09
@@ -0,0 +1,215 @@
 # 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`.
 ## 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="mean"`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
 ### GRPO (Group Relative Policy Optimization)
 On-policy PPO with group-normalized advantages:
 $$
 \text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
 $$
 $$
 L_{\text{GRPO}} = -\mathbb{E}\left[\min\left(\frac{\pi_\theta}{\pi_{\text{ref}}}A,\; \text{clip}\left(\frac{\pi_\theta}{\pi_{\text{ref}}}, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}\left[(\log\pi_\theta - \log\pi_{\text{ref}})^2\right]
 $$
 Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`, `sync_interval=200`, `reduction="mean"`.
 Keys: `prompts`, `responses`, `masks`, `rewards`.
 ## LR Schedulers
 | Type | Class | Description |
 |------|-------|-------------|
 | Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
 | SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
 | WSD | `WSDScheduler` | Warmup-Stable-Decay with sqrt cooldown |
 Created by `SchedulerFactory.create(schedule_type, optimizer, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`, `"wsd"`. Omit to use no scheduler.
 ## Gradient Checkpointing
 Trades compute for memory by recomputing activations during backward pass. Specify module types via `gradient_checkpointing_modules`:
 ```python
 from astrai.model.components.decoder_block import DecoderBlock
 config = TrainConfig(..., gradient_checkpointing_modules=[DecoderBlock])
 ```
 Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoint(use_reentrant=False)`, compatible with `torch.compile`. Uses `nn.Module.apply()` for traversal — works through DDP wrappers without manual unwrap. Empty list (default) means no-op.
 ## Checkpoint
 ```
 Checkpoint(state_dict, epoch, consumed_samples, extra, meta, config)
  ├── save(save_dir)    rank-0 only: meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
  └── load(save_dir, broadcast=False)    loads from local disk; set broadcast=True to broadcast metadata from rank-0
 ```
 Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
 Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.
 ## TrainContextBuilder (Builder Pattern)
 ```python
 context = (
    TrainContextBuilder(config)
        .with_resume_dir(resume_dir)
        .build()
 )
 # Returns TrainContext with model, strategy, optimizer, scheduler, dataloader, checkpoint
 ```
 - Loads checkpoint weights if provided
 - Creates executor via `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)`
 - Calls `executor.prepare(model, optimizer, dataloader, scheduler)` for model distribution (e.g. DDP) + gradient accumulation wrappers
 - Creates `ResumableDistributedSampler` for shuffle+resume
 - Builds strategy via `StrategyFactory.create(train_type, model, device, **kwargs)`
 ## Training CLI
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 nohup python scripts/tools/train.py \
    --nprocs=4 \
    --parallel_mode=ddp \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --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-09
@@ -1,32 +1,98 @@
-__version__ = "1.3.5"
+__version__ = "1.3.9"
 __author__ = "ViperEkura"
 from astrai.config import (
-    ModelConfig,
+    AutoRegressiveLMConfig,
    BaseModelConfig,
    ConfigFactory,
    EncoderConfig,
    PipelineConfig,
    TrainConfig,
 )
-from astrai.dataset import DatasetFactory
+from astrai.dataset import (
    BaseDataset,
    DatasetFactory,
    ResumableDistributedSampler,
    Store,
    StoreFactory,
 )
 from astrai.factory import BaseFactory
 from astrai.inference import (
    GenerationRequest,
    InferenceEngine,
    ProtocolHandler,
    SamplingPipeline,
    get_app,
    run_server,
    sample,
 )
 from astrai.model import (
    AutoModel,
    AutoRegressiveLM,
    EmbeddingEncoder,
    LoRAConfig,
    inject_lora,
 )
 from astrai.parallel import (
    ExecutorFactory,
    get_rank,
    get_world_size,
    only_on_rank,
    spawn_parallel_fn,
 )
 from astrai.preprocessing import Pipeline, filter_by_length
 from astrai.serialization import Checkpoint
 from astrai.tokenize import AutoTokenizer, ChatTemplate
 from astrai.trainer import (
    BaseScheduler,
    BaseStrategy,
    CallbackFactory,
    SchedulerFactory,
    StrategyFactory,
    TrainCallback,
    Trainer,
 )
 from astrai.model import AutoModel, Transformer
 from astrai.tokenize import AutoTokenizer
 from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer
 __all__ = [
-    "Transformer",
+    "AutoRegressiveLM",
-    "ModelConfig",
+    "AutoRegressiveLMConfig",
-    "TrainConfig",
+    "AutoModel",
    "DatasetFactory",
    "AutoTokenizer",
    "BaseDataset",
    "BaseFactory",
    "BaseModelConfig",
    "BaseScheduler",
    "BaseStrategy",
    "CallbackFactory",
    "ChatTemplate",
    "Checkpoint",
    "ConfigFactory",
    "DatasetFactory",
    "EmbeddingEncoder",
    "EncoderConfig",
    "ExecutorFactory",
    "GenerationRequest",
    "InferenceEngine",
-    "Trainer",
+    "LoRAConfig",
-    "CallbackFactory",
+    "Pipeline",
-    "StrategyFactory",
+    "PipelineConfig",
    "ProtocolHandler",
    "ResumableDistributedSampler",
    "SamplingPipeline",
    "SchedulerFactory",
-    "BaseFactory",
+    "Store",
-    "AutoModel",
+    "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",
 ]
@@ -1,8 +1,25 @@
-from astrai.config.model_config import ModelConfig
+from astrai.config.model_config import (
    AutoRegressiveLMConfig,
    BaseModelConfig,
    ConfigFactory,
    EncoderConfig,
 )
 from astrai.config.preprocess_config import (
    InputConfig,
    OutputConfig,
    PipelineConfig,
    ProcessingConfig,
 )
 from astrai.config.train_config import TrainConfig
 __all__ = [
-    # Model configuration
+    "BaseModelConfig",
-    "ModelConfig",
+    "AutoRegressiveLMConfig",
    "EncoderConfig",
    "ConfigFactory",
    "TrainConfig",
    "InputConfig",
    "OutputConfig",
    "PipelineConfig",
    "ProcessingConfig",
 ]
@@ -0,0 +1,98 @@
 import json
 from dataclasses import MISSING, dataclass, fields
 from pathlib import Path
 from typing import Any, Dict, Optional, Self, Union, get_type_hints
@dataclass
 class BaseConfig:
    def to_dict(self) -> Dict[str, Any]:
        d = {}
        for fld in fields(self):
            v = getattr(self, fld.name)
            if isinstance(v, (str, int, float, bool)):
                d[fld.name] = v
            elif v is None:
                d[fld.name] = None
            elif isinstance(v, (dict, list, tuple)):
                try:
                    val = list(v) if isinstance(v, tuple) else v
                    json.dumps(val)
                    d[fld.name] = val
                except (TypeError, ValueError):
                    pass
            elif isinstance(v, BaseConfig):
                d[fld.name] = v.to_dict()
            elif hasattr(v, "__dataclass_fields__"):
                sub = {}
                for f in fields(v):
                    a = getattr(v, f.name)
                    sub[f.name] = list(a) if isinstance(a, tuple) else a
                d[fld.name] = sub
        return d
    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> Self:
        hints = get_type_hints(cls)
        inst = cls.__new__(cls)
        for fld in fields(cls):
            if fld.name in d:
                v = d[fld.name]
                target = cls._unwrap_optional(hints.get(fld.name))
                if target is not None:
                    try:
                        v = cls._coerce(v, target)
                    except (TypeError, ValueError):
                        pass
                object.__setattr__(inst, fld.name, v)
            elif fld.default is not MISSING:
                object.__setattr__(inst, fld.name, fld.default)
            elif fld.default_factory is not MISSING:
                object.__setattr__(inst, fld.name, fld.default_factory())
            else:
                object.__setattr__(inst, fld.name, None)
        return inst
    @staticmethod
    def _unwrap_optional(tp) -> Optional[type]:
        if tp is None:
            return None
        origin = getattr(tp, "__origin__", None)
        if origin is not None:
            args = getattr(tp, "__args__", ())
            non_none = [a for a in args if a is not type(None)]
            return non_none[0] if non_none else None
        return tp
    @staticmethod
    def _coerce(value: Any, target_type: type) -> Any:
        if target_type is bool and isinstance(value, bool):
            return value
        if (
            target_type is int
            and isinstance(value, (int, float))
            and not isinstance(value, bool)
        ):
            return int(value)
        if (
            target_type is float
            and isinstance(value, (int, float))
            and not isinstance(value, bool)
        ):
            return float(value)
        if target_type is str and isinstance(value, str):
            return value
        if isinstance(value, target_type):
            return value
        if isinstance(value, dict) and issubclass(target_type, BaseConfig):
            return target_type.from_dict(value)
        raise TypeError
    @classmethod
    def from_file(cls, path: Union[str, Path]) -> Self:
        with open(path, "r", encoding="utf-8") as f:
            return cls.from_dict(json.load(f))
    def to_file(self, path: Union[str, Path]):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
@@ -1,42 +1,82 @@
-import json
+from dataclasses import dataclass
-from dataclasses import asdict, dataclass
+from typing import Any, Dict, Optional
-from typing import Optional, Self
+
 from astrai.config.base import BaseConfig
 from astrai.factory import BaseFactory
 class ConfigFactory(BaseFactory[BaseConfig]):
    """Factory that dispatches config classes by ``model_type``."""
    @classmethod
    def load(cls, raw: Dict[str, Any]) -> BaseConfig:
        model_type = raw.get("model_type") or "autoregressive_lm"
        config_cls = cls.get_component_class(model_type)
        return config_cls.from_dict(raw)
@dataclass
-class ModelConfig:
+class BaseModelConfig(BaseConfig):
-    # basic config
+    """Base config with ``model_type`` dispatch and file I/O."""
    model_type: Optional[str] = None
    neftune_alpha: float = 0.0
@dataclass
@ConfigFactory.register("autoregressive_lm")
 class AutoRegressiveLMConfig(BaseModelConfig):
    """Configuration for autoregressive language model."""
    vocab_size: Optional[int] = None
    dim: Optional[int] = None
    n_layers: Optional[int] = None
    norm_eps: Optional[float] = None
    dim_ffn: Optional[int] = None
    tie_weight: Optional[bool] = None
    # RoPE
    max_len: Optional[int] = None
    rope_theta: Optional[float] = None
    rope_scaling: Optional[dict] = None
-    # GQA
+    attn_type: str = "gqa"
    n_heads: Optional[int] = None
    n_kv_heads: Optional[int] = None
    use_qk_norm: Optional[bool] = None
    use_gated_attention: Optional[bool] = None
-    def load(self, config_path: str) -> Self:
+    kv_lora_rank: Optional[int] = None
-        config = {}
+    qk_nope_head_dim: Optional[int] = None
-        with open(config_path, "r") as f:
+    qk_rope_head_dim: Optional[int] = None
            config.update(json.load(f))
-        for key, value in config.items():
+    ffn_type: str = "mlp"
-            if hasattr(self, key):
+    n_routed_experts: Optional[int] = None
-                setattr(self, key, value)
+    n_shared_experts: Optional[int] = None
    n_activated_experts: Optional[int] = None
    topk_method: Optional[str] = None
        return self
-    def save(self, config_path: str):
+@dataclass
-        config_dict = {k: v for k, v in asdict(self).items() if v is not None}
+@ConfigFactory.register("embedding")
-        with open(config_path, "w") as f:
+class EncoderConfig(BaseModelConfig):
-            json.dump(config_dict, f, indent=4)
+    """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
@@ -0,0 +1,109 @@
 """Pipeline configuration for JSONL preprocessing.
 Supports single-sequence (SFT/pretrain) and multi-output (DPO/GRPO)
 modes, both driven declaratively through ``input.sections`` or
 ``input.sources``.
 """
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional
 from astrai.config.base import BaseConfig
@dataclass
 class InputConfig(BaseConfig):
    """Declarative input mapping.
    Single-output mode (backward-compatible)::
        {"input": {"sections": [{"field": "messages", ...}]}}
    Multi-output mode (DPO / GRPO)::
        {"input": {"sources": {
            "chosen": {"sections": [{"field": "chosen", ...}]},
            "rejected": {"sections": [{"field": "rejected", ...}]},
        }}}
    """
    sections: Optional[List[Dict]] = None
    sources: Optional[Dict[str, Dict]] = None
@dataclass
 class ProcessingConfig(BaseConfig):
    """Processing configuration.
    Parameters
    ----------
    max_seq_len : int
        Maximum sequence length (default: 2048).
    min_chars : int
        Minimum number of characters to keep (default: 50).
    max_chars : int
        Maximum number of characters to keep (default: 2_000_000).
    max_items : Optional[int]
        Maximum number of items to process (default: None, unlimited).
    packing_strategy : str
        How to pack sequences into a contiguous stream.
        - ``"simple"``: sequential concatenation (default, backward compatible).
        - ``"bfd"``: best-fit decreasing bin packing, minimises wasted tokens.
        - ``"bfd_split"``: BFD with over-length sequences split into chunks.
    max_packed_len : int
        Maximum length of a packed bin. Sequences longer than this are
        truncated or split depending on ``packing_strategy`` (default: 8192).
    truncation_mode : str
        How to truncate sequences longer than ``max_packed_len``.
        - ``"keep_start"``: keep the first ``max_packed_len`` tokens (default).
        - ``"keep_end"``: keep the last ``max_packed_len`` tokens.
    """
    max_seq_len: int = 2048
    min_chars: int = 50
    max_chars: int = 2_000_000
    max_items: Optional[int] = None
    packing_strategy: str = "simple"
    max_packed_len: int = 8192
    truncation_mode: str = "keep_start"
@dataclass
 class OutputConfig(BaseConfig):
    """Output configuration.
    Parameters
    ----------
    domain_key : Optional[str]
        Domain key for the output store (default: None).
    storage_format : str
        Storage format, one of ``"bin"``, ``"jsonl"`` (default: ``"bin"``).
    max_tokens_per_shard : int
        Maximum tokens per shard before splitting (default: 100_000_000).
    dtype : Dict[str, str]
        Per-key dtype overrides, e.g. ``{"input_ids": "int32"}`` (default: {}).
    position_ids_mode : Optional[str]
        How to compute position_ids in packed sequences.
         - ``"none"``: do not generate (default).
        - ``"doc_reset"``: reset to 0 at each document boundary.
        - ``"continuous"``: sequential 0, 1, 2, ... (pretrain, single doc).
    """
    domain_key: Optional[str] = None
    storage_format: str = "bin"
    max_tokens_per_shard: int = 100_000_000
    dtype: Dict[str, str] = field(default_factory=dict)
    position_ids_mode: str = "doc_reset"
@dataclass
 class PipelineConfig(BaseConfig):
    version: int = 1
    input: InputConfig = field(default_factory=InputConfig)
    mask: Dict[str, str] = field(default_factory=dict)
    mask_default: str = "mask"
    preprocessing: ProcessingConfig = field(default_factory=ProcessingConfig)
    output: OutputConfig = field(default_factory=OutputConfig)
@@ -1,43 +1,79 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
-from typing import Callable, Optional
+from typing import Any, Callable, Dict, List, Optional
 import torch.nn as nn
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import Dataset
 from astrai.config.base import BaseConfig
 from astrai.model.components.lora import LoRAConfig
 def required(**kw):
    return {"required": True, **kw}
@dataclass
-class TrainConfig:
+class TrainConfig(BaseConfig):
    # basic setting
-    model: nn.Module = field(default=None, metadata={"help": "Model for training."})
+    model_fn: Callable[[], nn.Module] = field(
-    strategy: str = field(default=None, metadata={"help": "Training strategy."})
+        default=None, metadata=required(help="Model factory for training.")
-    dataset: Dataset = field(default=None, metadata={"help": "Dataset for training."})
+    )
    strategy: str = field(default=None, metadata=required(help="Training strategy."))
    dataset: Dataset = field(
        default=None, metadata=required(help="Dataset for training.")
    )
    optimizer_fn: Callable[[nn.Module], Optimizer] = field(
-        default=None, metadata={"help": "Optimizer factory for training."}
+        default=None, metadata=required(help="Optimizer factory for training.")
    )
    scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
-        default=None, metadata={"help": "Scheduler factory for training."}
+        default=None, metadata=required(help="Scheduler factory for training.")
    )
    n_epoch: int = field(default=1, metadata={"help": "Number of epochs for training."})
-    batch_size: int = field(default=4, metadata={"help": "Batch size for training."})
+    batch_per_device: int = field(
-    accumulation_steps: int = field(
+        default=4, metadata={"help": "Batch size per device."}
    )
    grad_accum_steps: int = field(
        default=1, metadata={"help": "Number of iterations between steps."}
    )
    max_grad_norm: float = field(
        default=1.0, metadata={"help": "Maximum gradient norm."}
    )
    gradient_checkpointing_modules: List[str] = field(
        default_factory=list,
        metadata={"help": "Module types to enable activation checkpointing for."},
    )
    # checkpoint setting
    start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
-    start_batch: int = field(
+    start_samples: int = field(
-        default=0, metadata={"help": "Start batch iteration for training."}
+        default=0,
        metadata={
            "help": "Start samples count (per rank). Superseded by checkpoint consumed_samples."
        },
    )
    ckpt_dir: str = field(
        default="./checkpoint", metadata={"help": "Checkpoint directory."}
    )
    ckpt_interval: int = field(
-        default=5000, metadata={"help": "Number of iterations between checkpoints."}
+        default=5000,
        metadata={"help": "Number of optimizer steps between checkpoints."},
    )
    # lora setting
    lora: Optional[LoRAConfig] = field(
        default=None,
        metadata={"help": "LoRA config. None means full fine-tuning."},
    )
    # metric setting
    log_dir: str = field(
        default="./checkpoint/logs", metadata={"help": "Directory for metric logs."}
    )
    metrics: List[str] = field(
        default_factory=lambda: ["loss", "lr", "grad_norm"],
        metadata={"help": "Metrics to record during training."},
    )
    # dataloader setting
@@ -66,18 +102,42 @@ class TrainConfig:
    master_port: str = field(
        default="29500", metadata={"help": "Master port for distributed training."}
    )
-    parallel_wrapper: Optional[Callable] = field(
+    parallel_mode: str = field(
-        default=None, metadata={"help": "Parallel function for training."}
+        default="none",
        metadata={"help": "Parallel strategy: none, ddp, fsdp."},
    )
-    state_dict_fn: Optional[Callable] = field(
+    start_method: str = field(
-        default=None, metadata={"help": "Parallel function for state  dict saving."}
+        default="spawn",
        metadata={"help": "Multiprocessing start method (spawn/fork/forkserver)."},
    )
    # others
    device_type: str = field(
        default="cuda", metadata={"help": "Device type for distributed training."}
    )
-    extra_kwargs: dict = field(
+    val_dataset: Optional[Dataset] = field(
        default=None, metadata={"help": "Dataset for validation."}
    )
    val_split: Optional[float] = field(
        default=None,
        metadata={
            "help": "Ratio to split from training dataset for validation (e.g. 0.05). Ignored if val_dataset is set."
        },
    )
    val_step: int = field(
        default=1000,
        metadata={"help": "Number of optimizer steps between validation runs."},
    )
    neftune_alpha: float = field(
        default=0.0,
        metadata={"help": "NEFTune noise alpha (0=disabled, typical: 5.0)."},
    )
    executor_kwargs: Dict[str, Any] = field(
        default_factory=dict,
        metadata={"help": "Extra kwargs passed to ExecutorFactory.create()."},
    )
    extra_kwargs: Dict[str, Any] = field(
        default_factory=dict, metadata={"help": "Other arguments."}
    )
@@ -85,14 +145,6 @@ class TrainConfig:
        self.validate()
    def validate(self):
-        required_fields = [
+        for fld in fields(self):
-            "model",
+            if fld.metadata.get("required") and getattr(self, fld.name) is None:
-            "strategy",
+                raise ValueError(f"TrainConfig.{fld.name} is required but got None.")
            "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.")
@@ -4,34 +4,32 @@ from astrai.dataset.dataset import (
 )
 from astrai.dataset.sampler import ResumableDistributedSampler
 from astrai.dataset.storage import (
-    BaseSegmentFetcher,
+    H5Store,
-    BaseStorage,
+    JsonlStore,
-    H5Storage,
+    MmapStore,
-    JSONStorage,
+    Store,
-    MultiSegmentFetcher,
+    StoreFactory,
    available_storage_types,
    create_storage,
    detect_format,
 )
 from astrai.serialization import (
    load_bin,
    load_h5,
-    load_json,
+    save_bin,
    save_h5,
    save_json,
 )
 __all__ = [
    "BaseDataset",
    "DatasetFactory",
-    "BaseSegmentFetcher",
+    "Store",
-    "MultiSegmentFetcher",
+    "StoreFactory",
-    "BaseStorage",
+    "H5Store",
-    "H5Storage",
+    "MmapStore",
-    "JSONStorage",
+    "JsonlStore",
    "create_storage",
    "detect_format",
    "available_storage_types",
    "save_h5",
    "load_h5",
-    "save_json",
+    "save_bin",
-    "load_json",
+    "load_bin",
    "ResumableDistributedSampler",
 ]
@@ -8,8 +8,8 @@ from torch import Tensor
 from torch.utils.data import Dataset
 from astrai.dataset.storage import (
-    BaseStorage,
+    Store,
-    create_storage,
+    StoreFactory,
    detect_format,
 )
 from astrai.factory import BaseFactory
@@ -26,33 +26,49 @@ class BaseDataset(Dataset, ABC):
        super().__init__()
        self.window_size = window_size
        self.stride = stride
-        self.storage: Optional[BaseStorage] = None
+        self.storage: Optional[Store] = None
-    def load(self, load_path: str, storage_type: Optional[str] = None, tokenizer=None):
+    @property
    def required_keys(self) -> List[str]:
        """Return required storage keys for this dataset type.
        Subclasses should override to specify expected keys.
        """
        return []
    def _validate_keys(self):
        if not self.required_keys:
            return
        actual_keys = set(self.storage.keys)
        missing = [k for k in self.required_keys if k not in actual_keys]
        if missing:
            raise KeyError(
                f"Dataset {type(self).__name__} requires keys {self.required_keys}, "
                f"but storage at {self._load_path} only has {sorted(actual_keys)}. "
                f"Missing: {missing}"
            )
    def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
        """Load dataset from the given path.
        Auto-detects the storage format if not specified.
        Args:
            load_path: Path to the data directory or file
-            storage_type: Force a specific storage type ("h5", "json"),
+            storage_type: Force a specific storage type ("h5", "bin", "jsonl"),
                          or None for auto-detection
-            tokenizer: Callable str -> List[int], used to tokenize raw text
+            **kwargs: Extra arguments forwarded to the store constructor and
-                       in JSON files. Ignored for HDF5.
+                      to ``store.load()``.
        Raises:
            KeyError: If the loaded storage is missing required keys.
        """
        if storage_type is None:
            storage_type = detect_format(load_path)
-        self.storage = create_storage(storage_type)
+        self.storage = StoreFactory.create(storage_type, **kwargs)
-        self.storage.load(load_path, tokenizer=tokenizer)
+        self._load_path = load_path
-
+        self.storage.load(load_path, **kwargs)
-    def load_json(self, load_path: str, tokenizer=None):
+        self._validate_keys()
        """Load dataset from JSON files explicitly.
        Args:
            load_path: Path to the JSON data file or directory
            tokenizer: Optional tokenizer callable for raw text JSON.
        """
        self.load(load_path, storage_type="json", tokenizer=tokenizer)
    @property
    def count(self) -> int:
@@ -122,26 +138,6 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
        dataset = DatasetFactory.create("custom", window_size, stride)
    """
    @classmethod
    def _validate_component(cls, dataset_cls: type) -> None:
        """Validate that the dataset class inherits from BaseDataset."""
        if not issubclass(dataset_cls, BaseDataset):
            raise TypeError(f"{dataset_cls.__name__} must inherit from BaseDataset")
    @classmethod
    def create(cls, train_type: str, window_size: int, stride: int) -> "BaseDataset":
        """Create a dataset instance.
        Args:
            train_type: Type of training ("seq", "sft", "dpo", "grpo")
            window_size: Window size for data sampling
            stride: Stride between consecutive samples
        Returns:
            Dataset instance
        """
        return super().create(train_type, window_size, stride)
    @classmethod
    def load(
        cls,
@@ -150,7 +146,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
        window_size: int,
        stride: Optional[int] = None,
        storage_type: Optional[str] = None,
-        tokenizer=None,
+        **kwargs,
    ) -> "BaseDataset":
        """Create and load a dataset in one step.
@@ -159,8 +155,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
            load_path: Path to the data file
            window_size: Window size for data sampling
            stride: Stride between consecutive samples (default: same as window_size)
-            storage_type: Storage type ("h5", "json") or None for auto-detection
+            storage_type: Storage type ("h5", "bin", "jsonl") or None for auto-detection
-            tokenizer: Callable str -> List[int] for raw text JSON tokenization
+            **kwargs: Extra arguments forwarded to ``dataset.load()``.
        Returns:
            Loaded dataset instance
@@ -169,22 +165,18 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
            stride = window_size
        dataset = cls.create(train_type, window_size, stride)
-        dataset.load(load_path, storage_type=storage_type, tokenizer=tokenizer)
+        dataset.load(load_path, storage_type=storage_type, **kwargs)
        return dataset
    @classmethod
    def available_types(cls) -> list:
        """Return list of registered dataset type names."""
        return cls.list_registered()
@DatasetFactory.register("seq")
 class SEQDataset(BaseDataset):
    """Dataset for sequential next-token prediction training."""
-    def __init__(self, window_size: int, stride: int):
+    @property
-        super().__init__(window_size, stride)
+    def required_keys(self) -> List[str]:
        return ["sequence"]
    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
        return self.storage.fetch(begin_idx, end_idx, "sequence")
@@ -202,8 +194,9 @@ class SEQDataset(BaseDataset):
 class SFTDataset(BaseDataset):
    """Dataset for supervised fine-tuning with loss masking."""
-    def __init__(self, window_size: int, stride: int):
+    @property
-        super().__init__(window_size, stride)
+    def required_keys(self) -> List[str]:
        return ["sequence", "loss_mask", "position_ids"]
    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
        return self.storage.fetch(begin_idx, end_idx, key)
@@ -211,23 +204,26 @@ class SFTDataset(BaseDataset):
    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)
+        x = self._fetch_data(begin_idx, end_idx, "sequence")
-        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(
+        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence")
-            dtype=torch.long
+        position_ids = self._fetch_data(begin_idx, end_idx, "position_ids")
-        )
+        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask")
        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}
+        return {
            "input_ids": x.to(dtype=torch.long),
            "target_ids": y.to(dtype=torch.long),
            "position_ids": position_ids.to(dtype=torch.long),
            "loss_mask": loss_mask.to(dtype=torch.bool),
        }
@DatasetFactory.register("dpo")
 class DPODataset(BaseDataset):
    """Dataset for Direct Preference Optimization training."""
-    def __init__(self, window_size: int, stride: int):
+    @property
-        super().__init__(window_size, stride)
+    def required_keys(self) -> List[str]:
        return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
        return self.storage.fetch(begin_idx, end_idx, key)
@@ -256,8 +252,9 @@ class DPODataset(BaseDataset):
 class GRPODataset(BaseDataset):
    """Dataset for Group Relative Policy Optimization training."""
-    def __init__(self, window_size: int, stride: int):
+    @property
-        super().__init__(window_size, stride)
+    def required_keys(self) -> List[str]:
        return ["prompts", "responses", "masks", "rewards"]
    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
        return self.storage.fetch(begin_idx, end_idx, key)
@@ -265,9 +262,11 @@ class GRPODataset(BaseDataset):
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        begin_idx, end_idx = self.get_index(index)
-        prompts = self._fetch_data(begin_idx, end_idx, "prompts")
+        prompts = self._fetch_data(begin_idx, end_idx, "prompts").to(dtype=torch.long)
-        responses = self._fetch_data(begin_idx, end_idx, "responses")
+        responses = self._fetch_data(begin_idx, end_idx, "responses").to(
-        masks = self._fetch_data(begin_idx, end_idx, "masks")
+            dtype=torch.long
        )
        masks = self._fetch_data(begin_idx, end_idx, "masks").to(dtype=torch.bool)
        rewards = self._fetch_data(begin_idx, end_idx, "rewards")
        return {
@@ -43,6 +43,7 @@ class ResumableDistributedSampler(Sampler[int]):
        offset = 0 if drop_last else self.num_replicas - 1
        self.num_samples_per_replica = (self.num_samples + offset) // self.num_replicas
        self.total_size = self.num_samples_per_replica * self.num_replicas
        self.iter = self.iter % self.num_samples_per_replica
        self._indices = None
@@ -73,6 +74,12 @@ class ResumableDistributedSampler(Sampler[int]):
        self.epoch += 1
        self._indices = None
        self.iter = self.iter % self.num_samples_per_replica
    @property
    def _remaining(self):
        remaining = self.num_samples_per_replica - self.iter
        return max(remaining, 0)
    def __len__(self):
-        return self.num_samples_per_replica
+        return self._remaining
@@ -1,105 +1,44 @@
 """Storage backends for different data formats.
-Each storage handles format-specific loading (HDF5, JSON, etc.) and provides
+Layers:
-a uniform interface for data access and length observation via fetchers.
+  - I/O layer:       save_* / load_* functions, read/write raw files (HDF5/bin)
                      return Dict[str, List[Tensor]] — format-specific, no state
  - Store (ABC):     central abstraction, normalizes multi-segment into
                      Dict[str, List[Tensor]] per key via _normalize(),
                      fetch() uses bisect across segments — no forced concat
  - Dataset layer:   BaseDataset owns a Store, only calls store.fetch(begin, end, key)
 Key properties:
  - Multi-segment:   segments kept as-is, no forced concatenation — safe for
                      datasets larger than RAM
  - Explicit length: _length = min(total elements across keys), set at load,
                      __len__ returns O(1)
  - Zero-copy mmap:  MmapStore wraps np.memmap(mode="r"), all DataLoader
                       workers share OS page-cache pages
 """
 import bisect
 import glob
 import json
-import os
+import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Callable, Dict, List, Optional, Union
+from typing import Dict, List, Union
 import h5py
 import torch
 from torch import Tensor
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.factory import BaseFactory
 from astrai.preprocessing.builder import MaskBuilderFactory
 from astrai.preprocessing.position_id import PositionIdStrategyFactory
 from astrai.serialization import (
    load_bin,
    load_h5,
 )
 from astrai.tokenize import AutoTokenizer
-def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
+logger = logging.getLogger(__name__)
    os.makedirs(file_path, exist_ok=True)
    full_file_path = os.path.join(file_path, f"{file_name}.h5")
    with h5py.File(full_file_path, "w") as f:
        for key, tensors in tensor_group.items():
            grp = f.create_group(key)
            for idx, tensor in enumerate(tensors):
                arr = tensor.cpu().numpy()
                grp.create_dataset(f"data_{idx}", data=arr)
 def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
    tensor_group: Dict[str, List[Tensor]] = {}
    root_path = Path(file_path)
    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
    for h5_file in h5_files:
        with h5py.File(h5_file, "r") as f:
            for key in f.keys():
                grp = f[key]
                dsets = []
                for dset_name in grp.keys():
                    dset = grp[dset_name]
                    tensor = torch.from_numpy(dset[:])
                    if share_memory:
                        tensor = tensor.share_memory_()
                    dsets.append(tensor)
                if tensor_group.get(key) is None:
                    tensor_group[key] = []
                tensor_group[key].extend(dsets)
    return tensor_group
 def save_json(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
    os.makedirs(file_path, exist_ok=True)
    full_file_path = os.path.join(file_path, f"{file_name}.json")
    json_data = {}
    for key, tensors in tensor_group.items():
        json_data[key] = [tensor.tolist() for tensor in tensors]
    with open(full_file_path, "w", encoding="utf-8") as f:
        json.dump(json_data, f, ensure_ascii=False)
 def load_json(
    file_path: str,
    share_memory: bool = True,
    tokenizer: Optional[Callable[[str], List[int]]] = None,
 ) -> Dict[str, List[Tensor]]:
    """Load tensor data from JSON files.
    Supports two modes:
    - Pre-tokenized: JSON values are List[List[int]] (token IDs), loaded as-is.
    - Raw text: JSON values are List[str], tokenized via ``tokenizer`` callable
      at load time. A ``tokenizer`` receives a str and returns List[int].
    Non-data JSON files (e.g. config.json) with scalar/object values are
    silently skipped.
    """
    tensor_group: Dict[str, List[Tensor]] = {}
    root_path = Path(file_path)
    json_files = list(root_path.rglob("*.json")) + list(root_path.rglob("*.jsonl"))
    for json_file in json_files:
        with open(json_file, "r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, dict):
            continue
        for key, sequences in data.items():
            if not isinstance(sequences, list):
                continue
            tensors = []
            for seq in sequences:
                if tokenizer is not None and isinstance(seq, str):
                    seq = tokenizer(seq)
                tensor = torch.tensor(seq, dtype=torch.long)
                if share_memory:
                    tensor = tensor.share_memory_()
                tensors.append(tensor)
            if tensor_group.get(key) is None:
                tensor_group[key] = []
            tensor_group[key].extend(tensors)
    return tensor_group
 def detect_format(load_path: str) -> str:
@@ -109,7 +48,7 @@ def detect_format(load_path: str) -> str:
        load_path: Directory or file path
    Returns:
-        Format string ("h5" or "json")
+        Format string ("h5", "bin", or "jsonl")
    Raises:
        FileNotFoundError: If no supported data files are found
@@ -119,194 +58,255 @@ def detect_format(load_path: str) -> str:
        suffix = root.suffix.lower()
        if suffix in (".h5", ".hdf5"):
            return "h5"
-        if suffix in (".json", ".jsonl"):
+        if suffix == ".jsonl":
-            return "json"
+            return "jsonl"
        raise ValueError(f"Unsupported file format: {suffix}")
-    h5_files = list(root.rglob("*.h5")) + list(root.rglob("*.hdf5"))
+    h5_files = [
        Path(p)
        for pattern in ("*.h5", "*.hdf5")
        for p in glob.glob(str(root / "**" / pattern), recursive=True)
    ]
    if h5_files:
        return "h5"
-    json_files = list(root.rglob("*.json")) + list(root.rglob("*.jsonl"))
+    bin_files = [Path(p) for p in glob.glob(str(root / "**" / "*.bin"), recursive=True)]
-    if json_files:
+    if bin_files:
-        return "json"
+        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 BaseSegmentFetcher:
+class Store(ABC):
-    """Fetches data segments across multiple tensor segments.
+    """String keys -> segmented tensors with ``fetch(begin, end, keys)``.
-    Maintains cumulative lengths for efficient range queries across
+    Each key maps to one or more tensor segments (no forced concatenation).
-    multiple discontinuous segments.
+    ``len(store)`` returns ``self._length`` (explicit, O(1)), the minimum
-    """
+    total element count across all keys.
-    def __init__(self, segments: List[Tensor]):
+    Subclasses fill ``self._data`` and ``self._cum`` during ``load()``
-        self.segments = segments
+    via ``_normalize()``.
        self.cum_lengths = []
        total = 0
        for seg in segments:
            total += torch.numel(seg)
            self.cum_lengths.append(total)
        self.total_length = total
    def __len__(self) -> int:
        return self.total_length
    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
        """Fetch data in the range [begin_idx, end_idx)."""
        if not (
            0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length
        ):
            raise ValueError("begin_idx or end_idx out of bounds")
        if begin_idx >= end_idx:
            return torch.tensor([], dtype=torch.long)
        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
        result_segments = []
        for i in range(seg_start_idx, seg_end_idx + 1):
            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
            start = max(begin_idx - prev_cum, 0)
            end = min(end_idx - prev_cum, len(self.segments[i]))
            result_segments.append(self.segments[i][start:end])
        return torch.cat(result_segments, dim=0)
 class MultiSegmentFetcher:
    """Manages multiple segment fetchers for different data keys."""
    def __init__(self, multi_segments: Dict):
        self.multi_keys = list(multi_segments.keys())
        self.multi_fetchers = {
            key: BaseSegmentFetcher(segments)
            for key, segments in multi_segments.items()
        }
    def __len__(self) -> int:
        """Returns the minimum length across all fetchers."""
        if not self.multi_fetchers:
            return 0
        len_list = [len(seg) for seg in self.multi_fetchers.values()]
        return min(len_list)
    def key_fetch(
        self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]
    ) -> Dict:
        """Fetch data for specific keys."""
        fetch_dict = {}
        keys = [keys] if isinstance(keys, str) else keys
        for key in keys:
            fetcher = self.multi_fetchers[key]
            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
            fetch_dict[key] = fetch_tensor
        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
    def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
        """Fetch all keys."""
        return self.key_fetch(begin_idx, end_idx, self.multi_keys)
 class BaseStorage(ABC):
    """Abstract storage backend for loading and dispatching data.
    Storage encapsulates format-specific loading and provides a uniform
    interface for data access and length observation. Subclasses handle
    different data formats (HDF5, JSON, etc.) while exposing the same
    fetch interface.
    """
    def __init__(self):
-        self._fetcher: Optional[MultiSegmentFetcher] = None
+        self._data: Dict[str, List[Tensor]] = {}
        self._cum: Dict[str, List[int]] = {}
        self._length: int = 0
    @abstractmethod
-    def load(self, load_path: str, tokenizer=None) -> None:
+    def load(self, path: str) -> None:
        """Load data from the given path into internal fetcher."""
        raise NotImplementedError
    def __len__(self) -> int:
        """Total number of raw elements (tokens) in storage."""
        if self._fetcher is None:
            return 0
        return len(self._fetcher)
    def fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]):
        """Fetch data for the given keys and index range.
        Args:
            begin_idx: Starting index (inclusive)
            end_idx: Ending index (exclusive)
            keys: Single key or list of keys to fetch
        Returns:
            Tensor if single key, Dict[str, Tensor] if multiple keys
        """
        if self._fetcher is None:
            raise RuntimeError("Storage not loaded")
        return self._fetcher.key_fetch(begin_idx, end_idx, keys)
    @property
    def keys(self) -> List[str]:
-        """Return the data keys available in this storage."""
+        return list(self._data.keys())
-        if self._fetcher is None:
+
-            return []
+    def __len__(self) -> int:
-        return self._fetcher.multi_keys
+        return self._length
    def fetch(
        self,
        begin: int,
        end: int,
        keys: Union[str, List[str]],
    ):
        if not self._data:
            raise RuntimeError("Store not loaded")
        if not (0 <= begin < self._length and 0 <= end <= self._length):
            raise ValueError(
                f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
            )
        if isinstance(keys, str):
            return self._fetch_key(keys, begin, end)
        return {k: self._fetch_key(k, begin, end) for k in keys}
    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
        """Fetch slice [begin, end) across potentially multiple segments."""
        segments = self._data[key]
        cum = self._cum[key]
        seg_start = bisect.bisect_right(cum, begin)
        seg_end = bisect.bisect_left(cum, end)
        results = []
        for i in range(seg_start, seg_end + 1):
            prev = cum[i - 1] if i > 0 else 0
            s = max(begin - prev, 0)
            e = min(end - prev, segments[i].shape[0])
            results.append(segments[i][s:e])
        return results[0] if len(results) == 1 else torch.cat(results, dim=0)
    def _normalize(self, raw: Dict[str, List[Tensor]]):
        """Register segments and pre-compute cumulative lengths.
        Does NOT concatenate — segments are kept as-is to avoid OOM on
        large datasets.  Sets ``self._length`` to the minimum total
        element count across all keys.
        """
        for key, tensors in raw.items():
            self._data[key] = tensors
            cum = []
            total = 0
            for t in tensors:
                total += t.shape[0]
                cum.append(total)
            self._cum[key] = cum
        self._length = (
            min((cum[-1] if cum else 0) for cum in self._cum.values())
            if self._cum
            else 0
        )
-class H5Storage(BaseStorage):
+class StoreFactory(BaseFactory["Store"]):
    """Factory for creating Store instances by type name.
    Example::
        @StoreFactory.register("custom")
        class CustomStore(Store):
            ...
    """
@StoreFactory.register("h5")
 class H5Store(Store):
    """HDF5-based storage backend (pre-tokenized data)."""
-    def load(self, load_path: str, tokenizer=None) -> None:
+    def load(self, path: str):
-        segments = load_h5(load_path)
+        self._normalize(load_h5(path))
        self._fetcher = MultiSegmentFetcher(segments)
-class JSONStorage(BaseStorage):
+@StoreFactory.register("bin")
-    """JSON-based storage backend.
+class MmapStore(Store):
    """Memory-mapped binary storage backend.
-    Supports two modes:
+    Each key is a single .bin file backed by ``np.memmap(mode="r")``.
-    - Pre-tokenized: JSON values are List[List[int]], loaded as-is.
+    No per-process memory duplication — all DataLoader workers share the
-    - Raw text: JSON values are List[str], tokenized via ``tokenizer``
+    same OS page-cache pages.
-      callable (str -> List[int]) at load time.
+
    Format on disk::
        data_root/
          meta.json          # {key: {shape, dtype}, ...}
          <key>.bin          # raw numpy array, one per key
    """
-    def load(self, load_path: str, tokenizer=None) -> None:
+    def load(self, path: str):
-        segments = load_json(load_path, tokenizer=tokenizer)
+        self._mmap_refs = []
-        self._fetcher = MultiSegmentFetcher(segments)
+        root = Path(path)
        all_raw: Dict[str, List[Tensor]] = {}
        meta_paths = [
            Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)
        ]
        for meta_path in meta_paths:
            raw = load_bin(str(meta_path.parent))
            for key, tensors in raw.items():
                if key not in all_raw:
                    all_raw[key] = []
                all_raw[key].extend(tensors)
        if not meta_paths:
            raise FileNotFoundError(f"No meta.json found under {path}")
        self._normalize(all_raw)
        for tensors in self._data.values():
            self._mmap_refs.extend(tensors)
-_STORAGE_REGISTRY: Dict[str, type] = {
+@StoreFactory.register("jsonl")
-    "h5": H5Storage,
+class JsonlStore(Store):
-    "json": JSONStorage,
+    """On-the-fly tokenization store for raw JSONL files.
 }
-
+    A JSONL dataset directory contains ``*.jsonl`` files plus a
-def create_storage(storage_type: str) -> BaseStorage:
+    ``dataset_config.json`` file that follows the same schema as
-    """Create a storage instance by type name.
+    :class:`PipelineConfig` with an additional ``tokenizer_path`` field.
-
+    Records are tokenized when the store is loaded and concatenated into
-    Args:
+    segmented tensors matching the key layout expected by the dataset
-        storage_type: Storage type name ("h5", "json")
+    classes (``sequence``, ``loss_mask``, ``position_ids``, ...).
    Returns:
        Storage instance
    Raises:
        ValueError: If the storage type is unknown
    """
-    storage_cls = _STORAGE_REGISTRY.get(storage_type)
+
-    if storage_cls is None:
+    CONFIG_NAME = "dataset_config.json"
-        raise ValueError(
+
-            f"Unknown storage type: '{storage_type}'. "
+    def load(self, path: str):
-            f"Available: {sorted(_STORAGE_REGISTRY.keys())}"
+        root = Path(path)
        config_path = root / self.CONFIG_NAME
        if not config_path.exists():
            raise FileNotFoundError(
                f"JSONL dataset config not found: {config_path}. "
                f"Expected {self.CONFIG_NAME} alongside *.jsonl files."
            )
        with open(config_path, "r", encoding="utf-8") as f:
            raw_config = json.load(f)
        tokenizer_path = raw_config.pop("tokenizer_path", None)
        if tokenizer_path is None:
            raise ValueError(
                f"JSONL dataset config must specify 'tokenizer_path': {config_path}"
            )
        self.config = PipelineConfig.from_dict(raw_config)
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        mask_builder = MaskBuilderFactory.create("sectioned")
        position_strategy = PositionIdStrategyFactory.create(
            self.config.output.position_ids_mode
        )
    return storage_cls()
        raw: Dict[str, List[Tensor]] = {}
        doc_sequences: List[List[int]] = []
-def available_storage_types() -> List[str]:
+        for jsonl_path in sorted(root.glob("*.jsonl")):
-    """Return list of registered storage type names."""
+            with open(jsonl_path, "r", encoding="utf-8") as f:
-    return sorted(_STORAGE_REGISTRY.keys())
+                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        item = json.loads(line)
                    except json.JSONDecodeError:
                        logger.warning(
                            "Failed to parse JSON line in %s, skipping", jsonl_path
                        )
                        continue
                    result = mask_builder.build(item, self.config, tokenizer)
                    if result is None:
                        continue
                    result.pop("domain", None)
                    primary_ids = self._primary_ids(result)
                    if not primary_ids:
                        continue
                    doc_sequences.append(primary_ids)
                    for key, ids in result.items():
                        if key not in raw:
                            raw[key] = []
                        raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
        pos_ids = position_strategy.generate(doc_sequences)
        if pos_ids:
            raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
        self._normalize(raw)
    @staticmethod
    def _primary_ids(result: dict) -> List[int]:
        """Return the first integer list in *result* as the primary id sequence."""
        for val in result.values():
            if isinstance(val, list) and val and isinstance(val[0], int):
                return val
        return []
    @staticmethod
    def _infer_dtype(ids: List) -> torch.dtype:
        """Infer tensor dtype from the first element of a token/value list."""
        if ids and isinstance(ids[0], float):
            return torch.float32
        return torch.int32
@@ -0,0 +1,21 @@
 """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)
 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",
 ]
@@ -0,0 +1,36 @@
 """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)
@@ -0,0 +1,149 @@
 """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.
 Add new kernel wrappers here; split into per-variant files only if this file
 grows large.
 """
 import torch
 import torch.nn.functional as F
 from astrai.extension.loader import _available, _modules
 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 _torch_fallback(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: torch.Tensor | None,
    is_causal: bool,
    scale: float | None,
 ) -> torch.Tensor:
    """Reference attention via ``scaled_dot_product_attention``."""
    k, v = _expand_kv_heads(k, v, q.size(1))
    attn_mask = mask[:, None, None, :] if mask is not None else None
    return F.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale
    )
 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, n_kv_heads, kv_len, head_dim]
    """
    batch, max_pages = page_table.shape
    n_pages, 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}")
    k = k_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
    v = v_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
    for b in range(batch):
        for pos in range(kv_len):
            log_pg = pos // page_size
            pg_off = pos % page_size
            phys = int(page_table[b, log_pg].item())
            k[b, :, pos, :] = k_cache[phys, pg_off, :, :]
            v[b, :, pos, :] = v_cache[phys, pg_off, :, :]
    return k, v
 def attn_decode(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: torch.Tensor | None = None,
    is_causal: bool = False,
    causal_offset: int = 0,
    scale: float | None = None,
 ) -> torch.Tensor:
    if _available["attn_decode"]:
        return _modules["attn_decode"].attn_decode(
            q,
            k,
            v,
            mask=mask,
            is_causal=is_causal,
            causal_offset=causal_offset,
            scale=scale,
        )
    return _torch_fallback(q, k, v, mask, is_causal, scale)
 def attn_prefill(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: torch.Tensor | None = None,
    is_causal: bool = False,
    causal_offset: int = 0,
    scale: float | None = None,
 ) -> torch.Tensor:
    if _available["attn_prefill"]:
        return _modules["attn_prefill"].attn_prefill(
            q,
            k,
            v,
            mask=mask,
            is_causal=is_causal,
            causal_offset=causal_offset,
            scale=scale,
        )
    return _torch_fallback(q, k, v, mask, is_causal, scale)
 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,
    is_causal: bool = False,
    causal_offset: int = 0,
    scale: float | None = None,
 ) -> torch.Tensor:
    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,
            is_causal=is_causal,
            causal_offset=causal_offset,
            scale=scale,
        )
    k, v = _gather_kv_from_pages(page_table, k_cache, v_cache, page_size, kv_len)
    return _torch_fallback(q, k, v, mask, is_causal, scale)
@@ -1,210 +1,144 @@
-"""Base factory class for extensible component registration."""
+"""Base factory with decorator-based registration and kwarg-filtered instantiation."""
 import inspect
 import sys
 from abc import ABC
-from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+from typing import (
    Callable,
    Dict,
    ForwardRef,
    Generic,
    List,
    Optional,
    Type,
    TypeVar,
    Union,
 )
 from typing import get_args as _get_args
 from typing import get_origin as _get_origin
 T = TypeVar("T")
-class Registry:
+def _resolve_type(
-    """Flexible registry for component classes with category and priority support.
+    arg: Union[Type, str, ForwardRef], factory_cls: type
 ) -> Optional[Type]:
    """Resolve a generic type-arg (str forward-ref, ForwardRef, or class)."""
    if not isinstance(arg, (str, ForwardRef)):
        return arg
-    This registry stores component classes with optional metadata (category, priority).
+    name = arg if isinstance(arg, str) else arg.__forward_arg__
-    It provides methods for registration, retrieval, and listing with filtering.
+    if name == factory_cls.__name__:
-    """
+        return factory_cls
-    def __init__(self):
+    mod = sys.modules.get(factory_cls.__module__)
-        self._entries = {}  # name -> (component_cls, category, priority)
+    if mod is None:
        return None
    ns = vars(mod)
-    def register(
+    if isinstance(arg, ForwardRef):
-        self,
+        return arg._evaluate(ns, None, recursive_guard=frozenset())
        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)
-    def get(self, name: str) -> Type:
+    return ns.get(name)
        """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 class for component registration and creation.
+    """Generic factory with decorator-based component registration.
-    This base class provides a decorator-based registration pattern
+        class MyFactory(BaseFactory[MyBase]):
    for creating extensible component factories.
    Example usage:
        class MyFactory(BaseFactory[MyBaseClass]):
            pass
        @MyFactory.register("custom")
-        class CustomComponent(MyBaseClass):
+        class CustomComponent(MyBase):
            ...
-        component = MyFactory.create("custom", *args, **kwargs)
+        obj = MyFactory.create("custom", *args, **kwargs)
    ``create()`` filters kwargs to match the component's ``__init__``
    signature so components don't need ``**kwargs`` just to absorb
    unrelated parameters.
    """
-    _registry: Registry
+    _entries: Dict[str, Type[T]]
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
-        cls._registry = Registry()
+        for orig_base in getattr(cls, "__orig_bases__", ()):
            if _get_origin(orig_base) is BaseFactory:
                (arg,) = _get_args(orig_base)
                cls._entries = {}
                cls._component_base = _resolve_type(arg, cls)
                return
    @classmethod
-    def register(
+    def register(cls, name: str) -> Callable[[Type[T]], Type[T]]:
-        cls, name: str, category: Optional[str] = None, priority: int = 0
+        """Decorator to register a component class.
    ) -> Callable[[Type[T]], Type[T]]:
        """Decorator to register a component class with optional category and priority.
-        Args:
+        Validates that the decorated class inherits from the generic
-            name: Registration name for the component
+        type parameter ``T`` declared on the factory.
            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)
-            cls._registry.register(
+            if name in cls._entries:
-                name, component_cls, category=category, priority=priority
+                raise ValueError(f"Component '{name}' is already registered")
-            )
+            cls._entries[name] = component_cls
            return component_cls
        return decorator
    @classmethod
    def create(cls, name: str, *args, **kwargs) -> T:
-        """Create a component instance by name.
+        """Create a component instance by name, filtering kwargs to match
-
+        the component's ``__init__`` signature.
        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
        """
-        if not cls._registry.contains(name):
+        entry = cls._entries.get(name)
        if entry is None:
            raise ValueError(
-                f"Unknown component: '{name}'. "
+                f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
                f"Supported types: {sorted(cls._registry.list_names())}"
            )
-        component_cls = cls._registry.get(name)
+        component_cls = entry
        sig = inspect.signature(component_cls.__init__)
        has_var_kwargs = any(
            p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
        )
        if not has_var_kwargs:
            valid = {
                p.name
                for p in sig.parameters.values()
                if p.name != "self" and p.kind != inspect.Parameter.VAR_KEYWORD
            }
            kwargs = {k: v for k, v in kwargs.items() if k in valid}
        return component_cls(*args, **kwargs)
    @classmethod
-    def _validate_component(cls, component_cls: Type[T]) -> None:
+    def _validate_component(cls, component_cls: Type[T]):
-        """Validate that the component class is valid for this factory.
+        """Validate the decorated class inherits from the factory's base type.
-        Override this method in subclasses to add custom validation.
+        Override for custom validation beyond ``issubclass``.
        Args:
            component_cls: Component class to validate
        Raises:
            TypeError: If the component class is invalid
        """
-        pass
+        base = cls._component_base
        if base is not None and not issubclass(component_cls, base):
            raise TypeError(
                f"{component_cls.__name__} must inherit from {base.__name__}"
            )
    @classmethod
    def get_component_class(cls, name: str) -> Type[T]:
-        """Get the registered component class by name without instantiating it.
+        """Get the registered component class without instantiating it."""
-
+        entry = cls._entries.get(name)
-        Args:
+        if entry is None:
            name: Registered name of the component
        Returns:
            The component class itself
        Raises:
            ValueError: If the component name is not registered
        """
        if not cls._registry.contains(name):
            raise ValueError(
-                f"Unknown component: '{name}'. "
+                f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
                f"Supported types: {sorted(cls._registry.list_names())}"
            )
-        return cls._registry.get(name)
+        return entry
    @classmethod
-    def list_registered(cls) -> list:
+    def list_registered(cls) -> List[str]:
-        """List all registered component names.
+        """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.
+        """Check if a component name is registered."""
-
+        return name in cls._entries
        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,32 +1,43 @@
 """Inference module for continuous batching.
 Layers:
-  - core/:           Core inference loop (cache, executor, scheduler, task)
+  - core/:        Core inference loop (cache, executor, scheduler, task)
-  - api/:            HTTP protocol handlers (OpenAI, Anthropic)
+  - api/:         HTTP orchestration (ProtocolHandler, server)
-  - engine.py:       Facade (InferenceEngine), Value Object (GenerationRequest)
+  - protocols/:   Response builders (OpenAI, Anthropic)
-  - sample.py:       Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
+  - transport/:   SSE transport utilities
  - engine.py:    Facade (InferenceEngine), Value Object (GenerationRequest)
  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
 """
 from astrai.inference.api import (
    AnthropicHandler,
    AnthropicMessage,
    BaseToolParser,
    ChatCompletionRequest,
    ChatMessage,
    FunctionDef,
    GenContext,
    MessagesRequest,
    OpenAIHandler,
    ProtocolHandler,
    SimpleJsonToolParser,
    StopChecker,
-    StreamContext,
+    ToolDef,
-    app,
+    ToolParserFactory,
    get_app,
    run_server,
 )
 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,
-    KvcacheView,
+    PageCache,
    PageCacheView,
    PagePool,
    PrefixCache,
    Storage,
@@ -36,10 +47,7 @@ from astrai.inference.core import (
    TaskTable,
    page_hash,
 )
-from astrai.inference.engine import (
+from astrai.inference.engine import GenerationRequest, InferenceEngine
    GenerationRequest,
    InferenceEngine,
 )
 from astrai.inference.sample import (
    BaseSamplingStrategy,
    SamplingPipeline,
@@ -50,43 +58,46 @@ from astrai.inference.sample import (
 )
 __all__ = [
    # Engine / Requests
    "InferenceEngine",
    "GenerationRequest",
    # Core scheduler
    "InferenceScheduler",
    "Executor",
    "STOP",
    "Task",
    "TaskManager",
    "TaskStatus",
    # Core cache
    "Allocator",
    "CacheView",
    "KVCache",
-    "KvcacheView",
+    "ContiguousCache",
    "ContiguousCacheView",
    "PageCache",
    "PageCacheView",
    "PagePool",
    "PrefixCache",
    "Storage",
    "TaskTable",
    "page_hash",
    # Sampling (Strategy pattern)
    "sample",
    "BaseSamplingStrategy",
    "TemperatureStrategy",
    "TopKStrategy",
    "TopPStrategy",
    "SamplingPipeline",
    # Protocol
    "ProtocolHandler",
    "StopChecker",
-    "StreamContext",
+    "GenContext",
-    "AnthropicHandler",
+    "BaseToolParser",
-    "OpenAIHandler",
+    "SimpleJsonToolParser",
-    # Server
+    "ToolParserFactory",
    "OpenAIResponseBuilder",
    "AnthropicResponseBuilder",
    "ChatMessage",
    "ChatCompletionRequest",
    "FunctionDef",
    "ToolDef",
    "AnthropicMessage",
    "MessagesRequest",
-    "app",
+    "get_app",
    "run_server",
 ]
@@ -1,31 +1,39 @@
-"""Inference API: protocol handlers and FastAPI server."""
+"""Inference API: protocol handler, stop checker, tool parsers, and FastAPI server.
-from astrai.inference.api.protocol import (
+``app`` is no longer a module-level global. Use :func:`get_app` to access the
-    AnthropicHandler,
+lazy singleton FastAPI instance.
-    OpenAIHandler,
+"""
-    ProtocolHandler,
+
-    StopChecker,
+from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
    StreamContext,
 )
 from astrai.inference.api.server import (
    AnthropicMessage,
    ChatCompletionRequest,
    ChatMessage,
    FunctionDef,
    MessagesRequest,
-    app,
+    ToolDef,
    get_app,
    run_server,
 )
 from astrai.inference.api.tool_parser import (
    BaseToolParser,
    SimpleJsonToolParser,
    ToolParserFactory,
 )
 __all__ = [
    "AnthropicHandler",
    "OpenAIHandler",
    "ProtocolHandler",
    "StopChecker",
-    "StreamContext",
+    "GenContext",
    "BaseToolParser",
    "SimpleJsonToolParser",
    "ToolParserFactory",
    "AnthropicMessage",
    "ChatCompletionRequest",
    "ChatMessage",
    "FunctionDef",
    "ToolDef",
    "MessagesRequest",
-    "app",
+    "get_app",
    "run_server",
 ]
@@ -0,0 +1,142 @@
 """Anthropic message completion response builder."""
 import time
 import uuid
 from typing import Any, Dict, List, Tuple, Union
 from pydantic import BaseModel
 from astrai.inference.api.protocol import (
    GenContext,
    ResponseBuilder,
    StopInfo,
    sse_event,
 )
 from astrai.inference.engine import InferenceEngine
 def _extract_text(content: Union[str, List[Dict[str, Any]]]) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                return block.get("text", "")
    return ""
 class AnthropicResponseBuilder(ResponseBuilder):
    def prepare(
        self, request: BaseModel, engine: InferenceEngine
    ) -> Tuple[str, GenContext, List[str]]:
        messages: List[Dict[str, str]] = []
        system = getattr(request, "system", None)
        if system:
            messages.append({"role": "system", "content": system})
        for m in request.messages:
            text = _extract_text(m.content)
            if text:
                messages.append({"role": m.role, "content": text})
        prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
        ctx = GenContext(
            resp_id=f"msg_{uuid.uuid4().hex[:24]}",
            created=int(time.time()),
            model=request.model,
        )
        stop_sequences = getattr(request, "stop_sequences", None) or []
        return prompt, ctx, stop_sequences
    def format_stream_start(self, ctx: GenContext) -> List[str]:
        return [
            sse_event(
                {
                    "type": "message_start",
                    "message": {
                        "id": ctx.resp_id,
                        "type": "message",
                        "role": "assistant",
                        "model": ctx.model,
                        "content": [],
                        "usage": {"input_tokens": ctx.prompt_tokens},
                    },
                },
                event="message_start",
            ),
            sse_event(
                {
                    "type": "content_block_start",
                    "index": 0,
                    "content_block": {"type": "text", "text": ""},
                },
                event="content_block_start",
            ),
        ]
    def format_chunk(self, token: str, **kwargs) -> List[str]:
        return [
            sse_event(
                {
                    "type": "content_block_delta",
                    "index": 0,
                    "delta": {"type": "text_delta", "text": token},
                },
                event="content_block_delta",
            )
        ]
    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
        events: List[str] = []
        if stop.matched:
            trimmed = stop.body[: stop.body.rfind(stop.matched)]
            unyielded = trimmed[len(stop.yielded) :]
            if unyielded:
                events.append(
                    sse_event(
                        {
                            "type": "content_block_delta",
                            "index": 0,
                            "delta": {"type": "text_delta", "text": unyielded},
                        },
                        event="content_block_delta",
                    )
                )
        events.append(
            sse_event(
                {"type": "content_block_stop", "index": 0},
                event="content_block_stop",
            )
        )
        events.append(
            sse_event(
                {
                    "type": "message_delta",
                    "delta": {
                        "stop_reason": "stop_sequence" if stop.matched else "end_turn",
                        "stop_sequence": stop.matched,
                    },
                    "usage": {"output_tokens": ctx.completion_tokens},
                },
                event="message_delta",
            )
        )
        events.append(sse_event({"type": "message_stop"}, event="message_stop"))
        return events
    def format_response(
        self, ctx: GenContext, content: str, stop: StopInfo
    ) -> Dict[str, Any]:
        if stop.matched:
            content = content[: content.rfind(stop.matched)]
        return {
            "id": ctx.resp_id,
            "type": "message",
            "role": "assistant",
            "model": ctx.model,
            "content": [{"type": "text", "text": content}],
            "stop_reason": "stop_sequence" if stop.matched else "end_turn",
            "stop_sequence": stop.matched,
            "usage": {
                "input_tokens": ctx.prompt_tokens,
                "output_tokens": ctx.completion_tokens,
            },
        }
@@ -0,0 +1,278 @@
 """OpenAI chat completion response builder."""
 import logging
 import time
 import uuid
 from typing import Any, Dict, List, Optional, Tuple, Union
 from pydantic import BaseModel
 from astrai.inference.api.protocol import (
    GenContext,
    ResponseBuilder,
    StopInfo,
    sse_event,
 )
 from astrai.inference.api.tool_parser import BaseToolParser, ToolParserFactory
 from astrai.inference.engine import InferenceEngine
 logger = logging.getLogger(__name__)
 _UNSUPPORTED_PARAMS = (
    "n",
    "presence_penalty",
    "frequency_penalty",
    "logit_bias",
    "user",
 )
 def _resolve_tool_choice(
    request: BaseModel,
 ) -> Union[str, Dict[str, Any]]:
    tc = getattr(request, "tool_choice", None)
    if tc is None:
        return "auto"
    if isinstance(tc, str):
        return tc
    if isinstance(tc, dict):
        return tc
    return "auto"
 def _resolve_tools(request: BaseModel) -> Optional[List[Dict[str, Any]]]:
    raw = getattr(request, "tools", None)
    if not raw:
        return None
    if isinstance(raw, list):
        return [t.model_dump() if hasattr(t, "model_dump") else t for t in raw]
    return None
 class OpenAIResponseBuilder(ResponseBuilder):
    def prepare(
        self, request: BaseModel, engine: InferenceEngine
    ) -> Tuple[str, GenContext, List[str]]:
        messages = [{"role": m.role, "content": m.content} for m in request.messages]
        tools = _resolve_tools(request)
        prompt = engine.tokenizer.apply_chat_template(
            messages, tokenize=False, tools=tools or []
        )
        self._resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
        self._model = request.model
        for param in _UNSUPPORTED_PARAMS:
            value = getattr(request, param, None)
            fields = getattr(type(request), "model_fields", {})
            default = fields[param].default if param in fields else None
            if value is not None and value != default:
                logger.warning(
                    "ChatCompletionRequest param '%s'=%r is not supported"
                    " and will be ignored",
                    param,
                    value,
                )
        self._parser: Optional[BaseToolParser] = None
        if tools:
            tool_choice = _resolve_tool_choice(request)
            self._parser = ToolParserFactory.create(
                "simple_json", tools=tools, tool_choice=tool_choice
            )
        self._content_started = False
        ctx = GenContext(
            resp_id=self._resp_id,
            created=int(time.time()),
            model=self._model,
        )
        stop = request.stop
        stop_sequences = (
            [] if stop is None else [stop] if isinstance(stop, str) else stop
        )
        return prompt, ctx, stop_sequences
    def format_stream_start(self, ctx: GenContext) -> List[str]:
        return [
            sse_event(
                {
                    "id": self._resp_id,
                    "object": "chat.completion.chunk",
                    "created": ctx.created,
                    "model": self._model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"role": "assistant"},
                            "finish_reason": None,
                        }
                    ],
                }
            )
        ]
    def format_chunk(self, token: str, **kwargs) -> List[str]:
        body = kwargs.get("body", "")
        if self._parser is not None:
            return self._format_tool_chunk(body, **kwargs)
        return [
            sse_event(
                {
                    "id": self._resp_id,
                    "object": "chat.completion.chunk",
                    "created": 0,
                    "model": self._model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"content": token},
                            "finish_reason": None,
                        }
                    ],
                }
            )
        ]
    def _format_tool_chunk(self, body: str, **kwargs) -> List[str]:
        deltas = self._parser.feed(
            body,
            current_token_ids=kwargs.get("current_token_ids"),
            delta_token_ids=kwargs.get("delta_token_ids"),
        )
        events: List[str] = []
        for d in deltas:
            if "content" in d:
                if not self._content_started:
                    events.append(self._role_chunk())
                    self._content_started = True
                events.append(
                    sse_event(
                        {
                            "id": self._resp_id,
                            "object": "chat.completion.chunk",
                            "created": 0,
                            "model": self._model,
                            "choices": [
                                {
                                    "index": 0,
                                    "delta": {"content": d["content"]},
                                    "finish_reason": None,
                                }
                            ],
                        }
                    )
                )
            elif "tool_calls" in d:
                if not self._content_started:
                    events.append(self._role_chunk())
                    self._content_started = True
                events.append(
                    sse_event(
                        {
                            "id": self._resp_id,
                            "object": "chat.completion.chunk",
                            "created": 0,
                            "model": self._model,
                            "choices": [
                                {
                                    "index": 0,
                                    "delta": {"tool_calls": d["tool_calls"]},
                                    "finish_reason": None,
                                }
                            ],
                        }
                    )
                )
        return events
    def _role_chunk(self) -> str:
        return sse_event(
            {
                "id": self._resp_id,
                "object": "chat.completion.chunk",
                "created": 0,
                "model": self._model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {"role": "assistant"},
                        "finish_reason": None,
                    }
                ],
            }
        )
    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
        finish_reason = "stop"
        if self._parser is not None and self._parser.has_tool_calls:
            finish_reason = "tool_calls"
        return [
            sse_event(
                {
                    "id": self._resp_id,
                    "object": "chat.completion.chunk",
                    "created": ctx.created,
                    "model": self._model,
                    "choices": [
                        {"index": 0, "delta": {}, "finish_reason": finish_reason}
                    ],
                }
            ),
            sse_event(
                {
                    "prompt_tokens": ctx.prompt_tokens,
                    "completion_tokens": ctx.completion_tokens,
                    "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
                }
            ),
        ]
    def format_response(
        self, ctx: GenContext, content: str, stop: StopInfo
    ) -> Dict[str, Any]:
        if self._parser is not None:
            parsed = self._parser.parse_complete(content)
            if parsed and parsed.get("tool_calls"):
                return {
                    "id": self._resp_id,
                    "object": "chat.completion",
                    "created": ctx.created,
                    "model": self._model,
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": parsed.get("content"),
                                "tool_calls": parsed["tool_calls"],
                            },
                            "finish_reason": "tool_calls",
                        }
                    ],
                    "usage": {
                        "prompt_tokens": ctx.prompt_tokens,
                        "completion_tokens": ctx.completion_tokens,
                        "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
                    },
                }
        return {
            "id": self._resp_id,
            "object": "chat.completion",
            "created": ctx.created,
            "model": self._model,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": content},
                    "finish_reason": "stop",
                }
            ],
            "usage": {
                "prompt_tokens": ctx.prompt_tokens,
                "completion_tokens": ctx.completion_tokens,
                "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
            },
        }
@@ -1,15 +1,13 @@
-"""Protocol handlers for OpenAI and Anthropic chat completion APIs.
+"""Orchestration layer: ProtocolHandler, StopChecker, GenContext, StopInfo, ResponseBuilder, SSE utils.
-Template Method + Builder patterns eliminate the 45% code duplication between
+ProtocolHandler orchestrates the async generation loop and delegates
-stream/non-stream branches and across protocol adapters.
+protocol-specific formatting to a ResponseBuilder.
 """
 import json
 import time
 import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
@@ -17,7 +15,7 @@ from pydantic import BaseModel
 from astrai.inference.engine import InferenceEngine
-def _sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
+def sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
    lines: List[str] = []
    if event:
        lines.append(f"event: {event}")
@@ -26,22 +24,28 @@ def _sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
    return "\n".join(lines)
-def _sse_done() -> str:
+def sse_done() -> str:
    return "data: [DONE]\n\n"
@dataclass
-class StreamContext:
+class GenContext:
-    """Shared state across the streaming generation lifecycle."""
+    """Per-generation metadata passed to builder format methods."""
    resp_id: str
    created: int
    model: str
-    prompt_tokens: int
+    prompt_tokens: int = 0
    completion_tokens: int = 0
-    accumulated: str = ""
+
-    stop_matched: Optional[str] = None
+
-    last_yield_trimmed: str = ""
+@dataclass
 class StopInfo:
    """Stop-check result passed to format_stream_end / format_response."""
    matched: Optional[str] = None
    body: str = ""
    yielded: str = ""
 class StopChecker:
@@ -56,95 +60,67 @@ class StopChecker:
                return seq
        return None
    def trim(self, text: str, matched: str) -> str:
        idx = text.rfind(matched)
        return text[:idx] if idx != -1 else text
-    @property
+class ResponseBuilder(ABC):
-    def has_sequences(self) -> bool:
+    """Interface for protocol-specific response formatting.
        return len(self._sequences) > 0
-
+    A new protocol requires one concrete builder implementing 5 methods.
 class ProtocolHandler(ABC):
    """Template-method base for API protocol handlers.
    Subclasses implement format hooks; the base class orchestrates the
    generate-async loop and SSE/JSON response construction.
    Lifecycle::
        handle()
          ├─ build_prompt()          # protocol-specific prompt assembly
          ├─ create_response_id()    # unique response identifier
          ├─ [stream]
          │   ├─ format_stream_start()
          │   ├─ format_stream_token()  × N
          │   │   └─ on_token() hook for stop-sequence interception
          │   └─ format_stream_end()
          └─ [non-stream]
              ├─ (accumulate tokens)
              └─ format_non_stream_response()
    """
-    request_model: type[BaseModel]
+    @abstractmethod
    def prepare(
        self, request: BaseModel, engine: InferenceEngine
    ) -> Tuple[str, GenContext, List[str]]:
        """Return (prompt, ctx, stop_sequences) for a generation request."""
-    def __init__(self, request: BaseModel, engine: InferenceEngine):
+    @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
    @abstractmethod
    def build_prompt(self) -> str:
        """Build the full prompt string from the request messages."""
    @abstractmethod
    def create_response_id(self) -> str:
        """Generate a unique response ID following the protocol convention."""
    @abstractmethod
    def format_stream_start(self, ctx: StreamContext) -> List[str]:
        """Yield SSE events that open the stream (role marker, metadata)."""
    @abstractmethod
    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
        """Yield an SSE event for a single generated token."""
    @abstractmethod
    def format_stream_end(self, ctx: StreamContext) -> List[str]:
        """Yield SSE events that close the stream (finish reason, usage stats)."""
    @abstractmethod
    def format_non_stream_response(
        self, ctx: StreamContext, content: str
    ) -> Dict[str, Any]:
        """Build the JSON response body for non-streaming mode."""
    def get_stop_sequences(self) -> List[str]:
        return []
    def create_stop_checker(self) -> StopChecker:
        return StopChecker(self.get_stop_sequences())
    def on_token(
        self, ctx: StreamContext, token: str, stop_checker: StopChecker
    ) -> Optional[str]:
        """Hook after each token is appended to accumulated.
        Return a matched stop-sequence string to break the loop,
        or None to continue.
        """
        return None
    async def handle(self) -> Union[StreamingResponse, Dict[str, Any]]:
-        ctx = StreamContext(
+        prompt, ctx, stop_sequences = self.builder.prepare(self.request, self.engine)
-            resp_id=self.create_response_id(),
+        ctx.prompt_tokens = len(self.engine.tokenizer.encode(prompt))
            created=int(time.time()),
            model=self.request.model,
            prompt_tokens=self._count_prompt_tokens(),
        )
        agen = self.engine.generate_async(
-            prompt=self.build_prompt(),
+            prompt=prompt,
            max_tokens=self.request.max_tokens,
            temperature=self.request.temperature,
            top_p=self.request.top_p,
@@ -152,33 +128,47 @@ class ProtocolHandler(ABC):
        )
        if self.request.stream:
-            return self._handle_stream(agen, ctx)
+            return self._handle_stream(agen, ctx, stop_sequences)
        else:
-            return await self._handle_non_stream(agen, ctx)
+            return await self._handle_non_stream(agen, ctx, stop_sequences)
-    def _count_prompt_tokens(self) -> int:
+    def _handle_stream(
-        return len(self.engine.tokenizer.encode(self.build_prompt()))
+        self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
-
+    ) -> StreamingResponse:
-    def _handle_stream(self, agen, ctx: StreamContext) -> StreamingResponse:
+        checker = StopChecker(stop_sequences)
        stop_checker = self.create_stop_checker()
        async def event_stream():
-            for event in self.format_stream_start(ctx):
+            for event in self.builder.format_stream_start(ctx):
                yield event
            body = ""
            yielded = ""
            matched = None
            token_ids: List[int] = []
            async for token in agen:
-                ctx.completion_tokens += 1
+                body += token
                ctx.accumulated += token
-                matched = self.on_token(ctx, token, stop_checker)
+                new_ids = self.engine.tokenizer.encode(token)
                token_ids.extend(new_ids)
                matched = checker.check(body)
                if matched:
                    break
-                yield self.format_stream_token(ctx, token)
+                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
-            for event in self.format_stream_end(ctx):
+            stop = StopInfo(matched=matched, body=body, yielded=yielded)
            for event in self.builder.format_stream_end(ctx, stop):
                yield event
-            yield _sse_done()
+            yield sse_done()
        return StreamingResponse(
            event_stream(),
@@ -186,249 +176,24 @@ class ProtocolHandler(ABC):
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )
-    async def _handle_non_stream(self, agen, ctx: StreamContext) -> Dict[str, Any]:
+    async def _handle_non_stream(
-        stop_checker = self.create_stop_checker()
+        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:
            ctx.completion_tokens += 1
            ctx.accumulated += token
            chunks.append(token)
            body += token
-            matched = self.on_token(ctx, token, stop_checker)
+            matched = checker.check(body)
            if matched:
                break
            ctx.completion_tokens += 1
        content = "".join(chunks)
-        return self.format_non_stream_response(ctx, content)
+        stop = StopInfo(matched=matched, body=body)
-
+        return self.builder.format_response(ctx, content, stop)
 def _extract_text_content(content: Union[str, List[Dict[str, Any]]]) -> str:
    """Extract plain text from an Anthropic content block (string or list)."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                return block.get("text", "")
    return ""
 class OpenAIHandler(ProtocolHandler):
    """OpenAI-compatible /v1/chat/completions handler."""
    def build_prompt(self) -> str:
        messages = [
            {"role": m.role, "content": m.content} for m in self.request.messages
        ]
        return self.engine.tokenizer.apply_chat_template(messages, tokenize=False)
    def create_response_id(self) -> str:
        return f"chatcmpl-{uuid.uuid4().hex[:12]}"
    def format_stream_start(self, ctx: StreamContext) -> List[str]:
        return [
            _sse_event(
                {
                    "id": ctx.resp_id,
                    "object": "chat.completion.chunk",
                    "created": ctx.created,
                    "model": ctx.model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"role": "assistant"},
                            "finish_reason": None,
                        }
                    ],
                }
            )
        ]
    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
        return _sse_event(
            {
                "id": ctx.resp_id,
                "object": "chat.completion.chunk",
                "created": ctx.created,
                "model": ctx.model,
                "choices": [
                    {"index": 0, "delta": {"content": token}, "finish_reason": None}
                ],
            }
        )
    def format_stream_end(self, ctx: StreamContext) -> List[str]:
        return [
            _sse_event(
                {
                    "id": ctx.resp_id,
                    "object": "chat.completion.chunk",
                    "created": ctx.created,
                    "model": ctx.model,
                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
                }
            ),
            _sse_event(
                {
                    "prompt_tokens": ctx.prompt_tokens,
                    "completion_tokens": ctx.completion_tokens,
                    "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
                }
            ),
        ]
    def format_non_stream_response(
        self, ctx: StreamContext, content: str
    ) -> Dict[str, Any]:
        return {
            "id": ctx.resp_id,
            "object": "chat.completion",
            "created": ctx.created,
            "model": ctx.model,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": content},
                    "finish_reason": "stop",
                }
            ],
            "usage": {
                "prompt_tokens": ctx.prompt_tokens,
                "completion_tokens": ctx.completion_tokens,
                "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
            },
        }
 class AnthropicHandler(ProtocolHandler):
    """Anthropic-compatible /v1/messages handler."""
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._yielded = ""
    def build_prompt(self) -> str:
        messages: List[Dict[str, str]] = []
        system = getattr(self.request, "system", None)
        if system:
            messages.append({"role": "system", "content": system})
        for m in self.request.messages:
            content = _extract_text_content(m.content)
            if content:
                messages.append({"role": m.role, "content": content})
        return self.engine.tokenizer.apply_chat_template(messages, tokenize=False)
    def create_response_id(self) -> str:
        return f"msg_{uuid.uuid4().hex[:24]}"
    def get_stop_sequences(self) -> List[str]:
        return getattr(self.request, "stop_sequences", None) or []
    def on_token(
        self, ctx: StreamContext, token: str, stop_checker: StopChecker
    ) -> Optional[str]:
        matched = stop_checker.check(ctx.accumulated)
        if not matched:
            return None
        ctx.stop_matched = matched
        trimmed = ctx.accumulated[: ctx.accumulated.rfind(matched)]
        unyielded = trimmed[len(self._yielded) :]
        if unyielded:
            ctx.last_yield_trimmed = unyielded
        return matched
    def format_stream_start(self, ctx: StreamContext) -> List[str]:
        return [
            _sse_event(
                {
                    "type": "message_start",
                    "message": {
                        "id": ctx.resp_id,
                        "type": "message",
                        "role": "assistant",
                        "model": ctx.model,
                        "content": [],
                        "usage": {"input_tokens": ctx.prompt_tokens},
                    },
                },
                event="message_start",
            ),
            _sse_event(
                {
                    "type": "content_block_start",
                    "index": 0,
                    "content_block": {"type": "text", "text": ""},
                },
                event="content_block_start",
            ),
        ]
    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
        self._yielded += token
        return _sse_event(
            {
                "type": "content_block_delta",
                "index": 0,
                "delta": {"type": "text_delta", "text": token},
            },
            event="content_block_delta",
        )
    def format_stream_end(self, ctx: StreamContext) -> List[str]:
        matched = ctx.stop_matched
        events: List[str] = []
        last_yielded = ctx.last_yield_trimmed
        if last_yielded:
            events.append(
                _sse_event(
                    {
                        "type": "content_block_delta",
                        "index": 0,
                        "delta": {"type": "text_delta", "text": last_yielded},
                    },
                    event="content_block_delta",
                )
            )
        events.append(
            _sse_event(
                {"type": "content_block_stop", "index": 0},
                event="content_block_stop",
            )
        )
        events.append(
            _sse_event(
                {
                    "type": "message_delta",
                    "delta": {
                        "stop_reason": "stop_sequence" if matched else "end_turn",
                        "stop_sequence": matched,
                    },
                    "usage": {"output_tokens": ctx.completion_tokens},
                },
                event="message_delta",
            )
        )
        events.append(_sse_event({"type": "message_stop"}, event="message_stop"))
        return events
    def format_non_stream_response(
        self, ctx: StreamContext, content: str
    ) -> Dict[str, Any]:
        matched = ctx.stop_matched
        if matched:
            content = content[: content.rfind(matched)]
        return {
            "id": ctx.resp_id,
            "type": "message",
            "role": "assistant",
            "model": ctx.model,
            "content": [{"type": "text", "text": content}],
            "stop_reason": "stop_sequence" if matched else "end_turn",
            "stop_sequence": matched,
            "usage": {
                "input_tokens": ctx.prompt_tokens,
                "output_tokens": ctx.completion_tokens,
            },
        }
@@ -3,6 +3,9 @@ OpenAI / Anthropic-compatible chat completion server backed by continuous-batchi
 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
@@ -12,22 +15,37 @@ from typing import Any, Dict, List, Optional, Union
 import torch
 import uvicorn
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import APIRouter, FastAPI, HTTPException
 from pydantic import BaseModel, Field
-from astrai.inference.api.protocol import AnthropicHandler, OpenAIHandler
+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__)
-_project_root = Path(__file__).parent.parent.parent
+_app_instance: Optional[FastAPI] = None
 class ChatMessage(BaseModel):
    role: str
-    content: 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):
@@ -46,6 +64,8 @@ class ChatCompletionRequest(BaseModel):
    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):
@@ -67,14 +87,30 @@ class MessagesRequest(BaseModel):
    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: Optional[Path] = None,
+    param_path: Path,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    max_batch_size: int = 16,
 ) -> InferenceEngine:
    if param_path is None:
        param_path = _project_root / "params"
    if not param_path.exists():
        raise FileNotFoundError(f"Parameter directory not found: {param_path}")
@@ -92,67 +128,66 @@ def _create_engine(
    return engine
-@asynccontextmanager
+def get_app() -> FastAPI:
-async def lifespan(app: FastAPI):
+    """Return the singleton FastAPI instance (lazily created on first call)."""
-    config = app.state.server_config
+    global _app_instance
-    if not config.get("_test", False):
+    if _app_instance is None:
-        try:
+        _app_instance = FastAPI(
-            app.state.engine = _create_engine(**config)
+            title="AstrAI Inference Server",
-        except Exception as e:
+            version="0.2.0",
-            logger.error(f"Failed to load model: {e}")
+            lifespan=lifespan,
-            raise
+        )
-    yield
+        _app_instance.include_router(router)
-    if app.state.engine:
+        _app_instance.state.server_config = {}
-        app.state.engine.shutdown()
+        _app_instance.state.engine = None
-        logger.info("Inference engine shutdown complete")
+    return _app_instance
-app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
+def _get_engine() -> InferenceEngine:
-
+    engine = get_app().state.engine
 def _get_engine(request: Request) -> InferenceEngine:
    engine = request.app.state.engine
    if engine is None:
        raise HTTPException(status_code=503, detail="Engine not initialized")
    return engine
-@app.get("/health")
+@router.get("/health")
-async def health(request: Request):
+async def health():
    app = get_app()
    return {
        "status": "ok",
-        "model_loaded": request.app.state.engine is not None,
+        "model_loaded": app.state.engine is not None,
    }
-@app.get("/stats")
+@router.get("/stats")
-async def get_stats(request: Request):
+async def get_stats():
-    return _get_engine(request).get_stats()
+    return _get_engine().get_stats()
-@app.post("/v1/chat/completions")
+@router.post("/v1/chat/completions")
-async def chat_completion(request: ChatCompletionRequest, req: Request):
+async def chat_completion(request: ChatCompletionRequest):
-    engine = _get_engine(req)
+    engine = _get_engine()
-    handler = OpenAIHandler(request, engine)
+    handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
    return await handler.handle()
-@app.post("/v1/messages")
+@router.post("/v1/messages")
-async def create_message(request: MessagesRequest, req: Request):
+async def create_message(request: MessagesRequest):
-    engine = _get_engine(req)
+    engine = _get_engine()
-    handler = AnthropicHandler(request, 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,
    param_path: Optional[Path] = None,
    max_batch_size: int = 16,
 ):
    app = get_app()
    app.state.server_config = {
        "device": device,
        "dtype": dtype,
@@ -163,4 +198,5 @@ def run_server(
        app,
        host=host,
        port=port,
        reload=reload,
    )
@@ -0,0 +1,325 @@
 """Tool call parsers for extracting structured tool calls from model output.
 Patterned after vLLM's ToolParser abstraction. Each parser knows how to
 detect and incrementally extract tool calls from raw generated text.
 Subclasses may optionally consume ``token_ids`` for token-level parsing
 (e.g. Harmony / VLM-style parsers).
 """
 import re
 import uuid
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional
 from astrai.factory import BaseFactory
 class BaseToolParser(ABC):
    """Abstract tool call parser — one instance per request.
    Maintains streaming state internally so that each call to :meth:`feed`
    can diff against previously emitted content.
    Parameters
    ----------
    tools : list of dict, optional
        Tool definitions from the request.
    tool_choice : str
        ``"auto"`` / ``"required"`` / ``"none"`` or a named tool choice
        dict.
    """
    def __init__(self, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
        self.tools = tools or []
        self.tool_choice = tool_choice
    @abstractmethod
    def feed(
        self,
        body: str,
        current_token_ids: Optional[List[int]] = None,
        delta_token_ids: Optional[List[int]] = None,
    ) -> List[Dict]:
        """Feed the *full* accumulated text each step.
        Returns a list of delta dicts to emit. Each delta is one of:
        - ``{"content": "text"}``      — plain text delta
        - ``{"tool_calls": [...]}``    — tool-call delta (OpenAI format)
        Returns an empty list when nothing new should be emitted.
        Parameters
        ----------
        body : str
            The complete accumulated generated text so far.
        current_token_ids : list of int, optional
            All token IDs decoded into *body* (cumulative).
        delta_token_ids : list of int, optional
            Only the token IDs for this chunk.
        """
    @abstractmethod
    def parse_complete(self, body: str) -> Optional[Dict]:
        """Parse the *complete* generated text after generation ends.
        Returns ``None`` when no tool calls were found, otherwise a dict
        with ``content`` (str or None) and ``tool_calls`` (list of dicts).
        """
    @property
    @abstractmethod
    def has_tool_calls(self) -> bool:
        """True if the parser detected at least one tool call in the stream."""
 class ToolParserFactory(BaseFactory["BaseToolParser"]):
    pass
 _TOOL_CALL_HEAD_RE = re.compile(r'\{\s*"name"\s*:')
 def _scan_json(text: str, start: int = 0):
    """Scan for a complete JSON object starting at *start*.
    Returns ``(end, complete)`` where *end* is one-past the closing
    brace (or ``len(text)`` if unclosed), and *complete* is a bool.
    """
    depth = 0
    in_string = False
    escape = False
    for i in range(start, len(text)):
        c = text[i]
        if escape:
            escape = False
            continue
        if c == "\\":
            escape = True
            continue
        if c == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if c == "{":
            depth += 1
        elif c == "}":
            depth -= 1
            if depth == 0:
                return i + 1, True
    return len(text), False
 def _parse_tool_call_json(json_str: str, complete: bool):
    """Extract *name* and *arguments* from a tool-call JSON string.
    Returns ``(name, args, valid)``.
    """
    name_match = re.search(r'"name"\s*:\s*"([^"]*)"', json_str)
    if not name_match:
        return None, "", False
    name = name_match.group(1)
    args_match = re.search(r'"arguments"\s*:\s*(.*)', json_str, re.DOTALL)
    if not args_match:
        return name, "", True
    raw = args_match.group(1).rstrip()
    if complete and raw.endswith("}"):
        raw = raw[:-1].rstrip()
    if raw.startswith("{"):
        inner = raw[1:].rstrip()
        if inner.endswith("}"):
            inner = inner[:-1].rstrip()
        raw = inner
    return name, raw, True
 def _find_tool_calls(text: str, start_pos: int = 0):
    """Find all complete ``{...}`` tool-call objects in *text*.
    Returns a list of dicts with keys *start*, *end*, *name*, *args*,
    *complete*.
    """
    results = []
    pos = start_pos
    while True:
        brace = text.find("{", pos)
        if brace == -1:
            break
        end, complete = _scan_json(text, brace)
        if not complete:
            break
        json_str = text[brace:end]
        if not _TOOL_CALL_HEAD_RE.search(json_str):
            pos = end
            continue
        name, args, valid = _parse_tool_call_json(json_str, complete=True)
        if not valid or name is None:
            pos = end
            continue
        results.append(
            {
                "start": brace,
                "end": end,
                "name": name,
                "args": args,
                "complete": True,
            }
        )
        pos = end
    return results
 def _find_partial_tool_call(text: str, start_pos: int = 0):
    """Find one incomplete (still-generating) tool-call JSON object."""
    brace = text.find("{", start_pos)
    if brace == -1:
        return None
    json_str = text[brace:]
    if not _TOOL_CALL_HEAD_RE.search(json_str):
        return None
    name, args, valid = _parse_tool_call_json(json_str, complete=False)
    if not valid or name is None:
        return None
    return {
        "start": brace,
        "name": name,
        "args": args,
        "complete": False,
    }
@ToolParserFactory.register("simple_json")
 class SimpleJsonToolParser(BaseToolParser):
    """Parser for models that output tool calls as plain JSON objects.
    Detects ``{"name": "<func>", "arguments": {...}}`` anywhere in the
    generated text.  Handles single and (non-overlapping) multiple tool
    calls.  Text preceding the first tool call is emitted as plain
    ``content`` deltas.
    """
    def __init__(self, tools=None, tool_choice="auto"):
        super().__init__(tools, tool_choice)
        self._emitted_content_len = 0
        self._tc_state: List[Dict] = []
        self._has_tool_calls = False
    # -------------------------------------------------------------- feed
    def feed(
        self,
        body: str,
        current_token_ids: Optional[List[int]] = None,
        delta_token_ids: Optional[List[int]] = None,
    ) -> List[Dict]:
        deltas: List[Dict] = []
        completed = _find_tool_calls(body)
        if not completed:
            partial = _find_partial_tool_call(body)
            if not partial:
                return self._emit_plain_content(body, deltas)
            all_tcs = [partial]
        else:
            all_tcs = completed
            partial = _find_partial_tool_call(body, completed[-1]["end"])
            if partial:
                all_tcs = completed + [partial]
        first_start = all_tcs[0]["start"]
        if first_start > self._emitted_content_len:
            content = body[self._emitted_content_len : first_start]
            self._emitted_content_len = first_start
            if content:
                deltas.append({"content": content})
        for i, tc in enumerate(all_tcs):
            if i >= len(self._tc_state):
                self._tc_state.append(
                    {
                        "id": f"call_{uuid.uuid4().hex[:12]}",
                        "name_emitted": False,
                        "args_emitted_len": 0,
                    }
                )
                self._has_tool_calls = True
            st = self._tc_state[i]
            if not st["name_emitted"]:
                st["name_emitted"] = True
                deltas.append(
                    {
                        "tool_calls": [
                            {
                                "index": i,
                                "id": st["id"],
                                "type": "function",
                                "function": {"name": tc["name"], "arguments": ""},
                            }
                        ]
                    }
                )
            new_args = tc["args"]
            if len(new_args) > st["args_emitted_len"]:
                diff = new_args[st["args_emitted_len"] :]
                st["args_emitted_len"] = len(new_args)
                deltas.append(
                    {
                        "tool_calls": [
                            {
                                "index": i,
                                "function": {"arguments": diff},
                            }
                        ]
                    }
                )
        return deltas
    def _emit_plain_content(self, body: str, deltas: List[Dict]) -> List[Dict]:
        new_content = body[self._emitted_content_len :]
        if new_content:
            self._emitted_content_len = len(body)
            deltas.append({"content": new_content})
        return deltas
    # -------------------------------------------------------- complete
    def parse_complete(self, body: str) -> Optional[Dict]:
        completed = _find_tool_calls(body)
        if not completed:
            return None
        content = body[: completed[0]["start"]].strip() or None
        tool_calls = []
        for i, tc in enumerate(completed):
            tool_calls.append(
                {
                    "id": f"call_{uuid.uuid4().hex[:12]}",
                    "type": "function",
                    "function": {
                        "name": tc["name"],
                        "arguments": tc["args"],
                    },
                }
            )
        return {"content": content, "tool_calls": tool_calls}
    @property
    def has_tool_calls(self) -> bool:
        return self._has_tool_calls
@@ -2,8 +2,12 @@
 from astrai.inference.core.cache import (
    Allocator,
    CacheView,
    ContiguousCache,
    ContiguousCacheView,
    KVCache,
-    KvcacheView,
+    PageCache,
    PageCacheView,
    PagePool,
    PrefixCache,
    Storage,
@@ -16,8 +20,12 @@ from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 __all__ = [
    "Allocator",
    "CacheView",
    "KVCache",
-    "KvcacheView",
+    "ContiguousCache",
    "ContiguousCacheView",
    "PageCache",
    "PageCacheView",
    "PagePool",
    "PrefixCache",
    "Storage",
@@ -1,4 +1,5 @@
 import threading
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from typing import Callable, Dict, List, Optional, Tuple
@@ -42,7 +43,7 @@ class Allocator:
                return idx
            return -1
-    def free(self, idx: int, keep_cached: bool = False) -> None:
+    def free(self, idx: int, keep_cached: bool = False):
        with self._lock:
            self._refs[idx] -= 1
            if self._refs[idx] == 0:
@@ -51,7 +52,7 @@ class Allocator:
                else:
                    self._free_mask |= 1 << idx
-    def inc_ref(self, idx: int) -> None:
+    def inc_ref(self, idx: int):
        with self._lock:
            self._refs[idx] += 1
            self._lru.pop(idx, None)
@@ -60,9 +61,10 @@ class Allocator:
        with self._lock:
            return self._refs[idx]
-    def touch(self, idx: int) -> None:
+    def touch(self, idx: int):
        with self._lock:
-            self._lru.move_to_end(idx)
+            if idx in self._lru:
                self._lru.move_to_end(idx)
 class PrefixCache:
@@ -74,7 +76,7 @@ class PrefixCache:
        self._hash_to_page: Dict[int, int] = {}
        self._lock = threading.Lock()
-    def evict(self, idx: int) -> None:
+    def evict(self, idx: int):
        with self._lock:
            h = self._page_to_hash.pop(idx, None)
            if h is not None:
@@ -96,9 +98,7 @@ class PrefixCache:
                hits.append(p)
            return hits
-    def record(
+    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
        self, page_idx: int, token_ids: List[int], logical_page_idx: int
    ) -> None:
        with self._lock:
            h = page_hash(token_ids, logical_page_idx, self._page_size)
            old_h = self._page_to_hash.pop(page_idx, None)
@@ -127,13 +127,13 @@ class PagePool:
    def alloc(self) -> int:
        return self._alloc.alloc()
-    def free(self, idx: int) -> None:
+    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) -> None:
+    def inc_ref(self, idx: int):
        self._alloc.inc_ref(idx)
    def lookup(self, token_ids: List[int]) -> List[int]:
@@ -142,9 +142,7 @@ class PagePool:
            self._alloc.touch(p)
        return hits
-    def record(
+    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
        self, page_idx: int, token_ids: List[int], logical_page_idx: int
    ) -> None:
        self._prefix.record(page_idx, token_ids, logical_page_idx)
@@ -157,7 +155,7 @@ class TaskTable:
        self._cached: Dict[str, int] = {}
        self._lock = threading.Lock()
-    def set(self, task_id: str, page_table: List[int], cached: int) -> None:
+    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
@@ -220,7 +218,7 @@ class Storage:
        start_pos: int,
        k: Tensor,
        v: Tensor,
-    ) -> None:
+    ):
        seq_len = k.size(1)
        if seq_len == 0:
            return
@@ -278,7 +276,42 @@ class Storage:
        return k, v
-class KvcacheView:
+class CacheView(ABC):
    """Abstract view passed to attention layers for KV-cache I/O."""
    @abstractmethod
    def write(self, layer_id: int, k: Tensor, v: Tensor): ...
    @abstractmethod
    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: ...
 class KVCache(ABC):
    """Abstract KV-cache facade for scheduler/executor."""
    @abstractmethod
    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: ...
    @abstractmethod
    def task_free(self, task_id: str): ...
    @abstractmethod
    def task_extend(self, task_id: str, pos: int) -> bool: ...
    @abstractmethod
    def bind_tasks(
        self, task_ids: List[str], total_len: int, device: torch.device
    ) -> CacheView: ...
    def task_cached(self, task_id: str) -> int:
        return 0
    def task_record_hashes(
        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
    ): ...
 class PageCacheView(CacheView):
    """Bundles Storage + page_table + total_len for attention layers."""
    def __init__(self, storage: Storage, page_table: Tensor, total_len: int = 0):
@@ -286,7 +319,7 @@ class KvcacheView:
        self._page_table = page_table
        self._total_len = total_len
-    def write(self, layer_id: int, k: Tensor, v: Tensor) -> None:
+    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)
@@ -294,8 +327,8 @@ class KvcacheView:
        return self._storage.gather(layer_id, self._page_table, self._total_len)
-class KVCache:
+class PageCache(KVCache):
-    """Facade: page management + KV-cache I/O for continuous batching."""
+    """Paged KV-cache with prefix sharing."""
    def __init__(
        self,
@@ -339,7 +372,7 @@ class KVCache:
        self._table.set(task_id, hits + new_pages, cached)
        return True
-    def task_free(self, task_id: str) -> None:
+    def task_free(self, task_id: str):
        page_table, _ = self._table.pop(task_id)
        for idx in page_table:
            self._pool.free(idx)
@@ -359,14 +392,108 @@ class KVCache:
    def task_record_hashes(
        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
-    ) -> None:
+    ):
        page_table = self._table.get(task_id)
        full_pages = len(prompt_ids) // self.page_size
        for i in range(start_logical_page, full_pages):
            self._pool.record(page_table[i], prompt_ids, i)
-    def make_table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor:
+    def bind_tasks(
-        return self._table.table_tensor(task_ids, device)
+        self, task_ids: List[str], total_len: int, device: torch.device
    ) -> PageCacheView:
        page_table = self._table.table_tensor(task_ids, device)
        return PageCacheView(self._storage, page_table, total_len)
-    def bind(self, page_table: Tensor, total_len: int = 0) -> KvcacheView:
+
-        return KvcacheView(self._storage, page_table, total_len)
+class ContiguousCacheView(CacheView):
    """Contiguous KV-cache view for attention layers."""
    def __init__(
        self, cache: "ContiguousCache", batch_indices: Tensor, total_len: int = 0
    ):
        self._cache = cache
        self._batch_indices = batch_indices
        self._total_len = total_len
    def write(self, layer_id: int, k: Tensor, v: Tensor):
        seq_len = k.size(1)
        start_pos = self._total_len - seq_len
        indices = self._batch_indices
        self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k
        self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v
        new_len = start_pos + seq_len
        for s in indices.tolist():
            cur = self._cache._slot_len.get(s, 0)
            if new_len > cur:
                self._cache._slot_len[s] = new_len
    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
        max_len = max(
            self._cache._slot_len.get(int(s), 0) for s in self._batch_indices.tolist()
        )
        indices = self._batch_indices
        k = self._cache.k[layer_id, indices, :max_len]
        v = self._cache.v[layer_id, indices, :max_len]
        return k, v
 class ContiguousCache(KVCache):
    """Contiguous per-slot KV cache (default implementation)."""
    def __init__(
        self,
        n_layers: int,
        max_batch_size: int,
        max_seq_len: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.max_seq_len = max_seq_len
        self.k = torch.zeros(
            n_layers,
            max_batch_size,
            max_seq_len,
            n_kv_heads,
            head_dim,
            device=device,
            dtype=dtype,
        )
        self.v = torch.zeros(
            n_layers,
            max_batch_size,
            max_seq_len,
            n_kv_heads,
            head_dim,
            device=device,
            dtype=dtype,
        )
        self._slot_len: Dict[int, int] = {}
        self._task_slot: Dict[str, int] = {}
        self._free_slots = list(range(max_batch_size))
        self._device = device
    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
        if not self._free_slots:
            return False
        slot = self._free_slots.pop(0)
        self._task_slot[task_id] = slot
        self._slot_len[slot] = 0
        return True
    def task_free(self, task_id: str):
        slot = self._task_slot.pop(task_id, None)
        if slot is not None:
            self._slot_len.pop(slot, None)
            self._free_slots.append(slot)
    def task_extend(self, task_id: str, pos: int) -> bool:
        return pos < self.max_seq_len
    def bind_tasks(
        self, task_ids: List[str], total_len: int, device: torch.device
    ) -> ContiguousCacheView:
        slots = [self._task_slot[tid] for tid in task_ids]
        batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
        return ContiguousCacheView(self, batch_indices, total_len)
@@ -19,19 +19,17 @@ class Executor:
        self,
        model: AutoModel,
        tokenizer: AutoTokenizer,
-        page_cache: KVCache,
+        kv_cache: KVCache,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
-        self.page_cache = page_cache
+        self.kv_cache = kv_cache
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype
-    def execute_prefill(
+    def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
    ) -> None:
        if start_pos >= prompt_len:
            return
@@ -45,7 +43,6 @@ class Executor:
        )
        task_ids = [t.task_id for t in tasks]
        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
        with torch.inference_mode():
            self.model(
@@ -55,7 +52,7 @@ class Executor:
                )
                .unsqueeze(0)
                .expand(batch_sz, -1),
-                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
+                paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
            )
    def execute_decode(self, tasks: List[Task]) -> List[int]:
@@ -74,7 +71,6 @@ class Executor:
        total_len = position_ids.max().item() + 1
        task_ids = [t.task_id for t in tasks]
        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
@@ -83,7 +79,7 @@ class Executor:
        with torch.inference_mode():
            outputs = self.model(
                input_ids.unsqueeze(1),
-                paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
+                paged_cache=self.kv_cache.bind_tasks(task_ids, total_len, self.device),
                position_ids=position_ids.unsqueeze(1),
            )
            logits = outputs["logits"][:, -1, :]
@@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple
 import torch
-from astrai.inference.core.cache import KVCache
+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
@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 class InferenceScheduler:
-    """Four-phase continuous batching loop: cleanup -> refill -> prefill -> decode."""
+    """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
    def __init__(
        self,
@@ -22,30 +22,39 @@ class InferenceScheduler:
        tokenizer: AutoTokenizer,
        max_batch_size: int = 16,
        max_seq_len: Optional[int] = None,
-        max_prompt_len: int = 512,
+        max_prompt_len: int = 2048,
        page_size: int = 64,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
        cache: Optional[KVCache] = None,
    ):
        config = model.config
-        self.max_seq_len = max_seq_len or config.max_len
+        if max_seq_len is not None:
            self.max_seq_len = max_seq_len
        elif config.max_len is not None:
            self.max_seq_len = config.max_len
        else:
            raise ValueError(
                "max_seq_len must be provided either as argument "
                "or in model config (config.max_len)"
            )
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype
-        n_pages = (
+        head_dim = config.dim // config.n_heads
            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
        ) // page_size
-        self._page_cache = KVCache(
+        if cache is not None:
-            config.n_layers,
+            self._cache = cache
-            n_pages,
+        else:
-            page_size,
+            self._cache = ContiguousCache(
-            config.n_kv_heads,
+                config.n_layers,
-            config.dim // config.n_heads,
+                max_batch_size,
-            self.device,
+                self.max_seq_len,
-            self.dtype,
+                config.n_kv_heads,
-        )
+                head_dim,
                self.device,
                self.dtype,
            )
        self._task_mgr = TaskManager(
            tokenizer=tokenizer,
@@ -57,30 +66,32 @@ class InferenceScheduler:
        self._executor = Executor(
            model=model,
            tokenizer=tokenizer,
-            page_cache=self._page_cache,
+            kv_cache=self._cache,
            device=self.device,
            dtype=self.dtype,
        )
-        self._running = False
+        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) -> None:
+    def remove_task(self, task_id: str):
        for task in self._task_mgr.remove_task(task_id):
-            self._page_cache.task_free(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) -> None:
+    def _run_generation_loop(self):
        stop_ids = self._task_mgr.tokenizer.stop_ids
        cache = self._cache
        try:
-            while self._running:
+            while not self._stop_event.is_set():
                finished = self._task_mgr.remove_finished_tasks(stop_ids)
                for task in finished:
-                    self._page_cache.task_free(task.task_id)
+                    cache.task_free(task.task_id)
                active = self._task_mgr.get_active_tasks()
                available = self._task_mgr.max_batch_size - len(active)
@@ -88,7 +99,7 @@ class InferenceScheduler:
                    candidates = self._task_mgr.pull_candidates(available)
                    failed = []
                    for task in candidates:
-                        if self._page_cache.task_alloc(task.task_id, task.prompt_ids):
+                        if cache.task_alloc(task.task_id, task.prompt_ids):
                            self._task_mgr.activate(task)
                        else:
                            failed.append(task)
@@ -100,7 +111,10 @@ class InferenceScheduler:
                    continue
                to_prefill = [
-                    t for t in self._task_mgr.get_active_tasks() if t.output_tokens == 0
+                    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:
@@ -110,36 +124,34 @@ class InferenceScheduler:
                    for t in to_prefill:
                        key = (
                            len(t.prompt_ids),
-                            self._page_cache.task_cached(t.task_id),
+                            cache.task_cached(t.task_id),
                        )
                        groups.setdefault(key, []).append(t)
                    for (prompt_len, start_pos), group in groups.items():
                        self._executor.execute_prefill(group, prompt_len, start_pos)
-                        start_logical_page = start_pos // self._page_cache.page_size
+                        start_logical_page = start_pos // getattr(
                            cache, "page_size", 64
                        )
                        for t in group:
-                            self._page_cache.task_record_hashes(
+                            cache.task_record_hashes(
-                                t.task_id,
+                                t.task_id, t.prompt_ids, start_logical_page
                                t.prompt_ids,
                                start_logical_page=start_logical_page,
                            )
                pos_groups: Dict[int, List[Task]] = {}
                for t in self._task_mgr.get_active_tasks():
                    pos_groups.setdefault(t.next_pos, []).append(t)
-                if pos_groups:
+                for next_pos in sorted(pos_groups.keys()):
-                    best_key = max(pos_groups, key=lambda k: len(pos_groups[k]))
+                    group = sorted(pos_groups[next_pos], key=lambda t: t.task_id)
                    group = sorted(pos_groups[best_key], key=lambda t: t.task_id)
                    valid: List[Task] = []
                    for t in group:
-                        if self._page_cache.task_extend(t.task_id, t.next_pos):
+                        if cache.task_extend(t.task_id, t.next_pos):
                            valid.append(t)
                        else:
                            t.status = TaskStatus.ABORTED
-                            if t.stream_callback:
+                            self._task_mgr.invoke_callback(t.task_id, STOP)
                                t.stream_callback(STOP)
                    if valid:
                        next_tokens = self._executor.execute_decode(valid)
@@ -147,41 +159,44 @@ class InferenceScheduler:
                        for t, ntok in zip(valid, next_tokens):
                            t.output_ids.append(ntok)
                            t.output_tokens += 1
-                            pos = t.input_tokens + t.output_tokens
+                            self._task_mgr.invoke_callback(
-                            self._page_cache.task_extend(t.task_id, pos)
+                                t.task_id,
-                            if t.stream_callback:
+                                self._task_mgr.tokenizer.decode([ntok]),
-                                t.stream_callback(
+                            )
                                    self._task_mgr.tokenizer.decode([ntok])
                                )
                        for t in valid:
                            if t.is_finished(stop_ids):
-                                if t.stream_callback:
+                                self._task_mgr.invoke_callback(t.task_id, STOP)
                                    t.stream_callback(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():
-                if task.stream_callback:
+                self._task_mgr.invoke_callback(task.task_id, STOP)
-                    task.stream_callback(STOP)
+                cache.task_free(task.task_id)
-                self._page_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()
            raise
-    def start(self) -> None:
+    def start(self):
-        if not self._running:
+        if self._loop_thread is not None and self._loop_thread.is_alive():
-            self._running = True
+            return
-            t = threading.Thread(target=self._run_generation_loop, daemon=True)
+        self._stop_event.clear()
-            t.start()
+        t = threading.Thread(target=self._run_generation_loop, daemon=True)
-            self._loop_thread = t
+        t.start()
        self._loop_thread = t
-    def stop(self) -> None:
+    def stop(self):
-        self._running = False
+        self._stop_event.set()
        self._task_mgr.wake()
-        if hasattr(self, "_loop_thread"):
+        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._page_cache.task_free(task.task_id)
+            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()
@@ -33,7 +33,6 @@ class Task:
        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
@@ -48,7 +47,6 @@ class Task:
        self.output_tokens: int = 0
        self.arrival_time = time.time()
        self.finish_time: Optional[float] = None
        self.stream_callback = stream_callback
    @property
    def next_pos(self) -> int:
@@ -79,6 +77,7 @@ class TaskManager:
        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()
@@ -117,12 +116,13 @@ class TaskManager:
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            stream_callback=stream_callback,
        )
        with self._lock:
            self.waiting_queue.append(task)
            self._total_tasks += 1
            if stream_callback:
                self._callbacks[task_id] = stream_callback
        self._task_event.set()
        return task_id
@@ -134,8 +134,14 @@ class TaskManager:
                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,
@@ -172,12 +178,12 @@ class TaskManager:
                to_add.append(self.waiting_queue.popleft())
        return to_add
-    def activate(self, task: Task) -> None:
+    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]) -> None:
+    def return_to_waiting(self, tasks: List[Task]):
        with self._lock:
            for task in reversed(tasks):
                self.waiting_queue.appendleft(task)
@@ -185,18 +191,26 @@ class TaskManager:
    def has_work(self) -> bool:
        return bool(self.active_tasks or self.waiting_queue)
-    def wait_for_tasks(self, timeout: float = 1.0) -> None:
+    def wait_for_tasks(self, timeout: float = 1.0):
-        self._task_event.clear()
+        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 clear_queues(self) -> None:
+    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) -> None:
+    def wake(self):
        self._task_event.set()
@@ -8,22 +8,12 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple,
 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.tokenize import AutoTokenizer
 def _validate_sampling_params(
    top_k: int, top_p: float, temperature: float, max_tokens: Optional[int] = None
 ):
    if not (isinstance(top_k, int) and top_k >= 0):
        raise ValueError("top_k must be a non-negative integer")
    if not (0.0 <= top_p <= 1.0):
        raise ValueError("top_p must be a float between 0.0 and 1.0")
    if not (isinstance(temperature, (int, float)) and temperature >= 0):
        raise ValueError("temperature must be a non-negative number")
 class GenerateResult:
    """Thread-safe token accumulator for streaming and non-streaming modes."""
@@ -59,7 +49,7 @@ class GenerateResult:
    def wait(self, timeout: Optional[float] = None) -> bool:
        return self._event.wait(timeout=timeout)
-    def wait_completion(self, timeout: float = 300.0) -> None:
+    def wait_completion(self, timeout: float = 300.0):
        with self._cond:
            if not self._cond.wait_for(
                lambda: self._completed >= self._total, timeout=timeout
@@ -86,7 +76,12 @@ class GenerationRequest:
        max_tokens: Optional[int] = None,
        stream: bool = False,
    ):
-        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
+        if not (isinstance(top_k, int) and top_k >= 0):
            raise ValueError("top_k must be a non-negative integer")
        if not (0.0 <= top_p <= 1.0):
            raise ValueError("top_p must be a float between 0.0 and 1.0")
        if not (isinstance(temperature, (int, float)) and temperature > 0):
            raise ValueError("temperature must be a positive number")
        self.messages = messages
        self.top_k = top_k
@@ -107,6 +102,7 @@ class InferenceEngine:
        max_seq_len: Optional[int] = None,
        max_prompt_len: int = 2048,
        page_size: int = 128,
        cache: Optional[KVCache] = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
@@ -116,7 +112,7 @@ class InferenceEngine:
            max_batch_size=max_batch_size,
            max_seq_len=max_seq_len,
            max_prompt_len=max_prompt_len,
-            page_size=page_size,
+            cache=cache,
        )
        self.scheduler.start()
@@ -137,7 +133,6 @@ class InferenceEngine:
        top_p: float = 1.0,
        top_k: int = 50,
    ) -> Union[Generator, str, List[str]]:
        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
        is_batch = isinstance(prompt, list)
        prompts = prompt if is_batch else [prompt]
@@ -158,7 +153,6 @@ class InferenceEngine:
        top_p: float = 1.0,
        top_k: int = 50,
    ) -> AsyncGenerator[str, None]:
        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
        sync_gen = self._generate_streaming(
            [prompt], False, max_tokens, temperature, top_p, top_k
        )
@@ -289,7 +283,7 @@ class InferenceEngine:
    def get_stats(self) -> Dict[str, Any]:
        return self.scheduler.get_stats()
-    def shutdown(self) -> None:
+    def shutdown(self):
        self.scheduler.stop()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
@@ -29,6 +29,7 @@ class BaseSamplingStrategy(ABC):
        Returns:
            Transformed logits tensor.
        """
        raise NotImplementedError
 class TemperatureStrategy(BaseSamplingStrategy):
@@ -41,13 +42,15 @@ class TemperatureStrategy(BaseSamplingStrategy):
    def __init__(self, temperature: Union[float, Tensor] = 1.0):
        self.temperature = temperature
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
        t = self.temperature
        if isinstance(t, Tensor):
            t = t.to(logits.device, non_blocking=True).view(-1, 1)
            t = torch.clamp(t, min=1e-8)
            if (t != 1.0).any():
-                logits = logits / t.to(logits.device, non_blocking=True).view(-1, 1)
+                logits = logits / t
        elif t != 1.0:
-            logits = logits / t
+            logits = logits / max(t, 1e-8)
        return logits
@@ -61,7 +64,7 @@ class TopKStrategy(BaseSamplingStrategy):
    def __init__(self, top_k: Union[int, Tensor] = 0):
        self.top_k = top_k
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
        tk = self.top_k
        if isinstance(tk, Tensor):
            tk = tk.to(logits.device, non_blocking=True).long().clamp(min=0)
@@ -98,7 +101,9 @@ class TopPStrategy(BaseSamplingStrategy):
    def __init__(self, top_p: Union[float, Tensor] = 1.0):
        self.top_p = top_p
-    def _apply(self, logits, top_p, filter_value):
+    def _apply(
        self, logits: Tensor, top_p: Union[float, Tensor], filter_value: float
    ) -> Tensor:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        remove = cum_probs > top_p
@@ -109,7 +114,7 @@ class TopPStrategy(BaseSamplingStrategy):
        logits[mask] = filter_value
        return logits
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
        tp = self.top_p
        if isinstance(tp, Tensor):
            tp = tp.to(logits.device, non_blocking=True)
@@ -140,7 +145,7 @@ class SamplingPipeline(BaseSamplingStrategy):
    def __init__(self, strategies: List[BaseSamplingStrategy]):
        self.strategies = strategies
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
        for strategy in self.strategies:
            logits = strategy.apply(logits, filter_value)
        return logits
@@ -1,12 +1,18 @@
 from astrai.model.automodel import AutoModel
-from astrai.model.module import (
+from astrai.model.components.attention import GQA
-    GQA,
+from astrai.model.components.decoder_block import DecoderBlock
-    MLP,
+from astrai.model.components.linear import Linear
-    DecoderBlock,
+from astrai.model.components.lora import (
-    Linear,
+    LoRAConfig,
-    RMSNorm,
+    inject_lora,
    load_lora,
    merge_lora,
    save_lora,
 )
-from astrai.model.transformer import Transformer
+from astrai.model.components.mlp import MLP
 from astrai.model.components.norm import RMSNorm
 from astrai.model.encoder import EmbeddingEncoder
 from astrai.model.transformer import AutoRegressiveLM
 __all__ = [
    # Modules
@@ -16,6 +22,13 @@ __all__ = [
    "GQA",
    "DecoderBlock",
    # Models
-    "Transformer",
+    "AutoRegressiveLM",
    "EmbeddingEncoder",
    "AutoModel",
    # LoRA
    "LoRAConfig",
    "inject_lora",
    "merge_lora",
    "save_lora",
    "load_lora",
 ]
@@ -6,16 +6,20 @@ from contextlib import contextmanager
 from pathlib import Path
 from typing import Self, Union
 import safetensors.torch as st
 import torch.nn as nn
-from astrai.config import ModelConfig
+from astrai.config.model_config import BaseModelConfig, ConfigFactory
 from astrai.factory import BaseFactory
 from astrai.serialization import load_model_config, load_model_weights, save_model
@contextmanager
 def _disable_random_init(enable: bool = True):
-    init_functions = [
+    if not enable:
        yield
        return
    names = (
        "xavier_normal_",
        "xavier_uniform_",
        "kaiming_normal_",
@@ -25,18 +29,15 @@ def _disable_random_init(enable: bool = True):
        "constant_",
        "normal_",
        "uniform_",
-    ]
+    )
-    original_funcs = {}
+    orig = {n: getattr(nn.init, n) for n in names if hasattr(nn.init, n)}
-    for name in init_functions:
+    for n in orig:
-        if enable and hasattr(nn.init, name):
+        setattr(nn.init, n, lambda *a, **kw: None)
            original_funcs[name] = getattr(nn.init, name)
            setattr(nn.init, name, lambda *args, **kwargs: None)
    try:
        yield
    finally:
-        if enable:
+        for n, fn in orig.items():
-            for name, orig_func in original_funcs.items():
+            setattr(nn.init, n, fn)
                setattr(nn.init, name, orig_func)
 class AutoModel(BaseFactory["AutoModel"], nn.Module):
@@ -45,7 +46,7 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
    Provides model loading/saving, registration, and generation.
    """
-    def __init__(self, config: ModelConfig):
+    def __init__(self, config: BaseModelConfig):
        super().__init__()
        self.config = config
@@ -59,24 +60,22 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
        model_path = Path(path)
        # Load config
        config = ModelConfig()
        config_path = model_path / "config.json"
-        if config_path.exists():
+        if not config_path.exists():
            config.load(str(config_path))
        else:
            raise FileNotFoundError(f"Config file not found: {config_path}")
-        model_type = config.model_type or "transformer"
+        raw = load_model_config(str(model_path))
        config = ConfigFactory.load(raw)
        model_type = config.model_type or "autoregressive_lm"
        actual_cls = AutoModel.get_component_class(model_type)
        with _disable_random_init(enable=disable_random_init):
            model = actual_cls(config)
        # Load weights
        weights_path = model_path / "model.safetensors"
        if weights_path.exists():
-            state_dict = st.load_file(str(weights_path))
+            state_dict = load_model_weights(str(model_path))
            model.load_state_dict(state_dict, strict=strict)
        return model
@@ -84,15 +83,12 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
    def save_pretrained(
        self,
        save_directory: Union[str, Path],
-    ) -> None:
+    ):
-        save_path = Path(save_directory)
+        save_model(
-        save_path.mkdir(parents=True, exist_ok=True)
+            config=self.config.to_dict(),
-
+            state_dict=self.state_dict(),
-        # Save config
+            save_directory=str(save_directory),
-        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."""
@@ -0,0 +1,25 @@
 from astrai.model.components.attention import GQA, MLA, repeat_kv
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
 from astrai.model.components.mlp import MLP
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import (
    RotaryEmbedding,
    apply_rotary_emb,
    get_rotary_emb,
 )
 __all__ = [
    "Linear",
    "RMSNorm",
    "MLP",
    "Embedding",
    "GQA",
    "MLA",
    "DecoderBlock",
    "RotaryEmbedding",
    "apply_rotary_emb",
    "get_rotary_emb",
    "repeat_kv",
 ]
@@ -5,11 +5,14 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
-from astrai.inference.core.cache import KvcacheView
+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:
    """Repeat KV heads n_rep times for GQA."""
    bs, slen, n_heads, head_dim = x.shape
    if n_rep == 1:
        return x
@@ -20,88 +23,11 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
    )
-def get_rotary_emb(
+class AttnFactory(BaseFactory[nn.Module]):
-    dim: int,
+    pass
    max_len: int,
    base: float = 10000,
    device: Optional[torch.device] = None,
 ) -> Tensor:
    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
    freqs = torch.outer(t, theta).float()
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    return torch.complex(cos, sin)
 def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
    dtype = x.dtype
    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
    x_complex = torch.view_as_complex(x_)
    freqs_cis = freqs_cis.unsqueeze(2)
    x_rotated = x_complex * freqs_cis
    x_out = torch.view_as_real(x_rotated).flatten(-2)
    return x_out.to(dtype)
 class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_len: int, base: int = 10000):
        super().__init__()
        self.dim = dim
        self.max_len = max_len
        self.base = base
        self._set_rotary_buffer(self.max_len)
    def _set_rotary_buffer(self, max_len: int):
        rotary_emb = get_rotary_emb(self.dim, max_len, self.base)
        freqs_cis = torch.view_as_real(rotary_emb)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
        if position_ids is None:
            position_ids = (
                torch.arange(x.size(1), device=x.device)
                .unsqueeze(0)
                .expand(x.size(0), -1)
            )
        position_freq_cis = self.freqs_cis[position_ids].float()
        return torch.view_as_complex(position_freq_cis)
 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
@AttnFactory.register("gqa")
 class GQA(nn.Module):
    def __init__(
        self,
@@ -112,6 +38,7 @@ class GQA(nn.Module):
        norm_eps: float,
        use_gated_attention: bool,
        layer_id: int,
        n_layers: int = 1,
    ):
        super().__init__()
        assert dim % n_heads == 0
@@ -129,7 +56,7 @@ class GQA(nn.Module):
        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)
+        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)
@@ -148,11 +75,10 @@ class GQA(nn.Module):
        x: Tensor,
        rotary_emb: Tensor,
        attn_mask: Tensor = None,
-        paged_cache: Optional[KvcacheView] = None,
+        paged_cache: Optional[CacheView] = None,
    ) -> Tensor:
        is_causal = attn_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)
@@ -167,7 +93,6 @@ class GQA(nn.Module):
        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, attn_mask, is_causal=is_causal)
@@ -183,6 +108,7 @@ class GQA(nn.Module):
        return out
@AttnFactory.register("mla")
 class MLA(nn.Module):
    def __init__(
        self,
@@ -193,8 +119,10 @@ class MLA(nn.Module):
        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
@@ -206,19 +134,25 @@ class MLA(nn.Module):
        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)
        # 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),
+            n_kv_heads * (2 * self.head_dim),
        )
-        self.o_proj = Linear(dim, dim, bias=False)
+        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)
@@ -228,7 +162,7 @@ class MLA(nn.Module):
        x: Tensor,
        rotary_emb: Tensor,
        attn_mask: Tensor = None,
-        paged_cache: Optional[KvcacheView] = None,
+        paged_cache: Optional[CacheView] = None,
    ) -> Tensor:
        bsz, seq_len, _ = x.size()
        is_causal = attn_mask is None
@@ -248,7 +182,7 @@ class MLA(nn.Module):
        q_nope, q_rope = (
            q[..., : self.qk_nope_head_dim],
-            q[..., self.qk_rope_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)
@@ -256,6 +190,10 @@ class MLA(nn.Module):
        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)
@@ -274,57 +212,3 @@ class MLA(nn.Module):
        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: Tensor,
        attention_mask: Optional[Tensor] = None,
        paged_cache: Optional[KvcacheView] = None,
    ) -> Tensor:
        attn_output = self.attention(
            self.input_norm(x),
            rotary_emb,
            attention_mask,
            paged_cache,
        )
        x = attn_output + x
        x = self.mlp(self.post_attention_norm(x)) + x
        return x
 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)
@@ -0,0 +1,39 @@
 from dataclasses import asdict
 from typing import Optional
 import torch.nn as nn
 from torch import Tensor
 from astrai.inference.core.cache import CacheView
 from astrai.model.components.attention import AttnFactory
 from astrai.model.components.mlp import FFNFactory
 from astrai.model.components.norm import RMSNorm
 class DecoderBlock(nn.Module):
    def __init__(self, config, layer_id: int):
        super().__init__()
        cfg = asdict(config)
        cfg["down_init_std"] = 0.02 / (2 * config.n_layers) ** 0.5
        self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
        self.input_norm = RMSNorm(config.dim, config.norm_eps)
        self.post_attention_norm = RMSNorm(config.dim, config.norm_eps)
        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
    def forward(
        self,
        x: Tensor,
        rotary_emb: Tensor,
        attention_mask: Optional[Tensor] = None,
        paged_cache: Optional[CacheView] = None,
    ) -> Tensor:
        attn_output = self.attention(
            self.input_norm(x),
            rotary_emb,
            attention_mask,
            paged_cache,
        )
        x = attn_output + x
        x = self.mlp(self.post_attention_norm(x)) + x
        return x
@@ -0,0 +1,26 @@
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 class Embedding(nn.Module):
    def __init__(self, vocab_size: int, embedding_dim: int, neftune_alpha: float = 0.0):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
        self.neftune_noise_alpha = neftune_alpha
    def set_neftune_alpha(self, alpha: float):
        self.neftune_noise_alpha = alpha
    def reset_parameters(self):
        nn.init.normal_(self.weight, mean=0.0, std=0.02)
    def forward(self, x: Tensor) -> Tensor:
        out = F.embedding(x, self.weight)
        if self.training and self.neftune_noise_alpha > 0.0:
            eps = self.neftune_noise_alpha / math.sqrt(out.size(1))
            out = out + eps * torch.randn_like(out)
        return out
@@ -0,0 +1,24 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 class Linear(nn.Module):
    def __init__(
        self, in_dim: int, out_dim: int, bias: bool = False, init_std: float = 0.02
    ):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
        self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
        self.init_std = init_std
    def reset_parameters(self):
        nn.init.normal_(self.weight, mean=0.0, std=self.init_std)
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / (fan_in**0.5)
            nn.init.uniform_(self.bias, -bound, bound)
    def forward(self, x: Tensor) -> Tensor:
        return F.linear(x, self.weight, self.bias)
@@ -0,0 +1,194 @@
 import logging
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Optional, Set
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from astrai.model.components.linear import Linear
 from astrai.serialization import (
    load_json,
    load_safetensors,
    save_json,
    save_safetensors,
 )
 logger = logging.getLogger(__name__)
 TARGET_MODULES_ATTN = {"q_proj", "k_proj", "v_proj", "o_proj"}
 TARGET_MODULES_FFN = {"up", "gate", "down"}
@dataclass
 class LoRAConfig:
    r: int = 16
    alpha: int = 32
    target_modules: tuple = ("q_proj", "v_proj")
 class LoRALinear(nn.Module):
    def __init__(self, base: Linear, r: int = 16, alpha: int = 32):
        super().__init__()
        self.register_parameter("weight", base.weight)
        self.weight.requires_grad_(False)
        self.bias = base.bias
        if self.bias is not None:
            self.bias.requires_grad_(False)
        self.r = r
        self.scaling = alpha / r
        self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1]) / r)
        self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r))
        self._merged = False
    def forward(self, x):
        out = F.linear(x, self.weight, self.bias)
        if not self._merged:
            out += (F.linear(x, self.lora_A) @ self.lora_B.T) * self.scaling
        return out
    def merge(self):
        if self._merged:
            return
        self.weight.data += (self.lora_B @ self.lora_A) * self.scaling
        self._merged = True
        del self.lora_A
        del self.lora_B
 def _collect_lora_info(model: nn.Module) -> dict:
    names = {}
    for n, m in model.named_modules():
        if isinstance(m, Linear):
            _, _, child = n.rpartition(".")
            names.setdefault(child, []).append(n)
    return names
 def _get_lora_count(model: nn.Module) -> int:
    return sum(1 for m in model.modules() if isinstance(m, LoRALinear))
 def inject_lora(
    model: nn.Module,
    r: int = 16,
    alpha: int = 32,
    target_modules: Optional[Set[str]] = None,
 ) -> LoRAConfig:
    if target_modules is None:
        target_modules = TARGET_MODULES_ATTN
    available = _collect_lora_info(model)
    injected = 0
    for name, module in list(model.named_modules()):
        if not isinstance(module, Linear):
            continue
        parent_name, _, child_name = name.rpartition(".")
        if child_name not in target_modules:
            continue
        parent = model.get_submodule(parent_name) if parent_name else model
        setattr(parent, child_name, LoRALinear(module, r=r, alpha=alpha))
        injected += 1
    if injected == 0:
        logger.warning(
            "No LoRA layers injected. Available Linear child names: %s. "
            "target_modules: %s. Check model type and target_modules.",
            sorted(available),
            sorted(target_modules),
        )
    else:
        logger.info("LoRA injected: %d layers (r=%d, alpha=%d)", injected, r, alpha)
    return LoRAConfig(r=r, alpha=alpha, target_modules=tuple(target_modules))
 def merge_lora(model: nn.Module):
    n = 0
    for module in model.modules():
        if isinstance(module, LoRALinear):
            module.merge()
            n += 1
    if n == 0:
        logger.warning("No LoRA layers to merge.")
    else:
        logger.info("Merged %d LoRA layers", n)
 def save_lora(model: nn.Module, save_dir: str, config: LoRAConfig):
    lora_sd = {
        k: v
        for k, v in model.state_dict().items()
        if k.endswith((".lora_A", ".lora_B"))
    }
    if not lora_sd:
        raise RuntimeError(
            "No LoRA parameters found in model. "
            "The model may not have been injected or was already merged."
        )
    path = Path(save_dir)
    path.mkdir(parents=True, exist_ok=True)
    save_safetensors(lora_sd, path / "adapter_model.safetensors")
    save_json(asdict(config), path / "adapter_config.json")
    logger.info("LoRA adapter saved to %s (%d keys)", save_dir, len(lora_sd))
 def load_lora(model: nn.Module, load_dir: str) -> LoRAConfig:
    path = Path(load_dir)
    raw = load_json(path / "adapter_config.json")
    config = LoRAConfig(
        r=raw["r"], alpha=raw["alpha"], target_modules=tuple(raw["target_modules"])
    )
    existing = _get_lora_count(model)
    if existing > 0:
        logger.warning(
            "Model already has %d LoRA layers. Skipping injection, "
            "loading weights onto existing layers only.",
            existing,
        )
    else:
        inject_lora(
            model,
            r=config.r,
            alpha=config.alpha,
            target_modules=set(config.target_modules),
        )
    weights = load_safetensors(path / "adapter_model.safetensors")
    try:
        missing, unexpected = model.load_state_dict(weights, strict=False)
    except RuntimeError as e:
        msg = str(e)
        if "size mismatch" in msg:
            raise RuntimeError(
                f"LoRA weight shapes do not match the model. "
                f"The adapter config (r={config.r}) may not match the injected layers. "
                f"Original error: {msg}"
            ) from e
        raise
    injected = _get_lora_count(model)
    if injected == 0:
        raise RuntimeError(
            "No LoRA layers found after loading. "
            "Inject LoRA before calling load_lora, or check the adapter config."
        )
    if missing:
        lora_missing = [k for k in missing if "lora" in k]
        if lora_missing:
            raise RuntimeError(
                f"LoRA weight keys not found in model: {lora_missing}. "
                f"The adapter config (r={config.r}) may not match the model."
            )
        logger.debug("LoRA load: %d missing base-weight keys (expected)", len(missing))
    if unexpected:
        logger.warning("LoRA load: %d unexpected keys", len(unexpected))
    logger.info("LoRA adapter loaded from %s", load_dir)
    return config
@@ -0,0 +1,100 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from astrai.factory import BaseFactory
 from astrai.model.components.linear import Linear
 class FFNFactory(BaseFactory[nn.Module]):
    pass
@FFNFactory.register("mlp")
 class MLP(nn.Module):
    def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02):
        super().__init__()
        self.up = Linear(dim, dim_ffn)
        self.gate = Linear(dim, dim_ffn)
        self.down = Linear(dim_ffn, dim, init_std=down_init_std)
    def forward(self, x: Tensor) -> Tensor:
        gated = self.up(x) * F.silu(self.gate(x))
        out = self.down(gated)
        return out
@FFNFactory.register("moe")
 class DeepSeekMoE(nn.Module):
    def __init__(
        self,
        dim: int,
        dim_ffn: int,
        n_routed_experts: int,
        n_shared_experts: int = 1,
        n_activated_experts: int = 2,
        topk_method: str = "greedy",
        n_layers: int = 1,
    ):
        super().__init__()
        self.dim = dim
        self.n_routed_experts = n_routed_experts
        self.n_shared_experts = n_shared_experts
        self.n_activated_experts = n_activated_experts
        self.topk_method = topk_method
        self.router = Linear(dim, n_routed_experts, bias=False)
        moe_scale = 1 / max(n_shared_experts, 1) + 1 / n_activated_experts
        down_init_std = 0.02 / (2 * n_layers * moe_scale) ** 0.5
        self.shared_experts = nn.ModuleList(
            [
                MLP(dim, dim_ffn, down_init_std=down_init_std)
                for _ in range(n_shared_experts)
            ]
        )
        self.routed_experts = nn.ModuleList(
            [
                MLP(dim, dim_ffn, down_init_std=down_init_std)
                for _ in range(n_routed_experts)
            ]
        )
    def forward(self, x: Tensor) -> Tensor:
        bsz, seq_len, dim = x.shape
        x_flat = x.view(-1, dim)
        shared_out = self._shared_forward(x_flat)
        routed_out = self._routed_forward(x_flat)
        out = (shared_out + routed_out).view(bsz, seq_len, dim)
        return out
    def _shared_forward(self, x: Tensor) -> Tensor:
        if self.n_shared_experts == 0:
            return torch.zeros_like(x)
        return sum(e(x) for e in self.shared_experts) / self.n_shared_experts
    def _routed_forward(self, x: Tensor) -> Tensor:
        N, D = x.shape
        K = self.n_activated_experts
        router_logits = self.router(x)
        router_probs = torch.softmax(router_logits.float(), dim=-1).to(x.dtype)
        topk_weights, topk_indices = torch.topk(router_probs, K, dim=-1)
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
        output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
        for expert_idx in range(self.n_routed_experts):
            expert_mask = topk_indices == expert_idx
            token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
            if token_idx.numel() == 0:
                continue
            expert_input = x[token_idx]
            expert_output = self.routed_experts[expert_idx](expert_input)
            weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
            output.index_add_(0, token_idx, expert_output * weights)
        return output
@@ -0,0 +1,15 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 class RMSNorm(nn.Module):
    def __init__(self, dim, norm_eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.normalized_shape = (dim,)
        self.norm_eps = norm_eps
    def forward(self, x: Tensor) -> Tensor:
        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)
@@ -0,0 +1,71 @@
 from typing import Dict, Optional
 import torch
 import torch.nn as nn
 from torch import Tensor
 def get_rotary_emb(
    dim: int,
    max_len: int,
    base: float = 10000,
    device: Optional[torch.device] = None,
 ) -> Tensor:
    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
    freqs = torch.outer(t, theta).float()
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    return torch.complex(cos, sin)
 def ntk_base(base: float, dim: int, factor: float) -> float:
    return base * (factor ** (dim / (dim - 2)))
 def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
    dtype = x.dtype
    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
    x_complex = torch.view_as_complex(x_)
    freqs_cis = freqs_cis.unsqueeze(2)
    x_rotated = x_complex * freqs_cis
    x_out = torch.view_as_real(x_rotated).flatten(-2)
    return x_out.to(dtype)
 class RotaryEmbedding(nn.Module):
    def __init__(
        self,
        dim: int,
        max_len: int,
        base: float = 10000,
        rope_scaling: Optional[Dict] = None,
    ):
        super().__init__()
        self.dim = dim
        self.max_len = max_len
        self.base = base
        self.rope_scaling = rope_scaling
        if rope_scaling is not None:
            scaling_type = rope_scaling.get("type", "ntk")
            factor = rope_scaling.get("factor", 1.0)
            if scaling_type == "ntk":
                self.base = ntk_base(base, dim, factor)
        self._set_rotary_buffer(self.max_len)
    def _set_rotary_buffer(self, max_len: int):
        rotary_emb = get_rotary_emb(self.dim, max_len, self.base)
        freqs_cis = torch.view_as_real(rotary_emb)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
        if position_ids is None:
            position_ids = (
                torch.arange(x.size(1), device=x.device)
                .unsqueeze(0)
                .expand(x.size(0), -1)
            )
        position_freq_cis = self.freqs_cis[position_ids].float()
        return torch.view_as_complex(position_freq_cis)
@@ -0,0 +1,89 @@
 from typing import Any, Mapping, Optional
 import torch
 import torch.nn as nn
 from torch import Tensor
 from astrai.config.model_config import EncoderConfig
 from astrai.model.automodel import AutoModel
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import RotaryEmbedding
 from astrai.model.transformer import process_attention_mask
@AutoModel.register("embedding")
 class EmbeddingEncoder(AutoModel):
    def __init__(self, config: EncoderConfig):
        super().__init__(config)
        self.config = config
        rope_dim = config.dim // config.n_heads
        rope_base = config.rope_theta if config.rope_theta is not None else 10000
        self.rotary_embedding = RotaryEmbedding(
            rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
        )
        self.embed_tokens = Embedding(
            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
        )
        self.layers = nn.ModuleList(
            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
        )
        self.norm = RMSNorm(config.dim, config.norm_eps)
        self.pooling_type = config.pooling_type or "mean"
        self.normalize_embeddings = config.normalize_embeddings or False
        self.apply(self._init_weights)
    def _init_weights(self, module):
        if hasattr(module, "reset_parameters"):
            module.reset_parameters()
    def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
        state_dict = dict(state_dict)
        state_dict.pop("lm_head.weight", None)
        return super().load_state_dict(state_dict, strict=strict, assign=assign)
    def forward(
        self,
        input_ids: Tensor,
        input_mask: Optional[Tensor] = None,
        position_ids: Optional[Tensor] = None,
    ) -> Tensor:
        assert input_ids.ndim == 2
        B, S = input_ids.shape
        x = self.embed_tokens(input_ids)
        rotary_emb = self.rotary_embedding(x, position_ids)
        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=False)
        for layer in self.layers:
            x = layer(x, rotary_emb, attn_mask, paged_cache=None)
        hidden_states = self.norm(x)
        if self.pooling_type == "cls":
            pooled = hidden_states[:, 0]
        elif self.pooling_type == "last":
            if input_mask is not None:
                lengths = input_mask.sum(dim=1) - 1
                pooled = hidden_states[torch.arange(B, device=x.device), lengths]
            else:
                pooled = hidden_states[:, -1]
        else:
            if input_mask is not None:
                mask = input_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
                pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(
                    min=1.0
                )
            else:
                pooled = hidden_states.mean(dim=1)
        if self.normalize_embeddings:
            pooled = torch.nn.functional.normalize(pooled, p=2, dim=-1)
        return pooled
@@ -1,19 +1,17 @@
-from typing import Any, Mapping, Optional
+from typing import Any, Dict, Mapping, Optional
 import torch
 import torch.nn as nn
 from torch import Tensor
-from astrai.config.model_config import ModelConfig
+from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.core.cache import KvcacheView
+from astrai.inference.core.cache import CacheView
 from astrai.model.automodel import AutoModel
-from astrai.model.module import (
+from astrai.model.components.decoder_block import DecoderBlock
-    DecoderBlock,
+from astrai.model.components.embedding import Embedding
-    Embedding,
+from astrai.model.components.linear import Linear
-    Linear,
+from astrai.model.components.norm import RMSNorm
-    RMSNorm,
+from astrai.model.components.rope import RotaryEmbedding
    RotaryEmbedding,
 )
 def process_attention_mask(
@@ -28,66 +26,58 @@ def process_attention_mask(
        return input_mask
    device = input_tensor.device
-    dtype = input_tensor.dtype
+    B = input_tensor.size(0)
    B, S = input_tensor.size()[:2]
    T = position_ids.max().item() + 1
    if input_mask is None:
        if position_ids.min().item() == 0 and is_causal:
            return None
-        pad = torch.ones(B, T, dtype=torch.bool, device=device)
+        attend = torch.ones(B, 1, T, dtype=torch.bool, device=device)
    else:
-        pad = input_mask[:, :T].to(device=device, dtype=torch.bool)
+        attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1)
    attend = pad.view(B, 1, T).expand(B, S, T).clone()
    if is_causal:
-        attend &= position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
+        causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
        attend = attend & causal
-    return torch.full(
+    return attend.unsqueeze(1)
        (B, 1, S, T), -torch.finfo(dtype).max / 2, dtype=dtype, device=device
    ).masked_fill_(attend.unsqueeze(1), 0.0)
-@AutoModel.register("transformer")
+@AutoModel.register("autoregressive_lm")
-class Transformer(AutoModel):
+class AutoRegressiveLM(AutoModel):
-    """Transformer language model with paged KV cache."""
+    """Autoregressive language model with paged KV cache."""
-    def __init__(self, config: ModelConfig):
+    def __init__(self, config: AutoRegressiveLMConfig):
        super().__init__(config)
        self.config = config
        rope_dim = (
            config.qk_rope_head_dim
            if config.attn_type == "mla"
            else config.dim // config.n_heads
        )
        rope_base = config.rope_theta if config.rope_theta is not None else 10000
        self.rotary_embedding = RotaryEmbedding(
-            config.dim // config.n_heads, config.max_len
+            rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
        )
        self.embed_tokens = Embedding(
            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
        )
        self.embed_tokens = Embedding(config.vocab_size, config.dim)
        self.layers = nn.ModuleList(
-            [
+            [DecoderBlock(config, 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:
+        if self.config.tie_weight is True:
            self.lm_head.weight = self.embed_tokens.weight
-        self._init_weights()
+        self.apply(self._init_weights)
-    def _init_weights(self):
+    def _init_weights(self, module):
-        for param in self.parameters():
+        if hasattr(module, "reset_parameters"):
-            if param.dim() > 1:
+            module.reset_parameters()
                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"
@@ -95,7 +85,7 @@ class Transformer(AutoModel):
        state_dict = dict(state_dict)
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
            # same tensor for embed and lm_head
            if embed_key in state_dict:
                state_dict[lm_head_key] = state_dict[embed_key]
@@ -111,7 +101,7 @@ class Transformer(AutoModel):
            destination=destination, prefix=prefix, keep_vars=keep_vars
        )
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
            lm_head_key = prefix + "lm_head.weight"
            if lm_head_key in state_dict:
                del state_dict[lm_head_key]
@@ -122,9 +112,9 @@ class Transformer(AutoModel):
        self,
        input_ids: Tensor,
        input_mask: Optional[Tensor] = None,
-        paged_cache: Optional[KvcacheView] = None,
+        paged_cache: Optional[CacheView] = None,
        position_ids: Optional[Tensor] = None,
-    ) -> Tensor:
+    ) -> Dict[str, Tensor]:
        assert input_ids.ndim == 2
        x = self.embed_tokens(input_ids)
@@ -1,3 +1,13 @@
 from astrai.parallel.executor import (
    AccumOptimizer,
    AccumScheduler,
    BaseExecutor,
    DDPExecutor,
    ExecutorFactory,
    FSDPExecutor,
    GradientState,
    NoneExecutor,
 )
 from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
 from astrai.parallel.setup import (
    get_current_device,
@@ -17,4 +27,12 @@ __all__ = [
    "spawn_parallel_fn",
    "RowParallelLinear",
    "ColumnParallelLinear",
    "ExecutorFactory",
    "BaseExecutor",
    "GradientState",
    "AccumOptimizer",
    "AccumScheduler",
    "NoneExecutor",
    "DDPExecutor",
    "FSDPExecutor",
 ]
@@ -0,0 +1,286 @@
 """Unified training executor — parallel strategy + gradient accumulation."""
 import contextlib
 import logging
 import os
 from contextlib import contextmanager
 from typing import Optional, Tuple
 import torch
 import torch.nn as nn
 from torch.distributed.fsdp import FullStateDictConfig, StateDictType
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 from astrai.factory import BaseFactory
 from astrai.parallel.setup import get_rank, get_world_size
 logger = logging.getLogger(__name__)
 class GradientState:
    def __init__(self, grad_accum_steps: int = 1):
        self.num_steps = max(grad_accum_steps, 1)
        self._step: int = 0
        self._sync_gradients: bool = True
    @property
    def sync_gradients(self) -> bool:
        return self._sync_gradients
    def _do_sync(self):
        self._step += 1
        self._sync_gradients = self._step % self.num_steps == 0
 class AccumOptimizer:
    def __init__(self, optimizer: Optimizer, gradient_state: GradientState):
        self.optimizer = optimizer
        self.gradient_state = gradient_state
    def step(self, closure=None):
        if self.gradient_state.sync_gradients:
            self.optimizer.step(closure)
    def zero_grad(self):
        if self.gradient_state.sync_gradients:
            self.optimizer.zero_grad()
    @property
    def param_groups(self):
        return self.optimizer.param_groups
    def state_dict(self):
        return self.optimizer.state_dict()
    def load_state_dict(self, d):
        self.optimizer.load_state_dict(d)
 class AccumScheduler:
    def __init__(self, scheduler: LRScheduler, gradient_state: GradientState):
        self.scheduler = scheduler
        self.gradient_state = gradient_state
    def step(self):
        if self.gradient_state.sync_gradients:
            self.scheduler.step()
    def state_dict(self):
        return self.scheduler.state_dict()
    def load_state_dict(self, d):
        self.scheduler.load_state_dict(d)
    def get_last_lr(self):
        return self.scheduler.get_last_lr()
 class BaseExecutor:
    def __init__(self, grad_accum_steps: int = 1):
        self.gradient_state = GradientState(grad_accum_steps)
    def prepare(
        self,
        model: nn.Module,
        optimizer: Optional[Optimizer] = None,
        dataloader: Optional[DataLoader] = None,
        scheduler: Optional[LRScheduler] = None,
    ) -> Tuple[
        nn.Module, Optional[Optimizer], Optional[DataLoader], Optional[LRScheduler]
    ]:
        model = self._prepare_model(model)
        if optimizer is not None:
            optimizer = AccumOptimizer(optimizer, self.gradient_state)
        if scheduler is not None:
            scheduler = AccumScheduler(scheduler, self.gradient_state)
        return model, optimizer, dataloader, scheduler
    def _prepare_model(self, model: nn.Module) -> nn.Module:
        return model
    def _no_sync(self, model: nn.Module):
        return contextlib.nullcontext()
    @contextmanager
    def accumulate(self, model: nn.Module):
        self.gradient_state._do_sync()
        if not self.gradient_state.sync_gradients:
            with self._no_sync(model):
                yield
        else:
            yield
    def backward(self, loss: torch.Tensor):
        loss.backward()
    def unwrap_model(self, model: nn.Module):
        return model.state_dict()
    @property
    def use_distributed(self) -> bool:
        return get_world_size() > 1
    @property
    def sync_gradients(self) -> bool:
        return self.gradient_state.sync_gradients
    @property
    def grad_accum_steps(self) -> int:
        return self.gradient_state.num_steps
    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
        total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        if isinstance(total_norm, torch.Tensor):
            return total_norm.item()
        return total_norm
 class ExecutorFactory(BaseFactory[BaseExecutor]):
    pass
@ExecutorFactory.register("none")
 class NoneExecutor(BaseExecutor):
    pass
@ExecutorFactory.register("ddp")
 class DDPExecutor(BaseExecutor):
    def __init__(
        self,
        grad_accum_steps: int = 1,
        dim: int = 0,
        broadcast_buffers: bool = True,
        init_sync: bool = True,
        process_group=None,
        bucket_cap_mb: int = 25,
        find_unused_parameters: bool = False,
        check_reduction: bool = False,
        gradient_as_bucket_view: bool = False,
        static_graph: bool = False,
        delay_all_reduce_named_params=None,
        param_to_hook_all_reduce=None,
        mixed_precision=None,
        device_mesh=None,
    ):
        super().__init__(grad_accum_steps=grad_accum_steps)
        self._ddp_kwargs = dict(
            dim=dim,
            broadcast_buffers=broadcast_buffers,
            init_sync=init_sync,
            process_group=process_group,
            bucket_cap_mb=bucket_cap_mb,
            find_unused_parameters=find_unused_parameters,
            check_reduction=check_reduction,
            gradient_as_bucket_view=gradient_as_bucket_view,
            static_graph=static_graph,
            delay_all_reduce_named_params=delay_all_reduce_named_params,
            param_to_hook_all_reduce=param_to_hook_all_reduce,
            mixed_precision=mixed_precision,
            device_mesh=device_mesh,
        )
    def _prepare_model(self, model: nn.Module) -> nn.Module:
        if not self.use_distributed:
            logger.warning("DDP backend selected but world_size=1, model not wrapped")
            return model
        local_rank = int(os.environ.get("LOCAL_RANK", get_rank()))
        model = DDP(
            model,
            device_ids=[local_rank],
            output_device=local_rank,
            **self._ddp_kwargs,
        )
        logger.info("Model wrapped with DDP (world_size=%d)", get_world_size())
        return model
    def _no_sync(self, model: nn.Module):
        if isinstance(model, DDP):
            return model.no_sync()
        return contextlib.nullcontext()
    def unwrap_model(self, model: nn.Module):
        if isinstance(model, DDP):
            return model.module.state_dict()
        return model.state_dict()
@ExecutorFactory.register("fsdp")
 class FSDPExecutor(BaseExecutor):
    def __init__(
        self,
        grad_accum_steps: int = 1,
        process_group=None,
        sharding_strategy=None,
        cpu_offload=None,
        auto_wrap_policy=None,
        backward_prefetch=None,
        mixed_precision=None,
        ignored_modules=None,
        param_init_fn=None,
        sync_module_states: bool = False,
        forward_prefetch: bool = False,
        limit_all_gathers: bool = True,
        ignored_states=None,
        device_mesh=None,
    ):
        super().__init__(grad_accum_steps=grad_accum_steps)
        self._fsdp_kwargs = {
            k: v
            for k, v in dict(
                process_group=process_group,
                sharding_strategy=sharding_strategy,
                cpu_offload=cpu_offload,
                auto_wrap_policy=auto_wrap_policy,
                backward_prefetch=backward_prefetch,
                mixed_precision=mixed_precision,
                ignored_modules=ignored_modules,
                param_init_fn=param_init_fn,
                sync_module_states=sync_module_states,
                forward_prefetch=forward_prefetch,
                limit_all_gathers=limit_all_gathers,
                use_orig_params=True,
                ignored_states=ignored_states,
                device_mesh=device_mesh,
            ).items()
            if v is not None
        }
        self._original_model: Optional[nn.Module] = None
    def _prepare_model(self, model: nn.Module) -> nn.Module:
        if not self.use_distributed:
            logger.warning("FSDP backend selected but world_size=1, model not wrapped")
            return model
        self._original_model = model
        device_id = torch.device("cuda", get_rank())
        model = FSDP(model, device_id=device_id, **self._fsdp_kwargs)
        logger.info("Model wrapped with FSDP (world_size=%d)", get_world_size())
        return model
    def _no_sync(self, model: nn.Module):
        if isinstance(model, FSDP):
            return model.no_sync()
        return contextlib.nullcontext()
    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
        if isinstance(model, FSDP) and self.use_distributed:
            total_norm = model.clip_grad_norm_(max_norm)
            if isinstance(total_norm, torch.Tensor):
                return total_norm.item()
            return total_norm
        return super().clip_grad_norm(model, max_norm)
    def unwrap_model(self, model: nn.Module):
        if isinstance(model, FSDP) and self.use_distributed:
            with FSDP.state_dict_type(
                model,
                StateDictType.FULL_STATE_DICT,
                FullStateDictConfig(offload_to_cpu=True, rank0_only=False),
            ):
                return model.state_dict()
        return model.state_dict()
@@ -1,13 +1,21 @@
 import os
 import socket
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from functools import wraps
-from typing import Callable
+from typing import Callable, Optional
 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"]
@@ -30,6 +38,7 @@ def get_rank() -> int:
 def setup_parallel(
    rank: int,
    world_size: int,
    local_rank: int,
    backend: str = "nccl",
    master_addr: str = "localhost",
    master_port: str = "29500",
@@ -41,20 +50,26 @@ def setup_parallel(
        return
    if world_size <= 1:
        device_id = torch.device(device_type, local_rank)
        os.environ["LOCAL_RANK"] = str(local_rank)
        os.environ["WORLD_SIZE"] = "1"
        os.environ["LOCAL_DEVICE"] = str(device_id)
        yield None
        return
-    device_id = torch.device(device_type, rank)
+    device_id = torch.device(device_type, local_rank)
    os.environ["MASTER_ADDR"] = master_addr
    os.environ["MASTER_PORT"] = master_port
-    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(local_rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["LOCAL_DEVICE"] = str(device_id)
-    dist.init_process_group(
+    pg_kwargs = dict(rank=rank, world_size=world_size, backend=backend)
-        rank=rank, world_size=world_size, backend=backend, device_id=device_id
+    if backend in ("nccl", "ccl"):
-    )
+        pg_kwargs["device_id"] = device_id
    dist.init_process_group(**pg_kwargs)
    try:
        if backend == "nccl" and torch.cuda.is_available():
@@ -90,7 +105,7 @@ def only_on_rank(rank, sync=False):
    return decorator
-def wrapper_spawn_func(
+def _run_single_rank(
    rank: int,
    world_size: int,
    backend: str,
@@ -100,20 +115,108 @@ def wrapper_spawn_func(
    func: Callable,
    kwargs: dict,
 ):
-    try:
+    with setup_parallel(
        rank=rank,
        world_size=world_size,
        local_rank=rank,
        backend=backend,
        master_addr=master_addr,
        master_port=master_port,
        device_type=device_type,
    ):
        func(**kwargs)
 class LaunchStrategy(ABC):
    """Strategy for launching a function in a distributed context."""
    def __init__(
        self,
        world_size: int,
        backend: str,
        master_addr: str,
        master_port: str,
        device_type: str,
        start_method: str,
    ):
        self.world_size = world_size
        self.backend = backend
        self.master_addr = master_addr
        self.master_port = master_port
        self.device_type = device_type
        self.start_method = start_method
    @abstractmethod
    def launch(self, func: Callable, **kwargs):
        raise NotImplementedError
 class TorchrunStrategy(LaunchStrategy):
    """External orchestrator (torchrun, SLURM, K8s) — env vars pre-set."""
    def launch(self, func: Callable, **kwargs):
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ.get("LOCAL_RANK", rank))
        with setup_parallel(
            rank=rank,
            world_size=world_size,
-            backend=backend,
+            local_rank=local_rank,
-            master_addr=master_addr,
+            backend=self.backend,
-            master_port=master_port,
+            master_addr=os.environ.get("MASTER_ADDR", self.master_addr),
-            device_type=device_type,
+            master_port=os.environ.get("MASTER_PORT", self.master_port),
            device_type=self.device_type,
        ):
            func(**kwargs)
-    except Exception as e:
+
-        print(f"Error in rank {rank}: {e}")
+class LocalStrategy(LaunchStrategy):
-        raise
+    """Local launcher — single-process or mp.start_processes."""
    def launch(self, func: Callable, **kwargs):
        args = (
            self.world_size,
            self.backend,
            self.master_addr,
            self.master_port,
            self.device_type,
            func,
            kwargs,
        )
        if self.world_size == 1:
            _run_single_rank(0, *args)
            return
        ctx = mp.start_processes(
            _run_single_rank,
            args=args,
            nprocs=self.world_size,
            start_method=self.start_method,
            join=False,
        )
        try:
            while not ctx.join():
                pass
        except BaseException:
            for p in ctx.processes:
                p.terminate()
            ctx.join()
            raise
 def _detect_launcher() -> str:
    """Detect the distributed launcher from environment.
    Returns one of: "torchelastic", "torchrun", "external", "local".
    """
    if dist.is_torchelastic_launched():
        return "torchelastic"
    if "LOCAL_WORLD_SIZE" in os.environ:
        return "torchrun"
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        return "external"
    return "local"
 def spawn_parallel_fn(
@@ -121,41 +224,20 @@ def spawn_parallel_fn(
    world_size: int,
    backend: str = "nccl",
    master_addr: str = "localhost",
-    master_port: str = "29500",
+    master_port: Optional[str] = None,
    device_type: str = "cuda",
    start_method: str = "spawn",
    **kwargs,
 ):
-    # clear environment variables
+    if master_port is None:
-    for key in [
+        master_port = find_free_port()
-        "MASTER_ADDR",
+    launcher = _detect_launcher()
-        "MASTER_PORT",
+    if launcher in ("torchelastic", "torchrun", "external"):
-        "RANK",
+        strategy = TorchrunStrategy(
-        "WORLD_SIZE",
+            world_size, backend, master_addr, master_port, device_type, start_method
-        "LOCAL_RANK",
+        )
-        "LOCAL_DEVICE",
+    else:
-    ]:
+        strategy = LocalStrategy(
-        if key in os.environ:
+            world_size, backend, master_addr, master_port, device_type, start_method
-            del os.environ[key]
+        )
-
+    strategy.launch(func, **kwargs)
    if world_size == 1:
        device_id = torch.device(device_type, 0)
        os.environ["LOCAL_RANK"] = "0"
        os.environ["WORLD_SIZE"] = "1"
        os.environ["LOCAL_DEVICE"] = str(device_id)
        func(**kwargs)
        return
    wrapper_spawn_func_args = (
        world_size,
        backend,
        master_addr,
        master_port,
        device_type,
        func,
        kwargs,
    )
    mp.spawn(
        wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
    )
@@ -0,0 +1,36 @@
 from astrai.preprocessing.builder import (
    BaseMaskBuilder,
    MaskBuilderFactory,
    MultiOutputMaskBuilder,
    SectionedMaskBuilder,
    SingleOutputMaskBuilder,
 )
 from astrai.preprocessing.packing import (
    PackingStrategy,
    PackingStrategyFactory,
 )
 from astrai.preprocessing.pipeline import Pipeline, filter_by_length
 from astrai.preprocessing.position_id import (
    PositionIdStrategy,
    PositionIdStrategyFactory,
 )
 from astrai.preprocessing.writer import (
    StoreWriter,
    StoreWriterFactory,
 )
 __all__ = [
    "BaseMaskBuilder",
    "MaskBuilderFactory",
    "MultiOutputMaskBuilder",
    "PackingStrategy",
    "PackingStrategyFactory",
    "Pipeline",
    "PositionIdStrategy",
    "PositionIdStrategyFactory",
    "SectionedMaskBuilder",
    "SingleOutputMaskBuilder",
    "StoreWriter",
    "StoreWriterFactory",
    "filter_by_length",
 ]
@@ -0,0 +1,324 @@
 """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):
        all_ids: list[int] = []
        loss_mask: list[int] = []
        for sec in sections:
            field = sec["field"]
            action = sec["action"]
            use_template = sec.get("template", False)
            values = item.get(field)
            if not isinstance(values, list):
                continue
            for val in values:
                if use_template:
                    if isinstance(val, list):
                        wrapper = {field: val}
                        self._append_template(
                            wrapper,
                            field,
                            action,
                            tokenizer,
                            config,
                            all_ids,
                            loss_mask,
                        )
                else:
                    wrapper = {field: str(val)}
                    self._append_text(
                        wrapper,
                        field,
                        action,
                        tokenizer,
                        False,
                        False,
                        config,
                        all_ids,
                        loss_mask,
                    )
        max_len = config.preprocessing.max_seq_len
        all_ids = all_ids[:max_len]
        loss_mask = loss_mask[: len(all_ids)]
        if not all_ids:
            return None, None
        return all_ids, loss_mask
    @staticmethod
    def is_value_section(sections: list) -> bool:
        return len(sections) == 1 and sections[0].get("action") == "value"
    @staticmethod
    def extract_raw_value(item: dict, sections: list):
        sec = sections[0]
        field = sec["field"]
        raw = item.get(field)
        if raw is None:
            return None
        if isinstance(raw, list):
            return [float(v) for v in raw]
        return [float(raw)]
    def _append_template(
        self, item, field, action, tokenizer, config, all_ids, loss_mask
    ):
        messages = item.get(field)
        if not isinstance(messages, list) or not messages:
            return False
        for msg in messages:
            role = msg.get("role", "")
            act = _resolve_action(action, role, config)
            rendered = tokenizer.apply_chat_template(
                [msg], tokenize=False, add_generation_prompt=False
            )
            ids = tokenizer.encode(rendered, add_special_tokens=False)
            all_ids.extend(ids)
            val = 1 if act == "train" else 0
            loss_mask.extend([val] * len(ids))
        return True
    def _append_text(
        self,
        item,
        field,
        action,
        tokenizer,
        add_special,
        is_text_config,
        config,
        all_ids,
        loss_mask,
    ):
        text = str(item.get(field, ""))
        if not text.strip():
            return False
        if is_text_config:
            pp = config.preprocessing
            if pp.min_chars > 0 and len(text) < pp.min_chars:
                return False
            if len(text) > pp.max_chars:
                return False
        ids = tokenizer.encode(text, add_special_tokens=add_special)
        all_ids.extend(ids)
        val = 1 if action == "train" else 0
        loss_mask.extend([val] * len(ids))
        return True
 class BaseMaskBuilder(ABC):
    """Convert a JSONL item into token ids and optional loss_mask."""
    @abstractmethod
    def build(self, item: dict, config, tokenizer) -> Optional[dict]: ...
 class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
    pass
@MaskBuilderFactory.register("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
                )
            else:
                ids, mask = self.renderer.process_sections(
                    item, sections, config, tokenizer, is_top_level=True
                )
            if ids is None:
                continue
            result[output_key] = ids
            if not all(m == 1 for m in mask):
                result[mask_key] = mask
            elif "mask_key" in spec:
                result[mask_key] = mask
            any_output = True
        if not any_output:
            return None
        result["domain"] = _extract_domain(item, config.output.domain_key)
        return result
@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)
@@ -0,0 +1,121 @@
 """Sequence packing strategies for shard-level reordering and truncation.
 Each strategy receives the accumulated ``{key: [list of token lists]}``
 dict for a shard and returns a reordered / truncated version.  The
 pipeline later flattens the result into contiguous tensors.
 """
 from abc import ABC, abstractmethod
 from typing import Dict, List
 from astrai.factory import BaseFactory
 def _truncate(seq: List[int], max_len: int, mode: str) -> List[int]:
    if len(seq) <= max_len:
        return seq
    if mode == "keep_end":
        return seq[-max_len:]
    return seq[:max_len]
 class PackingStrategy(ABC):
    """Reorder and truncate sequences within a shard."""
    @abstractmethod
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
        max_packed_len: int,
        truncation_mode: str,
    ) -> Dict[str, List[List[int]]]:
        raise NotImplementedError
 class PackingStrategyFactory(BaseFactory["PackingStrategy"]):
    pass
@PackingStrategyFactory.register("simple")
 class SimplePacking(PackingStrategy):
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
        max_packed_len: int,
        truncation_mode: str,
    ) -> Dict[str, List[List[int]]]:
        return {
            k: [_truncate(v, max_packed_len, truncation_mode) for v in vals]
            for k, vals in keys.items()
        }
@PackingStrategyFactory.register("bfd")
 class BFDPacking(PackingStrategy):
    """Best-Fit Decreasing bin packing.
    Assigns sequences to bins using a best-fit heuristic (sorted by
    decreasing length) and concatenates sequences within each bin into
    a single packed sequence.  Packed sequences are truncated to
    *max_packed_len* so that each packed bin fits within one context
    window during training.
    """
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
        max_packed_len: int,
        truncation_mode: str,
    ) -> Dict[str, List[List[int]]]:
        sequences = keys.get("sequence", [])
        if not sequences:
            return keys
        bins = self._plan(sequences, max_packed_len, truncation_mode)
        packed: Dict[str, List[List[int]]] = {}
        for k, vals in keys.items():
            packed[k] = [
                _truncate(
                    self._concat_bin(vals, bin_indices),
                    max_packed_len,
                    truncation_mode,
                )
                for bin_indices in bins
            ]
        return packed
    @staticmethod
    def _concat_bin(vals: List[List[int]], indices: List[int]) -> List[int]:
        result: List[int] = []
        for i in indices:
            result.extend(vals[i])
        return result
    @staticmethod
    def _plan(
        sequences: List[List[int]], max_packed_len: int, truncation_mode: str
    ) -> List[List[int]]:
        n = len(sequences)
        order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
        bins: List[List[int]] = []
        bin_lengths: List[int] = []
        for orig_idx in order:
            seq_len = len(
                _truncate(sequences[orig_idx], max_packed_len, truncation_mode)
            )
            best_bin = None
            best_remain = max_packed_len + 1
            for i, bl in enumerate(bin_lengths):
                remain = max_packed_len - bl
                if seq_len <= remain < best_remain:
                    best_remain = remain
                    best_bin = i
            if best_bin is not None:
                bins[best_bin].append(orig_idx)
                bin_lengths[best_bin] += seq_len
            else:
                bins.append([orig_idx])
                bin_lengths.append(seq_len)
        return bins
@@ -0,0 +1,192 @@
 """Config-driven JSONL preprocessing pipeline.
 Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
 sharding and flush to ``.h5`` / ``.bin`` storage.  Packing, position-id
 generation and storage writing are each delegated to pluggable strategies,
 dispatched by configuration keys.
 """
 import json
 import logging
 import os
 from collections import defaultdict
 from itertools import chain
 from typing import Dict, List, Optional
 import torch
 import tqdm
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.builder import MaskBuilderFactory
 from astrai.preprocessing.packing import PackingStrategyFactory
 from astrai.preprocessing.position_id import PositionIdStrategyFactory
 from astrai.preprocessing.writer import StoreWriterFactory
 from astrai.tokenize import AutoTokenizer
 logger = logging.getLogger(__name__)
 _STR_TO_DTYPE: dict[str, torch.dtype] = {
    "bool": torch.bool,
    "uint8": torch.uint8,
    "int8": torch.int8,
    "int16": torch.int16,
    "int32": torch.int32,
    "int64": torch.int64,
    "float16": torch.float16,
    "float32": torch.float32,
    "float64": torch.float64,
 }
 def filter_by_length(text: str, min_len: int = 50, max_len: int = 2_000_000) -> bool:
    return min_len <= len(text) <= max_len
 class Pipeline:
    """Tokenization pipeline driven by a declarative :class:`PipelineConfig`.
    Usage::
        config = PipelineConfig.from_file("sft_pipeline.json")
        Pipeline(config, ["data.jsonl"], output_dir="out", tokenizer_path="params").run()
    """
    def __init__(
        self,
        config: PipelineConfig,
        input_paths: list[str],
        output_dir: str,
        tokenizer_path: str,
    ):
        os.makedirs(output_dir, exist_ok=True)
        self.config = config
        self.paths = input_paths
        self.output_dir = output_dir
        self.tokenizer_path = tokenizer_path
        self.mask_builder = MaskBuilderFactory.create("sectioned")
        self._packer = PackingStrategyFactory.create(
            config.preprocessing.packing_strategy
        )
        self._position_id = PositionIdStrategyFactory.create(
            config.output.position_ids_mode
        )
        self._writer = StoreWriterFactory.create(config.output.storage_format)
    def transform(self, item: dict) -> Optional[dict]:
        return self.mask_builder.build(item, self.config, self._tokenizer)
    def run(self):
        self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path)
        domains: dict = defaultdict(lambda: defaultdict(list))
        total_tokens = 0
        shard_idx: dict[str, int] = defaultdict(int)
        count = 0
        pp = self.config.preprocessing
        for item in tqdm.tqdm(
            self._iter_items(), desc="Tokenizing", unit="docs", mininterval=0.5
        ):
            if pp.max_items and count >= pp.max_items:
                break
            try:
                result = self.transform(item)
            except Exception:
                logger.warning(
                    "Failed to process item #%d, skipping", count + 1, exc_info=True
                )
                continue
            if result is None:
                continue
            domain = result.pop("domain", "__default__")
            is_multi = bool(getattr(self.config.input, "sources", None))
            if is_multi:
                ids = self._primary_ids(result)
            else:
                ids = result.pop("sequence")
                result["sequence"] = ids
            if not ids:
                continue
            bucket = domains[domain]
            self._align_bucket(bucket, result, ids)
            for key, val in result.items():
                bucket[key].append(val)
            count += 1
            total_tokens += len(ids)
            if total_tokens >= self.config.output.max_tokens_per_shard:
                self._flush(domains, shard_idx)
                domains.clear()
                total_tokens = 0
        if total_tokens > 0:
            self._flush(domains, shard_idx)
    @staticmethod
    def _primary_ids(result: dict) -> list:
        """Return the first list-valued entry in *result* as the primary id
        sequence for token counting."""
        for val in result.values():
            if isinstance(val, list) and val and isinstance(val[0], int):
                return val
        return []
    @staticmethod
    def _align_bucket(bucket: dict, result: dict, ids: list):
        """Pad previously-accumulated keys that are missing from *result*."""
        for key in list(bucket.keys()):
            if key in result:
                continue
            bucket[key].append([0] * len(ids))
    def _iter_items(self):
        for path in self.paths:
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    yield json.loads(line)
    def _flush(self, domains, shard_idx):
        for domain, keys in domains.items():
            idx = shard_idx[domain]
            pp = self.config.preprocessing
            original_sequences = keys.get("sequence", [])
            mode = self.config.output.position_ids_mode
            if mode == "doc_reset" and original_sequences:
                keys["position_ids"] = [list(range(len(s))) for s in original_sequences]
            keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode)
            tensors: Dict[str, List[torch.Tensor]] = {}
            for key, ids_list in keys.items():
                dt = _STR_TO_DTYPE.get(
                    self.config.output.dtype.get(key, "int32"), torch.int32
                )
                tensors[key] = [
                    torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
                ]
            if mode == "continuous" and original_sequences:
                pos_ids = self._position_id.generate(keys.get("sequence", []))
                if pos_ids:
                    tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
            self._writer.save(self.output_dir, domain, idx, tensors)
            shard_idx[domain] = idx + 1
            first_key = "sequence" if "sequence" in tensors else next(iter(tensors))
            tqdm.tqdm.write(
                f"  saved {domain}/shard_{idx:04d}  "
                f"({tensors[first_key][0].numel():,} tokens)"
            )
@@ -0,0 +1,46 @@
 """Position-id generation strategies for packed sequences.
 Each strategy takes the list of per-document token sequences after packing
 and returns a flat list of position ids (same total length as all
 sequences combined).  The pipeline wraps the result into a tensor and
 attaches it as ``position_ids``.
 """
 from abc import ABC, abstractmethod
 from typing import List
 from astrai.factory import BaseFactory
 class PositionIdStrategy(ABC):
    """Generate ``position_ids`` for packed sequences."""
    @abstractmethod
    def generate(self, sequences: List[List[int]]) -> List[int]:
        raise NotImplementedError
 class PositionIdStrategyFactory(BaseFactory["PositionIdStrategy"]):
    pass
@PositionIdStrategyFactory.register("none")
 class NoPositionId(PositionIdStrategy):
    def generate(self, sequences: List[List[int]]) -> List[int]:
        return []
@PositionIdStrategyFactory.register("doc_reset")
 class DocResetPositionId(PositionIdStrategy):
    def generate(self, sequences: List[List[int]]) -> List[int]:
        pos_ids = []
        for seq in sequences:
            pos_ids.extend(range(len(seq)))
        return pos_ids
@PositionIdStrategyFactory.register("continuous")
 class ContinuousPositionId(PositionIdStrategy):
    def generate(self, sequences: List[List[int]]) -> List[int]:
        total = sum(len(seq) for seq in sequences)
        return list(range(total))
@@ -0,0 +1,75 @@
 """Storage writer strategies for pipeline output.
 The :class:`StoreWriter` abstraction decouples the pipeline from the
 concrete storage format (bin / h5).  The pipeline builds a ``{key:
 List[Tensor]}`` dict and delegates the write to the writer selected
 by ``output.storage_format``.
 """
 import logging
 import os
 import shutil
 from abc import ABC, abstractmethod
 from typing import Dict, List
 import torch
 from astrai.factory import BaseFactory
 from astrai.serialization import save_bin, save_h5
 logger = logging.getLogger(__name__)
 class StoreWriter(ABC):
    """Write pre-tokenized tensors to disk in a format-specific way."""
    @abstractmethod
    def save(
        self,
        output_dir: str,
        domain: str,
        shard_idx: int,
        tensors: Dict[str, List[torch.Tensor]],
    ) -> None: ...
 class StoreWriterFactory(BaseFactory["StoreWriter"]):
    pass
@StoreWriterFactory.register("bin")
 class BinWriter(StoreWriter):
    def save(self, output_dir, domain, shard_idx, tensors):
        shard_path = os.path.join(output_dir, domain, f"shard_{shard_idx:04d}")
        try:
            save_bin(shard_path, tensors)
        except Exception:
            if os.path.exists(shard_path):
                shutil.rmtree(shard_path, ignore_errors=True)
            logger.error(
                "Failed to write shard %s/%s_%04d, cleaned up partial output",
                domain,
                "shard",
                shard_idx,
                exc_info=True,
            )
            raise
@StoreWriterFactory.register("h5")
 class H5Writer(StoreWriter):
    def save(self, output_dir, domain, shard_idx, tensors):
        chunk_dir = os.path.join(output_dir, domain)
        file_path = os.path.join(chunk_dir, f"data_{shard_idx:04d}.h5")
        try:
            save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
        except Exception:
            if os.path.exists(file_path):
                os.remove(file_path)
            logger.error(
                "Failed to write shard %s/data_%04d.h5, cleaned up partial output",
                domain,
                shard_idx,
                exc_info=True,
            )
            raise
@@ -0,0 +1,21 @@
 """Training component protocols — structural subtyping for optimizer/scheduler wrappers."""
 from typing import Any, Protocol, runtime_checkable
@runtime_checkable
 class OptimizerProtocol(Protocol):
    def step(self, closure=None): ...
    def zero_grad(self): ...
    @property
    def param_groups(self) -> Any: ...
    def state_dict(self) -> dict: ...
    def load_state_dict(self, d: dict): ...
@runtime_checkable
 class SchedulerProtocol(Protocol):
    def step(self): ...
    def state_dict(self) -> dict: ...
    def load_state_dict(self, d: dict): ...
    def get_last_lr(self): ...
@@ -1,77 +0,0 @@
 import json
 from pathlib import Path
 from typing import Any, Dict, Optional
 import safetensors.torch as st
 import torch
 import torch.distributed as dist
 from astrai.parallel.setup import get_rank
 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,
        )
@@ -0,0 +1,43 @@
 """Serialization utilities for models and datasets.
 This package re-exports checkpoint helpers and dataset storage helpers so
 that existing imports from ``astrai.serialization`` continue to work.
 """
 from astrai.serialization.checkpoint import (
    Checkpoint,
    load_json,
    load_model_config,
    load_model_weights,
    load_safetensors,
    load_state_dict,
    load_torch,
    save_json,
    save_model,
    save_safetensors,
    save_torch,
 )
 from astrai.serialization.dataset import (
    load_bin,
    load_h5,
    save_bin,
    save_h5,
 )
 __all__ = [
    "Checkpoint",
    "load_json",
    "load_model_config",
    "load_model_weights",
    "load_safetensors",
    "load_state_dict",
    "load_torch",
    "save_json",
    "save_model",
    "save_safetensors",
    "save_torch",
    "load_bin",
    "load_h5",
    "save_bin",
    "save_h5",
 ]
@@ -0,0 +1,204 @@
 """Model checkpoint serialization helpers."""
 import io
 import json
 import os
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 import safetensors.torch as st
 import torch
 import torch.distributed as dist
 from astrai.parallel.setup import get_rank
 _META_FILE = "meta.json"
 _CONFIG_FILE = "config.json"
 _WEIGHTS_FILE = "model.safetensors"
 def save_safetensors(state_dict: dict, path: Union[str, Path]):
    st.save_file(state_dict, str(path))
 def load_safetensors(path: Union[str, Path], broadcast: bool = False) -> dict:
    if not broadcast or not dist.is_initialized():
        return st.load_file(str(path))
    rank = get_rank()
    if rank == 0:
        state_dict = st.load_file(str(path))
    else:
        state_dict = {}
    tmp = [state_dict]
    dist.broadcast_object_list(tmp, src=0)
    return tmp[0]
 def save_json(data: dict, path: Union[str, Path]):
    with open(str(path), "w") as f:
        json.dump(data, f, indent=2)
 def load_json(path: Union[str, Path], broadcast: bool = False) -> dict:
    if not broadcast or not dist.is_initialized():
        with open(str(path), "r") as f:
            return json.load(f)
    rank = get_rank()
    if rank == 0:
        with open(str(path), "r") as f:
            data = json.load(f)
    else:
        data = {}
    tmp = [data]
    dist.broadcast_object_list(tmp, src=0)
    return tmp[0]
 def save_torch(obj: Any, path: Union[str, Path]):
    torch.save(obj, str(path))
 def load_torch(path: Union[str, Path], broadcast: bool = False) -> Any:
    if not broadcast or not dist.is_initialized():
        return torch.load(str(path), map_location="cpu", weights_only=False)
    path = Path(path)
    rank = get_rank()
    if rank == 0:
        with open(path, "rb") as f:
            raw = f.read()
        data_tensor = torch.frombuffer(bytearray(raw), dtype=torch.uint8)
        num_bytes = torch.tensor([len(raw)], dtype=torch.long)
    else:
        num_bytes = torch.tensor([0], dtype=torch.long)
    dist.broadcast(num_bytes, src=0)
    if rank != 0:
        data_tensor = torch.empty(num_bytes.item(), dtype=torch.uint8)
    dist.broadcast(data_tensor, src=0)
    buf = io.BytesIO(data_tensor.numpy().tobytes())
    return torch.load(buf, map_location="cpu", weights_only=False)
 def save_model(config: dict, state_dict: dict, save_directory: str):
    save_path = Path(save_directory)
    save_path.mkdir(parents=True, exist_ok=True)
    save_json(config, save_path / _CONFIG_FILE)
    save_safetensors(state_dict, save_path / _WEIGHTS_FILE)
 def load_model_config(save_directory: str) -> dict:
    return load_json(Path(save_directory) / _CONFIG_FILE)
 def load_model_weights(save_directory: str) -> dict:
    return load_state_dict(Path(save_directory) / _WEIGHTS_FILE)
 def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict:
    path = Path(path)
    if not broadcast or not dist.is_initialized():
        return load_safetensors(path)
    rank = get_rank()
    if rank == 0:
        state_dict = load_safetensors(path)
        specs = [
            (k, list(state_dict[k].shape), str(state_dict[k].dtype).split(".")[-1])
            for k in sorted(state_dict)
        ]
    else:
        state_dict = {}
        specs = []
    specs_list = [specs]
    dist.broadcast_object_list(specs_list, src=0)
    specs = specs_list[0]
    for key, shape, dtype_name in specs:
        dtype = getattr(torch, dtype_name)
        if rank != 0:
            tensor = torch.empty(shape, dtype=dtype, device="cpu")
        else:
            tensor = state_dict[key].contiguous().cpu()
        dist.broadcast(tensor, src=0)
        if rank != 0:
            state_dict[key] = tensor
    return state_dict
@dataclass
 class Checkpoint:
    state_dict: Dict[str, Any] = field(default_factory=dict)
    epoch: int = 0
    consumed_samples: int = 0
    extra: Dict[str, Any] = field(default_factory=dict)
    meta: Dict[str, Any] = field(default_factory=dict)
    config: Dict[str, Any] = field(default_factory=dict)
    def save(self, save_dir: str):
        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)
        if get_rank() != 0:
            return
        meta = {
            "epoch": self.epoch,
            "consumed_samples": self.consumed_samples,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            **self.meta,
        }
        save_json(meta, save_path / _META_FILE)
        save_json(self.config, save_path / _CONFIG_FILE)
        save_safetensors(self.state_dict, save_path / _WEIGHTS_FILE)
        for key, value in self.extra.items():
            save_torch(value, save_path / f"{key}.pt")
    @classmethod
    def load(cls, save_dir: str, broadcast: bool = False) -> "Checkpoint":
        save_path = Path(save_dir)
        meta = load_json(save_path / _META_FILE, broadcast)
        config = load_json(save_path / _CONFIG_FILE, broadcast)
        state_dict = load_state_dict(save_path / _WEIGHTS_FILE, broadcast=broadcast)
        extra = {}
        for f in sorted(save_path.iterdir()):
            if f.suffix == ".pt":
                extra[f.stem] = load_torch(f, broadcast=broadcast)
        return cls(
            state_dict=state_dict,
            epoch=meta.get("epoch", 0),
            consumed_samples=meta.get("consumed_samples", 0),
            extra=extra,
            config=config,
        )
    @classmethod
    def load_any(cls, save_dir: str, broadcast: bool = False) -> Optional["Checkpoint"]:
        save_path = Path(save_dir)
        meta_path = save_path / _META_FILE
        weights_path = save_path / _WEIGHTS_FILE
        if meta_path.exists():
            return cls.load(save_dir, broadcast=broadcast)
        if weights_path.exists():
            state_dict = load_state_dict(weights_path, broadcast=broadcast)
            config = {}
            config_path = save_path / _CONFIG_FILE
            if config_path.exists():
                config = load_json(config_path, broadcast)
            return cls(state_dict=state_dict, config=config)
        return None
@@ -0,0 +1,76 @@
 """Dataset storage serialization helpers (HDF5 / memory-mapped binary)."""
 import json
 import os
 from pathlib import Path
 from typing import Dict, List
 import h5py
 import numpy as np
 import torch
 from torch import Tensor
 def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
    os.makedirs(file_path, exist_ok=True)
    full_file_path = os.path.join(file_path, f"{file_name}.h5")
    with h5py.File(full_file_path, "w") as f:
        for key, tensors in tensor_group.items():
            grp = f.create_group(key)
            for idx, tensor in enumerate(tensors):
                arr = tensor.cpu().numpy()
                grp.create_dataset(f"data_{idx}", data=arr)
 def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
    tensor_group: Dict[str, List[Tensor]] = {}
    root_path = Path(file_path)
    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]]):
    os.makedirs(file_path, exist_ok=True)
    meta = {}
    for key, tensors in tensor_group.items():
        cat = torch.cat(tensors, dim=0)
        meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]}
        np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
    with open(os.path.join(file_path, "meta.json"), "w") as f:
        json.dump(meta, f)
 def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
    with open(os.path.join(file_path, "meta.json"), "r") as f:
        meta = json.load(f)
    segments: Dict[str, List[Tensor]] = {}
    for key, info in meta.items():
        arr = np.memmap(
            os.path.join(file_path, f"{key}.bin"),
            dtype=info["dtype"],
            mode="r+",
            shape=tuple(info["shape"]),
        )
        segments[key] = [torch.from_numpy(arr)]
    return segments
@@ -1,13 +1,10 @@
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 from jinja2 import Template
 # Message type for chat messages
 type MessageType = Dict[str, Any]
@dataclass
 class ChatTemplate:
    """A chat template with Jinja2 rendering support.
@@ -15,23 +12,24 @@ class ChatTemplate:
        name: Unique identifier for the template.
        template_str: Jinja2 template string.
        description: Optional description.
-        default_variables: Optional dictionary of default variable values
+        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.
    """
-    name: str
+    def __init__(
-    template_str: str
+        self,
-    description: str = ""
+        name: str = "",
-    default_variables: Dict[str, Any] = None
+        template_str: str = "",
-    special_tokens: Dict[str, str] = None
+        description: str = "",
-
+        default_variables: Optional[Dict[str, Any]] = None,
-    def __post_init__(self):
+        special_tokens: Optional[Dict[str, str]] = None,
-        if self.default_variables is None:
+    ):
-            self.default_variables = {}
+        self.name = name
-        if self.special_tokens is None:
+        self.template_str = template_str
-            self.special_tokens = {}
+        self.description = description
        self.default_variables = default_variables or {}
        self.special_tokens = special_tokens or {}
        self._compiled: Template = Template(template_str)
    @classmethod
    def from_string(
@@ -43,7 +41,7 @@ class ChatTemplate:
    ) -> "ChatTemplate":
        """Create a ChatTemplate instance directly from a template string."""
        return cls(
-            name="",  # empty name for ad‑hoc templates
+            name="",
            template_str=template_str,
            description=description,
            default_variables=default_variables,
@@ -73,5 +71,4 @@ class ChatTemplate:
        if system_prompt is not None:
            variables["system_prompt"] = system_prompt
-        jinja_template = Template(self.template_str)
+        return self._compiled.render(**variables)
        return jinja_template.render(**variables)
@@ -51,9 +51,26 @@ class AutoTokenizer:
                self.set_chat_template(config["chat_template"])
    @classmethod
-    def from_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoTokenizer":
+    def from_pretrained(cls, path: Union[str, Path]) -> "AutoTokenizer":
-        """Load tokenizer from pretrained directory."""
+        """Load tokenizer from pretrained directory.
        Raises:
            FileNotFoundError: If tokenizer.json is missing.
            RuntimeError: If tokenizer failed to initialize.
        """
        path = Path(path)
        tokenizer_file = path / "tokenizer.json"
        if not tokenizer_file.exists():
            raise FileNotFoundError(
                f"Tokenizer file not found: {tokenizer_file}. "
                "A valid tokenizer.json is required."
            )
        instance = cls(path)
        if instance._tokenizer is None:
            raise RuntimeError(
                f"Failed to load tokenizer from {path}. "
                "The tokenizer.json may be corrupted or incompatible."
            )
        return instance
    def save_pretrained(self, save_path: str):
@@ -1,42 +1,25 @@
-from typing import Any, Callable, Dict
+from typing import Dict
 import torch
 import torch.nn as nn
-def _grad_stat(
+def grad_norm(model: nn.Module, per_param: bool = False) -> float | Dict[str, float]:
-    model: nn.Module, fn: Callable[[torch.Tensor], Any], default: Any
+    grads = [p.grad.detach() for p in model.parameters() if p.grad is not None]
-) -> dict:
+    if not grads:
-    results = {}
+        return 0.0
    for name, param in model.named_parameters():
        results[name] = default
        if param.grad is not None:
            results[name] = fn(param.grad.data)
    return results
-
+    total_sq = torch.stack([g.pow(2).sum() for g in grads]).sum()
-def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+    if per_param:
-    return _grad_stat(model, lambda g: g.norm(norm_type).item(), 0.0)
+        norms = {}
-
+        for name, param in model.named_parameters():
-
+            if param.grad is not None:
-def grad_std(model: nn.Module) -> Dict[str, float]:
+                norms[name] = param.grad.norm(2).item()
-    return _grad_stat(model, lambda g: g.std().item(), 0.0)
+            else:
-
+                norms[name] = 0.0
-
+        norms["total"] = total_sq.sqrt().item()
-def grad_max(model: nn.Module) -> Dict[str, float]:
+        return norms
-    return _grad_stat(model, lambda g: g.max().item(), -float("inf"))
+    return total_sq.sqrt().item()
 def grad_min(model: nn.Module) -> Dict[str, float]:
    return _grad_stat(model, lambda g: g.min().item(), float("inf"))
 def grad_mean(model: nn.Module) -> Dict[str, float]:
    return _grad_stat(model, lambda g: g.mean().item(), 0.0)
 def grad_nan_num(model: nn.Module) -> Dict[str, int]:
    return _grad_stat(model, lambda g: g.isnan().sum().item(), 0)
 def ctx_get_loss(ctx):
@@ -47,25 +30,9 @@ def ctx_get_lr(ctx):
    return ctx.optimizer.param_groups[-1]["lr"]
 def ctx_get_val_loss(ctx):
    return ctx.val_loss
 def ctx_get_grad_norm(ctx):
-    return grad_norm(ctx.model)
+    return ctx.grad_norm
 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, Type
+from typing import Any, Dict, List
 from torch.optim.lr_scheduler import LRScheduler
@@ -31,7 +31,6 @@ class SchedulerFactory(BaseFactory["BaseScheduler"]):
    """Factory class for creating learning rate schedulers.
    Supports decorator-based registration for extensible scheduler types.
    Also supports creation from ScheduleConfig objects.
    Example usage:
        @SchedulerFactory.register("custom")
@@ -41,33 +40,6 @@ class SchedulerFactory(BaseFactory["BaseScheduler"]):
        scheduler = SchedulerFactory.create("custom", optimizer, **kwargs)
    """
    @classmethod
    def _validate_component(cls, scheduler_cls: Type[BaseScheduler]) -> None:
        """Validate that the scheduler class inherits from BaseScheduler."""
        if not issubclass(scheduler_cls, BaseScheduler):
            raise TypeError(f"{scheduler_cls.__name__} must inherit from BaseScheduler")
    @classmethod
    def create(
        cls, optimizer, schedule_type: str = "none", **kwargs
    ) -> "BaseScheduler":
        """Create a scheduler instance by type name.
        Args:
            optimizer: PyTorch optimizer
            schedule_type: Type of scheduler ("cosine", "sgdr")
            **kwargs: Arguments passed to the scheduler constructor
        Returns:
            Scheduler instance
        """
        return super().create(schedule_type, optimizer, **kwargs)
    @classmethod
    def available_types(cls) -> list:
        """Return list of registered scheduler type names."""
        return cls.list_registered()
 # ----------- Scheduler implementations -----------
@@ -81,7 +53,7 @@ class CosineScheduler(BaseScheduler):
        optimizer,
        warmup_steps: int,
        lr_decay_steps: int,
-        min_rate: float = 0.05,
+        min_rate: float = 0.01,
        last_epoch: int = -1,
    ):
        self.warmup_steps = warmup_steps
@@ -93,11 +65,15 @@ class CosineScheduler(BaseScheduler):
    def get_lr(self) -> List[float]:
        # warmup
        if self.last_epoch < self.warmup_steps:
-            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
+            warmup_factor = max(
                self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
            )
            return [base_lr * warmup_factor for base_lr in self.base_lrs]
        # cosine decay
-        decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
+        decay_progress = (self.last_epoch - self.warmup_steps) / max(
            self.lr_decay_steps, 1
        )
        decay_progress = min(decay_progress, 1.0)
        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
        decay_factor = max(self.min_rate, cosine_decay)
@@ -132,7 +108,7 @@ class SGDRScheduler(BaseScheduler):
        optimizer,
        warmup_steps: int,
        cycle_length: int,
-        min_rate: float = 0.05,
+        min_rate: float = 0.01,
        t_mult: int = 2,
        last_epoch: int = -1,
    ):
@@ -146,7 +122,9 @@ class SGDRScheduler(BaseScheduler):
    def get_lr(self):
        # warmup
        if self.last_epoch < self.warmup_steps:
-            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
+            warmup_factor = max(
                self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
            )
            return [base_lr * warmup_factor for base_lr in self.base_lrs]
        # SGDR
@@ -192,3 +170,66 @@ class SGDRScheduler(BaseScheduler):
        self.min_rate = state_dict.pop("min_rate")
        self.t_mult = state_dict.pop("t_mult")
        super().load_state_dict(state_dict)
@SchedulerFactory.register("wsd")
 class WSDScheduler(BaseScheduler):
    """WSD (Warmup-Stable-Decay) scheduler with sqrt cooldown.
    warmup_steps:  linear warmup from min_rate to 1.0
    stable_steps:   constant at base_lr
    decay_steps:    sqrt decay from base_lr to min_rate
    min_rate:       minimum lr as fraction of base_lr (default 0.0)
    """
    def __init__(
        self,
        optimizer,
        warmup_steps: int,
        stable_steps: int,
        decay_steps: int,
        min_rate: float = 0.01,
        last_epoch: int = -1,
    ):
        self.warmup_steps = warmup_steps
        self.stable_steps = stable_steps
        self.decay_steps = decay_steps
        self.min_rate = min_rate
        self.total_steps = warmup_steps + stable_steps + decay_steps
        super().__init__(optimizer, last_epoch)
    def get_lr(self) -> List[float]:
        if self.last_epoch < self.warmup_steps:
            factor = max(self.min_rate, self.last_epoch / max(self.warmup_steps, 1))
            return [base_lr * factor for base_lr in self.base_lrs]
        offset = self.last_epoch - self.warmup_steps
        if offset < self.stable_steps:
            return list(self.base_lrs)
        decay_ratio = (offset - self.stable_steps) / max(self.decay_steps, 1)
        decay_ratio = min(decay_ratio, 1.0)
        factor = (1.0 - self.min_rate) * (1.0 - decay_ratio) ** 2 + self.min_rate
        return [base_lr * factor for base_lr in self.base_lrs]
    def state_dict(self):
        state = super().state_dict()
        state.update(
            {
                "warmup_steps": self.warmup_steps,
                "stable_steps": self.stable_steps,
                "decay_steps": self.decay_steps,
                "min_rate": self.min_rate,
                "total_steps": self.total_steps,
            }
        )
        return state
    def load_state_dict(self, state_dict):
        self.warmup_steps = state_dict.pop("warmup_steps")
        self.stable_steps = state_dict.pop("stable_steps")
        self.decay_steps = state_dict.pop("decay_steps")
        self.min_rate = state_dict.pop("min_rate")
        self.total_steps = state_dict.pop("total_steps")
        super().load_state_dict(state_dict)
@@ -1,39 +1,28 @@
 """Training strategy implementations with factory pattern."""
 import copy
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Union
+from typing import Callable, Dict, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.parallel import DistributedDataParallel as DDP
 from astrai.factory import BaseFactory
-def unwrap_model(model: nn.Module) -> nn.Module:
+def create_ref_model(
-    """Unwrap DDP wrapper if present to get the original model."""
+    model_fn: Callable[[], nn.Module], state_dict: Dict[str, Tensor]
-    if isinstance(model, DDP):
+) -> nn.Module:
-        return model.module
+    """Create a frozen reference model from model_fn + full state dict."""
-    return model
+    ref_model = model_fn()
-
+    ref_model.load_state_dict(state_dict)
 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) -> Any:
+def move_to_device(batch: Dict[str, Tensor], device: str) -> Dict[str, Tensor]:
    """Move batch tensors to specified device with non-blocking transfer."""
    return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
@@ -43,7 +32,7 @@ def get_logprobs(
    input_ids: Tensor,
    mask: Tensor,
    reduction: str,
-):
+) -> Tensor:
    """Compute token-wise log probabilities from model outputs.
    Args:
@@ -81,14 +70,35 @@ def get_logprobs(
        return token_logprobs * shifted_mask
 def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
    S = position_ids.size(1)
    device = position_ids.device
    boundaries = position_ids[:, 1:] <= position_ids[:, :-1]
    doc_ids = torch.cat(
        [
            torch.zeros(position_ids.size(0), 1, dtype=torch.long, device=device),
            boundaries.long().cumsum(dim=1),
        ],
        dim=1,
    )
    same_doc = doc_ids.unsqueeze(-1) == doc_ids.unsqueeze(-2)
    causal = torch.tril(torch.ones(S, S, dtype=torch.bool, device=device))
    return (same_doc & causal).unsqueeze(1)
 class BaseStrategy(ABC):
    """Abstract base class for training strategies."""
    def __init__(
-        self, model: Union[Callable[..., Dict[str, Tensor]]], device: str, **kwargs
+        self,
        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
        device: str,
        **kwargs,
    ):
        self.model = model
        self.device = device
        self.executor = kwargs.pop("executor", None)
        self.model_fn = kwargs.pop("model_fn", None)
        self.extra_kwargs = kwargs
    @abstractmethod
@@ -122,32 +132,6 @@ class StrategyFactory(BaseFactory["BaseStrategy"]):
        strategy = StrategyFactory.create("custom", model, device)
    """
    @classmethod
    def _validate_component(cls, strategy_cls: type) -> None:
        """Validate that the strategy class inherits from BaseStrategy."""
        if not issubclass(strategy_cls, BaseStrategy):
            raise TypeError(f"{strategy_cls.__name__} must inherit from BaseStrategy")
    @classmethod
    def create(cls, train_type: str, model, device: str, **kwargs) -> "BaseStrategy":
        """Create a strategy instance based on training type.
        Args:
            train_type: Type of training ("seq", "sft", "dpo", "grpo")
            model: Model instance for the strategy
            device: Device to run the strategy on
            **kwargs: Additional arguments passed to strategy constructor
        Returns:
            Strategy instance
        """
        return super().create(train_type, model, device, **kwargs)
    @classmethod
    def available_strategies(cls) -> list:
        """Return list of registered strategy names."""
        return cls.list_registered()
 # ============== Strategy Classes ==============
 # All strategies are registered at class definition time using the decorator
@@ -160,7 +144,13 @@ class SEQStrategy(BaseStrategy):
    Computes cross-entropy loss for next token prediction.
    """
-    def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
+    def __init__(
        self,
        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
        device: str,
        label_smoothing: float = 0.0,
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
        self.label_smoothing = label_smoothing
@@ -185,21 +175,31 @@ class SFTStrategy(BaseStrategy):
    Applies cross-entropy loss only to tokens where loss_mask is True.
    """
-    def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
+    def __init__(
        self,
        model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
        device: str,
        label_smoothing: float = 0.0,
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
        self.label_smoothing = label_smoothing
    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
        batch = move_to_device(batch, self.device)
-        input_ids, target_ids, loss_mask = (
+        input_ids, target_ids, position_ids, loss_mask = (
            batch["input_ids"],
            batch["target_ids"],
            batch["position_ids"],
            batch["loss_mask"],
        )
        ignore_index = -100
-        logits = self.model(input_ids=input_ids)["logits"]
+        input_mask = make_doc_boundary_mask(position_ids)
-        target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
+        target_ids = target_ids.masked_fill(~loss_mask, ignore_index)
        logits = self.model(
            input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
        )["logits"]
        loss = F.cross_entropy(
            input=logits.flatten(0, 1).float(),
@@ -228,7 +228,9 @@ class DPOStrategy(BaseStrategy):
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(model)
+        self.ref_model = create_ref_model(
            self.model_fn, self.executor.unwrap_model(model)
        ).to(device=self.device)
        self.beta = beta
        self.reduction = reduction
@@ -282,7 +284,9 @@ class GRPOStrategy(BaseStrategy):
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(model)
+        self.ref_model = create_ref_model(
            self.model_fn, self.executor.unwrap_model(model)
        ).to(device=self.device)
        self.clip_eps = clip_eps
        self.kl_coef = kl_coef
        self.group_size = group_size
@@ -292,8 +296,7 @@ class GRPOStrategy(BaseStrategy):
    def sync_ref_model(self):
        """Copy current model weights to ref model."""
-        ref_state = self.model.state_dict()
+        self.ref_model.load_state_dict(self.executor.unwrap_model(self.model))
        self.ref_model.load_state_dict(ref_state)
    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
        self._step += 1
@@ -1,28 +1,31 @@
 import json
 import logging
 import os
 import sys
 import time
 from pathlib import Path
-from typing import Callable, List, Optional, Protocol, runtime_checkable
+from typing import IO, Callable, List, Optional, Protocol, runtime_checkable
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from torch.nn.utils import clip_grad_norm_
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
 from tqdm import tqdm
 from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
 from astrai.parallel.setup import get_current_device, get_rank
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
    ctx_get_grad_max,
    ctx_get_grad_mean,
    ctx_get_grad_min,
    ctx_get_grad_nan_num,
    ctx_get_grad_norm,
    ctx_get_grad_std,
    ctx_get_loss,
    ctx_get_lr,
    ctx_get_val_loss,
 )
 from astrai.trainer.train_context import TrainContext
 logger = logging.getLogger(__name__)
@runtime_checkable
 class TrainCallback(Protocol):
@@ -42,18 +45,15 @@ class TrainCallback(Protocol):
    def on_epoch_end(self, context: TrainContext):
        """Called at the end of each epoch."""
    def on_step_begin(self, context: TrainContext):
        """Called at the beginning of each step."""
    def on_step_end(self, context: TrainContext):
        """Called at the end of each step."""
    def on_batch_begin(self, context: TrainContext):
        """Called at the beginning of each batch."""
    def on_batch_end(self, context: TrainContext):
        """Called at the end of each batch."""
    def on_optimizer_step(self, context: TrainContext):
        """Called on every optimizer step (sync step only)."""
    def on_error(self, context: TrainContext):
        """Called when an error occurs during training."""
@@ -79,9 +79,45 @@ class GradientClippingCallback(TrainCallback):
    def __init__(self, max_grad_norm: float):
        self.max_grad_norm = max_grad_norm
-    def on_step_end(self, context: TrainContext):
+    def on_optimizer_step(self, context: TrainContext):
-        _ = context
+        context.grad_norm = context.executor.clip_grad_norm(
-        clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
+            context.model, self.max_grad_norm
        )
@CallbackFactory.register("gradient_checkpointing")
 class GradientCheckpointingCallback(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.
    """
    def __init__(self, modules: Optional[List[type]] = None):
        self.modules = tuple(modules) if modules else ()
    def _enable(self, module: nn.Module):
        if self.modules and isinstance(module, self.modules):
            fn = module.forward
            module._original_forward = fn
            module.forward = lambda *a, **kw: torch_checkpoint(
                fn, *a, use_reentrant=False, **kw
            )
    @staticmethod
    def _disable(module: nn.Module):
        if hasattr(module, "_original_forward"):
            module.forward = module._original_forward
            del module._original_forward
    def on_train_begin(self, context: TrainContext):
        context.model.apply(self._enable)
        logger.info("Gradient checkpointing enabled")
    def on_train_end(self, context: TrainContext):
        context.model.apply(self._disable)
@CallbackFactory.register("checkpoint")
@@ -90,54 +126,62 @@ class CheckpointCallback(TrainCallback):
    Checkpoint callback for trainer.
    """
    extra_keys = ("optimizer", "scheduler")
    def __init__(
        self,
        save_dir: str,
        interval: int,
        weight_only: bool = False,
        state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
        save_extra_fn: Optional[Callable[["TrainContext"], dict]] = None,
    ):
        self.save_dir = save_dir
        self.interval = interval
        self.weight_only = weight_only
-        self.state_dict_fn = state_dict_fn
+        self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
-        self.save_extra_fn = save_extra_fn
+        self.last_ckpt_step = 0
        self.last_ckpt_iter = 0
    @only_on_rank(0)
    def _save_checkpoint(self, context: TrainContext):
-        save_path = os.path.join(
+        state_dict = context.executor.unwrap_model(context.model)
-            self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
+        self.last_ckpt_step = context.optimizer_step
        )
        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
+        if get_rank() == 0:
-        context.checkpoint = Checkpoint(
+            save_path = os.path.join(
-            state_dict=state_dict,
+                self.save_dir,
-            epoch=context.epoch,
+                f"epoch_{context.epoch}_step_{context.optimizer_step}",
-            iteration=context.iteration,
+            )
-            extra=extra,
+            extra = self.save_extra_fn(context)
-        )
+            meta = context.config.to_dict()
-
+            context.checkpoint = Checkpoint(
-        context.checkpoint.save(save_path)
+                state_dict=state_dict,
-        self.last_ckpt_iter = context.iteration
+                epoch=context.epoch,
                consumed_samples=context.consumed_samples,
                config=context.model_config,
                extra=extra,
                meta=meta,
            )
            context.checkpoint.save(save_path)
    def on_batch_end(self, context: TrainContext):
-        if context.iteration - self.last_ckpt_iter >= self.interval:
+        if context.optimizer_step - self.last_ckpt_step >= self.interval:
            self._save_checkpoint(context)
    def on_train_end(self, context: TrainContext):
-        if context.iteration != self.last_ckpt_iter:
+        if context.optimizer_step != self.last_ckpt_step:
            self._save_checkpoint(context)
    def on_error(self, context: TrainContext):
        self._save_checkpoint(context)
    @staticmethod
    def save_extra(context: TrainContext) -> dict:
        extra = {}
        for name in CheckpointCallback.extra_keys:
            obj = getattr(context, name, None)
            if obj:
                extra[name] = obj.state_dict()
        return extra
@CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):
@@ -145,26 +189,36 @@ class ProgressBarCallback(TrainCallback):
    Progress bar callback for trainer.
    """
-    def __init__(self, num_epoch: int):
+    def __init__(
        self, num_epoch: int, log_interval: int = 100, file: Optional[IO[str]] = None
    ):
        self.num_epoch = num_epoch
        self.log_interval = log_interval
        self.file = file
        self.progress_bar: tqdm = None
    @only_on_rank(0)
    def on_epoch_begin(self, context: TrainContext):
        total_steps = len(context.dataloader) // context.executor.grad_accum_steps
        self.progress_bar = tqdm(
-            context.dataloader,
+            total=total_steps,
            desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
            dynamic_ncols=True,
            file=self.file or sys.stdout,
        )
    @only_on_rank(0)
-    def on_batch_end(self, context: TrainContext):
+    def on_optimizer_step(self, context: TrainContext):
-        self.progress_bar.set_postfix(
+        postfix = {
-            {
+            "step": context.optimizer_step,
-                "loss": f"{context.loss:.4f}",
+            "loss": f"{context.loss:.4f}",
-                "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
+            "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)
@@ -174,19 +228,20 @@ class ProgressBarCallback(TrainCallback):
            self.progress_bar.close()
-@CallbackFactory.register("metric_logger")
+@CallbackFactory.register("metric")
-class MetricLoggerCallback(TrainCallback):
+class MetricCallback(TrainCallback):
    def __init__(
        self,
        log_dir: str,
        save_interval: int,
        log_interval: int = 10,
        metrics: List[str] = None,
        val_step: int = 0,
    ):
-        self.last_log_iter = 0
+        self.last_log_flush_step = 0
        self.save_interval = save_interval
        self.log_interval = log_interval
        self.metrics = metrics or ["loss", "lr"]
        self.val_step = val_step
        self._next_val_step = 0
        self.log_dir = Path(log_dir) if log_dir else Path.cwd() / "logs"
        self.log_dir.mkdir(parents=True, exist_ok=True)
@@ -196,46 +251,84 @@ class MetricLoggerCallback(TrainCallback):
        self._metric_funcs = {
            "loss": ctx_get_loss,
            "lr": ctx_get_lr,
            "val_loss": ctx_get_val_loss,
            "grad_norm": ctx_get_grad_norm,
            "grad_std": ctx_get_grad_std,
            "grad_max": ctx_get_grad_max,
            "grad_min": ctx_get_grad_min,
            "grad_mean": ctx_get_grad_mean,
            "grad_nan_num": ctx_get_grad_nan_num,
        }
-    def _get_log_data(self, context: TrainContext):
+    def _metrics(self, context: TrainContext, names):
        return {
-            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+            m: self._metric_funcs[m](context)
-            "epoch": context.epoch,
+            for m in names
-            "iter": context.iteration,
+            if self._metric_funcs[m](context) is not None
            **{m: self._metric_funcs[m](context) for m in self.metrics},
        }
    @only_on_rank(0)
-    def _add_log(self, log_data):
+    def _append(self, event_type: str, context: TrainContext, **extra):
-        self.log_cache.append(log_data)
+        entry = {
            "type": event_type,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "epoch": context.epoch,
            "step": context.optimizer_step,
            "consumed_samples": context.consumed_samples,
            **extra,
        }
        self.log_cache.append(entry)
    def _run_validation(self, context: TrainContext) -> float:
        context.model.eval()
        total_loss = 0.0
        num_batches = 0
        with torch.no_grad():
            for batch in context.val_dataloader:
                loss = context.strategy(batch)
                total_loss += loss.item()
                num_batches += 1
        if context.world_size > 1 and dist.is_initialized():
            stats = torch.tensor(
                [total_loss, float(num_batches)], device=get_current_device()
            )
            dist.all_reduce(stats, op=dist.ReduceOp.SUM)
            avg_loss = (stats[0] / stats[1]).item()
        else:
            avg_loss = total_loss / max(num_batches, 1)
        context.model.train()
        return avg_loss
    @only_on_rank(0)
-    def _save_log(self, epoch, iter):
+    def _flush(self, epoch, step):
-        log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
+        log_file = self.log_dir / f"epoch_{epoch}_step_{step}_metric.jsonl"
-
+        log_file.parent.mkdir(parents=True, exist_ok=True)
        with open(log_file, "w") as f:
            for log in self.log_cache:
                f.write(json.dumps(log) + "\n")
-    def on_batch_end(self, context):
+    def on_optimizer_step(self, context):
-        if context.iteration % self.log_interval == 0:
+        if (
-            log_data = self._get_log_data(context)
+            context.val_dataloader is not None
-            self._add_log(log_data)
+            and self.val_step > 0
            and context.optimizer_step >= self._next_val_step
        ):
            context.val_loss = self._run_validation(context)
            self._next_val_step = context.optimizer_step + self.val_step
            self._append("validation", context, val_loss=context.val_loss)
-        if context.iteration - self.last_log_iter >= self.save_interval:
+        step_metrics = [m for m in self.metrics if m != "val_loss"]
-            self._save_log(context.epoch, context.iteration)
+        self._append("step", context, **self._metrics(context, step_metrics))
-            self.last_log_iter = context.iteration
+
        if context.optimizer_step - self.last_log_flush_step >= self.save_interval:
            self._flush(context.epoch, context.optimizer_step)
            self.last_log_flush_step = context.optimizer_step
    def on_epoch_end(self, context):
        self._append("epoch", context)
    def on_train_end(self, context):
-        if context.iteration != self.last_log_iter:
+        if context.optimizer_step != self.last_log_flush_step:
-            self._save_log(context.epoch, context.iteration)
+            self._flush(context.epoch, context.optimizer_step)
    def on_error(self, context):
-        self._save_log(context.epoch, context.iteration)
+        self._flush(context.epoch, context.optimizer_step)
@@ -1,15 +1,18 @@
 from dataclasses import dataclass, field
-from typing import Callable, Optional, Self
+from pathlib import Path
 from typing import Any, Dict, Optional, Self
 import torch
 import torch.nn as nn
-from torch.optim import Optimizer
+from torch.utils.data import DataLoader, random_split
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import ResumableDistributedSampler
 from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
-from astrai.serialization import Checkpoint
+from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import Checkpoint, load_json
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
@@ -18,84 +21,169 @@ class TrainContext:
    model: nn.Module = field(default=None)
    strategy: BaseStrategy = field(default=None)
    dataloader: DataLoader = field(default=None)
-    optimizer: Optimizer = field(default=None)
+    optimizer: OptimizerProtocol = field(default=None)
-    scheduler: LRScheduler = field(default=None)
+    scheduler: SchedulerProtocol = field(default=None)
    checkpoint: Checkpoint = field(default=None)
    config: TrainConfig = field(default=None)
    model_config: dict = field(default_factory=dict)
    executor: BaseExecutor = field(default=None)
    epoch: int = field(default=0)
-    iteration: int = field(default=0)
+    consumed_samples: int = field(default=0)
    loss: float = field(default=0.0)
    grad_norm: Optional[float] = field(default=None)
    val_dataloader: Optional[DataLoader] = field(default=None)
    val_loss: Optional[float] = field(default=None)
    world_size: int = field(default=1)
    rank: int = field(default=0)
-    kwargs: dict = field(default_factory=dict)
+    kwargs: Dict[str, Any] = field(default_factory=dict)
    @property
    def optimizer_step(self) -> int:
        return self.consumed_samples // (
            self.config.batch_per_device
            * self.world_size
            * self.config.grad_accum_steps
        )
 class TrainContextBuilder:
    def __init__(
        self,
        config: TrainConfig,
        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
    ):
        self.config = config
-        self._checkpoint: Optional[Checkpoint] = None
+        self._resume_dir: Optional[str] = None
        self._load_extra_fn = load_extra_fn
-    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+    def with_resume_dir(self, resume_dir: Optional[str]) -> Self:
-        self._checkpoint = checkpoint
+        self._resume_dir = resume_dir
        return self
    def build(self) -> TrainContext:
-        context = TrainContext(
+        cfg = self.config
-            model=self.config.model,
+        device = get_current_device()
-            world_size=get_world_size(),
+
-            rank=get_rank(),
+        executor = ExecutorFactory.create(
            cfg.parallel_mode,
            grad_accum_steps=cfg.grad_accum_steps,
            **cfg.executor_kwargs,
        )
-        device = get_current_device()
+        model = cfg.model_fn()
-        context.model = context.model.to(device=device)
+        model = model.to(device=device)
-        if self.config.nprocs > 1 and self.config.parallel_wrapper:
+        model_config = {}
-            context.model = self.config.parallel_wrapper(context.model)
+        if self._resume_dir:
            config_path = Path(self._resume_dir) / "config.json"
            if config_path.exists():
                model_config = load_json(config_path)
-        if self._checkpoint is not None:
+        if not model_config and hasattr(model, "config"):
-            context.epoch = max(self._checkpoint.epoch, self.config.start_epoch)
+            model_config = model.config.to_dict()
-            context.iteration = max(self._checkpoint.iteration, self.config.start_batch)
+
-            context.model.load_state_dict(self._checkpoint.state_dict)
+        context = TrainContext(
-            context.checkpoint = self._checkpoint
+            model=model,
-        else:
+            world_size=get_world_size(),
-            context.checkpoint = Checkpoint(
+            rank=get_rank(),
-                state_dict=context.model.state_dict(),
+            config=cfg,
            model_config=model_config,
            executor=executor,
        )
        if self._resume_dir:
            checkpoint = Checkpoint.load_any(self._resume_dir)
            if checkpoint is not None:
                model.load_state_dict(checkpoint.state_dict, strict=False)
                if checkpoint.config:
                    context.model_config = checkpoint.config
                context.epoch = checkpoint.epoch or cfg.start_epoch
                if checkpoint.consumed_samples > 0:
                    context.consumed_samples = checkpoint.consumed_samples
                else:
                    context.consumed_samples = cfg.start_samples * context.world_size
                context.checkpoint = checkpoint
        if cfg.lora is not None:
            inject_lora(
                model,
                r=cfg.lora.r,
                alpha=cfg.lora.alpha,
                target_modules=set(cfg.lora.target_modules),
            )
-        context.optimizer = self.config.optimizer_fn(context.model)
+        context.optimizer = cfg.optimizer_fn(model)
-        context.scheduler = self.config.scheduler_fn(context.optimizer)
+        context.scheduler = cfg.scheduler_fn(context.optimizer)
-        if self._checkpoint and self._checkpoint.extra and self._load_extra_fn:
+        train_dataset = cfg.dataset
-            self._load_extra_fn(self._checkpoint.extra, context)
+        val_dataset = cfg.val_dataset
-        cfg = self.config
+        if val_dataset is None and cfg.val_split is not None:
-        sampler_offset = context.iteration * cfg.batch_size
+            n_total = len(cfg.dataset)
            n_val = max(1, int(n_total * cfg.val_split))
            n_train = n_total - n_val
            generator = torch.Generator().manual_seed(cfg.random_seed)
            train_dataset, val_dataset = random_split(
                cfg.dataset, [n_train, n_val], generator=generator
            )
        sampler_offset = context.consumed_samples // context.world_size
        sampler = ResumableDistributedSampler(
-            data_source=cfg.dataset,
+            data_source=train_dataset,
            start_epoch=context.epoch,
            start_iter=sampler_offset,
            seed=cfg.random_seed,
        )
        context.dataloader = DataLoader(
-            cfg.dataset,
+            train_dataset,
-            batch_size=cfg.batch_size,
+            batch_size=cfg.batch_per_device,
            sampler=sampler,
            num_workers=cfg.num_workers,
            pin_memory=cfg.pin_memory,
            prefetch_factor=cfg.prefetch_factor,
        )
        if val_dataset is not None:
            val_sampler = ResumableDistributedSampler(
                data_source=val_dataset,
                start_epoch=0,
                start_iter=0,
                seed=cfg.random_seed,
                shuffle=False,
            )
            context.val_dataloader = DataLoader(
                val_dataset,
                batch_size=cfg.batch_per_device,
                sampler=val_sampler,
                num_workers=cfg.num_workers,
                pin_memory=cfg.pin_memory,
                prefetch_factor=cfg.prefetch_factor,
            )
        context.model, context.optimizer, context.dataloader, context.scheduler = (
            executor.prepare(
                model,
                context.optimizer,
                context.dataloader,
                context.scheduler,
            )
        )
        if context.checkpoint and context.checkpoint.extra:
            extra = context.checkpoint.extra
            for name in ("optimizer", "scheduler"):
                if name in extra:
                    obj = getattr(context, name, None)
                    if obj is not None:
                        obj.load_state_dict(extra[name])
        context.strategy = StrategyFactory.create(
            cfg.strategy,
            model=context.model,
            train_type=self.config.strategy,
            device=device,
-            **self.config.extra_kwargs,
+            executor=executor,
            model_fn=cfg.model_fn,
            **cfg.extra_kwargs,
        )
        return context
@@ -1,10 +1,8 @@
 import logging
 from itertools import batched
 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,
@@ -26,17 +24,27 @@ class Trainer:
    def _get_default_callbacks(self) -> List[TrainCallback]:
        cfg = self.train_config
-        return [
+        callbacks = [
            CallbackFactory.create(
                "gradient_checkpointing",
                modules=cfg.gradient_checkpointing_modules,
            ),
            CallbackFactory.create(
                "checkpoint",
                cfg.ckpt_dir,
                cfg.ckpt_interval,
            ),
            CallbackFactory.create(
                "metric",
                log_dir=cfg.log_dir,
                save_interval=cfg.ckpt_interval,
                metrics=cfg.metrics,
                val_step=cfg.val_step,
            ),
            CallbackFactory.create("progress_bar", cfg.n_epoch),
            CallbackFactory.create("checkpoint", cfg.ckpt_dir, cfg.ckpt_interval),
            CallbackFactory.create("metric_logger", cfg.ckpt_dir, cfg.ckpt_interval),
            CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
        ]
-
+        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:
@@ -44,56 +52,58 @@ class Trainer:
            if method:
                method(context)
-    def train(self, checkpoint: Optional[Checkpoint] = None):
+    def _trainer_loop(self, resume_dir: Optional[str] = None):
-        config = self.train_config
+        context = (
-        spawn_parallel_fn(
+            TrainContextBuilder(self.train_config).with_resume_dir(resume_dir).build()
            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()
            accumulation_steps = max(self.train_config.accumulation_steps, 1)
-            for epoch in range(context.epoch, self.train_config.n_epoch):
+            for epoch in range(context.epoch, context.config.n_epoch):
                context.epoch = epoch
                self._call_callbacks("on_epoch_begin", context)
-                for steps in batched(context.dataloader, accumulation_steps):
+                for batch in context.dataloader:
-                    self._call_callbacks("on_step_begin", context)
+                    with executor.accumulate(context.model):
                    step_batch_nums = len(steps)
                    for batch in steps:
                        self._call_callbacks("on_batch_begin", context)
                        loss = context.strategy(batch)
                        context.loss = loss.item()
-                        context.iteration += 1
+                        stand_loss = loss / executor.grad_accum_steps
-
+                        executor.backward(stand_loss)
-                        stand_loss = loss / step_batch_nums
+                        context.consumed_samples += (
-                        stand_loss.backward()
+                            context.config.batch_per_device * context.world_size
                        )
                        self._call_callbacks("on_batch_end", context)
-                    self._call_callbacks("on_step_end", context)
+                        if executor.sync_gradients:
-                    context.optimizer.step()
+                            self._call_callbacks("on_optimizer_step", context)
-                    context.optimizer.zero_grad()
+                            context.optimizer.step()
                            context.optimizer.zero_grad()
-                    if context.scheduler:
+                            if context.scheduler:
-                        context.scheduler.step()
+                                context.scheduler.step()
                self._call_callbacks("on_epoch_end", context)
        except Exception as e:
-            logger.error(f"Training failed: {str(e)}", exc_info=True)
+            logger.error("Training failed: %s", str(e), exc_info=True)
            self._call_callbacks("on_error", context)
            raise
        finally:
            self._call_callbacks("on_train_end", context)
    def train(self, resume_dir: Optional[str] = None):
        cfg = self.train_config
        spawn_parallel_fn(
            self._trainer_loop,
            backend=cfg.backend,
            world_size=cfg.nprocs,
            master_addr=cfg.master_addr,
            master_port=cfg.master_port,
            device_type=cfg.device_type,
            start_method=cfg.start_method,
            resume_dir=resume_dir,
        )
@@ -0,0 +1,2 @@
 # Source directory for CUDA kernels — build-time only.
 # Compiled .so files live in astrAI/_ext/.
@@ -0,0 +1,48 @@
 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")
@@ -0,0 +1,54 @@
 #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 is_causal;
    int causal_offset;
    int num_splits;
    float scale;
    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 is_causal;
    int causal_offset;
    float scale;
    int num_splits;
    int page_size;
    int max_pages;
    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;
 };
@@ -0,0 +1,113 @@
 #include "attn_decode_split_kv.cuh"
 #include "attn_entry_utils.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "attn_decode_split_kv_mma.cuh"
 #endif
 static int decode_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)));
 }
 // 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 = decode_num_splits(p.batch * p.kv_head, chunks_total);
    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 = o_part.data_ptr<float>();
    p.ml_part = ml_part.data_ptr<float>();
    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 = decode_num_splits(p.batch * p.kv_head, tiles_total);
    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 = o_part.data_ptr<float>();
    p.ml_part = ml_part.data_ptr<float>();
    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 (!p.use_mask && 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,
    bool is_causal = false,
    int64_t causal_offset = 0,
    c10::optional<double> scale = c10::nullopt
 ) {
    AttentionParams<bf16> p;
    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, 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");
    auto O = torch::empty_like(q);
    p.o = (bf16*)O.data_ptr();
    switch (p.head_dim) {
        case 32:
            dispatch_decode<32>(p);
            break;
        case 64:
            dispatch_decode<64>(p);
            break;
        case 128:
            dispatch_decode<128>(p);
            break;
        case 256:
            dispatch_decode<256>(p);
            break;
        default:
            TORCH_CHECK(false, "decode: unsupported head_dim ", p.head_dim,
                        " (supported: 32, 64, 128, 256)");
    }
    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("is_causal") = false,
        py::arg("causal_offset") = 0,
        py::arg("scale") = py::none(),
        "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
 }
@@ -0,0 +1,116 @@
 #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;
    float q_reg[8];
    int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
    for (int i = 0; i < hd_per_thread; i++)
        q_reg[i] = __bfloat162float(p.q[q_off + i]);
    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim;
    int mask_base = batch * p.kv_len;
    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);
        int total = this_chunk * p.head_dim;
        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y)
            k_smem[i] = p.k[kv_base + chunk_start * p.head_dim + i];
        __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;
            if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s])
                partial = -FLT_MAX;
            if (p.is_causal && (chunk_start + s) > 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 v_off = kv_base + (chunk_start + s) * p.head_dim + lane * hd_per_thread;
            for (int i = 0; i < hd_per_thread; i++)
                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i]) * 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;
    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;
    p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
 }
@@ -0,0 +1,200 @@
 #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.
 //
 // Partial layout (float, contiguous):
 //   o_part : [batch, q_head, num_splits, HEAD_DIM]
 //   ml_part: [batch, q_head, num_splits, 2]  (m, l)
 //
 // Optimizations:
 //   - cp.async global→shared for K/V (bypasses registers, cuts instruction count)
 //   - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts
 //   - Q loaded directly from global into mma A-operand registers (no sQ staging,
 //     no prologue syncwarp) — frees shared memory for double-buffering
 //   - Double-buffered KV (STAGES=2): next tile's cp.async overlaps current
 //     tile's MMA compute — hides global load latency / boosts bandwidth
 //     utilization for small-batch (low-occupancy) decode
 //   - Predicated cp.async (cp_async_16_pred) for full AND partial tiles on one
 //     uniform path — eliminates the scalar fallback branch
 //
 // Smem footprint (BC=32): STAGES=2 → 2*(sK+sV) = 2*2*32*HEAD_DIM*2 bytes.
 //   D=128: 16 KB (fits 48 KB static cap).  D=256: 32 KB (also fits).
 // STAGES=1 fallback (4/8 KB) for smem-constrained configs.
 template <int HEAD_DIM, int BC, int STAGES = 2>
 __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
    constexpr int BR = 16;
    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 ----
    // Same layout as prefill: frag[0]/[2] = row gid, frag[1]/[3] = row gid+8
    // cols kt*16 + tid4*2 + {0,1} / +{8,9}. pau[0]=cols c,c+1; pau[4]=c+8,c+9.
    const int q_base = (batch * p.q_head + q_head0) * HEAD_DIM;
    const int qra = gid;
    const int qrb = gid + 8;
    const bool va = qra < G, vb = qrb < G;
    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*>(
            &p.q[q_base + qra * HEAD_DIM + c]);
        const unsigned* pbu = reinterpret_cast<const unsigned*>(
            &p.q[q_base + qrb * HEAD_DIM + c]);
        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;
    }
    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 kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM;
    const int mask_base = batch * p.kv_len;
    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);
            cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid);
            cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + 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;
        // 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;
        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
                                    mask_base, 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;
        }
    }
 }
@@ -0,0 +1,43 @@
 #pragma once
 #include <torch/extension.h>
 #include "attn_common.h"
 template<typename T>
 inline void attn_pack_params(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    c10::optional<torch::Tensor> mask,
    bool is_causal,
    int64_t causal_offset,
    c10::optional<double> scale,
    AttentionParams<T>& p
 ) {
    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);
    p.batch = (int)q.size(0);
    p.q_head = (int)q.size(1);
    p.kv_head = (int)k.size(1);
    p.q_len = (int)q.size(2);
    p.kv_len = (int)k.size(2);
    p.head_dim = (int)q.size(3);
    p.use_mask = mask.has_value() ? 1 : 0;
    p.is_causal = is_causal ? 1 : 0;
    p.causal_offset = (int)causal_offset;
    p.scale = scale.has_value() ? (float)scale.value() : 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();
    if (p.use_mask) {
        TORCH_CHECK(mask.value().dtype() == torch::kBool);
        TORCH_CHECK(mask.value().dim() == 2);
        TORCH_CHECK(mask.value().size(0) == p.batch);
        TORCH_CHECK(mask.value().size(1) == p.kv_len);
        p.mask = mask.value().data_ptr<bool>();
    } else {
        p.mask = nullptr;
    }
 }
@@ -0,0 +1,252 @@
 #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));
 }
 // ---------------------------------------------------------------------------
 // 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).
 // 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 mask_base,
    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;
    #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_base + cc]);
        bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base + c1]);
        bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base + cc]);
        bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base + 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]);
        }
    }
 }
@@ -0,0 +1,149 @@
 #include "attn_paged_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "attn_paged_decode_split_kv_mma.cuh"
 #endif
 #include <torch/extension.h>
 #include <c10/cuda/CUDAGuard.h>
 static int paged_decode_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)));
 }
 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 = paged_decode_num_splits(p.batch * p.kv_head, chunks_total);
    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 = o_part.data_ptr<float>();
    p.ml_part = ml_part.data_ptr<float>();
    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 = paged_decode_num_splits(p.batch * p.kv_head, tiles_total);
    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 = o_part.data_ptr<float>();
    p.ml_part = ml_part.data_ptr<float>();
    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 (!p.use_mask && 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,
    bool is_causal = false,
    int64_t causal_offset = 0,
    c10::optional<double> scale = c10::nullopt
 ) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
    int batch = q.size(0);
    int q_head = q.size(1);
    int head_dim = q.size(3);
    int kv_head = k_cache.size(2);
    int max_pages = page_table.size(1);
    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(q.size(2) == 1, "Q seq_len must be 1 (decode)");
    TORCH_CHECK(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);
    TORCH_CHECK(k_cache.size(0) >= 0, "k_cache must have at least 0 pages");
    float scale_val = scale.has_value()
        ? static_cast<float>(scale.value())
        : 1.0f / std::sqrt(static_cast<float>(head_dim));
    auto O = torch::empty_like(q);
    PagedAttentionParams<bf16, float> p;
    p.batch = batch;
    p.q_head = q_head;
    p.kv_head = kv_head;
    p.q_len = static_cast<int>(q.size(2));
    p.kv_len = static_cast<int>(kv_len);
    p.head_dim = head_dim;
    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
    p.is_causal = is_causal ? 1 : 0;
    p.causal_offset = static_cast<int>(causal_offset);
    p.scale = scale_val;
    p.page_size = static_cast<int>(page_size);
    p.max_pages = max_pages;
    p.page_table = page_table.data_ptr<int64_t>();
    p.k_cache = reinterpret_cast<const bf16*>(k_cache.data_ptr());
    p.v_cache = reinterpret_cast<const bf16*>(v_cache.data_ptr());
    p.q = reinterpret_cast<const bf16*>(q.data_ptr());
    p.mask = p.use_mask ? mask.value().data_ptr<bool>() : nullptr;
    p.o = reinterpret_cast<bf16*>(O.data_ptr());
    p.o_part = nullptr;
    p.ml_part = nullptr;
    switch (p.head_dim) {
        case 32:  dispatch_paged_decode<32>(p);  break;
        case 64:  dispatch_paged_decode<64>(p);  break;
        case 128: dispatch_paged_decode<128>(p); break;
        case 256: dispatch_paged_decode<256>(p); break;
        default:
            TORCH_CHECK(false, "paged_decode: unsupported head_dim ", p.head_dim,
                        " (supported: 32, 64, 128, 256)");
    }
    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("is_causal") = false,
        py::arg("causal_offset") = 0,
        py::arg("scale") = py::none(),
        "Paged GQA decode — split-KV with direct page-table access.");
 }
@@ -0,0 +1,140 @@
 #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;
    float q_reg[8];
    int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
    #pragma unroll
    for (int i = 0; i < hd_per_thread; i++)
        q_reg[i] = __bfloat162float(p.q[q_off + i]);
    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.kv_len;
    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;
            if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s])
                partial = -FLT_MAX;
            if (p.is_causal && (chunk_start + s) > 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;
    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;
    p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
 }
@@ -0,0 +1,182 @@
 #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.
 //
 // Optimizations mirror attn_decode_split_kv_mma_kernel:
 //   - Q loaded directly from global into mma A-operand registers (no sQ)
 //   - Double-buffered KV (STAGES=2) for D<=128, single-buffer for D=256
 //   - Predicated cp.async for unified full/partial tile path
 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 BR = 16;
    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_idx = 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_idx * 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_head + q_head0) * HEAD_DIM;
    const int qra = gid;
    const int qrb = gid + 8;
    const bool va = qra < G, vb = qrb < G;
    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*>(
            &p.q[q_base + qra * HEAD_DIM + c]);
        const unsigned* pbu = reinterpret_cast<const unsigned*>(
            &p.q[q_base + qrb * HEAD_DIM + c]);
        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;
    }
    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 mask_base = batch * p.kv_len;
    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_idx * 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;
        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
                                    mask_base, 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;
        }
    }
 }
@@ -0,0 +1,79 @@
 #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,
    bool is_causal = false,
    int64_t causal_offset = 0,
    c10::optional<double> scale = c10::nullopt
 ) {
    AttentionParams<bf16> p;
    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
    TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
    auto O = torch::empty_like(q);
    p.o = (bf16*)O.data_ptr();
    switch (p.head_dim) {
        case 32:
            dispatch_prefill<32>(p);
            break;
        case 64:
            dispatch_prefill<64>(p);
            break;
        case 128:
            dispatch_prefill<128>(p);
            break;
        case 256:
            dispatch_prefill<256>(p);
            break;
        default:
            TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
                        " (supported: 32,64,128,256)");
    }
    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("is_causal") = false,
        py::arg("causal_offset") = 0,
        py::arg("scale") = py::none(),
        "GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
 }
@@ -0,0 +1,140 @@
 #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];
    float qreg[DPT];
    if (q_row < p.q_len) {
        int q_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
 #pragma unroll
        for (int i = 0; i < DPT; i++)
            qreg[i] = __bfloat162float(p.q[q_off + i]) * 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;
    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM;
    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);
        for (int i = lid; i < tlen * HEAD_DIM; i += tt) {
            int gidx = kv_base + (kv0 + i / HEAD_DIM) * HEAD_DIM + (i % HEAD_DIM);
            sK[i] = p.k[gidx];
            sV[i] = p.v[gidx];
        }
        __syncthreads();
        int lim = tlen;
        if (p.is_causal && 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;
        }
        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);
            if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s])
                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) {
        int o_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
        float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
 #pragma unroll
        for (int i = 0; i < DPT; i++)
            p.o[o_off + i] = __float2bfloat16(acc[i] * rl);
    }
 }
@@ -0,0 +1,204 @@
 #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 — sized by template parameters at compile time.
    // 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 the Q fragments straight from global into the mma A-operand layout
    // (m16n8k16, row-major): no sQ staging area and no serialized per-warp
    // prologue barriers. Each lane reads exactly the 8 Q elements ldmatrix
    // would have produced, pre-scaled by the attention scale. Kept resident in
    // registers across the tile loop.
    //   frag[0]/[2]: row = qrow0 + gid ;  frag[1]/[3]: row = qrow0 + gid + 8
    //   frag[0]/[1]: cols kt*16 + tid4*2 + {0,1} ; frag[2]/[3]: + 8
    const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
    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];
 #pragma unroll
    for (int kt = 0; kt < KD; kt++) {
        int c = kt * 16 + tid4 * 2;
        const unsigned* pau = reinterpret_cast<const unsigned*>(
            &p.q[q_base + qra * HEAD_DIM + c]);
        const unsigned* pbu = reinterpret_cast<const unsigned*>(
            &p.q[q_base + qrb * HEAD_DIM + c]);
        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;
    }
    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 kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM;
    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 is_causal == 0)
    const int use_skip = p.is_causal;
    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;
    const int mb = batch * p.kv_len;
    // 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;
    // 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);
            cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid);
            cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + d], valid);
        }
        cp_async_commit();
    };
    // Prologue: kick off the first tile's 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.is_causal ? min(p.kv_len, qr0 + p.causal_offset + 1)
                                : p.kv_len;
        int maxc1 = p.is_causal ? min(p.kv_len, qr1 + p.causal_offset + 1)
                                : p.kv_len;
        mma_softmax_tile<NC8, DN8>(kv0, maxc0, maxc1,
                                    mb, 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;
    const int o_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
 #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 * HEAD_DIM + 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 * HEAD_DIM + d]) = v;
        }
    }
 }
@@ -0,0 +1,199 @@
 /*
 Pure-C test:
 nvcc -I csrc -arch=sm_89 -O3 \
    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
    csrc/tests/attn_decode_test.cu -o test && ./test
 */
 #include "test_utils.cuh"
 #include "../kernels/attn_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "../kernels/attn_decode_split_kv_mma.cuh"
 #endif
 // Split-K scratch (torch-free): the production launcher allocates these from
 // torch; here we pass pre-allocated device buffers so the bench loop doesn't
 // pay a cudaMalloc per iteration. Size for the maximum split count (32).
 struct DecodeScratch {
    float* o_part = nullptr;
    float* ml_part = nullptr;
 };
 // Launch the production decode path (tensor-core head-packing MMA on sm_80+,
 // scalar fallback otherwise), mirroring dispatch_decode() in attn_decode.cu.
 #ifndef ASTRAI_NO_MMA
 static bool decode_use_mma(const AttentionParams<bf16>& p) {
    int G = p.q_head / p.kv_head;
    return !p.use_mask && G > 1 && G <= 16;
 }
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
    int tiles_total = (p.kv_len + BC - 1) / BC;
    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
    p.o_part = sc.o_part;
    p.ml_part = sc.ml_part;
    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
 static void launch_scalar_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
    int gs = 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);
    p.o_part = sc.o_part;
    p.ml_part = sc.ml_part;
    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, gs), smem>>>(p);
    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 template <int HEAD_DIM>
 static void dispatch_decode_t(AttentionParams<bf16>& p, DecodeScratch& sc) {
 #ifndef ASTRAI_NO_MMA
    if (decode_use_mma(p)) { launch_mma_decode<HEAD_DIM, 32>(p, sc); return; }
 #endif
    launch_scalar_decode(p, sc);
 }
 static void dispatch_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
    dispatch_by_head_dim(p.head_dim, [&]<int D>() { dispatch_decode_t<D>(p, sc); });
 }
 // Warmed-up, CUDA-event timed sweep over the production decode MMA path.
 static void bench() {
    const int cfgs[][5] = {
        {1, 32, 4, 512, 128},    // B, Hq, Hk, kv_len, D
        {1, 32, 4, 1024, 128},
        {1, 32, 4, 2048, 128},
        {1, 32, 4, 4096, 128},
        {16, 32, 4, 2048, 128},
        {32, 32, 4, 1024, 128},
    };
    const int WARMUP = 10, ITERS = 100;
    printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
    print_bench_header();
    for (int ci = 0; ci < 6; ci++) {
        int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
        int sl = cfgs[ci][3], D = cfgs[ci][4];
        size_t nQ = (size_t)B * Hq * D;
        size_t nKV = (size_t)B * Hk * sl * D;
        bf16 *dQ, *dK, *dV, *dO;
        cudaMalloc(&dQ, nQ*2); cudaMalloc(&dK, nKV*2);
        cudaMalloc(&dV, nKV*2); cudaMalloc(&dO, nQ*2);
        size_t big = nQ > nKV ? nQ : nKV; bf16* tmp = new bf16[big];
        for (size_t i = 0; i < nQ; i++)  tmp[i] = f2bf(randf());
        cudaMemcpy(dQ, tmp, nQ*2, cudaMemcpyHostToDevice);
        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
        cudaMemcpy(dK, tmp, nKV*2, cudaMemcpyHostToDevice);
        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
        cudaMemcpy(dV, tmp, nKV*2, cudaMemcpyHostToDevice);
        delete[] tmp;
        AttentionParams<bf16> p;
        p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
        p.head_dim = D; p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
        p.scale = 1.0f / sqrtf((float)D);
        p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO;
        DecodeScratch sc;
        cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
        cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
        auto launch = [&]() { dispatch_decode(p, sc); };
        double flops = 4.0 * B * Hq * (double)sl * D;
        double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
        BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
        char cfg[64];
        snprintf(cfg, sizeof(cfg),
                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
                 B, Hq, Hk, 1, sl, D, 0);
        print_bench_row(cfg, r);
        cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
        cudaFree(sc.o_part); cudaFree(sc.ml_part);
    }
 }
 int main() {
    const int configs[][5] = {
        {1, 2, 1, 64, 32},    // B,Hq,Hk,seq_len,D
        {1, 32, 4, 512, 128},
        {1, 32, 4, 1024, 128},
    };
    int n_cfgs = sizeof(configs) / sizeof(configs[0]);
    for (int ci = 0; ci < n_cfgs; ci++) {
        int B = configs[ci][0], Hq = configs[ci][1], Hk = configs[ci][2];
        int sl = configs[ci][3], D = configs[ci][4], gs = Hq / Hk;
        printf("=== B=%d Hq=%d Hk=%d seq=%d D=%d gs=%d ===\n", B,Hq,Hk,sl,D,gs);
        size_t nQ = B*Hq*1*D, nKV = B*Hk*sl*D;
        float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
        for (size_t i=0;i<nQ;i++) hQ[i]=randf();
        for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
        bool* hMask=new bool[B*sl];
        for (int i=0;i<B*sl;i++) hMask[i]=true;
        bf16 *dQ,*dK,*dV,*dO,*tmp;
        bool* dMask;
        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
        cudaMalloc(&dMask,B*sl);
        tmp=new bf16[max(nQ,nKV)];
        for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
        cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
        AttentionParams<bf16> p;
        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
        p.use_mask=0; p.is_causal=0; p.causal_offset=0;
        p.scale=1.0f/sqrtf((float)D);
        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
        // Split-K scratch (max 32 splits), sized for the production MMA path.
        DecodeScratch sc;
        cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
        cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
        double t0=now_ms();
        dispatch_decode(p, sc);
        cudaDeviceSynchronize();
        double kms=now_ms()-t0;
        cudaError_t err=cudaGetLastError();
        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
        bf16* hOut=new bf16[nQ];
        cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
        float* ref=new float[nQ];
        cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, 0, 0);
        float max_err=0;
        for (size_t i=0;i<nQ;i++){
            float d=fabsf(bf2f(hOut[i])-ref[i]);
            if(d>max_err) max_err=d;
        }
        printf("kernel: %.3f ms  max_err: %.6e\n\n",kms,max_err);
        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
        cudaFree(sc.o_part);cudaFree(sc.ml_part);
        delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
    }
    printf("All tests passed!\n");
    bench();
    return 0;
 }
@@ -0,0 +1,330 @@
 // Compile:
 //   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
 //        --extra-device-vectorization csrc/tests/attn_paged_decode_test.cu \
 //        -o /tmp/test_paged && /tmp/test_paged
 #include <cstring>
 #include "test_utils.cuh"
 #include "../kernels/attn_paged_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "../kernels/attn_paged_decode_split_kv_mma.cuh"
 #endif
 // Copy contiguous K/V from page pool (reference gather)
 static void gather_kv_cpu(
    const bf16* h_k_pool, const bf16* h_v_pool,
    const int64_t* h_pt, int B, int Hkv, int kv_len,
    int page_size, int head_dim,
    bf16* h_k, bf16* h_v)
 {
    int max_pages = (kv_len + page_size - 1) / page_size;
    size_t page_stride = (size_t)page_size * Hkv * head_dim;
    for (int b = 0; b < B; b++) {
        for (int pos = 0; pos < kv_len; pos++) {
            int log_pg = pos / page_size;
            int pg_off = pos % page_size;
            int phys = (int)h_pt[b * max_pages + log_pg];
            for (int h = 0; h < Hkv; h++) {
                size_t src_base = (size_t)phys * page_stride
                                + (size_t)pg_off * Hkv * head_dim
                                + h * head_dim;
                size_t dst_base = ((size_t)b * Hkv + h) * kv_len * head_dim + (size_t)pos * head_dim;
                memcpy(h_k + dst_base, h_k_pool + src_base, head_dim * sizeof(bf16));
                memcpy(h_v + dst_base, h_v_pool + src_base, head_dim * sizeof(bf16));
            }
        }
    }
 }
 template <int HEAD_DIM>
 static void launch_paged_decode(PagedAttentionParams<bf16, float>& p) {
 #ifndef ASTRAI_NO_MMA
    int G_check = p.q_head / p.kv_head;
    bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32;
    if (use_mma) {
        constexpr int STAGES = (HEAD_DIM <= 128) ? 2 : 1;
        int tiles_total = (p.kv_len + 32 - 1) / 32;
        p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
        paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32, STAGES>
            <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
    } else
 #endif
    {
        int group_sz = 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);
        size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
        paged_attn_decode_split_kv_kernel<<<
            dim3(p.batch * p.kv_head, 1, p.num_splits),
            dim3(32, group_sz), smem>>>(p);
    }
    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 template <int HEAD_DIM>
 static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed) {
    printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d ... ", B, Hq, Hkv, kv_len, page_size, HEAD_DIM);
    fflush(stdout);
    int max_pages = (kv_len + page_size - 1) / page_size;
    int n_phys_pages = B * max_pages;
    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
    size_t sz_o  = sz_q;
    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
    int max_splits = 32;
    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
    bf16 *d_q, *d_o_paged, *d_o_ref;
    bf16 *d_k_pool, *d_v_pool;
    int64_t* d_pt;
    float *d_op, *d_ml;
    cudaMalloc(&d_q, sz_q);
    cudaMalloc(&d_o_paged, sz_o);
    cudaMalloc(&d_o_ref, sz_o);
    cudaMalloc(&d_k_pool, sz_kv);
    cudaMalloc(&d_v_pool, sz_kv);
    cudaMalloc(&d_pt, sz_pt);
    cudaMalloc(&d_op, sz_op);
    cudaMalloc(&d_ml, sz_ml);
    srand(seed);
    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
    bf16* h_q = (bf16*)malloc(sz_q);
    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
        h_q[i] = __float2bfloat16(rnd());
    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
    bf16* h_k_pool = (bf16*)malloc(sz_kv);
    bf16* h_v_pool = (bf16*)malloc(sz_kv);
    size_t ps = (size_t)page_size * Hkv * HEAD_DIM;
    for (int pg = 0; pg < n_phys_pages; pg++) {
        for (int off = 0; off < page_size; off++) {
            for (int h = 0; h < Hkv; h++) {
                for (int d = 0; d < HEAD_DIM; d++) {
                    float v = sinf((float)(pg * 7919 + off * 1049 + h * 331 + d));
                    size_t idx = (size_t)pg * ps + (size_t)off * Hkv * HEAD_DIM + h * HEAD_DIM + d;
                    h_k_pool[idx] = __float2bfloat16(v);
                    h_v_pool[idx] = __float2bfloat16(v * 0.3f);
                }
            }
        }
    }
    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
    int64_t* h_pt = (int64_t*)malloc(sz_pt);
    int next_pg = 0;
    for (int b = 0; b < B; b++)
        for (int p = 0; p < max_pages; p++)
            h_pt[b * max_pages + p] = next_pg++;
    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
    bf16* h_k_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
    bf16* h_v_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
    gather_kv_cpu(h_k_pool, h_v_pool, h_pt, B, Hkv, kv_len, page_size, HEAD_DIM, h_k_cont, h_v_cont);
    float* h_q_f = (float*)malloc((size_t)B * Hq * HEAD_DIM * sizeof(float));
    float* h_k_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
    float* h_v_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
    for (int i = 0; i < B * kv_len * Hkv * HEAD_DIM; i++) {
        h_k_f[i] = bf2f(h_k_cont[i]);
        h_v_f[i] = bf2f(h_v_cont[i]);
    }
    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
    cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv, 1, kv_len, HEAD_DIM, 0, 0);
    float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
    PagedAttentionParams<bf16, float> p;
    p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1;
    p.kv_len = kv_len; p.head_dim = HEAD_DIM;
    p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
    p.num_splits = 1; p.scale = scale_val;
    p.page_size = page_size; p.max_pages = max_pages;
    p.page_table = d_pt;
    p.k_cache = d_k_pool; p.v_cache = d_v_pool;
    p.q = d_q; p.mask = nullptr; p.o = d_o_paged;
    p.o_part = d_op; p.ml_part = d_ml;
    launch_paged_decode<HEAD_DIM>(p);
    cudaDeviceSynchronize();
    bf16* h_o_bf16 = (bf16*)malloc(sz_o);
    cudaMemcpy(h_o_bf16, d_o_paged, sz_o, cudaMemcpyDeviceToHost);
    float* h_o_paged = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
        h_o_paged[i] = __bfloat162float(h_o_bf16[i]);
    float max_err = 0.0f;
    int bad_idx = -1;
    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
        float e = fabsf(h_o_paged[i] - h_o_ref[i]);
        if (e > max_err) { max_err = e; bad_idx = i; }
    }
    bool pass = max_err < 0.02f;
    if (pass) {
        printf("PASS (max_abs_err=%.4e)\n", max_err);
    } else {
        int b = bad_idx / (Hq * HEAD_DIM);
        int h = (bad_idx / HEAD_DIM) % Hq;
        int d = bad_idx % HEAD_DIM;
        printf("FAIL (max_abs_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n",
               max_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]);
        printf("  ref[0..7]:");
        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
            printf(" %.4f", h_o_ref[i]);
        printf("\n  got[0..7]:");
        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
            printf(" %.4f", h_o_paged[i]);
        printf("\n");
    }
    free(h_q); free(h_k_pool); free(h_v_pool); free(h_pt);
    free(h_k_cont); free(h_v_cont);
    free(h_q_f); free(h_k_f); free(h_v_f);
    free(h_o_ref); free(h_o_bf16); free(h_o_paged);
    cudaFree(d_q); cudaFree(d_o_paged); cudaFree(d_o_ref);
    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
    cudaFree(d_op); cudaFree(d_ml);
    return pass ? 0 : 1;
 }
 struct TestCase {
    int head_dim;
    int B, Hq, Hkv, kv_len, page_size, seed;
 };
 static const TestCase TESTS[] = {
    {128, 1, 1, 1, 8, 128, 1},
    {128, 1, 4, 4, 128, 128, 2},
    {128, 2, 4, 4, 256, 128, 3},
    {128, 1, 4, 1, 64, 64, 4},
    {128, 1, 8, 2, 64, 128, 5},
    {128, 2, 16, 4, 128, 128, 6},
    {64, 1, 4, 2, 32, 128, 7},
    {256, 1, 2, 1, 16, 128, 8},
    {32, 1, 4, 2, 32, 64, 9},
    {128, 3, 8, 2, 256, 128, 10},
    {128, 2, 32, 8, 512, 128, 11},
 #ifndef ASTRAI_NO_MMA
    {128, 1, 16, 2, 256, 128, 12},
    {128, 2, 32, 4, 512, 128, 13},
 #endif
 };
 static int dispatch_test(const TestCase& tc) {
    bool matched = false;
    int r = 0;
    dispatch_by_head_dim(tc.head_dim, [&]<int D>() {
        matched = true;
        r = run_test<D>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
    });
    return matched ? r : 1;
 }
 // Warmed-up, CUDA-event timed sweep over paged decode configs.
 // Bytes = K + V read through page table (B*Hk*kv*D each), bf16.
 template <int HEAD_DIM>
 static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
    int max_pages = (kv_len + page_size - 1) / page_size;
    int n_phys_pages = B * max_pages;
    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
    int max_splits = 32;
    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
    int64_t* d_pt;
    float *d_op, *d_ml;
    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
    cudaMalloc(&d_pt, sz_pt);
    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
    int64_t* h_pt = (int64_t*)malloc(sz_pt);
    int next_pg = 0;
    for (int b = 0; b < B; b++)
        for (int p = 0; p < max_pages; p++)
            h_pt[b * max_pages + p] = next_pg++;
    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
    free(h_pt);
    float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
    PagedAttentionParams<bf16, float> pa;
    pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1;
    pa.kv_len = kv_len; pa.head_dim = HEAD_DIM;
    pa.use_mask = 0; pa.is_causal = 0; pa.causal_offset = 0;
    pa.num_splits = 1; pa.scale = scale_val;
    pa.page_size = page_size; pa.max_pages = max_pages;
    pa.page_table = d_pt;
    pa.k_cache = d_k_pool; pa.v_cache = d_v_pool;
    pa.q = d_q; pa.mask = nullptr; pa.o = d_o;
    pa.o_part = d_op; pa.ml_part = d_ml;
    const int WARMUP = 10, ITERS = 100;
    auto launch = [&]() { launch_paged_decode<HEAD_DIM>(pa); };
    double flops = 4.0 * B * Hq * (double)kv_len * HEAD_DIM;
    size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
    double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
    BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
    char cfg[64];
    snprintf(cfg, sizeof(cfg),
             "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d page=%3d",
             B, Hq, Hkv, 1, kv_len, HEAD_DIM, page_size);
    print_bench_row(cfg, r);
    free(tmp);
    cudaFree(d_q); cudaFree(d_o);
    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
    cudaFree(d_op); cudaFree(d_ml);
 }
 static void bench() {
    printf("\n===== PAGED DECODE BENCH =====\n");
    print_bench_header();
    bench_config<128>(1, 32, 4, 512, 128);
    bench_config<128>(1, 32, 4, 1024, 128);
    bench_config<128>(1, 32, 4, 2048, 128);
    bench_config<128>(1, 32, 4, 4096, 128);
    bench_config<128>(16, 32, 4, 2048, 128);
    bench_config<128>(32, 32, 4, 1024, 128);
 }
 int main() {
    int n = sizeof(TESTS) / sizeof(TESTS[0]);
    int fail = 0;
    printf("=== Paged Decode vs CPU reference (%d cases) ===\n\n", n);
    for (int i = 0; i < n; i++) {
        fail += dispatch_test(TESTS[i]);
        if (fail) break;
    }
    if (fail) {
        printf("\nFAILED (%d/%d tests failed)\n", fail, n);
        return fail;
    }
    printf("\nAll %d tests passed!\n", n);
    bench();
    return 0;
 }
@@ -0,0 +1,176 @@
 /*
 Pure-C test:
 nvcc -I csrc -arch=sm_89 -O3 \
    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
    csrc/tests/attn_prefill_test.cu -o test && ./test
 */
 #include "test_utils.cuh"
 #include "../kernels/attn_prefill_split_q.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "../kernels/attn_prefill_split_q_mma.cuh"
 #endif
 // Launch the production prefill path (tensor-core MMA on sm_80+, else the
 // scalar fallback), mirroring dispatch_prefill() in attn_prefill.cu.
 template <int HEAD_DIM>
 static void launch_prefill(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
    constexpr int WARPS = 4, BR = 16;
    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);
    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
 }
 static void dispatch_prefill(AttentionParams<bf16>& p) {
    switch (p.head_dim) {
        case 64:  launch_prefill<64>(p);  break;
        case 128: launch_prefill<128>(p); break;
        default:  printf("bench: unsupported D=%d\n", p.head_dim);
    }
 }
 // Warmed-up, CUDA-event timed throughput sweep over the production MMA path.
 // Reports per-call latency and effective tensor-core TFLOP/s (2 matmuls:
 // QK^T and P@V, each 2*B*Hq*ql*kl*D flops; halved for causal).
 static void bench() {
    const int cfgs[][7] = {
        {1,32,4,512,512,128,0},
        {1,32,4,1024,1024,128,0},
        {1,32,4,2048,2048,128,0},
        {1,32,4,2048,2048,128,1},
        {4,32,4,2048,2048,128,1},
        {1,32,4,4096,4096,128,1},
    };
    int n = sizeof(cfgs)/sizeof(cfgs[0]);
    const int WARMUP = 10, ITERS = 50;
    printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
    printf("%-46s | %10s | %10s | %10s\n",
           "config", "latency", "bandwidth", "throughput");
    printf("---------------------------------------------------------------"
           "----------------------------\n");
    for (int ci = 0; ci < n; ci++) {
        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
        int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
        size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*D;
        bf16 *dQ,*dK,*dV,*dO,*tmp;
        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
        AttentionParams<bf16> p;
        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
        p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
        p.scale=1.0f/sqrtf((float)D);
        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
        for (int i=0;i<WARMUP;i++) dispatch_prefill(p);
        cudaDeviceSynchronize();
        cudaError_t err=cudaGetLastError();
        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
        cudaEventRecord(s);
        for (int i=0;i<ITERS;i++) dispatch_prefill(p);
        cudaEventRecord(e); cudaEventSynchronize(e);
        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
        double flops = 4.0*B*Hq*(double)ql*kl*D;
        if (causal) flops *= 0.5;
        double tflops = flops/(ms*1e-3)/1e12;
        // HBM traffic: Q + O (B*Hq*ql*D each) + K + V (B*Hk*kl*D each), bf16.
        double bytes = 2.0 * (2.0*nQ + 2.0*nKV);
        double gbps = bytes/(ms*1e-3)/1e9;
        char cfg[64];
        snprintf(cfg, sizeof(cfg),
                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
                 B,Hq,Hk,ql,kl,D,causal);
        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
               cfg, ms, gbps, tflops);
        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
    }
 }
 int main() {
    const int configs[][7] = {
        {1,2,1,64,128,64,0},     // tiny: B,Hq,Hk,q,kv,D,causal
        {1,32,4,512,512,128,0},  // standard
        {1,32,4,128,256,128,0},  // medium
        {1,4,2,256,256,128,1},   // causal
    };
    int n_configs = sizeof(configs) / sizeof(configs[0]);
    for (int ci = 0; ci < n_configs; ci++) {
        int B=configs[ci][0], Hq=configs[ci][1], Hk=configs[ci][2];
        int ql=configs[ci][3], kl=configs[ci][4], D=configs[ci][5];
        int causal=configs[ci][6];
        printf("=== B=%d Hq=%d Hk=%d q=%d kv=%d D=%d causal=%d ===\n",
               B,Hq,Hk,ql,kl,D,causal);
        size_t nQ = B*Hq*ql*D, nKV = B*Hk*kl*D;
        float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
        for (size_t i=0;i<nQ;i++) hQ[i]=randf();
        for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
        bf16 *dQ,*dK,*dV,*dO,*tmp;
        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
        tmp=new bf16[max(nQ,nKV)];
        for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
        AttentionParams<bf16> p;
        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
        p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
        p.scale=1.0f/sqrtf((float)D);
        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
        double t0=now_ms();
        dispatch_prefill(p);
        cudaDeviceSynchronize();
        double kms=now_ms()-t0;
        cudaError_t err=cudaGetLastError();
        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
        bf16* hOut=new bf16[nQ];
        cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
        float* ref=new float[nQ];
        cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal, 0);
        float max_err=0;
        for (size_t i=0;i<nQ;i++) {
            float d=fabsf(bf2f(hOut[i])-ref[i]);
            if(d>max_err) max_err=d;
        }
        printf("kernel: %.3f ms  max_err: %.6e\n\n",kms,max_err);
        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
        delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
    }
    printf("All tests passed!\n");
    bench();
    return 0;
 }
@@ -0,0 +1,154 @@
 #pragma once
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>
 #include <chrono>
 #include <cuda_bf16.h>
 using bf16 = __nv_bfloat16;
 inline bf16 f2bf(float x) { return __float2bfloat16(x); }
 inline float bf2f(bf16 x) { return __bfloat162float(x); }
 inline float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
 inline double now_ms() {
    using namespace std::chrono;
    return duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
 }
 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;
    if (n > tiles_total) n = tiles_total;
    if (n > 32) n = 32;
    if (n < 1) n = 1;
    return n;
 }
 #define CUDA_CHECK(call) \
    do { \
        cudaError_t _e = (call); \
        if (_e != cudaSuccess) { \
            printf("CUDA error %s at %s:%d\n", cudaGetErrorString(_e), __FILE__, __LINE__); \
            exit(1); \
        } \
    } while (0)
 struct BenchResult {
    float ms;
    double gbps;
    double tflops;
 };
 template <typename Fn>
 BenchResult bench_kernel(Fn launch, int warmup, int iters,
                         double flops, double bytes) {
    for (int i = 0; i < warmup; i++) launch();
    cudaDeviceSynchronize();
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        printf("CUDA error before bench: %s\n", cudaGetErrorString(err));
        return {0, 0, 0};
    }
    cudaEvent_t s, e;
    cudaEventCreate(&s); cudaEventCreate(&e);
    cudaEventRecord(s);
    for (int i = 0; i < iters; i++) launch();
    cudaEventRecord(e); cudaEventSynchronize(e);
    float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= iters;
    cudaEventDestroy(s); cudaEventDestroy(e);
    return {ms, bytes / (ms * 1e-3) / 1e9, flops / (ms * 1e-3) / 1e12};
 }
 inline void print_bench_header() {
    printf("%-46s | %10s | %10s | %10s\n",
           "config", "latency", "bandwidth", "throughput");
    printf("---------------------------------------------------------------"
           "----------------------------\n");
 }
 inline void print_bench_row(const char* cfg, const BenchResult& r) {
    printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
           cfg, r.ms, r.gbps, r.tflops);
 }
 template <int... Ds>
 struct _HeadSwitch;
 template <int D>
 struct _HeadSwitch<D> {
    template <typename Fn>
    static void call(int hd, Fn&& fn) { if (hd == D) fn.template operator()<D>(); }
 };
 template <int D, int... Rest>
 struct _HeadSwitch<D, Rest...> {
    template <typename Fn>
    static void call(int hd, Fn&& fn) {
        if (hd == D) fn.template operator()<D>();
        else _HeadSwitch<Rest...>::call(hd, fn);
    }
 };
 // Default set: 32, 64, 128, 256
 template <typename Fn>
 void dispatch_by_head_dim(int head_dim, Fn&& fn) {
    _HeadSwitch<32, 64, 128, 256>::call(head_dim, fn);
 }
 // Generic CPU reference for multi-query / grouped-query attention.
 // Tensor shapes (all float*):
 //   Q : [B, Hq, q_len, D]
 //   K : [B, Hk, kv_len, D]
 //   V : [B, Hk, kv_len, D]
 //   O : [B, Hq, q_len, D]
 // mask: if q_len == 1, shape is [B, kv_len]; otherwise mask is not supported.
 static void cpu_attention_ref(
    const float* Q, const float* K, const float* V, const bool* mask,
    float* O, int B, int Hq, int Hk, int q_len, int kv_len, int D,
    int is_causal, int causal_offset
 ) {
    float scale = 1.0f / sqrtf((float)D);
    int n_rep = Hq / Hk;
    for (int b = 0; b < B; b++) {
        for (int h = 0; h < Hq; h++) {
            int kv_h = h / n_rep;
            for (int qi = 0; qi < q_len; qi++) {
                float mv = -INFINITY, sv = 0.0f;
                float accum[256] = {0.0f};
                int lim = kv_len;
                if (is_causal) {
                    int c = qi + causal_offset + 1;
                    lim = (c < kv_len) ? c : kv_len;
                }
                for (int kj = 0; kj < lim; kj++) {
                    if (mask != nullptr && q_len == 1) {
                        if (!mask[b * kv_len + kj]) continue;
                    }
                    float dot = 0.0f;
                    size_t q_idx = ((size_t)b * Hq + h) * q_len + qi;
                    size_t kv_idx = ((size_t)b * Hk + kv_h) * kv_len + kj;
                    for (int d = 0; d < D; d++)
                        dot += Q[q_idx * D + d] * K[kv_idx * D + d];
                    dot *= scale;
                    float nm = fmaxf(mv, dot);
                    float a = expf(mv - nm);
                    float b_exp = expf(dot - nm);
                    sv = sv * a + b_exp;
                    for (int d = 0; d < D; d++)
                        accum[d] = accum[d] * a + V[kv_idx * D + d] * b_exp;
                    mv = nm;
                }
                float inv = 1.0f / sv;
                size_t o_idx = ((size_t)b * Hq + h) * q_len + qi;
                for (int d = 0; d < D; d++)
                    O[o_idx * D + d] = accum[d] * inv;
            }
        }
    }
 }
@@ -1,12 +1,13 @@
 services:
  server:
-    build: .
+    build:
-    image: astrai:latest
+      context: .
      dockerfile: Dockerfile
    user: "${UID:-1000}:${GID:-1000}"
    ports:
      - "8000:8000"
    volumes:
      - ./params:/app/params:ro
      - ./checkpoints:/app/checkpoints
    command: python -m scripts.tools.server --port 8000 --device cuda
    deploy:
      resources:
@@ -25,13 +26,14 @@ services:
  server-cpu:
    profiles: [cpu]
-    build: .
+    build:
-    image: astrai:latest
+      context: .
      dockerfile: Dockerfile
    user: "${UID:-1000}:${GID:-1000}"
    ports:
      - "8000:8000"
    volumes:
      - ./params:/app/params:ro
      - ./checkpoints:/app/checkpoints
    command: python -m scripts.tools.server --port 8000 --device cpu
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
@@ -9,8 +9,8 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
    "h5py==3.15.1",
-    "numpy==2.3.2",
+    "numpy==2.4.4",
-    "torch==2.7.1",
+    "torch==2.11.0",
    "tokenizers==0.21.4",
    "tqdm==4.67.1",
    "safetensors==0.5.3",
@@ -37,7 +37,7 @@ dev = ["pytest==9.0.2", "ruff"]
 where = ["."]
 [tool.pip]
-extra-index-url = "https://download.pytorch.org/whl/cu126"
+extra-index-url = "https://download.pytorch.org/whl/cu128"
 [tool.setuptools.dynamic]
 version = { attr = "astrai.__version__" }
@@ -11,7 +11,6 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 def generate_text():
    # Load model from pretrained
    model = AutoModel.from_pretrained(PARAMETER_ROOT)
    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
    model.to(device="cuda", dtype=torch.bfloat16)
@@ -22,16 +21,15 @@ def generate_text():
        model=model,
        tokenizer=tokenizer,
    )
-    response = engine.generate(
+    for token in engine.generate(
        prompt=query,
-        stream=False,
+        stream=True,
        max_tokens=2048,
        temperature=0.8,
        top_p=0.95,
        top_k=50,
-    )
+    ):
-
+        print(token, end="", flush=True)
    print(response)
 if __name__ == "__main__":
--- a/Show More
+++ b/Show More
		`@@ -0,0 +1,2 @@`
							`# Source directory for CUDA kernels — build-time only.`
							`# Compiled .so files live in astrAI/_ext/.`